From 109aee7ff12271a15c04a13c56408d71cf1550b2 Mon Sep 17 00:00:00 2001 From: Nodir Temirkhodjaev Date: Mon, 7 Feb 2022 14:46:11 +0300 Subject: [PATCH] Driver: PsTree: Enumerate processes on startup --- src/driver/fortdev.c | 5 + src/driver/fortps.c | 206 +++++++++++++++++++++++++++++++++++++- src/driver/fortps.h | 2 + src/driver/fortutl.c | 35 +++++-- src/driver/fortutl.h | 2 + src/driver/fortwrk.c | 11 +- src/driver/fortwrk.h | 2 + src/driver/wdm/um_ntifs.c | 32 ++++++ src/driver/wdm/um_ntifs.h | 34 +++++++ src/driver/wdm/um_wdm.c | 26 +++++ src/driver/wdm/um_wdm.h | 22 ++++ 11 files changed, 362 insertions(+), 15 deletions(-) diff --git a/src/driver/fortdev.c b/src/driver/fortdev.c index a27c067b..601ad514 100644 --- a/src/driver/fortdev.c +++ b/src/driver/fortdev.c @@ -350,6 +350,11 @@ FORT_API NTSTATUS fort_device_load(PDEVICE_OBJECT device) status = fort_syscb_time_register(); } + /* Enumerate processes in background */ + if (NT_SUCCESS(status)) { + fort_worker_queue(&fort_device()->worker, FORT_WORKER_PSTREE, &fort_pstree_enum_processes); + } + return status; } diff --git a/src/driver/fortps.c b/src/driver/fortps.c index 0d69a7ae..ed8533e7 100644 --- a/src/driver/fortps.c +++ b/src/driver/fortps.c @@ -9,6 +9,7 @@ #define FORT_PSTREE_POOL_TAG 'PwfF' #define FORT_SVCHOST_PREFIX L"\\svchost\\" +#define FORT_SVCHOST_EXE L"svchost.exe" #define FORT_PSTREE_NAME_LEN_MAX (120 * sizeof(WCHAR)) #define FORT_PSTREE_NAMES_POOL_SIZE (4 * 1024) @@ -59,8 +60,53 @@ typedef struct _SYSTEM_PROCESSES #endif #if defined(FORT_DRIVER) + +typedef struct _PEB_LDR_DATA +{ + BYTE Reserved1[8]; + PVOID Reserved2[3]; + LIST_ENTRY InMemoryOrderModuleList; +} PEB_LDR_DATA, *PPEB_LDR_DATA; + +typedef struct _RTL_USER_PROCESS_PARAMETERS +{ + BYTE Reserved1[16]; + PVOID Reserved2[10]; + UNICODE_STRING ImagePathName; + UNICODE_STRING CommandLine; +} RTL_USER_PROCESS_PARAMETERS, *PRTL_USER_PROCESS_PARAMETERS; + +typedef VOID(NTAPI *PPS_POST_PROCESS_INIT_ROUTINE)(VOID); + +typedef struct _PEB +{ + BYTE Reserved1[2]; + BYTE BeingDebugged; + BYTE Reserved2[1]; + PVOID Reserved3[2]; + PPEB_LDR_DATA Ldr; + PRTL_USER_PROCESS_PARAMETERS ProcessParameters; + PVOID Reserved4[3]; + PVOID AtlThunkSListPtr; + PVOID Reserved5; + ULONG Reserved6; + PVOID Reserved7; + ULONG Reserved8; + ULONG AtlThunkSListPtr32; + PVOID Reserved9[45]; + BYTE Reserved10[96]; + PPS_POST_PROCESS_INIT_ROUTINE PostProcessInitRoutine; + BYTE Reserved11[128]; + PVOID Reserved12[1]; + ULONG SessionId; +} PEB, *PPEB; + NTSTATUS NTAPI ZwQuerySystemInformation(ULONG systemInformationClass, PVOID systemInformation, ULONG systemInformationLength, PULONG returnLength); + +NTSTATUS NTAPI ZwQueryInformationProcess(HANDLE processHandle, ULONG processInformationClass, + PVOID processInformation, ULONG processInformationLength, PULONG returnLength); + #endif #define fort_pstree_get_proc(ps_tree, index) \ @@ -87,12 +133,12 @@ static void fort_pstree_name_del(PFORT_PSTREE ps_tree, PFORT_PSNAME ps_name) static BOOL fort_pstree_svchost_check( PCUNICODE_STRING path, PCUNICODE_STRING commandLine, PUNICODE_STRING serviceName) { - const USHORT svchostSize = sizeof(L"svchost.exe") - sizeof(WCHAR); /* skip terminating zero */ + const USHORT svchostSize = sizeof(FORT_SVCHOST_EXE) - sizeof(WCHAR); /* skip terminating zero */ const USHORT svchostCount = svchostSize / sizeof(WCHAR); const USHORT sys32Size = path->Length - svchostSize; const USHORT sys32Count = sys32Size / sizeof(WCHAR); - if (_wcsnicmp(path->Buffer + sys32Count, L"svchost.exe", svchostCount) != 0) + if (_wcsnicmp(path->Buffer + sys32Count, FORT_SVCHOST_EXE, svchostCount) != 0) return FALSE; PCUNICODE_STRING sys32Path = fort_system32_path(); @@ -261,8 +307,11 @@ static void NTAPI fort_pstree_notify( #endif if (proc == NULL) { - fort_pstree_handle_new_proc(ps_tree, createInfo->ImageFileName, createInfo->CommandLine, - pid_hash, processId, parentProcessId); + UNICODE_STRING path = *createInfo->ImageFileName; + fort_path_prefix_adjust(&path); + + fort_pstree_handle_new_proc( + ps_tree, &path, createInfo->CommandLine, pid_hash, processId, parentProcessId); } } @@ -337,3 +386,152 @@ FORT_API PFORT_PSNAME fort_pstree_get_proc_name( return ps_name; } + +static HANDLE OpenProcessById(DWORD processId) +{ + if (processId == 0 || processId == 4) + return NULL; // skip "System Idle Process" (0) and "System" (4) processes + + NTSTATUS status; + + PEPROCESS peProcess = NULL; + status = PsLookupProcessByProcessId((HANDLE) (ptrdiff_t) processId, &peProcess); + if (!NT_SUCCESS(status) || peProcess == NULL) { + LOG("PsTree: Lookup Process Error: %x\n", status); + return NULL; + } + + HANDLE processHandle = NULL; + status = ObOpenObjectByPointer(peProcess, 0, NULL, 0, 0, KernelMode, &processHandle); + + ObDereferenceObject(peProcess); + + if (!NT_SUCCESS(status) || processHandle == NULL) { + LOG("PsTree: Open Process Object Error: %x\n", status); + return NULL; + } + + return processHandle; +} + +static PWCHAR GetUnicodeStringBuffer(PCUNICODE_STRING string, PRTL_USER_PROCESS_PARAMETERS params) +{ +#ifdef _X86_ + if ((PCHAR) string->Buffer < (PCHAR) processParams) { + return (PWCHAR) ((PCHAR) params + (DWORD) (ptrdiff_t) string->Buffer); + } +#else + UNUSED(params); +#endif + + return string->Buffer; +} + +static NTSTATUS GetCurrentProcessPathArgs(PUNICODE_STRING path, PUNICODE_STRING commandLine) +{ + NTSTATUS status; + + PROCESS_BASIC_INFORMATION procBasicInfo; + status = ZwQueryInformationProcess(ZwCurrentProcess(), ProcessBasicInformation, &procBasicInfo, + sizeof(PROCESS_BASIC_INFORMATION), NULL); + if (!NT_SUCCESS(status)) { + LOG("PsTree: Query Process Error: %x\n", status); + return status; + } + + if (procBasicInfo.PebBaseAddress == NULL) { + LOG("PsTree: Query Process Error: PebBaseAddress\n"); + return STATUS_INVALID_ADDRESS; + } + + PRTL_USER_PROCESS_PARAMETERS params = procBasicInfo.PebBaseAddress->ProcessParameters; + + path->Length = params->ImagePathName.Length; + path->MaximumLength = params->ImagePathName.Length; + path->Buffer = GetUnicodeStringBuffer(¶ms->ImagePathName, params); + + commandLine->Length = params->CommandLine.Length; + commandLine->MaximumLength = params->CommandLine.Length; + commandLine->Buffer = GetUnicodeStringBuffer(¶ms->CommandLine, params); + + return STATUS_SUCCESS; +} + +static void fort_pstree_attach_process(PSYSTEM_PROCESSES processEntry, HANDLE processHandle) +{ + NTSTATUS status; + + PRKPROCESS process; + status = ObReferenceObjectByHandle( + processHandle, 0, *PsProcessType, KernelMode, (PVOID *) &process, NULL); + if (!NT_SUCCESS(status)) { + LOG("PsTree: Attach Process Error: %x\n", status); + return; + } + + KAPC_STATE apcState; + KeStackAttachProcess(process, &apcState); + { + UNICODE_STRING path; + UNICODE_STRING commandLine; + + status = GetCurrentProcessPathArgs(&path, &commandLine); + if (NT_SUCCESS(status)) { + const HANDLE processId = (HANDLE) (ptrdiff_t) processEntry->ProcessId; + const HANDLE parentProcessId = (HANDLE) (ptrdiff_t) processEntry->ParentProcessId; + + PS_CREATE_NOTIFY_INFO createInfo = { .ParentProcessId = parentProcessId, + .ImageFileName = &path, + .CommandLine = &commandLine }; + + fort_pstree_notify(/*process=*/NULL, processId, &createInfo); + } + } + KeUnstackDetachProcess(&apcState); + + ObDereferenceObject(process); +} + +static void fort_pstree_enum_processes_loop(PSYSTEM_PROCESSES processEntry) +{ + for (;;) { + const DWORD processId = (DWORD) processEntry->ProcessId; + + const HANDLE processHandle = OpenProcessById(processId); + if (processHandle != NULL) { + fort_pstree_attach_process(processEntry, processHandle); + + ZwClose(processHandle); + } + + if (processEntry->NextEntryOffset == 0) + break; + + processEntry = (PSYSTEM_PROCESSES) ((PUCHAR) processEntry + processEntry->NextEntryOffset); + } +} + +FORT_API void NTAPI fort_pstree_enum_processes(void) +{ + NTSTATUS status; + + ULONG bufferSize; + status = ZwQuerySystemInformation(SystemProcessInformation, NULL, 0, &bufferSize); + if (status != STATUS_INFO_LENGTH_MISMATCH) + return; + + bufferSize *= 3; /* for possibly new created processes/threads */ + + PVOID buffer = fort_mem_alloc(bufferSize, FORT_PSTREE_POOL_TAG); + if (buffer == NULL) + return; + + status = ZwQuerySystemInformation(SystemProcessInformation, buffer, bufferSize, &bufferSize); + if (NT_SUCCESS(status)) { + fort_pstree_enum_processes_loop(buffer); + } else { + LOG("PsTree: Enum Processes Error: %x\n", status); + } + + fort_mem_free(buffer, FORT_PSTREE_POOL_TAG); +} diff --git a/src/driver/fortps.h b/src/driver/fortps.h index 9fe48fda..9d6f86e1 100644 --- a/src/driver/fortps.h +++ b/src/driver/fortps.h @@ -35,6 +35,8 @@ FORT_API void fort_pstree_close(PFORT_PSTREE ps_tree); FORT_API PFORT_PSNAME fort_pstree_get_proc_name( PFORT_PSTREE ps_tree, DWORD processId, PUNICODE_STRING path); +FORT_API void NTAPI fort_pstree_enum_processes(void); + #ifdef __cplusplus } // extern "C" #endif diff --git a/src/driver/fortutl.c b/src/driver/fortutl.c index 914d5b06..b4dcd618 100644 --- a/src/driver/fortutl.c +++ b/src/driver/fortutl.c @@ -77,23 +77,35 @@ FORT_API NTSTATUS fort_driver_path( return status; } -static NTSTATUS fort_system32_path_set(PCUNICODE_STRING path, USHORT charCount) +static NTSTATUS fort_system32_path_set(PCUNICODE_STRING path) { - if (charCount >= sizeof(g_system32PathBuffer) / sizeof(WCHAR)) + if (path->Length >= sizeof(g_system32PathBuffer) - sizeof(WCHAR)) return STATUS_BUFFER_OVERFLOW; - const USHORT len = charCount * sizeof(WCHAR); + RtlCopyMemory(g_system32PathBuffer, path->Buffer, path->Length); + g_system32PathBuffer[path->Length / sizeof(WCHAR)] = L'\0'; - RtlCopyMemory(g_system32PathBuffer, path->Buffer, len); - g_system32PathBuffer[charCount] = L'\0'; - - g_system32Path.Length = len; + g_system32Path.Length = path->Length; g_system32Path.MaximumLength = sizeof(g_system32PathBuffer); g_system32Path.Buffer = g_system32PathBuffer; return STATUS_SUCCESS; } +FORT_API void fort_path_prefix_adjust(PUNICODE_STRING path) +{ + if (path->Length < 7) + return; + + PCWCHAR p = path->Buffer; + + if (p[0] == '\\' && p[1] == '?' && p[2] == '?' && p[3] == '\\' && p[5] == ':') { + path->Buffer += 4; + path->Length -= 4 * sizeof(WCHAR); + path->MaximumLength -= 4 * sizeof(WCHAR); + } +} + FORT_API NTSTATUS fort_system32_path_init(PDRIVER_OBJECT driver, PUNICODE_STRING regPath) { NTSTATUS status; @@ -116,10 +128,13 @@ FORT_API NTSTATUS fort_system32_path_init(PDRIVER_OBJECT driver, PUNICODE_STRING } if (sp != NULL) { - const USHORT charCount = (USHORT) (sp + 1 /* include the separator */ - - driverPath.Buffer); + UNICODE_STRING sys32Path = driverPath; + sys32Path.Length = (USHORT) ((PCHAR) (sp + 1) /* include the separator */ + - (PCHAR) (driverPath.Buffer)); - status = fort_system32_path_set(&driverPath, charCount); + fort_path_prefix_adjust(&sys32Path); + + status = fort_system32_path_set(&sys32Path); } else { status = STATUS_OBJECT_PATH_INVALID; } diff --git a/src/driver/fortutl.h b/src/driver/fortutl.h index e4ac1efd..e6d891b9 100644 --- a/src/driver/fortutl.h +++ b/src/driver/fortutl.h @@ -10,6 +10,8 @@ extern "C" { FORT_API NTSTATUS fort_driver_path( PDRIVER_OBJECT driver, PUNICODE_STRING regPath, PUNICODE_STRING outPath); +FORT_API void fort_path_prefix_adjust(PUNICODE_STRING path); + FORT_API NTSTATUS fort_system32_path_init(PDRIVER_OBJECT driver, PUNICODE_STRING regPath); FORT_API PUNICODE_STRING fort_system32_path(); diff --git a/src/driver/fortwrk.c b/src/driver/fortwrk.c index 4d5cfe92..f4dd4922 100644 --- a/src/driver/fortwrk.c +++ b/src/driver/fortwrk.c @@ -15,14 +15,23 @@ static void NTAPI fort_worker_callback(PVOID device, PVOID context, PIO_WORKITEM if (id_bits & FORT_WORKER_REAUTH) { worker->reauth_func(); } + + if (id_bits & FORT_WORKER_PSTREE) { + worker->pstree_func(); + } } FORT_API void fort_worker_queue(PFORT_WORKER worker, UCHAR work_id, FORT_WORKER_FUNC worker_func) { const UCHAR id_bits = InterlockedOr8(&worker->id_bits, work_id); - if (work_id == FORT_WORKER_REAUTH) { + switch (work_id) { + case FORT_WORKER_REAUTH: { worker->reauth_func = worker_func; + } break; + case FORT_WORKER_PSTREE: { + worker->pstree_func = worker_func; + } break; } if (id_bits == 0) { diff --git a/src/driver/fortwrk.h b/src/driver/fortwrk.h index cfe1b4bd..326f4467 100644 --- a/src/driver/fortwrk.h +++ b/src/driver/fortwrk.h @@ -4,6 +4,7 @@ #include "fortdrv.h" #define FORT_WORKER_REAUTH 0x01 +#define FORT_WORKER_PSTREE 0x02 typedef void (NTAPI *FORT_WORKER_FUNC)(void); @@ -12,6 +13,7 @@ typedef struct fort_worker UCHAR volatile id_bits; FORT_WORKER_FUNC reauth_func; + FORT_WORKER_FUNC pstree_func; PIO_WORKITEM item; } FORT_WORKER, *PFORT_WORKER; diff --git a/src/driver/wdm/um_ntifs.c b/src/driver/wdm/um_ntifs.c index bf095f68..173ddf80 100644 --- a/src/driver/wdm/um_ntifs.c +++ b/src/driver/wdm/um_ntifs.c @@ -8,3 +8,35 @@ NTSTATUS RtlDowncaseUnicodeString(PUNICODE_STRING destinationString, PCUNICODE_S UNUSED(allocateDestinationString); return STATUS_SUCCESS; } + +NTSTATUS PsLookupProcessByProcessId(HANDLE processId, PEPROCESS *process) +{ + UNUSED(processId); + UNUSED(process); + return STATUS_SUCCESS; +} + +NTSTATUS ObOpenObjectByPointer(PVOID object, ULONG handleAttributes, + PACCESS_STATE passedAccessState, ACCESS_MASK desiredAccess, POBJECT_TYPE objectType, + KPROCESSOR_MODE accessMode, PHANDLE handle) +{ + UNUSED(object); + UNUSED(handleAttributes); + UNUSED(passedAccessState); + UNUSED(desiredAccess); + UNUSED(objectType); + UNUSED(accessMode); + UNUSED(handle); + return STATUS_SUCCESS; +} + +void KeStackAttachProcess(PRKPROCESS process, PRKAPC_STATE apcState) +{ + UNUSED(process); + UNUSED(apcState); +} + +void KeUnstackDetachProcess(PRKAPC_STATE apcState) +{ + UNUSED(apcState); +} diff --git a/src/driver/wdm/um_ntifs.h b/src/driver/wdm/um_ntifs.h index edef1d20..665d81e7 100644 --- a/src/driver/wdm/um_ntifs.h +++ b/src/driver/wdm/um_ntifs.h @@ -7,9 +7,43 @@ extern "C" { #endif +typedef struct _KAPC_STATE +{ + LIST_ENTRY ApcListHead[MaximumMode]; + struct _KPROCESS *Process; + union { + UCHAR InProgressFlags; + struct + { + BOOLEAN KernelApcInProgress : 1; + BOOLEAN SpecialApcInProgress : 1; + }; + }; + + BOOLEAN KernelApcPending; + union { + BOOLEAN UserApcPendingAll; + struct + { + BOOLEAN SpecialUserApcPending : 1; + BOOLEAN UserApcPending : 1; + }; + }; +} KAPC_STATE, *PKAPC_STATE, *PRKAPC_STATE; + FORT_API NTSTATUS RtlDowncaseUnicodeString(PUNICODE_STRING destinationString, PCUNICODE_STRING sourceString, BOOLEAN allocateDestinationString); +FORT_API NTSTATUS PsLookupProcessByProcessId(HANDLE processId, PEPROCESS *process); + +FORT_API NTSTATUS ObOpenObjectByPointer(PVOID object, ULONG handleAttributes, + PACCESS_STATE passedAccessState, ACCESS_MASK desiredAccess, POBJECT_TYPE objectType, + KPROCESSOR_MODE accessMode, PHANDLE handle); + +FORT_API VOID KeStackAttachProcess(PRKPROCESS process, PRKAPC_STATE apcState); + +FORT_API VOID KeUnstackDetachProcess(PRKAPC_STATE apcState); + #ifdef __cplusplus } // extern "C" #endif diff --git a/src/driver/wdm/um_wdm.c b/src/driver/wdm/um_wdm.c index 390f524b..6bea8499 100644 --- a/src/driver/wdm/um_wdm.c +++ b/src/driver/wdm/um_wdm.c @@ -372,3 +372,29 @@ NTSTATUS ZwQuerySystemInformation(ULONG systemInformationClass, PVOID systemInfo UNUSED(returnLength); return STATUS_SUCCESS; } + +NTSTATUS ZwQueryInformationProcess(HANDLE processHandle, ULONG processInformationClass, + PVOID processInformation, ULONG processInformationLength, PULONG returnLength) +{ + UNUSED(processHandle); + UNUSED(processInformationClass); + UNUSED(processInformation); + UNUSED(processInformationLength); + UNUSED(returnLength); + return STATUS_SUCCESS; +} + +POBJECT_TYPE *PsProcessType = NULL; + +NTSTATUS ObReferenceObjectByHandle(HANDLE handle, ACCESS_MASK desiredAccess, + POBJECT_TYPE objectType, KPROCESSOR_MODE accessMode, PVOID *object, + POBJECT_HANDLE_INFORMATION handleInformation) +{ + UNUSED(handle); + UNUSED(desiredAccess); + UNUSED(objectType); + UNUSED(accessMode); + UNUSED(object); + UNUSED(handleInformation); + return STATUS_SUCCESS; +} diff --git a/src/driver/wdm/um_wdm.h b/src/driver/wdm/um_wdm.h index 601375a9..9a8e16e2 100644 --- a/src/driver/wdm/um_wdm.h +++ b/src/driver/wdm/um_wdm.h @@ -101,6 +101,16 @@ typedef struct KIRQL CancelIrql; } IRP, *PIRP; +typedef struct _ACCESS_STATE *PACCESS_STATE; +typedef struct _KPROCESS *PKPROCESS, *PRKPROCESS, *PEPROCESS; +typedef struct _OBJECT_TYPE *POBJECT_TYPE; + +typedef struct _OBJECT_HANDLE_INFORMATION +{ + ULONG HandleAttributes; + ACCESS_MASK GrantedAccess; +} OBJECT_HANDLE_INFORMATION, *POBJECT_HANDLE_INFORMATION; + struct _DRIVER_OBJECT; typedef struct _DEVICE_OBJECT @@ -207,6 +217,9 @@ FORT_API ULONG DbgPrintEx(ULONG componentId, ULONG level, PCSTR format, ...); #define NT_ASSERT(cond) assert((cond)) +#define NtCurrentProcess() ((HANDLE) (LONG_PTR) -1) +#define ZwCurrentProcess() NtCurrentProcess() + #define NonPagedPool 0 #define NonPagedPoolExecute NonPagedPool FORT_API PVOID ExAllocatePoolWithTag(PVOID type, SIZE_T size, ULONG tag); @@ -313,6 +326,15 @@ FORT_API NTSTATUS ZwQuerySymbolicLinkObject( FORT_API NTSTATUS ZwQuerySystemInformation(ULONG systemInformationClass, PVOID systemInformation, ULONG systemInformationLength, PULONG returnLength); +FORT_API NTSTATUS ZwQueryInformationProcess(HANDLE processHandle, ULONG processInformationClass, + PVOID processInformation, ULONG processInformationLength, PULONG returnLength); + +extern POBJECT_TYPE *PsProcessType; + +FORT_API NTSTATUS ObReferenceObjectByHandle(HANDLE handle, ACCESS_MASK desiredAccess, + POBJECT_TYPE objectType, KPROCESSOR_MODE accessMode, PVOID *object, + POBJECT_HANDLE_INFORMATION handleInformation); + #ifdef __cplusplus } // extern "C" #endif