Files
conquest/src/agent/core/process.nim
Jakob Friedl 4907639848 Small changes.
2025-11-06 16:48:06 +01:00

101 lines
4.0 KiB
Nim

import winim/lean
import tables
import ../utils/io
import ../../common/utils
import token
type
ProcessInfo* = object
pid*: DWORD
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.}
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.}
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
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 true:
var
status: NTSTATUS
hToken: HANDLE
hProcess: HANDLE
oa: OBJECT_ATTRIBUTES
clientId: CLIENT_ID
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](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:
result[pid].user = hToken.getTokenUser().username
defer:
discard pNtClose(hProcess)
discard pNtClose(hToken)
# Move to next process
if sysProcessInfo.NextEntryOffset == 0:
break
sysProcessInfo = cast[PSYSTEM_PROCESS_INFORMATION](cast[ULONG_PTR](sysProcessInfo) + sysProcessInfo.NextEntryOffset)