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.
This commit is contained in:
Jakob Friedl
@@ -1,6 +1,6 @@
|
||||
import winim, tables
|
||||
import ../../server/types
|
||||
export Task, TaskCommand, TaskResult, TaskStatus
|
||||
export Task, CommandType, TaskResult, TaskStatus
|
||||
|
||||
type
|
||||
ProductType* = enum
|
||||
|
||||
@@ -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()}")
|
||||
@@ -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]) =
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user