Updated 'ps' command implementation.

docs/7-MODULES.md

@@ -334,6 +334,8 @@ Usage    : ps
 Example  : ps
 ```
 
+![Ps command](../assets/modules-10.png)
+
 ### env
 Display environment variables.
 ```

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.}
-
+    NtClose = proc(handle: HANDLE): NTSTATUS {.stdcall.}
-const PROCESS_QUERY_LIMITED_INFORMATION = 0x00001000'i32
 
 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)
+    var sysProcessInfo = processSnapshot() 
-    if hSnapshot == INVALID_HANDLE_VALUE: 
+    defer: LocalFree(cast[HLOCAL](sysProcessInfo))
-        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)
     
     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(
+        var 
-            pid: pe32.th32ProcessID,
+            pid = cast[DWORD](sysProcessInfo.UniqueProcessId)
-            ppid: pe32.th32ParentProcessID,
+            ppid = cast[DWORD](sysProcessInfo.InheritedFromUniqueProcessId)
-            name: $cast[WideCString](addr pe32.szExeFile[0]),
+
+        # 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[pid].user = hToken.getTokenUser().username
+        defer:
+            discard pNtClose(hProcess)
+            discard pNtClose(hToken)
 
-        result[pe32.th32ProcessID] = procInfo
+        # Move to next process
+        if sysProcessInfo.NextEntryOffset == 0: 
+            break
             
-    for pid, procInfo in result.mpairs():
+        sysProcessInfo = cast[PSYSTEM_PROCESS_INFORMATION](cast[ULONG_PTR](sysProcessInfo) + sysProcessInfo.NextEntryOffset)
-        if result.contains(procInfo.ppid):
-            result[procInfo.ppid].children.add(pid)

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")]
+            let headers = @[
-            output &= fmt"{headers[0]:<10}{headers[1]:<10}{headers[2]:<40}{headers[3]}" & "\n"
+                protect("PID"), 
-            output &= "-".repeat(len(headers[0])).alignLeft(10) & "-".repeat(len(headers[1])).alignLeft(10) & "-".repeat(len(headers[2])).alignLeft(40) & "-".repeat(len(headers[3])) & "\n"
+                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.session:<10}{process.user}" & "\n"
-                output &= fmt"{process.pid:<10}{process.ppid:<10}{processName:<40}{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)