Cleaned up agent command parsing

agents/monarch/nim.cfg

@@ -5,4 +5,4 @@
 -d:Octet3="0"
 -d:Octet4="1"
 -d:ListenerPort=5555
--d:SleepDelay=1
+-d:SleepDelay=5

server/agent/interact.nim

@@ -1,6 +1,6 @@
 import argparse, times, strformat, terminal, nanoid, tables, json, sequtils
 import ./taskDispatcher
-import ../[types]
+import ../types
 
 #[
     Agent Argument parsing
@@ -101,8 +101,48 @@ proc initAgentCommands*(): Table[CommandType, Command] =
 
     return commands 
 
+let commands = initAgentCommands() 
+
+proc getCommandFromTable(cmd: string, commands: Table[CommandType, Command]): (CommandType, Command) =
+    let commandType = parseEnum[CommandType](cmd.toLowerAscii())
+    let command = commands[commandType]
+    (commandType, command)
+    
+proc parseAgentCommand(input: string): seq[string] = 
+    var i = 0
+    while i < input.len:
+
+        # Skip whitespaces/tabs
+        while i < input.len and input[i] in {' ', '\t'}: 
+            inc i
+        if i >= input.len: 
+            break
+        
+        var arg = ""
+        if input[i] == '"':
+            # Parse quoted argument
+            inc i # Skip opening quote
+
+            # Add parsed argument when quotation is closed
+            while i < input.len and input[i] != '"': 
+                arg.add(input[i]) 
+                inc i
+            
+            if i < input.len: 
+                inc i # Skip closing quote
+        
+        else:
+            while i < input.len and input[i] notin {' ', '\t'}: 
+                arg.add(input[i])
+                inc i
+        
+        # Add argument to returned result
+        if arg.len > 0: result.add(arg)
+
+
 proc displayHelp(cq: Conquest, commands: Table[CommandType, Command]) = 
     cq.writeLine("Available commands:")
+    cq.writeLine(" * back")
     for key, cmd in commands: 
         cq.writeLine(fmt" * {cmd.name:<15}{cmd.description}")
     cq.writeLine()
@@ -123,122 +163,90 @@ Usage   : {usage}
 
     if command.arguments.len > 0:
         cq.writeLine("Arguments:")
+
+        let header = @["Name", "Type", "", "Description"]
+        cq.writeLine(fmt"   {header[0]:<15} {header[1]:<8}{header[2]:<10} {header[3]}")
+        cq.writeLine(fmt"   {'-'.repeat(15)} {'-'.repeat(18)} {'-'.repeat(20)}")
+        
         for arg in command.arguments: 
-            let requirement = if arg.isRequired: "REQUIRED" else: "OPTIONAL"
+            let requirement = if arg.isRequired: "(REQUIRED)" else: "(OPTIONAL)"
-            cq.writeLine(fmt" * {arg.name:<15} {requirement}    {arg.description}")
+            cq.writeLine(fmt" * {arg.name:<15} {($arg.argumentType).toUpperAscii():<8}{requirement:<10} {arg.description}")
 
         cq.writeLine()
 
-proc parseAgentCommand(input: string): seq[string] = 
+proc handleHelp(cq: Conquest, parsed: seq[string], commands: Table[CommandType, Command]) = 
-    var i = 0
+    try: 
+        # Try parsing the first argument passed to 'help' as a command
+        let (commandType, command) = getCommandFromTable(parsed[1],  commands)
+        cq.displayCommandHelp(command)
+    except IndexDefect:
+        # 'help' command is called without additional parameters
+        cq.displayHelp(commands)
+    except ValueError: 
+        # Command was not found
+        cq.writeLine(fgRed, styleBright, fmt"[-] The command '{parsed[1]}' does not exist." & '\n')
 
-    while i < input.len:
+proc packageArguments(cq: Conquest, command: Command, arguments: seq[string]): JsonNode = 
-        # Skip whitespaces/tabs
-        while i < input.len and input[i] in {' ', '\t'}: 
-            inc i
-        if i >= input.len: 
-            break
 
-        var arg = ""
+    # Construct a JSON payload with argument names and values 
-        if input[i] == '"':
+    result = newJObject()
-            # Parse quoted argument
+    let parsedArgs = if arguments.len > 1: arguments[1..^1] else: @[] # Remove first element from sequence to only handle arguments
-            inc i # (Skip opening quote)
 
-            while i < input.len and input[i] != '"': 
+    # Check if the correct amount of parameters are passed
-                # Add parsed argument when quotation is closed
+    if parsedArgs.len < command.arguments.filterIt(it.isRequired).len: 
-                arg.add(input[i]) 
+        cq.displayCommandHelp(command)
-                inc i
+        raise newException(ValueError, "Missing required arguments.")
 
-            if i < input.len: 
+    for i, argument in command.arguments: 
-                inc i # (Skip closing quote)
         
+        # Argument provided - convert to the corresponding data type
+        if i < parsedArgs.len:
+            case argument.argumentType:
+            of Int:
+                result[argument.name] = %parseUInt(parsedArgs[i])
+            of Binary: 
+                # Read file into memory and convert it into a base64 string
+                result[argument.name] = %""
+            else:
+                # The last optional argument is joined together
+                # This is required for non-quoted input with infinite length, such as `shell mv arg1 arg2`
+                if i == command.arguments.len - 1 and not argument.isRequired:
+                    result[argument.name] = %parsedArgs[i..^1].join(" ")
+                else:
+                    result[argument.name] = %parsedArgs[i]
+        
+        # Argument not provided - set to empty string for optional args
         else:
