Implemented communication with custom binary structure instead of JSON requests

This commit is contained in:
Jakob Friedl
2025-07-19 16:49:27 +02:00
parent d22ad0bd0c
commit 99f55cc04f
19 changed files with 524 additions and 433 deletions

View File

@@ -2,6 +2,7 @@ import terminal, strformat, strutils, sequtils, tables, json, times, base64, sys
import ../[utils, globals]
import ../db/database
import ../task/packer
import ../../common/types
# Utility functions
@@ -40,6 +41,8 @@ proc getTasks*(listener, agent: string): seq[seq[byte]] =
{.cast(gcsafe).}:
var result: seq[seq[byte]]
# Check if listener exists
if not cq.dbListenerExists(listener.toUpperAscii):
cq.writeLine(fgRed, styleBright, fmt"[-] Task-retrieval request made to non-existent listener: {listener}.", "\n")
@@ -55,40 +58,47 @@ proc getTasks*(listener, agent: string): seq[seq[byte]] =
# if not cq.dbUpdateCheckin(agent.toUpperAscii, now().format("dd-MM-yyyy HH:mm:ss")):
# return nil
# Return tasks
return cq.agents[agent.toUpperAscii].tasks
# Return tasks
for task in cq.agents[agent.toUpperAscii].tasks:
let taskData = serializeTask(task)
result.add(taskData)
return result
proc handleResult*(listener, agent, task: string, taskResult: TaskResult) =
proc handleResult*(resultData: seq[byte]) =
{.cast(gcsafe).}:
let
taskResult = deserializeTaskResult(resultData)
taskId = uuidToString(taskResult.taskId)
agentId = uuidToString(taskResult.agentId)
listenerId = uuidToString(taskResult.listenerId)
let date: string = now().format("dd-MM-yyyy HH:mm:ss")
cq.writeLine(fgBlack, styleBright, fmt"[{date}] [*] ", resetStyle, fmt"{$resultData.len} bytes received.")
if taskResult.status == cast[uint8](STATUS_FAILED):
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgRed, styleBright, " [-] ", resetStyle, fmt"Task {task} failed.")
case cast[StatusType](taskResult.status):
of STATUS_COMPLETED:
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgGreen, " [+] ", resetStyle, fmt"Task {taskId} completed.")
if taskResult.data.len != 0:
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgRed, styleBright, " [-] ", resetStyle, "Output:")
of STATUS_FAILED:
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgRed, styleBright, " [-] ", resetStyle, fmt"Task {taskId} failed.")
case cast[ResultType](taskResult.resultType):
of RESULT_STRING:
if int(taskResult.length) > 0:
cq.writeLine(fgBlack, styleBright, fmt"[{date}] [*] ", resetStyle, "Output:")
# Split result string on newline to keep formatting
# for line in decode(taskResult.data).split("\n"):
# cq.writeLine(line)
else:
cq.writeLine()
for line in taskResult.data.toString().split("\n"):
cq.writeLine(line)
else:
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgGreen, " [+] ", resetStyle, fmt"Task {task} finished.")
if taskResult.data.len != 0:
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgGreen, " [+] ", resetStyle, "Output:")
of RESULT_BINARY:
# Write binary data to a file
cq.writeLine()
# Split result string on newline to keep formatting
# for line in decode(taskResult.data).split("\n"):
# cq.writeLine(line)
else:
cq.writeLine()
of RESULT_NO_OUTPUT:
cq.writeLine()
# Update task queue to include all tasks, except the one that was just completed
# cq.agents[agent].tasks = cq.agents[agent].tasks.filterIt(it.id != task)
return
cq.agents[agentId].tasks = cq.agents[agentId].tasks.filterIt(it.taskId != taskResult.taskId)

View File

