Refactored random byte generation functions.

conquest.nimble

@@ -26,4 +26,4 @@ requires "parsetoml >= 0.7.2"
 requires "nimcrypto >= 0.6.4"
 requires "tiny_sqlite >= 0.2.0"
 requires "prologue >= 0.6.6" 
-requires "winim >= 3.9.4"
+requires "winim >= 3.9.4"

src/agent/core/http.nim

@@ -1,7 +1,7 @@
 import httpclient, json, strformat, strutils, asyncdispatch, base64, tables, parsetoml, random
 
 import ../../common/[types, utils, profile]
-import sugar
+
 proc httpGet*(ctx: AgentCtx, heartbeat: seq[byte]): string = 
 
     let client = newAsyncHttpClient(userAgent = ctx.profile.getString("agent.user-agent"))

src/agent/protocol/heartbeat.nim

@@ -12,7 +12,7 @@ proc createHeartbeat*(ctx: AgentCtx): Heartbeat =
             size: 0'u32,
             agentId: string.toUuid(ctx.agentId),
             seqNr: 0'u32,  
-            iv: generateIV(),
+            iv: generateBytes(Iv),
             gmac: default(AuthenticationTag)
         ), 
         listenerId: string.toUuid(ctx.listenerId),

src/agent/protocol/registration.nim

@@ -203,7 +203,7 @@ proc collectAgentMetadata*(ctx: AgentCtx): AgentRegistrationData =
             size: 0'u32,
             agentId: string.toUuid(ctx.agentId),
             seqNr: nextSequence(string.toUuid(ctx.agentId)),                              
-            iv: generateIV(),
+            iv: generateBytes(Iv),
             gmac: default(AuthenticationTag)
         ), 
         agentPublicKey: ctx.agentPublicKey,

src/agent/protocol/result.nim

@@ -11,7 +11,7 @@ proc createTaskResult*(task: Task, status: StatusType, resultType: ResultType, r
             size: 0'u32,
             agentId: task.header.agentId,
             seqNr: nextSequence(task.header.agentId), 
-            iv: generateIV(),
+            iv: generateBytes(Iv),
             gmac: default(array[16, byte])
         ), 
         taskId: task.taskId,

src/common/crypto.nim

@@ -1,4 +1,4 @@
-import system
+import macros, system
 import nimcrypto
 
 import ./[utils, types]
@@ -7,18 +7,11 @@ import ./[utils, types]
     Symmetric AES256 GCM encryption for secure C2 traffic
     Ensures both confidentiality and integrity of the packet 
 ]#
-proc generateIV*(): Iv =
-    # Generate a random 98-bit (12-byte) initialization vector for AES-256 GCM mode
-    var iv: Iv
-    if randomBytes(iv) != sizeof(Iv): 
-        raise newException(CatchableError, "Failed to generate IV.")
-    return iv 
-
-proc generateKey*(): Key = 
-    var key: Key 
-    if randomBytes(key) != sizeof(Key):
-        raise newException(CatchableError, "Failed to generate IV.")
-    return key
+proc generateBytes*(T: typedesc[Key | Iv]): array =
+    var bytes: T
+    if randomBytes(bytes) != sizeof(T): 
+        raise newException(CatchableError, "Failed to generate byte array.")
+    return bytes
 
 proc encrypt*(key: Key, iv: Iv, data: seq[byte], sequenceNumber: uint32 = 0): (seq[byte], AuthenticationTag) =
     
@@ -97,7 +90,7 @@ proc wipeKey*(data: var openArray[byte]) =
 
 # Key pair generation
 proc generateKeyPair*(): KeyPair = 
-    let privateKey = generateKey() 
+    let privateKey = generateBytes(Key) 
     return KeyPair(
         privateKey: privateKey, 
         publicKey: getPublicKey(privateKey)
@@ -158,4 +151,4 @@ proc loadKeyPair*(keyFile: string): KeyPair =
     except IOError: 
         let keyPair = generateKeyPair() 
         writeKeyToDisk(keyFile, keyPair.privateKey)
-        return keyPair
+        return keyPair

src/server/core/builder.nim

@@ -24,8 +24,8 @@ proc serializeConfiguration(cq: Conquest, listener: Listener, sleep: int): seq[b
     packer.reset() 
 
     # Encrypt profile configuration data with a newly generated encryption key
-    var aesKey = generateKey() 
-    let iv = generateIV()
+    var aesKey = generateBytes(Key) 
+    let iv = generateBytes(Iv)
 
     let (encData, gmac) = encrypt(aesKey, iv, data)
 

src/server/core/logger.nim

@@ -2,7 +2,7 @@ import times, strformat, strutils, prompt, terminal
 import std/[dirs, paths]
 
 import ../globals
-import ../../common/[types, profile]
+import ../../common/types
 
 proc makeAgentLogDirectory*(cq: Conquest, agentId: string): bool = 
     try: 

src/server/core/server.nim

@@ -2,7 +2,7 @@ import prompt, terminal, argparse, parsetoml
 import strutils, strformat, system, tables
 
 import ./[agent, listener, builder]
-import ../[globals, utils]
+import ../globals
 import ../db/database
 import ../core/logger
 import ../../common/[types, crypto, profile]

src/server/globals.nim

@@ -1,4 +1,3 @@
-import os 
 import ../common/types
 
 # Global server context

src/server/protocol/parser.nim

@@ -106,7 +106,7 @@ proc createTask*(cq: Conquest, command: Command, arguments: seq[string]): Task =
     taskHeader.size = 0'u32
     taskHeader.agentId = string.toUuid(cq.interactAgent.agentId)
     taskHeader.seqNr = nextSequence(taskHeader.agentId)
-    taskHeader.iv = generateIV() # Generate a random IV for AES-256 GCM
+    taskHeader.iv = generateBytes(Iv) # Generate a random IV for AES-256 GCM
     taskHeader.gmac = default(AuthenticationTag)
 
     task.header = taskHeader