Updated task structure to use a JSON string instead of seq[string], making it possible to use multiple differently typed arguments

This commit is contained in:
Jakob Friedl
2025-07-09 14:38:30 +02:00
parent 71ff092975
commit bb56ed42f2
10 changed files with 124 additions and 73 deletions

View File

@@ -1,5 +1,5 @@
import argparse, times, strformat, terminal, nanoid
import ./task
import ./taskDispatcher
import ../[types]
#[
@@ -10,8 +10,8 @@ var parser = newParser:
command("shell"):
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])
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])
command("sleep"):
help("Update sleep delay configuration.")
@@ -39,6 +39,11 @@ var parser = newParser:
help("Remove directory.")
arg("directory", help="Relative or absolute path to the directory to delete.", nargs = -1)
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)
command("help"):
nohelpflag()
@@ -65,11 +70,7 @@ proc handleAgentCommand*(cq: Conquest, args: varargs[string]) =
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)
cq.taskExecuteShell(opts.shell.get.command, opts.shell.get.arguments)
of "sleep":
cq.taskExecuteSleep(parseInt(opts.sleep.get.delay))
@@ -92,6 +93,9 @@ proc handleAgentCommand*(cq: Conquest, args: varargs[string]) =
of "rmdir":
cq.taskRemoveDirectory(opts.rmdir.get.directory)
of "bof":
cq.taskExecuteBof(opts.bof.get.file, opts.bof.get.arguments)
# Handle help flag
except ShortCircuit as err:
if err.flag == "argparse_help":

View File

@@ -1,48 +0,0 @@
import nanoid, sequtils, strutils, strformat, terminal, times
import ../types
import ../db/database
# Generic task creation procedure
proc createTask(cq: Conquest, command: TaskCommand, args: seq[string], message: string) =
let
date = now().format("dd-MM-yyyy HH:mm:ss")
task = Task(
id: generate(alphabet=join(toSeq('A'..'Z'), ""), size=8),
agent: cq.interactAgent.name,
command: command,
args: args,
)
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
# Use the generic createTask function
createTask(cq, Sleep, @[$delay], "Tasked agent to update sleep settings.")
proc taskExecuteShell*(cq: Conquest, arguments: seq[string]) =
cq.createTask(ExecuteShell, arguments, "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]) =
cq.createTask(SetWorkingDirectory, arguments, "Tasked agent to change current working directory.")
proc taskListDirectory*(cq: Conquest, arguments: seq[string]) =
cq.createTask(ListDirectory, arguments, "Tasked agent to list files and directories.")
proc taskRemoveFile*(cq: Conquest, arguments: seq[string]) =
cq.createTask(RemoveFile, arguments, "Tasked agent to remove file.")
proc taskRemoveDirectory*(cq: Conquest, arguments: seq[string]) =
cq.createTask(RemoveDirectory, arguments, "Tasked agent to remove directory.")

View File

@@ -0,0 +1,75 @@
import nanoid, sequtils, strutils, strformat, terminal, times, json
import ../types
import ../db/database
# Generic task creation procedure
proc createTask(cq: Conquest, command: TaskCommand, args: string, message: string) =
let
date = now().format("dd-MM-yyyy HH:mm:ss")
task = Task(
id: generate(alphabet=join(toSeq('A'..'Z'), ""), size=8),
agent: cq.interactAgent.name,
command: command,
args: args,
)
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 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 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