From dfcafa9c24db9bae9c38b4f976a4c6ca2d7126c4 Mon Sep 17 00:00:00 2001 From: Jakob Friedl <71284620+jakobfriedl@users.noreply.github.com> Date: Fri, 1 Aug 2025 13:16:12 +0200 Subject: [PATCH] Implemented basic "ps" and "env" commands. --- src/common/types.nim | 5 ++ src/modules/environment.nim | 150 ++++++++++++++++++++++++++++++++++++ src/modules/filesystem.nim | 1 - src/modules/manager.nim | 4 +- src/modules/shell.nim | 1 - src/modules/sleep.nim | 1 - src/server/api/handlers.nim | 4 +- 7 files changed, 161 insertions(+), 5 deletions(-) create mode 100644 src/modules/environment.nim diff --git a/src/common/types.nim b/src/common/types.nim index f1c48d9..80fb19c 100644 --- a/src/common/types.nim +++ b/src/common/types.nim @@ -28,6 +28,7 @@ type FLAG_PLAINTEXT = 0'u16 FLAG_ENCRYPTED = 1'u16 FLAG_COMPRESSED = 2'u16 + FLAG_FRAGMENTED = 4'u16 CommandType* = enum CMD_SLEEP = 0'u16 @@ -39,10 +40,14 @@ type CMD_RMDIR = 6'u16 CMD_MOVE = 7'u16 CMD_COPY = 8'u16 + CMD_PS = 9'u16 + CMD_ENV = 10'u16 + CMD_WHOAMI = 11'u16 StatusType* = enum STATUS_COMPLETED = 0'u8 STATUS_FAILED = 1'u8 + STATUS_IN_PROGRESS = 2'u8 ResultType* = enum RESULT_STRING = 0'u8 diff --git a/src/modules/environment.nim b/src/modules/environment.nim new file mode 100644 index 0000000..5420578 --- /dev/null +++ b/src/modules/environment.nim @@ -0,0 +1,150 @@ +import ../common/[types, utils] + +# Declare function prototypes +proc executePs(config: AgentConfig, task: Task): TaskResult +proc executeEnv(config: AgentConfig, task: Task): TaskResult +proc executeWhoami(config: AgentConfig, task: Task): TaskResult + +# Command definitions +let commands*: seq[Command] = @[ + Command( + name: "ps", + commandType: CMD_PS, + description: "Display running processes.", + example: "ps", + arguments: @[], + execute: executePs + ), + Command( + name: "env", + commandType: CMD_ENV, + description: "Display environment variables.", + example: "env", + arguments: @[], + execute: executeEnv + ), + Command( + name: "whoami", + commandType: CMD_WHOAMI, + description: "Get user information.", + example: "whoami", + arguments: @[], + execute: executeWhoami + ) +] + +# Implement execution functions +when defined(server): + proc executePs(config: AgentConfig, task: Task): TaskResult = nil + proc executeEnv(config: AgentConfig, task: Task): TaskResult = nil + proc executeWhoami(config: AgentConfig, task: Task): TaskResult = nil + +when defined(agent): + + import winim + import os, strutils, sequtils, strformat, tables, algorithm + import ../agent/core/taskresult + + # TODO: Add user context to process information + type + ProcessInfo = object + pid: DWORD + ppid: DWORD + name: string + children: seq[DWORD] + + proc executePs(config: AgentConfig, task: Task): TaskResult = + + echo fmt" [>] Listing running processes." + + try: + + var processes: seq[DWORD] = @[] + var procMap = initTable[DWORD, ProcessInfo]() + var output: string = "" + + # Take a snapshot of running processes + let hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) + if hSnapshot == INVALID_HANDLE_VALUE: + raise newException(CatchableError, "Invalid permissions.\n") + + # Close handle after object is no longer used + defer: CloseHandle(hSnapshot) + + var pe32: PROCESSENTRY32 + pe32.dwSize = DWORD(sizeof(PROCESSENTRY32)) + + # Loop over processes to fill the map + if Process32First(hSnapshot, addr pe32) == FALSE: + raise newException(CatchableError, "Failed to get processes.\n") + + while true: + var procInfo = ProcessInfo( + pid: pe32.th32ProcessID, + ppid: pe32.th32ParentProcessID, + name: $cast[WideCString](addr pe32.szExeFile[0]), + children: @[] + ) + + procMap[pe32.th32ProcessID] = procInfo + + if Process32Next(hSnapshot, addr pe32) == FALSE: + break + + # Build child-parent relationship + for pid, procInfo in procMap.mpairs(): + if procMap.contains(procInfo.ppid): + procMap[procInfo.ppid].children.add(pid) + else: + processes.add(pid) + + # Format and print process + proc printProcess(pid: DWORD, indentSpaces: int = 0) = + if not procMap.contains(pid): + return + + var process = procMap[pid] + let indent = " ".repeat(indentSpaces) + + let tree = (indent & fmt"[{process.pid}]").alignLeft(30) + + output &= fmt"{tree}{process.name:<25}" & "\n" + + # Recursively print child processes with indentation + process.children.sort() + for childPid in process.children: + printProcess(childPid, indentSpaces + 2) + + # Iterate over root processes + processes.sort() + for pid in processes: + printProcess(pid) + + return createTaskResult(task, STATUS_COMPLETED, RESULT_STRING, output.toBytes()) + + except CatchableError as err: + return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes()) + + proc executeEnv(config: AgentConfig, task: Task): TaskResult = + + echo fmt" [>] Displaying environment variables." + + try: + var envVars: string = "" + for key, value in envPairs(): + envVars &= fmt"{key}: {value}" & '\n' + return createTaskResult(task, STATUS_COMPLETED, RESULT_STRING, envVars.toBytes()) + + except CatchableError as err: + return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes()) + + proc executeWhoami(config: AgentConfig, task: Task): TaskResult = + + echo fmt" [>] Getting user information." + + try: + echo "whoami" + + + except CatchableError as err: + return createTaskResult(task, STATUS_FAILED, RESULT_STRING, err.msg.toBytes()) \ No newline at end of file diff --git a/src/modules/filesystem.nim b/src/modules/filesystem.nim index a21331e..722f355 100644 --- a/src/modules/filesystem.nim +++ b/src/modules/filesystem.nim @@ -1,4 +1,3 @@ -import ./manager import ../common/[types, utils] # Define function prototypes diff --git a/src/modules/manager.nim b/src/modules/manager.nim index b78a544..bf1ee4f 100644 --- a/src/modules/manager.nim +++ b/src/modules/manager.nim @@ -5,7 +5,8 @@ import ../common/[types, utils] import shell, sleep, - filesystem + filesystem, + environment type ModuleManager* = object @@ -24,6 +25,7 @@ proc loadModules*() = registerCommands(shell.commands) registerCommands(sleep.commands) registerCommands(filesystem.commands) + registerCommands(environment.commands) proc getCommandByType*(cmdType: CommandType): Command = return manager.commandsByType[cmdType] diff --git a/src/modules/shell.nim b/src/modules/shell.nim index ee8e584..cabf59e 100644 --- a/src/modules/shell.nim +++ b/src/modules/shell.nim @@ -1,4 +1,3 @@ -import ./manager import ../common/[types, utils] # Define function prototype diff --git a/src/modules/sleep.nim b/src/modules/sleep.nim index a892a57..c51d7b9 100644 --- a/src/modules/sleep.nim +++ b/src/modules/sleep.nim @@ -1,4 +1,3 @@ -import ./manager import ../common/[types, utils] # Define function prototype diff --git a/src/server/api/handlers.nim b/src/server/api/handlers.nim index 257b0a1..98d5321 100644 --- a/src/server/api/handlers.nim +++ b/src/server/api/handlers.nim @@ -82,9 +82,11 @@ proc handleResult*(resultData: seq[byte]) = case cast[StatusType](taskResult.status): of STATUS_COMPLETED: cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgGreen, " [+] ", resetStyle, fmt"Task {taskId} completed.") - of STATUS_FAILED: cq.writeLine(fgBlack, styleBright, fmt"[{date}]", fgRed, styleBright, " [-] ", resetStyle, fmt"Task {taskId} failed.") + of STATUS_IN_PROGRESS: + discard + case cast[ResultType](taskResult.resultType): of RESULT_STRING: