Driver: PsTree: Enumerate processes on startup

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;
 }
 

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(&params->ImagePathName, params);
+
+    commandLine->Length = params->CommandLine.Length;
+    commandLine->MaximumLength = params->CommandLine.Length;
+    commandLine->Buffer = GetUnicodeStringBuffer(&params->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);
+}

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

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;
     }

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();
 

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) {

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;

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);
+}

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

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;
+}

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