Updated 'ps' command implementation.
This commit is contained in:
BIN
assets/modules-10.png
Normal file
BIN
assets/modules-10.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 301 KiB |
@@ -334,6 +334,8 @@ Usage : ps
|
||||
Example : ps
|
||||
```
|
||||
|
||||

|
||||
|
||||
### env
|
||||
Display environment variables.
|
||||
```
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user