Cleaned up parts of the serialization by removing redundant code.

src/server/api/handlers.nim

@@ -2,7 +2,7 @@ import terminal, strformat, strutils, sequtils, tables, json, times, base64, sys
 
 import ../[utils, globals]
 import ../db/database
-import ../task/packer
+import ../message/packer
 import ../../common/[types, utils]
 
 #[
@@ -58,11 +58,9 @@ proc getTasks*(checkinData: seq[byte]): seq[seq[byte]] =
 
         # Update the last check-in date for the accessed agent
         cq.agents[agentId].latestCheckin = cast[int64](timestamp).fromUnix().local()
-        # if not cq.dbUpdateCheckin(agent.toUpperAscii, now().format("dd-MM-yyyy HH:mm:ss")):
-        #    return nil
 
         # Return tasks
-        for task in cq.agents[agentId].tasks.mitems: # Iterate over mutable items in order to modify GMAC
+        for task in cq.agents[agentId].tasks.mitems: # Iterate over agents as mutable items in order to modify GMAC tag
             let taskData = cq.serializeTask(task)
             result.add(taskData)
         

src/server/core/agent.nim

@@ -1,7 +1,7 @@
 import terminal, strformat, strutils, tables, times, system, osproc, streams, base64
 
+import ./task
 import ../utils
-import ../task/dispatcher
 import ../db/database
 import ../../common/[types, utils]
 

src/server/core/task.nim

@@ -1,6 +1,7 @@
 import times, strformat, terminal, tables, json, sequtils, strutils
-import ./[parser]
+
 import ../utils
+import ../message/parser
 import ../../modules/manager
 import ../../common/[types, utils]
 
@@ -72,7 +73,7 @@ proc handleAgentCommand*(cq: Conquest, input: string) =
     try: 
         let 
             command = getCommandByName(parsedArgs[0])
-            task = cq.parseTask(command, parsedArgs[1..^1])
+            task = cq.createTask(command, parsedArgs[1..^1])
 
         # Add task to queue
         cq.interactAgent.tasks.add(task)

src/server/message/packer.nim

@@ -27,7 +27,7 @@ proc serializeTask*(cq: Conquest, task: var Task): seq[byte] =
     task.header.gmac = gmac
 
     # Serialize header 
-    let header = packer.packHeader(task.header, uint32(payload.len))
+    let header = packer.serializeHeader(task.header, uint32(payload.len))
 
     return header & encData
 
@@ -35,28 +35,15 @@ proc deserializeTaskResult*(cq: Conquest, resultData: seq[byte]): TaskResult =
 
     var unpacker = initUnpacker(resultData.toString)
 
-    let header = unpacker.unpackHeader()
+    let header = unpacker.deserializeHeader()
 
     # Packet Validation
-    if header.magic != MAGIC:
-        raise newException(CatchableError, "Invalid magic bytes.")
-
-    if header.packetType != cast[uint8](MSG_RESPONSE): 
-        raise newException(CatchableError, "Invalid packet type for task result, expected MSG_RESPONSE.")
-
-    # Validate sequence number 
-    if not validateSequence(header.agentId, header.seqNr, header.packetType): 
-        raise newException(CatchableError, "Invalid sequence number.")
+    validatePacket(header, cast[uint8](MSG_RESPONSE)) 
 
     # Decrypt payload 
     let payload = unpacker.getBytes(int(header.size))
+    let decData= validateDecryption(cq.agents[uuidToString(header.agentId)].sessionKey, header.iv, payload, header.seqNr, header)
 
-    let (decData, gmac) = decrypt(cq.agents[uuidToString(header.agentId)].sessionKey, header.iv, payload, header.seqNr)
-
-    # Verify that the authentication tags match, which ensures the integrity of the decrypted data and AAD
-    if gmac != header.gmac:
-        raise newException(CatchableError, "Invalid authentication tag (GMAC) for task result.")
-    
     # Deserialize decrypted data
     unpacker = initUnpacker(decData.toString)
 
@@ -86,18 +73,10 @@ proc deserializeNewAgent*(cq: Conquest, data: seq[byte]): Agent =
 
     var unpacker = initUnpacker(data.toString)
 
-    let header= unpacker.unpackHeader()
+    let header= unpacker.deserializeHeader()
 
     # Packet Validation
-    if header.magic != MAGIC: 
-        raise newException(CatchableError, "Invalid magic bytes.")
-
-    if header.packetType != cast[uint8](MSG_REGISTER): 
-        raise newException(CatchableError, "Invalid packet type for agent registration, expected MSG_REGISTER.")
-
-    # Validate sequence number 
-    if not validateSequence(header.agentId, header.seqNr, header.packetType): 
-        raise newException(CatchableError, "Invalid sequence number.")
+    validatePacket(header, cast[uint8](MSG_REGISTER)) 
 
     # Key exchange
     let agentPublicKey = unpacker.getKey()
@@ -105,11 +84,7 @@ proc deserializeNewAgent*(cq: Conquest, data: seq[byte]): Agent =
     
     # Decrypt payload 
     let payload = unpacker.getBytes(int(header.size)) 
-    let (decData, gmac) = decrypt(sessionKey, header.iv, payload, header.seqNr)
-
-    # Verify that the authentication tags match, which ensures the integrity of the decrypted data and AAD
-    if gmac != header.gmac:
-        raise newException(CatchableError, "Invalid authentication tag (GMAC) for agent registration.")
+    let decData= validateDecryption(sessionKey, header.iv, payload, header.seqNr, header)
 
     # Deserialize decrypted data
     unpacker = initUnpacker(decData.toString)
@@ -148,26 +123,14 @@ proc deserializeHeartbeat*(cq: Conquest, data: seq[byte]): Heartbeat =
 
     var unpacker = initUnpacker(data.toString)
 
-    let header = unpacker.unpackHeader()
+    let header = unpacker.deserializeHeader()
 
     # Packet Validation
-    if header.magic != MAGIC: 
-        raise newException(CatchableError, "Invalid magic bytes.")
-
-    if header.packetType != cast[uint8](MSG_HEARTBEAT):
-        raise newException(CatchableError, "Invalid packet type for checkin request, expected MSG_HEARTBEAT.")
-
-    # Validate sequence number 
-    if not validateSequence(header.agentId, header.seqNr, header.packetType): 
-        raise newException(CatchableError, "Invalid sequence number.")
+    validatePacket(header, cast[uint8](MSG_HEARTBEAT)) 
 
     # Decrypt payload
     let payload = unpacker.getBytes(int(header.size))
-    let (decData, gmac) = decrypt(cq.agents[uuidToString(header.agentId)].sessionKey, header.iv, payload, header.seqNr)
-
-    # Verify that the authentication tags match, which ensures the integrity of the decrypted data and AAD
-    if gmac != header.gmac:
-        raise newException(CatchableError, "Invalid authentication tag (GMAC) for heartbeat.")
+    let decData= validateDecryption(cq.agents[uuidToString(header.agentId)].sessionKey, header.iv, payload, header.seqNr, header)
 
     # Deserialize decrypted data
     unpacker = initUnpacker(decData.toString)

src/server/message/parser.nim

@@ -1,4 +1,5 @@
 import strutils, strformat, times
+
 import ../utils
 import ../../common/[types, utils, sequence, crypto]
 
@@ -72,7 +73,7 @@ proc parseArgument*(argument: Argument, value: string): TaskArg =
     
     return result
 
-proc parseTask*(cq: Conquest, command: Command, arguments: seq[string]): Task = 
+proc createTask*(cq: Conquest, command: Command, arguments: seq[string]): Task = 
 
     # Construct the task payload prefix
     var task: Task