From b8b276d887c128cc9e4b615b69c517c4054a8632 Mon Sep 17 00:00:00 2001 From: Jakob Friedl <71284620+jakobfriedl@users.noreply.github.com> Date: Mon, 14 Jul 2025 22:14:27 +0200 Subject: [PATCH] Refactored agent command handling to remove redundant boiler-plate code. Commands are parsed dynamically based on a single definition. Command-specific actions might still need distinct implementations. --- agents/monarch/types.nim | 2 +- server/agent/interact.nim | 299 ++++++++++++++++++++++---------- server/agent/taskDispatcher.nim | 32 +--- server/server.nim | 8 +- server/types.nim | 28 ++- 5 files changed, 240 insertions(+), 129 deletions(-) diff --git a/agents/monarch/types.nim b/agents/monarch/types.nim index c0a3550..8a05b4e 100644 --- a/agents/monarch/types.nim +++ b/agents/monarch/types.nim @@ -1,6 +1,6 @@ import winim, tables import ../../server/types -export Task, TaskCommand, TaskResult, TaskStatus +export Task, CommandType, TaskResult, TaskStatus type ProductType* = enum diff --git a/server/agent/interact.nim b/server/agent/interact.nim index b3c75a8..b7750c3 100644 --- a/server/agent/interact.nim +++ b/server/agent/interact.nim @@ -1,123 +1,244 @@ -import argparse, times, strformat, terminal, nanoid +import argparse, times, strformat, terminal, nanoid, tables, json, sequtils import ./taskDispatcher import ../[types] #[ Agent Argument parsing ]# -var parser = newParser: - help("Conquest Command & Control") +proc initAgentCommands*(): Table[CommandType, Command] = + var commands = initTable[CommandType, Command]() - command("shell"): - help("Execute a shell command and retrieve the output.") - arg("command", help="Command to be executed.", nargs = 1) - arg("arguments", help="Arguments to be passed to the command.", nargs = -1) # Handle 0 or more command-line arguments (seq[string]) + commands[ExecuteShell] = Command( + name: "shell", + commandType: ExecuteShell, + description: "Execute a shell command and retrieve the output.", + example: "shell whoami /all", + arguments: @[ + Argument(name: "command", description: "Command to be executed.", argumentType: String, isRequired: true), + Argument(name: "arguments", description: "Arguments to be passed to the command.", argumentType: String, isRequired: false) + ] + ) - command("sleep"): - help("Update sleep delay configuration.") - arg("delay", help="Delay in seconds.", nargs = 1) + commands[Sleep] = Command( + name: "sleep", + commandType: Sleep, + description: "Update sleep delay configuration.", + example: "sleep 5", + arguments: @[ + Argument(name: "delay", description: "Delay in seconds.", argumentType: Int, isRequired: true) + ] + ) - command("info"): - help("Display agent information and current settings.") + commands[GetWorkingDirectory] = Command( + name: "pwd", + commandType: GetWorkingDirectory, + description: "Retrieve current working directory.", + example: "pwd", + arguments: @[] + ) - command("pwd"): - help("Retrieve current working directory.") + commands[SetWorkingDirectory] = Command( + name: "cd", + commandType: SetWorkingDirectory, + description: "Change current working directory.", + example: "cd C:\\Windows\\Tasks", + arguments: @[ + Argument(name: "directory", description: "Relative or absolute path of the directory to change to.", argumentType: String, isRequired: true) + ] + ) - command("cd"): - help("Change current working directory.") - arg("directory", help="Relative or absolute path of the directory to change to.", nargs = -1) + commands[ListDirectory] = Command( + name: "ls", + commandType: ListDirectory, + description: "List files and directories.", + example: "ls C:\\Users\\Administrator\\Desktop", + arguments: @[ + Argument(name: "directory", description: "Relative or absolute path. Default: current working directory.", argumentType: String, isRequired: false) + ] + ) - command("ls"): - help("List files and directories.") - arg("directory", help="Relative or absolute path. Default: current working directory.", nargs = -1) + commands[RemoveFile] = Command( + name: "rm", + commandType: RemoveFile, + description: "Remove a file.", + example: "rm C:\\Windows\\Tasks\\payload.exe", + arguments: @[ + Argument(name: "file", description: "Relative or absolute path to the file to delete.", argumentType: String, isRequired: true) + ] + ) - command("rm"): - help("Remove file.") - arg("file", help="Relative or absolute path to the file to delete.", nargs = -1) + commands[RemoveDirectory] = Command( + name: "rmdir", + commandType: RemoveDirectory, + description: "Remove a directory.", + example: "rm C:\\Payloads", + arguments: @[ + Argument(name: "directory", description: "Relative or absolute path to the directory to delete.", argumentType: String, isRequired: true) + ] + ) - command("rmdir"): - help("Remove directory.") - arg("directory", help="Relative or absolute path to the directory to delete.", nargs = -1) + commands[Move] = Command( + name: "move", + commandType: Move, + description: "Move a file or directory.", + example: "move source.exe C:\\Windows\\Tasks\\destination.exe", + arguments: @[ + Argument(name: "source", description: "Source file path.", argumentType: String, isRequired: true), + Argument(name: "destination", description: "Destination file path.", argumentType: String, isRequired: true) + ] + ) - command("move"): - help("Move a file or directory.") - option("-f", "--from", help="Original file name.", required=true) - option("-t", "--to", help="New file name.", required=true) + commands[Copy] = Command( + name: "copy", + commandType: Copy, + description: "Copy a file or directory.", + example: "copy source.exe C:\\Windows\\Tasks\\destination.exe", + arguments: @[ + Argument(name: "source", description: "Source file path.", argumentType: String, isRequired: true), + Argument(name: "destination", description: "Destination file path.", argumentType: String, isRequired: true) + ] + ) - command("copy"): - help("Copy a file or directory.") - option("-f", "--from", help="Original file name.", required=true) - option("-t", "--to", help="New file name.", required=true) + return commands - command("bof"): - help("Execute COFF or BOF file (.o) in memory.") - arg("file", help="Local path to object file.", nargs = 1) - arg("arguments", help="Arguments to be passed to the BOF.", nargs = -1) +proc displayHelp(cq: Conquest, commands: Table[CommandType, Command]) = + cq.writeLine("Available commands:") + for key, cmd in commands: + cq.writeLine(fmt" * {cmd.name:<15}{cmd.description}") + cq.writeLine() - command("help"): - nohelpflag() +proc displayCommandHelp(cq: Conquest, command: Command) = + var usage = command.name & " " & command.arguments.mapIt( + if it.isRequired: fmt"<{it.name}>" else: fmt"[{it.name}]" + ).join(" ") - command("back"): - nohelpflag() + if command.example != "": + usage &= "\nExample : " & command.example -proc handleAgentCommand*(cq: Conquest, args: varargs[string]) = + cq.writeLine(fmt""" +{command.description} + +Usage : {usage} +""") + + if command.arguments.len > 0: + cq.writeLine("Arguments:") + for arg in command.arguments: + let requirement = if arg.isRequired: "REQUIRED" else: "OPTIONAL" + cq.writeLine(fmt" * {arg.name:<15} {requirement} {arg.description}") + + cq.writeLine() + +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) + + while i < input.len and input[i] != '"': + # Add parsed argument when quotation is closed + 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 handleAgentCommand*(cq: Conquest, input: string) = + + let commands = initAgentCommands() + var + commandType: CommandType + command: Command # Return if no command (or just whitespace) is entered - if args[0].replace(" ", "").len == 0: return + 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, args[0]) + cq.writeLine(fgBlue, styleBright, fmt"[{date}] ", fgYellow, fmt"[{cq.interactAgent.name}] ", resetStyle, styleBright, input) - try: - let opts = parser.parse(args[0].split(" ").filterIt(it.len > 0)) - - case opts.command + # Split the user input, taking quotes into consideration + let parsed = parseAgentCommand(input) - of "back": # Return to management mode - discard + # Handle 'back' command + if parsed[0] == "back": + return - of "help": # Display help menu - cq.writeLine(parser.help()) + # Handle 'help' command + if parsed[0] == "help": + try: + # 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 - of "shell": - cq.taskExecuteShell(opts.shell.get.command, opts.shell.get.arguments) + cq.displayCommandHelp(command) + return - of "sleep": - cq.taskExecuteSleep(parseInt(opts.sleep.get.delay)) + # Following this, commands require actions on the agent and thus a task needs to be created + # Determine the command used by checking the first positional argument + try: + commandType = parseEnum[CommandType](parsed[0].toLowerAscii()) + command = commands[commandType] + except ValueError: + cq.writeLine(fgRed, styleBright, "[-] Unknown command.\n") + return - of "info": - discard + # TODO: Client/Server-side command specific actions (e.g. updating sleep, ...) - of "pwd": - cq.taskGetWorkingDirectory() + # 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 - of "cd": - cq.taskSetWorkingDirectory(opts.cd.get.directory) + 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] = %"" - of "ls": - cq.taskListDirectory(opts.ls.get.directory) - - of "rm": - cq.taskRemoveFile(opts.rm.get.file) - - of "rmdir": - cq.taskRemoveDirectory(opts.rmdir.get.directory) - - of "move": - cq.taskMove(opts.move.get.from, opts.move.get.to) - - of "copy": - cq.taskCopy(opts.copy.get.from, opts.copy.get.to) - - of "bof": - cq.taskExecuteBof(opts.bof.get.file, opts.bof.get.arguments) - - # Handle help flag - except ShortCircuit as err: - if err.flag == "argparse_help": - cq.writeLine(err.help) - - # Handle invalid arguments - except CatchableError: - cq.writeLine(fgRed, styleBright, "[-] ", getCurrentExceptionMsg(), "\n") + except CatchableError: + cq.writeLine(fgRed, styleBright, "[-] Invalid syntax.\n") + return + # Task creation + cq.createTask(commandType, $payload, fmt"Tasked agent to {command.description.toLowerAscii()}") \ No newline at end of file diff --git a/server/agent/taskDispatcher.nim b/server/agent/taskDispatcher.nim index 0488994..109d386 100644 --- a/server/agent/taskDispatcher.nim +++ b/server/agent/taskDispatcher.nim @@ -3,7 +3,7 @@ import ../types import ../db/database # Generic task creation procedure -proc createTask(cq: Conquest, command: TaskCommand, args: string, message: string) = +proc createTask*(cq: Conquest, command: CommandType, args: string, message: string) = let date = now().format("dd-MM-yyyy HH:mm:ss") task = Task( @@ -32,36 +32,6 @@ proc taskExecuteSleep*(cq: Conquest, delay: int) = # Use the generic createTask function createTask(cq, Sleep, $payload, "Tasked agent to update sleep settings.") -proc taskExecuteShell*(cq: Conquest, command: string, arguments: seq[string]) = - let payload = %*{ "command": command, "arguments": arguments.join(" ")} - cq.createTask(ExecuteShell, $payload, "Tasked agent to execute shell command.") - -proc taskGetWorkingDirectory*(cq: Conquest) = - cq.createTask(GetWorkingDirectory, "", "Tasked agent to get current working directory.") - -proc taskSetWorkingDirectory*(cq: Conquest, arguments: seq[string]) = - let payload = %*{ "directory": arguments.join(" ").replace("\"").replace("'")} - cq.createTask(SetWorkingDirectory, $payload, "Tasked agent to change current working directory.") - -proc taskListDirectory*(cq: Conquest, arguments: seq[string]) = - let payload = %*{ "directory": arguments.join(" ").replace("\"").replace("'")} - cq.createTask(ListDirectory, $payload, "Tasked agent to list files and directories.") - -proc taskRemoveFile*(cq: Conquest, arguments: seq[string]) = - let payload = %*{ "file": arguments.join(" ").replace("\"").replace("'")} - cq.createTask(RemoveFile, $payload, "Tasked agent to remove file.") - -proc taskRemoveDirectory*(cq: Conquest, arguments: seq[string]) = - let payload = %*{ "directory": arguments.join(" ").replace("\"").replace("'")} - cq.createTask(RemoveDirectory, $payload, "Tasked agent to remove directory.") - -proc taskMove*(cq: Conquest, oldPath, newPath: string) = - let payload = %*{ "from": oldPath, "to": newPath} - cq.createTask(Move, $payload, "Tasked agent to move a file or directory.") - -proc taskCopy*(cq: Conquest, oldPath, newPath: string) = - let payload = %*{ "from": oldPath, "to": newPath} - cq.createTask(Copy, $payload, "Tasked agent to copy a file or directory.") proc taskExecuteBof*(cq: Conquest, file: string, arguments: seq[string]) = diff --git a/server/server.nim b/server/server.nim index e705f17..67b9f30 100644 --- a/server/server.nim +++ b/server/server.nim @@ -60,16 +60,16 @@ var parser = newParser: command("exit"): nohelpflag() -proc handleConsoleCommand*(cq: Conquest, args: varargs[string]) = +proc handleConsoleCommand*(cq: Conquest, args: string) = # Return if no command (or just whitespace) is entered - if args[0].replace(" ", "").len == 0: return + if args.replace(" ", "").len == 0: return let date: string = now().format("dd-MM-yyyy HH:mm:ss") - cq.writeLine(fgBlue, styleBright, fmt"[{date}] ", resetStyle, styleBright, args[0]) + cq.writeLine(fgBlue, styleBright, fmt"[{date}] ", resetStyle, styleBright, args) try: - let opts = parser.parse(args[0].split(" ").filterIt(it.len > 0)) + let opts = parser.parse(args.split(" ").filterIt(it.len > 0)) case opts.command diff --git a/server/types.nim b/server/types.nim index 25eec0a..047a895 100644 --- a/server/types.nim +++ b/server/types.nim @@ -8,8 +8,7 @@ import terminal Agent types & procs ]# type - - TaskCommand* = enum + CommandType* = enum ExecuteShell = "shell" ExecuteBof = "bof" ExecuteAssembly = "dotnet" @@ -23,6 +22,27 @@ type Move = "move" Copy = "copy" + ArgumentType* = enum + String = 0 + Int = 1 + Long = 2 + Bool = 3 + Binary = 4 + + Argument* = object + name*: string + description*: string + argumentType*: ArgumentType + isRequired*: bool + + Command* = object + name*: string + commandType*: CommandType + description*: string + example*: string + arguments*: seq[Argument] + dispatchMessage*: string + TaskStatus* = enum Completed = "completed" Created = "created" @@ -39,7 +59,7 @@ type Task* = ref object id*: string agent*: string - command*: TaskCommand + command*: CommandType args*: string # Json string containing all the positional arguments # Example: """{"command": "whoami", "arguments": "/all"}""" @@ -181,7 +201,7 @@ template clear*(cq: Conquest) = cq.prompt.clear() # Overwrite withOutput function to handle function arguments -proc withOutput*(cq: Conquest, outputFunction: proc(cq: Conquest, args: varargs[string]), args: varargs[string]) = +proc withOutput*(cq: Conquest, outputFunction: proc(cq: Conquest, args: string), args: string) = cq.hidePrompt() outputFunction(cq, args) cq.showPrompt()