@@ -94,7 +94,7 @@ proc getTasks*(ctx: Context) {.async.} =
try:
var response: seq[byte]
let tasks = getTasks(listener, agent)
let tasks: seq[seq[byte]] = getTasks(listener, agent)
if tasks.len <= 0:
resp "", Http200
@@ -134,21 +134,15 @@ proc postResults*(ctx: Context) {.async.} =
task = ctx.getPathParams("task")
# Check headers
# If POST data is not JSON data, return 404 error code
if ctx.request.contentType != "application/json":
# If POST data is not binary data, return 404 error code
if ctx.request.contentType != "application/octet-stream":
resp "", Http404
return
try:
let
taskResultJson: JsonNode = parseJson(ctx.request.body)
taskResult: TaskResult = taskResultJson.to(TaskResult)
# Handle and display task result
handleResult(listener, agent, task, taskResult)
handleResult(ctx.request.body.toBytes())
except CatchableError:
# JSON data is invalid or does not match the expected format (described above)
resp "", Http404
return

View File

@@ -1,5 +1,5 @@
import times, strformat, terminal, tables, json, sequtils, strutils
import ./[parser, packer]
import ./[parser]
import ../utils
import ../../common/types
@@ -177,12 +177,9 @@ proc handleAgentCommand*(cq: Conquest, input: string) =
let
command = getCommandFromTable(parsedArgs[0], commands)
task = cq.parseTask(command, parsedArgs[1..^1])
taskData: seq[byte] = cq.serializeTask(task)
# cq.writeLine(taskData.toHexDump())
# Add task to queue
cq.interactAgent.tasks.add(taskData)
cq.interactAgent.tasks.add(task)
cq.writeLine(fgBlack, styleBright, fmt"[{date}] [*] ", resetStyle, fmt"Tasked agent to {command.description.toLowerAscii()}")
except CatchableError:

View File

@@ -3,9 +3,9 @@ import ../utils
import ../../common/types
import ../../common/serialize
proc serializeTask*(cq: Conquest, task: Task): seq[byte] =
proc serializeTask*(task: Task): seq[byte] =
var packer = initTaskPacker()
var packer = initPacker()
# Serialize payload
packer
@@ -39,3 +39,63 @@ proc serializeTask*(cq: Conquest, task: Task): seq[byte] =
# TODO: Calculate and patch HMAC
return header & payload
proc deserializeTaskResult*(resultData: seq[byte]): TaskResult =
var unpacker = initUnpacker(resultData.toString)
let
magic = unpacker.getUint32()
version = unpacker.getUint8()
packetType = unpacker.getUint8()
flags = unpacker.getUint16()
seqNr = unpacker.getUint32()
size = unpacker.getUint32()
hmacBytes = unpacker.getBytes(16)
# Explicit conversion from seq[byte] to array[16, byte]
var hmac: array[16, byte]
copyMem(hmac.addr, hmacBytes[0].unsafeAddr, 16)
# Packet Validation
if magic != MAGIC:
raise newException(CatchableError, "Invalid magic bytes.")
# TODO: Validate sequence number
# TODO: Validate HMAC
# TODO: Decrypt payload
# let payload = unpacker.getBytes(size)
let
taskId = unpacker.getUint32()
agentId = unpacker.getUint32()
listenerId = unpacker.getUint32()
timestamp = unpacker.getUint32()
command = unpacker.getUint16()
status = unpacker.getUint8()
resultType = unpacker.getUint8()
length = unpacker.getUint32()
data = unpacker.getBytes(int(length))
return TaskResult(
header: Header(
magic: magic,
version: version,
packetType: packetType,
flags: flags,
seqNr: seqNr,
size: size,
hmac: hmac
),
taskId: taskId,
agentId: agentId,
listenerId: listenerId,
timestamp: timestamp,
command: command,
status: status,
resultType: resultType,
length: length,
data: data
)

View File

@@ -31,6 +31,11 @@ proc toString*(data: seq[byte]): string =
for i, b in data:
result[i] = char(b)
proc toBytes*(data: string): seq[byte] =
result = newSeq[byte](data.len)
for i, c in data:
result[i] = byte(c.ord)
proc toHexDump*(data: seq[byte]): string =
for i, b in data:
result.add(b.toHex(2))