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:
@@ -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":
|
||||
|
||||
@@ -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.")
|
||||
75
server/agent/taskDispatcher.nim
Normal file
75
server/agent/taskDispatcher.nim
Normal 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
|
||||
Reference in New Issue
Block a user