-            while i < input.len and input[i] notin {' ', '\t'}: 
+            # If a required argument is not provided, display the help text
-                arg.add(input[i])
+            if argument.isRequired:
-                inc i
+                cq.displayCommandHelp(command)
-        
+                return
-        # Add argument to returned result
+            else:
-        if arg.len > 0: result.add(arg)
+                result[argument.name] = %""
 
 proc handleAgentCommand*(cq: Conquest, input: string) = 
-
-    let commands = initAgentCommands() 
-    var 
-        commandType: CommandType
-        command: Command
-
     # Return if no command (or just whitespace) is entered
     if input.replace(" ", "").len == 0: return
 
     let date: string = now().format("dd-MM-yyyy HH:mm:ss")
     cq.writeLine(fgBlue, styleBright, fmt"[{date}] ", fgYellow, fmt"[{cq.interactAgent.name}] ", resetStyle, styleBright, input)
 
-    # Split the user input, taking quotes into consideration 
+    let parsedArgs = parseAgentCommand(input)
-    let parsed = parseAgentCommand(input)
     
     # Handle 'back' command
-    if parsed[0] == "back": 
+    if parsedArgs[0] == "back": 
         return
 
     # Handle 'help' command 
-    if parsed[0] == "help": 
+    if parsedArgs[0] == "help": 
-        try: 
+        cq.handleHelp(parsedArgs, commands)
-            # Try parsing the first argument passed to 'help' as a command
-            commandType = parseEnum[CommandType](parsed[1].toLowerAscii())
-            command = commands[commandType]
-        except IndexDefect:
-            # 'help' command is called without additional parameters
-            cq.displayHelp(commands)
-            return
-        except ValueError: 
-            # Command was not found
-            cq.writeLine(fgRed, styleBright, fmt"[-] The command {parsed[1]} does not exist." & '\n')
-            return 
-
-        cq.displayCommandHelp(command)
         return
 
-    # Following this, commands require actions on the agent and thus a task needs to be created
+    # Handle commands with actions on the agent
-    # Determine the command used by checking the first positional argument
     try: 
-        commandType = parseEnum[CommandType](parsed[0].toLowerAscii())
+        let (commandType, command) = getCommandFromTable(parsedArgs[0], commands)
-        command = commands[commandType]
+        let payload = cq.packageArguments(command, parsedArgs)
-    except ValueError: 
+        cq.createTask(commandType, $payload, fmt"Tasked agent to {command.description.toLowerAscii()}")
-        cq.writeLine(fgRed, styleBright, "[-] Unknown command.\n")
+    except ValueError as err: 
+        cq.writeLine(fgRed, styleBright, fmt"[-] {err.msg}" & "\n")
         return
-
-    # TODO: Client/Server-side command specific actions (e.g. updating sleep, ...)
-
-    # Construct a JSON payload with argument names and values 
-    var payload = newJObject()
-    let parsedArgs = if parsed.len > 1: parsed[1..^1] else: @[] # Remove first element from sequence to only handle arguments
-
-    try: 
-        for i, argument in command.arguments: 
-            
-            if i < parsedArgs.len:
-                # Argument provided - convert to the corresponding data type
-                case argument.argumentType:
-                of Int:
-                    payload[argument.name] = %parseInt(parsedArgs[i])
-                of Binary: 
-                    # Read file into memory and convert it into a base64 string
-                    discard
-                else:
-                    payload[argument.name] = %parsedArgs[i]
-            
-            else:
-                # Argument not provided - set to empty string for optional args
-                # If a required argument is not provided, display the help text
-                if argument.isRequired:
-                    cq.displayCommandHelp(command)
-                    return
-                else:
-                    payload[argument.name] = %""
-
-    except CatchableError:
-        cq.writeLine(fgRed, styleBright, "[-] Invalid syntax.\n")
-        return
-
-    # Task creation
-    cq.createTask(commandType, $payload, fmt"Tasked agent to {command.description.toLowerAscii()}")

server/agent/taskDispatcher.nim

@@ -15,39 +15,3 @@ proc createTask*(cq: Conquest, command: CommandType, args: string, message: stri
     
     cq.interactAgent.tasks.add(task)
     cq.writeLine(fgBlack, styleBright, fmt"[{date}] [*] ", resetStyle, message)
-
-# Agent task functions
-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
-
-    # Construct payload 
-    let payload = %*{ "delay": delay }
-
-    # Use the generic createTask function
-    createTask(cq, Sleep, $payload, "Tasked agent to update sleep settings.")
-
-
-proc taskExecuteBof*(cq: Conquest, file: string, arguments: seq[string]) = 
-    
-    # Verify that the object file exists 
-    
-    # Read object file into memory and base64-encode it
-
-    # Create the payload package, consisting of base64-encoded object file and the arguments passed to it
-    # Best way would be a custom binary structure, but for the time being, a JSON string would work, which is deserialized and parsed by the agent
-    #[
-        let payload = %*
-        {
-            "file": "AAAA...AA=="
-            "arguments": "arg1 arg2 123"
-        }
-    ]#
-
-    # Create a new task 
-    discard

server/types.nim

@@ -23,11 +23,11 @@ type
         Copy = "copy"
 
     ArgumentType* = enum 
-        String = 0
+        String = "string"
-        Int = 1 
+        Int = "int"
-        Long = 2
+        Long = "long"
-        Bool = 3 
+        Bool = "bool"
-        Binary = 4
+        Binary = "binary"
 
     Argument* = object 
         name*: string