Implemented communication with custom binary structure instead of JSON requests
This commit is contained in:
@@ -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)
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user