Reworked module system. Modules can now be individually set to be included in the agent. For example, it is possible to compile an agent only capable of executing BOFs and nothing else.

src/modules/situationalAwareness.nim

@@ -0,0 +1,159 @@
+import ../common/[types, utils]
+
+# Declare function prototypes
+proc executePs(ctx: AgentCtx, task: Task): TaskResult
+proc executeEnv(ctx: AgentCtx, task: Task): TaskResult
+proc executeWhoami(ctx: AgentCtx, task: Task): TaskResult
+
+# Module definition
+let module* = Module(
+    name: protect("situational-awareness"),
+    description: protect("Retrieve information about the target system and environment."),
+    commands: @[
+        Command(
+            name: protect("ps"),
+            commandType: CMD_PS,
+            description: protect("Display running processes."),
+            example: protect("ps"),
+            arguments: @[],
+            execute: executePs
+        ),
+        Command(
+            name: protect("env"),
+            commandType: CMD_ENV,
+            description: protect("Display environment variables."),
+            example: protect("env"),
+            arguments: @[],
+            execute: executeEnv
+        ),
+        Command(
+            name: protect("whoami"),
+            commandType: CMD_WHOAMI,
+            description: protect("Get user information."),
+            example: protect("whoami"),
+            arguments: @[],
+            execute: executeWhoami
+        )
+    ]
+)
+
+# Implement execution functions
+when defined(server):
+    proc executePs(ctx: AgentCtx, task: Task): TaskResult = nil
+    proc executeEnv(ctx: AgentCtx, task: Task): TaskResult = nil
+    proc executeWhoami(ctx: AgentCtx, task: Task): TaskResult = nil
+
+when defined(agent): 
+
+    import winim
+    import os, strutils, sequtils, strformat, tables, algorithm
+    import ../agent/protocol/result
+    import ../common/utils
+
+    # TODO: Add user context to process information
+    type 
+        ProcessInfo = object 
+            pid: DWORD
+            ppid: DWORD 
+            name: string 
+            children: seq[DWORD]
+
+    proc executePs(ctx: AgentCtx, task: Task): TaskResult = 
+        
+        echo protect("   [>] 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, protect("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, protect("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)
+
+            # Add header row
+            let headers = @[protect("PID"), protect("PPID"), protect("Process")]
+            output &= fmt"{headers[0]:<10}{headers[1]:<10}{headers[2]:<25}" & "\n"
+            output &= "-".repeat(len(headers[0])).alignLeft(10) & "-".repeat(len(headers[1])).alignLeft(10) & "-".repeat(len(headers[2])).alignLeft(25) & "\n"
+
+            # 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) 
+
+                output &= fmt"{process.pid:<10}{process.ppid:<10}{indent}{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, string.toBytes(output))
+
+        except CatchableError as err: 
+            return createTaskResult(task, STATUS_FAILED, RESULT_STRING, string.toBytes(err.msg))
+
+    proc executeEnv(ctx: AgentCtx, task: Task): TaskResult = 
+
+        echo protect("   [>] Displaying environment variables.")
+
+        try: 
+            var output: string = ""
+            for key, value in envPairs(): 
+               output &= fmt"{key}: {value}" & '\n'
+               
+            return createTaskResult(task, STATUS_COMPLETED, RESULT_STRING, string.toBytes(output))
+
+        except CatchableError as err: 
+            return createTaskResult(task, STATUS_FAILED, RESULT_STRING, string.toBytes(err.msg))
+
+    proc executeWhoami(ctx: AgentCtx, task: Task): TaskResult = 
+
+        echo protect("   [>] Getting user information.") 
+
+        try: 
+
+            let output = protect("Not implemented") 
+            return createTaskResult(task, STATUS_FAILED, RESULT_STRING, string.toBytes(output))
+
+        except CatchableError as err: 
+            return createTaskResult(task, STATUS_FAILED, RESULT_STRING, string.toBytes(err.msg))