diff --git a/assets/modules-10.png b/assets/modules-10.png new file mode 100644 index 0000000..4e1f3cf Binary files /dev/null and b/assets/modules-10.png differ diff --git a/docs/7-MODULES.md b/docs/7-MODULES.md index 5eccf57..72d197a 100644 --- a/docs/7-MODULES.md +++ b/docs/7-MODULES.md @@ -334,6 +334,8 @@ Usage : ps Example : ps ``` +![Ps command](../assets/modules-10.png) + ### env Display environment variables. ``` diff --git a/src/agent/core/process.nim b/src/agent/core/process.nim index 9a4317b..8c65e0d 100644 --- a/src/agent/core/process.nim +++ b/src/agent/core/process.nim @@ -1,5 +1,4 @@ import winim/lean -import winim/inc/tlhelp32 import strutils, strformat, tables, algorithm import ../utils/io import ../../common/[types, utils] @@ -11,37 +10,56 @@ type ppid*: DWORD name*: string user*: string + session*: ULONG children*: seq[DWORD] - # NtQuerySystemInformation = proc(systemInformationClass: SYSTEM_INFORMATION_CLASS, systemInformation: PVOID, systemInformationLength: ULONG, returnLength: PULONG): NTSTATUS {.stdcall.} + NtQuerySystemInformation = proc(systemInformationClass: SYSTEM_INFORMATION_CLASS, systemInformation: PVOID, systemInformationLength: ULONG, returnLength: PULONG): NTSTATUS {.stdcall.} NtOpenProcess = proc(hProcess: PHANDLE, desiredAccess: ACCESS_MASK, oa: PCOBJECT_ATTRIBUTES, clientId: PCLIENT_ID): NTSTATUS {.stdcall.} NtOpenProcessToken = proc(processHandle: HANDLE, desiredAccess: ACCESS_MASK, tokenHandle: PHANDLE): NTSTATUS {.stdcall.} - -const PROCESS_QUERY_LIMITED_INFORMATION = 0x00001000'i32 + NtClose = proc(handle: HANDLE): NTSTATUS {.stdcall.} proc cmp*(x, y: ProcessInfo): int = return cmp(x.pid, y.pid) +#[ + Retrieve snapshot of all currently running processes using NtQuerySystemInformation +]# +proc processSnapshot*(): PSYSTEM_PROCESS_INFORMATION = + var + pSystemProcInfo: PSYSTEM_PROCESS_INFORMATION + status: NTSTATUS = 0 + returnLength: ULONG = 0 + + let pNtQuerySystemInformation = cast[NtQuerySystemInformation](GetProcAddress(GetModuleHandleA(protect("ntdll")), protect("NtQuerySystemInformation"))) + + # Retrieve returnLength and allocate sufficient memory + discard pNtQuerySystemInformation(systemProcessInformation, NULL, 0, addr returnLength) + pSystemProcInfo = cast[PSYSTEM_PROCESS_INFORMATION](LocalAlloc(LMEM_FIXED, returnLength)) + if pSystemProcInfo == NULL: + raise newException(CatchableError, "1.2" & GetLastError().getError()) + + # Retrieve system process information + status = pNtQuerySystemInformation(systemProcessInformation, cast[PVOID](pSystemProcInfo), returnLength, addr returnLength) + if status != STATUS_SUCCESS: + raise newException(CatchableError, "b" & status.getNtError()) + + return pSystemProcInfo + +#[ + Retrieve information about running processes +]# proc processList*(): Table[DWORD, ProcessInfo] = result = initTable[DWORD, ProcessInfo]() # Take a snapshot of running processes - let hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) - if hSnapshot == INVALID_HANDLE_VALUE: - raise newException(CatchableError, GetLastError().getError) - 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, GetLastError().getError) + var sysProcessInfo = processSnapshot() + defer: LocalFree(cast[HLOCAL](sysProcessInfo)) let pNtOpenProcess = cast[NtOpenProcess](GetProcAddress(GetModuleHandleA(protect("ntdll")), protect("NtOpenProcess"))) let pNtOpenProcessToken = cast[NtOpenProcessToken](GetProcAddress(GetModuleHandleA(protect("ntdll")), protect("NtOpenProcessToken"))) + let pNtClose = cast[NtClose](GetProcAddress(GetModuleHandleA(protect("ntdll")), protect("NtClose"))) - while Process32Next(hSnapshot, addr pe32): + while true: var status: NTSTATUS hToken: HANDLE @@ -49,26 +67,35 @@ proc processList*(): Table[DWORD, ProcessInfo] = oa: OBJECT_ATTRIBUTES clientId: CLIENT_ID - var procInfo = ProcessInfo( - pid: pe32.th32ProcessID, - ppid: pe32.th32ParentProcessID, - name: $cast[WideCString](addr pe32.szExeFile[0]), + var + pid = cast[DWORD](sysProcessInfo.UniqueProcessId) + ppid = cast[DWORD](sysProcessInfo.InheritedFromUniqueProcessId) + + # Retrieve process information + result[pid] = ProcessInfo( + pid: pid, + ppid: ppid, + name: $sysProcessInfo.ImageName.Buffer, + session: sysProcessInfo.SessionId, children: @[] ) # Retrieve user context InitializeObjectAttributes(addr oa, NULL, 0, 0, NULL) - clientId.UniqueProcess = cast[HANDLE](pe32.th32ProcessID) + clientId.UniqueProcess = cast[HANDLE](pid) clientId.UniqueThread = 0 status = pNtOpenProcess(addr hProcess, PROCESS_QUERY_INFORMATION, addr oa, addr clientId) if status == STATUS_SUCCESS and hProcess != 0: status = pNtOpenProcessToken(hProcess, TOKEN_QUERY, addr hToken) if status == STATUS_SUCCESS and hToken != 0: - procInfo.user = hToken.getTokenUser().username - - result[pe32.th32ProcessID] = procInfo + result[pid].user = hToken.getTokenUser().username + defer: + discard pNtClose(hProcess) + discard pNtClose(hToken) - for pid, procInfo in result.mpairs(): - if result.contains(procInfo.ppid): - result[procInfo.ppid].children.add(pid) + # Move to next process + if sysProcessInfo.NextEntryOffset == 0: + break + + sysProcessInfo = cast[PSYSTEM_PROCESS_INFORMATION](cast[ULONG_PTR](sysProcessInfo) + sysProcessInfo.NextEntryOffset) \ No newline at end of file diff --git a/src/modules/systeminfo.nim b/src/modules/systeminfo.nim index 1f7f255..d324f40 100644 --- a/src/modules/systeminfo.nim +++ b/src/modules/systeminfo.nim @@ -52,14 +52,24 @@ when defined(agent): var procMap = processList() + # Create child-parent process relationships for pid, procInfo in procMap.mpairs(): - if not procMap.contains(procInfo.ppid): + if procMap.contains(procInfo.ppid) and procInfo.ppid != 0: + procMap[procInfo.ppid].children.add(pid) + else: processes.add(pid) # Add header row - let headers = @[protect("PID"), protect("PPID"), protect("Process"), protect("User context")] - output &= fmt"{headers[0]:<10}{headers[1]:<10}{headers[2]:<40}{headers[3]}" & "\n" - output &= "-".repeat(len(headers[0])).alignLeft(10) & "-".repeat(len(headers[1])).alignLeft(10) & "-".repeat(len(headers[2])).alignLeft(40) & "-".repeat(len(headers[3])) & "\n" + let headers = @[ + protect("PID"), + protect("PPID"), + protect("Process"), + protect("Session"), + protect("User context") + ] + + output &= fmt"{headers[0]:<10}{headers[1]:<10}{headers[2]:<40}{headers[3]:<10}{headers[4]}" & "\n" + output &= "-".repeat(len(headers[0])).alignLeft(10) & "-".repeat(len(headers[1])).alignLeft(10) & "-".repeat(len(headers[2])).alignLeft(40) & "-".repeat(len(headers[3])).alignLeft(10) & "-".repeat(len(headers[4])) & "\n" # Format and print process proc printProcess(pid: DWORD, indentSpaces: int = 0) = @@ -68,15 +78,14 @@ when defined(agent): var process = procMap[pid] let processName = " ".repeat(indentSpaces) & process.name - - output &= fmt"{process.pid:<10}{process.ppid:<10}{processName:<40}{process.user}" & "\n" - + output &= fmt"{$process.pid:<10}{$process.ppid:<10}{processName:<40}{$process.session:<10}{process.user}" & "\n" + # Recursively print child processes with indentation process.children.sort() for childPid in process.children: printProcess(childPid, indentSpaces + 2) - - # Iterate over root processes + + # Iterate over root processes to construct the output processes.sort() for pid in processes: printProcess(pid)