106 lines
4.4 KiB
Nim
106 lines
4.4 KiB
Nim
import terminal, strformat, strutils, sequtils, tables, json, times, base64, system
|
|
|
|
import ../[utils, globals]
|
|
import ../db/database
|
|
import ../message/packer
|
|
import ../../common/[types, utils]
|
|
|
|
#[
|
|
Agent API
|
|
Functions relevant for dealing with the agent API, such as registering new agents, querying tasks and posting results
|
|
]#
|
|
proc register*(registrationData: seq[byte]): bool =
|
|
|
|
# The following line is required to be able to use the `cq` global variable for console output
|
|
{.cast(gcsafe).}:
|
|
|
|
let agent: Agent = cq.deserializeNewAgent(registrationData)
|
|
|
|
# Validate that listener exists
|
|
if not cq.dbListenerExists(agent.listenerId.toUpperAscii):
|
|
cq.writeLine(fgRed, styleBright, fmt"[-] {agent.ip} attempted to register to non-existent listener: {agent.listenerId}.", "\n")
|
|
return false
|
|
|
|
# # Store agent in database
|
|
if not cq.dbStoreAgent(agent):
|
|
cq.writeLine(fgRed, styleBright, fmt"[-] Failed to insert agent {agent.agentId} into database.", "\n")
|
|
return false
|
|
|
|
cq.agents[agent.agentId] = agent
|
|
|
|
let date = agent.firstCheckin.format("dd-MM-yyyy HH:mm:ss")
|
|
cq.writeLine(fgYellow, styleBright, fmt"[{date}] ", resetStyle, "Agent ", fgYellow, styleBright, agent.agentId, resetStyle, " connected to listener ", fgGreen, styleBright, agent.listenerId, resetStyle, ": ", fgYellow, styleBright, fmt"{agent.username}@{agent.hostname}", "\n")
|
|
|
|
return true
|
|
|
|
proc getTasks*(checkinData: seq[byte]): seq[seq[byte]] =
|
|
|
|
{.cast(gcsafe).}:
|
|
|
|
# Deserialize checkin request to obtain agentId and listenerId
|
|
let
|
|
request: Heartbeat = cq.deserializeHeartbeat(checkinData)
|
|
agentId = uuidToString(request.header.agentId)
|
|
listenerId = uuidToString(request.listenerId)
|
|
timestamp = request.timestamp
|
|
|
|
var result: seq[seq[byte]]
|
|
|
|
# Check if listener exists
|
|
if not cq.dbListenerExists(listenerId):
|
|
cq.writeLine(fgRed, styleBright, fmt"[-] Task-retrieval request made to non-existent listener: {listenerId}.", "\n")
|
|
raise newException(ValueError, "Invalid listener.")
|
|
|
|
# Check if agent exists
|
|
if not cq.dbAgentExists(agentId):
|
|
cq.writeLine(fgRed, styleBright, fmt"[-] Task-retrieval request made to non-existent agent: {agentId}.", "\n")
|
|
raise newException(ValueError, "Invalid agent.")
|
|
|
|
# Update the last check-in date for the accessed agent
|
|
cq.agents[agentId].latestCheckin = cast[int64](timestamp).fromUnix().local()
|
|
|
|
# Return tasks
|
|
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)
|
|
|
|
return result
|
|
|
|
proc handleResult*(resultData: seq[byte]) =
|
|
|
|
{.cast(gcsafe).}:
|
|
|
|
let
|
|
taskResult = cq.deserializeTaskResult(resultData)
|
|
taskId = uuidToString(taskResult.taskId)
|
|
agentId = uuidToString(taskResult.header.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.")
|
|
|
|
case cast[StatusType](taskResult.status):
|
|
of STATUS_COMPLETED:
|
|
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgGreen, " [+] ", resetStyle, fmt"Task {taskId} completed.")
|
|
of STATUS_FAILED:
|
|
cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgRed, styleBright, " [-] ", resetStyle, fmt"Task {taskId} failed.")
|
|
of STATUS_IN_PROGRESS:
|
|
discard
|
|
|
|
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 taskResult.data.toString().split("\n"):
|
|
cq.writeLine(line)
|
|
|
|
of RESULT_BINARY:
|
|
# Write binary data to a file
|
|
cq.writeLine()
|
|
|
|
of RESULT_NO_OUTPUT:
|
|
cq.writeLine()
|
|
|
|
# Update task queue to include all tasks, except the one that was just completed
|
|
cq.agents[agentId].tasks = cq.agents[agentId].tasks.filterIt(it.taskId != taskResult.taskId) |