From ea2278e7b3cde166354e56c9694c229a5da42f1e Mon Sep 17 00:00:00 2001 From: donnaskiez <54422048+donnaskiez@users.noreply.github.com> Date: Sun, 9 Jun 2024 17:22:22 +1000 Subject: [PATCH] internal refactoring (#14) - Refactor process list. New implementation consists of a hashmap. Each process entry then contains the associated user modules. - Implement user module integrity checks on timer callback --- driver/callbacks.c | 484 ++++++++++++++----- driver/callbacks.h | 47 +- driver/common.h | 65 ++- driver/crypt.c | 8 +- driver/driver.c | 80 ++- driver/driver.h | 15 +- driver/driver.vcxproj | 2 + driver/driver.vcxproj.filters | 6 + driver/hw.c | 2 +- driver/integrity.c | 184 ++++--- driver/integrity.h | 7 +- driver/io.c | 41 +- driver/map.c | 241 +++++++++ driver/map.h | 78 +++ driver/modules.c | 34 +- driver/pe.c | 52 ++ driver/pe.h | 13 + driver/pool.c | 26 +- driver/session.c | 44 +- driver/util.c | 64 +++ driver/util.h | 8 + module/dispatcher/dispatcher.cpp | 11 +- module/dispatcher/dispatcher.h | 6 +- module/kernel_interface/kernel_interface.cpp | 8 +- module/kernel_interface/kernel_interface.h | 31 +- module/module.cpp | 44 +- module/module.h | 8 + 27 files changed, 1276 insertions(+), 333 deletions(-) create mode 100644 driver/map.c create mode 100644 driver/map.h diff --git a/driver/callbacks.c b/driver/callbacks.c index 10d9de4..f0a61e5 100644 --- a/driver/callbacks.c +++ b/driver/callbacks.c @@ -10,6 +10,10 @@ #include "list.h" #include "session.h" #include "crypt.h" +#include "map.h" +#include "util.h" + +#define PROCESS_HASHMAP_BUCKET_COUNT 101 STATIC BOOLEAN @@ -27,17 +31,6 @@ EnumHandleCallback(_In_ PHANDLE_TABLE HandleTable, # pragma alloc_text(PAGE, ExUnlockHandleTableEntry) #endif -/* - * Its important on unload we dereference any objects to ensure the kernels - * reference count remains correct. - */ -VOID -CleanupProcessListFreeCallback(_In_ PPROCESS_LIST_ENTRY ProcessListEntry) -{ - ImpObDereferenceObject(ProcessListEntry->parent); - ImpObDereferenceObject(ProcessListEntry->process); -} - VOID CleanupThreadListFreeCallback(_In_ PTHREAD_LIST_ENTRY ThreadListEntry) { @@ -48,8 +41,8 @@ CleanupThreadListFreeCallback(_In_ PTHREAD_LIST_ENTRY ThreadListEntry) VOID UnregisterProcessCreateNotifyRoutine() { - PPROCESS_LIST_HEAD list = GetProcessList(); - InterlockedExchange(&list->active, FALSE); + PRTL_HASHMAP map = GetProcessHashmap(); + InterlockedExchange(&map->active, FALSE); ImpPsSetCreateProcessNotifyRoutine(ProcessCreateNotifyRoutine, TRUE); } @@ -69,24 +62,6 @@ UnregisterThreadCreateNotifyRoutine() ImpPsRemoveCreateThreadNotifyRoutine(ThreadCreateNotifyRoutine); } -/* - * While ExDeleteLookasideListEx already frees each item, we wanna allow - * ourselves to reduce the reference count to any objects we are referencing. - */ -VOID -CleanupProcessListOnDriverUnload() -{ - PPROCESS_LIST_HEAD list = GetProcessList(); - DEBUG_VERBOSE("Freeing process list"); - for (;;) { - if (!LookasideListFreeFirstEntry( - &list->start, &list->lock, CleanupProcessListFreeCallback)) { - ExDeleteLookasideListEx(&list->lookaside_list); - return; - } - } -} - VOID CleanupThreadListOnDriverUnload() { @@ -132,27 +107,6 @@ unlock: ImpKeReleaseGuardedMutex(&list->lock); } -VOID -EnumerateProcessListWithCallbackRoutine( - _In_ PROCESSLIST_CALLBACK_ROUTINE CallbackRoutine, _In_opt_ PVOID Context) -{ - PPROCESS_LIST_HEAD list = GetProcessList(); - ImpKeAcquireGuardedMutex(&list->lock); - - if (!CallbackRoutine) - goto unlock; - - PPROCESS_LIST_ENTRY entry = list->start.Next; - - while (entry) { - CallbackRoutine(entry, Context); - entry = (PPROCESS_LIST_ENTRY)entry->list.Next; - } - -unlock: - ImpKeReleaseGuardedMutex(&list->lock); -} - VOID EnumerateDriverListWithCallbackRoutine( _In_ DRIVERLIST_CALLBACK_ROUTINE CallbackRoutine, _In_opt_ PVOID Context) @@ -283,6 +237,79 @@ unlock: ImpKeReleaseGuardedMutex(&list->lock); } +STATIC +BOOLEAN +ProcessHashmapCompareFunction(_In_ PVOID Struct1, _In_ PVOID Struct2) +{ + HANDLE h1 = *((PHANDLE)Struct1); + HANDLE h2 = *((PHANDLE)Struct2); + + return h1 == h2 ? TRUE : FALSE; +} + +STATIC +UINT32 +ProcessHashmapHashFunction(_In_ UINT64 Key) +{ + return ((UINT32)Key) % PROCESS_HASHMAP_BUCKET_COUNT; +} + +STATIC +VOID +ImageLoadInsertNonSystemImageIntoProcessHashmap(_In_ PIMAGE_INFO ImageInfo, + _In_ HANDLE ProcessId, + _In_opt_ PUNICODE_STRING + FullImageName) +{ + NTSTATUS status = STATUS_UNSUCCESSFUL; + PEPROCESS process = NULL; + PRTL_HASHMAP map = GetProcessHashmap(); + PPROCESS_LIST_ENTRY entry = NULL; + PPROCESS_MAP_MODULE_ENTRY module = NULL; + PPROCESS_MODULE_MAP_CONTEXT context = NULL; + + if (!map->active) + return; + + status = PsLookupProcessByProcessId(ProcessId, &process); + + if (!NT_SUCCESS(status)) + return; + + KeAcquireGuardedMutex(&map->lock); + + /* the PEPROCESS is the first element and is the only thing compared, hence + * we can simply pass it in the context parameter.*/ + entry = RtlLookupEntryHashmap(GetProcessHashmap(), ProcessId, &ProcessId); + + /* critical error has occured */ + if (!entry) { + DEBUG_ERROR("RtlLookupEntryHashmap failed."); + goto end; + } + + context = (PPROCESS_MODULE_MAP_CONTEXT)map->context; + module = ExAllocateFromLookasideListEx(&context->pool); + + if (!module) + goto end; + + /* for now lets just do base and size */ + module->base = ImageInfo->ImageBase; + module->size = ImageInfo->ImageSize; + + /* We dont care if this errors. */ + if (FullImageName) + UnicodeToCharBufString( + FullImageName, module->path, sizeof(module->path)); + + InsertTailList(&entry->module_list, &module->entry); + entry->list_count++; + +end: + KeReleaseGuardedMutex(&map->lock); +} + VOID ImageLoadNotifyRoutineCallback(_In_opt_ PUNICODE_STRING FullImageName, _In_ HANDLE ProcessId, @@ -299,11 +326,15 @@ ImageLoadNotifyRoutineCallback(_In_opt_ PUNICODE_STRING FullImageName, if (InterlockedExchange(&list->active, list->active) == FALSE) return; - if (ImageInfo->SystemModeImage == FALSE) + if (ImageInfo->SystemModeImage == FALSE) { + ImageLoadInsertNonSystemImageIntoProcessHashmap( + ImageInfo, ProcessId, FullImageName); return; + } FindDriverEntryByBaseAddress(ImageInfo->ImageBase, &entry); + /* if we image exists, exit */ if (entry) return; @@ -322,23 +353,10 @@ ImageLoadNotifyRoutineCallback(_In_opt_ PUNICODE_STRING FullImageName, module.ImageSize = ImageInfo->ImageSize; if (FullImageName) { - status = RtlUnicodeStringToAnsiString(&ansi_path, FullImageName, TRUE); - - if (!NT_SUCCESS(status)) { - DEBUG_ERROR("RtlUnicodeStringToAnsiString failed with status %x", - status); - goto hash; - } - - if (ansi_path.Length > sizeof(module.FullPathName)) { - RtlFreeAnsiString(&ansi_path); - goto hash; - } - - RtlCopyMemory(module.FullPathName, ansi_path.Buffer, ansi_path.Length); - RtlCopyMemory(entry->path, ansi_path.Buffer, ansi_path.Length); - - RtlFreeAnsiString(&ansi_path); + UnicodeToCharBufString( + FullImageName, module.FullPathName, sizeof(module.FullPathName)); + RtlCopyMemory( + entry->path, module.FullPathName, sizeof(module.FullPathName)); } DEBUG_VERBOSE("New system image ansi: %s", entry->path); @@ -359,29 +377,172 @@ hash: ListInsert(&list->start, entry, &list->lock); } -NTSTATUS -InitialiseProcessList() +/* assumes map lock is held */ +VOID +FreeProcessEntryModuleList(_In_ PPROCESS_LIST_ENTRY Entry, + _In_opt_ PVOID Context) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - PPROCESS_LIST_HEAD list = GetProcessList(); + UNREFERENCED_PARAMETER(Context); - status = ExInitializeLookasideListEx(&list->lookaside_list, + PRTL_HASHMAP map = GetProcessHashmap(); + PLIST_ENTRY list = NULL; + PPROCESS_MAP_MODULE_ENTRY list_entry = NULL; + PPROCESS_MODULE_MAP_CONTEXT context = map->context; + + while (!IsListEmpty(&Entry->module_list)) { + list = RemoveTailList(&Entry->module_list); + list_entry = CONTAINING_RECORD(list, PROCESS_MAP_MODULE_ENTRY, entry); + + ExFreeToLookasideListEx(&context->pool, list_entry); + } +} + +VOID +EnumerateProcessModuleList(_In_ HANDLE ProcessId, + _In_ PROCESS_MODULE_CALLBACK Callback, + _In_opt_ PVOID Context) +{ + PRTL_HASHMAP map = GetProcessHashmap(); + BOOLEAN ret = FALSE; + PPROCESS_LIST_ENTRY entry = NULL; + PLIST_ENTRY list = NULL; + PPROCESS_MAP_MODULE_ENTRY module = NULL; + + if (!map->active) + return; + + KeAcquireGuardedMutex(&map->lock); + + entry = RtlLookupEntryHashmap(map, ProcessId, &ProcessId); + + if (!entry) + goto end; + + for (list = entry->module_list.Flink; list != &entry->module_list; + list = list->Flink) { + module = CONTAINING_RECORD(list, PROCESS_MAP_MODULE_ENTRY, entry); + + if (Callback(module, Context)) + goto end; + } + +end: + KeReleaseGuardedMutex(&map->lock); +} + +VOID +FindOurUserModeModuleEntry(_In_ PROCESS_MODULE_CALLBACK Callback, + _In_opt_ PVOID Context) +{ + PRTL_HASHMAP map = GetProcessHashmap(); + PPROCESS_LIST_ENTRY entry = NULL; + PACTIVE_SESSION session = GetActiveSession(); + PLIST_ENTRY list = NULL; + PPROCESS_MAP_MODULE_ENTRY module = NULL; + + if (!map->active) + return; + + KeAcquireGuardedMutex(&map->lock); + + entry = RtlLookupEntryHashmap(map, session->km_handle, &session->km_handle); + + if (!entry) + return; + + for (list = entry->module_list.Flink; list != &entry->module_list; + list = list->Flink) { + module = CONTAINING_RECORD(list, PROCESS_MAP_MODULE_ENTRY, entry); + + if (module->base == session->module.base_address && + module->size == session->module.size) { + Callback(module, Context); + goto end; + } + } + +end: + KeReleaseGuardedMutex(&map->lock); +} + +VOID +CleanupProcessHashmap() +{ + PRTL_HASHMAP map = GetProcessHashmap(); + PRTL_HASHMAP_ENTRY entry = NULL; + PRTL_HASHMAP_ENTRY temp = NULL; + PLIST_ENTRY list = NULL; + PPROCESS_MODULE_MAP_CONTEXT context = NULL; + + map->active = FALSE; + + KeAcquireGuardedMutex(&map->lock); + + /* First, free all module lists */ + RtlEnumerateHashmap(map, FreeProcessEntryModuleList, NULL); + + for (UINT32 index = 0; index < map->bucket_count; index++) { + entry = &map->buckets[index]; + + while (!IsListEmpty(&entry->entry)) { + list = RemoveHeadList(&entry->entry); + temp = CONTAINING_RECORD(list, RTL_HASHMAP_ENTRY, entry); + ExFreePoolWithTag(temp, POOL_TAG_HASHMAP); + } + } + + context = map->context; + + ExDeleteLookasideListEx(&context->pool); + ExFreePoolWithTag(map->context, POOL_TAG_HASHMAP); + RtlDeleteHashmap(map); + + KeReleaseGuardedMutex(&map->lock); +} + +NTSTATUS +InitialiseProcessHashmap() +{ + PAGED_CODE(); + + NTSTATUS status = STATUS_UNSUCCESSFUL; + PPROCESS_MODULE_MAP_CONTEXT context = NULL; + + context = ExAllocatePool2(POOL_FLAG_NON_PAGED, + sizeof(PROCESS_MODULE_MAP_CONTEXT), + POOL_TAG_HASHMAP); + + if (!context) + return STATUS_INSUFFICIENT_RESOURCES; + + status = ExInitializeLookasideListEx(&context->pool, NULL, NULL, - POOL_NX_ALLOCATION, + NonPagedPoolNx, 0, - sizeof(PROCESS_LIST_ENTRY), - POOL_TAG_PROCESS_LIST, + sizeof(PROCESS_MAP_MODULE_ENTRY), + POOL_TAG_MODULE_LIST, 0); if (!NT_SUCCESS(status)) { - DEBUG_ERROR("ExInitializeLookasideListEx failed with status %x", - status); + ExFreePoolWithTag(context, POOL_TAG_HASHMAP); + return status; + } + + status = RtlCreateHashmap(PROCESS_HASHMAP_BUCKET_COUNT, + sizeof(PROCESS_LIST_ENTRY), + ProcessHashmapHashFunction, + ProcessHashmapCompareFunction, + context, + GetProcessHashmap()); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("RtlCreateHashmap: %lx", status); + ExDeleteLookasideListEx(&context->pool); + ExFreePoolWithTag(context, POOL_TAG_HASHMAP); return status; } - InterlockedExchange(&list->active, TRUE); - ListInit(&list->start, &list->lock); return status; } @@ -411,28 +572,6 @@ InitialiseThreadList() return status; } -VOID -FindProcessListEntryByProcess(_In_ PKPROCESS Process, - _Out_ PPROCESS_LIST_ENTRY* Entry) -{ - PPROCESS_LIST_HEAD list = GetProcessList(); - ImpKeAcquireGuardedMutex(&list->lock); - *Entry = NULL; - - PPROCESS_LIST_ENTRY entry = (PPROCESS_LIST_ENTRY)list->start.Next; - - while (entry) { - if (entry->process == Process) { - *Entry = entry; - goto unlock; - } - - entry = entry->list.Next; - } -unlock: - ImpKeReleaseGuardedMutex(&list->lock); -} - VOID FindThreadListEntryByThreadAddress(_In_ PKTHREAD Thread, _Out_ PTHREAD_LIST_ENTRY* Entry) @@ -464,19 +603,66 @@ CanInitiateDeferredHashing(_In_ LPCSTR ProcessName, _In_ PDRIVER_LIST_HEAD Head) : FALSE; } +VOID +EnumerateAndPrintProcessHashmap() +{ + PRTL_HASHMAP map = GetProcessHashmap(); + PRTL_HASHMAP_ENTRY entry = NULL; + PPROCESS_LIST_ENTRY proc_entry = NULL; + PPROCESS_MAP_MODULE_ENTRY mod_entry = NULL; + PLIST_ENTRY list_head = NULL; + PLIST_ENTRY list_entry = NULL; + PLIST_ENTRY mod_list_entry = NULL; + + KeAcquireGuardedMutex(&map->lock); + + for (UINT32 index = 0; index < map->bucket_count; index++) { + list_head = &map->buckets[index]; + list_entry = list_head->Flink; + + DEBUG_VERBOSE("Bucket %u:\n", index); + + while (list_entry != list_head) { + entry = CONTAINING_RECORD(list_entry, RTL_HASHMAP_ENTRY, entry); + + if (entry->in_use == TRUE) { + proc_entry = (PPROCESS_LIST_ENTRY)entry->object; + DEBUG_VERBOSE(" -> process id: %lx", proc_entry->process_id); + DEBUG_VERBOSE(" -> process: %llx", proc_entry->process); + DEBUG_VERBOSE(" -> parent: %llx", proc_entry->parent); + + mod_list_entry = proc_entry->module_list.Flink; + while (mod_list_entry != &proc_entry->module_list) { + mod_entry = CONTAINING_RECORD( + mod_list_entry, PROCESS_MAP_MODULE_ENTRY, entry); + DEBUG_VERBOSE(" -> module base: %llx", mod_entry->base); + DEBUG_VERBOSE(" -> module size: %lx", mod_entry->size); + + mod_list_entry = mod_list_entry->Flink; + } + } + + list_entry = list_entry->Flink; + } + } + + KeReleaseGuardedMutex(&map->lock); +} + VOID ProcessCreateNotifyRoutine(_In_ HANDLE ParentId, _In_ HANDLE ProcessId, _In_ BOOLEAN Create) { - PPROCESS_LIST_ENTRY entry = NULL; + BOOLEAN new = FALSE; PKPROCESS parent = NULL; PKPROCESS process = NULL; - PPROCESS_LIST_HEAD list = GetProcessList(); PDRIVER_LIST_HEAD driver_list = GetDriverList(); LPCSTR process_name = NULL; + PRTL_HASHMAP map = GetProcessHashmap(); + PPROCESS_LIST_ENTRY entry = NULL; - if (!list->active) + if (!map->active) return; ImpPsLookupProcessByProcessId(ParentId, &parent); @@ -487,19 +673,21 @@ ProcessCreateNotifyRoutine(_In_ HANDLE ParentId, process_name = ImpPsGetProcessImageFileName(process); + KeAcquireGuardedMutex(&map->lock); + if (Create) { - entry = ExAllocateFromLookasideListEx(&list->lookaside_list); + entry = RtlInsertEntryHashmap(map, ProcessId); if (!entry) - return; + goto end; - ImpObfReferenceObject(parent); - ImpObfReferenceObject(process); + entry->process_id = ProcessId; + entry->process = process; + entry->parent = parent; - entry->parent = parent; - entry->process = process; + InitializeListHead(&entry->module_list); - ListInsert(&list->start, entry, &list->lock); + entry->list_count = 0; /* * Notify to our driver that we can hash x86 modules, and hash @@ -513,16 +701,22 @@ ProcessCreateNotifyRoutine(_In_ HANDLE ParentId, } } else { - FindProcessListEntryByProcess(process, &entry); + entry = RtlLookupEntryHashmap(map, ProcessId, &ProcessId); - if (!entry) - return; + if (!entry) { + DEBUG_ERROR("UNABLE TO FIND PROCESS NODE!!!"); + goto end; + } ImpObDereferenceObject(entry->parent); ImpObDereferenceObject(entry->process); - LookasideListRemoveEntry(&list->start, entry, &list->lock); + FreeProcessEntryModuleList(entry, NULL); + RtlDeleteEntryHashmap(map, ProcessId, &ProcessId); } + +end: + KeReleaseGuardedMutex(&map->lock); } VOID @@ -954,22 +1148,21 @@ end: } NTSTATUS -EnumerateProcessHandles(_In_ PPROCESS_LIST_ENTRY ProcessListEntry, - _In_opt_ PVOID Context) +EnumerateProcessHandles(_In_ PPROCESS_LIST_ENTRY Entry, _In_opt_ PVOID Context) { /* Handles are stored in pageable memory */ PAGED_CODE(); UNREFERENCED_PARAMETER(Context); - if (!ProcessListEntry) + if (!Entry) return STATUS_INVALID_PARAMETER; - if (ProcessListEntry->process == PsInitialSystemProcess) + if (Entry->process == PsInitialSystemProcess) return STATUS_SUCCESS; PHANDLE_TABLE handle_table = - *(PHANDLE_TABLE*)((uintptr_t)ProcessListEntry->process + + *(PHANDLE_TABLE*)((uintptr_t)Entry->process + EPROCESS_HANDLE_TABLE_OFFSET); if (!handle_table) @@ -990,14 +1183,43 @@ EnumerateProcessHandles(_In_ PPROCESS_LIST_ENTRY ProcessListEntry, #define REPEAT_TIME_10_SEC 10000 +STATIC +VOID +TimerObjectValidateProcessModuleCallback(_In_ PPROCESS_MAP_MODULE_ENTRY Entry, + _In_opt_ PVOID Context) +{ + CHAR hash[SHA_256_HASH_LENGTH] = {0}; + NTSTATUS status = STATUS_UNSUCCESSFUL; + PACTIVE_SESSION session = (PACTIVE_SESSION)Context; + + if (!ARGUMENT_PRESENT(Context)) + return; + + status = HashUserModule(Entry, hash, sizeof(hash)); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("HashUserModule: %x", status); + return; + } + + if (RtlCompareMemory(hash, session->module.module_hash, sizeof(hash)) != + sizeof(hash)) { + DEBUG_ERROR("User module hash not matching!! MODIFIED!"); + return; + } + + DEBUG_VERBOSE("User module hash valid."); +} + STATIC VOID TimerObjectWorkItemRoutine(_In_ PDEVICE_OBJECT DeviceObject, _In_opt_ PVOID Context) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - PTIMER_OBJECT timer = (PTIMER_OBJECT)Context; - PDRIVER_LIST_HEAD list = GetDriverList(); + NTSTATUS status = STATUS_UNSUCCESSFUL; + PTIMER_OBJECT timer = (PTIMER_OBJECT)Context; + PDRIVER_LIST_HEAD list = GetDriverList(); + PACTIVE_SESSION session = GetActiveSession(); UNREFERENCED_PARAMETER(DeviceObject); @@ -1018,6 +1240,16 @@ TimerObjectWorkItemRoutine(_In_ PDEVICE_OBJECT DeviceObject, if (!NT_SUCCESS(status)) DEBUG_ERROR("ValidateOurDriverImage failed with status %x", status); + KeAcquireGuardedMutex(&session->lock); + if (!session->is_session_active) { + KeReleaseGuardedMutex(&session->lock); + goto end; + } + + FindOurUserModeModuleEntry(TimerObjectValidateProcessModuleCallback, + session); + + KeReleaseGuardedMutex(&session->lock); end: InterlockedExchange(&timer->state, FALSE); } diff --git a/driver/callbacks.h b/driver/callbacks.h index 2f098d9..0ca9150 100644 --- a/driver/callbacks.h +++ b/driver/callbacks.h @@ -9,9 +9,6 @@ typedef void (*THREADLIST_CALLBACK_ROUTINE)( _In_ PTHREAD_LIST_ENTRY ThreadListEntry, _In_opt_ PVOID Context); -typedef void (*PROCESSLIST_CALLBACK_ROUTINE)( - _In_ PPROCESS_LIST_ENTRY ProcessListEntry, _In_opt_ PVOID Context); - #define DRIVER_PATH_LENGTH 0x100 #define SHA_256_HASH_LENGTH 32 @@ -36,6 +33,9 @@ typedef struct _DRIVER_LIST_ENTRY { typedef void (*DRIVERLIST_CALLBACK_ROUTINE)( _In_ PDRIVER_LIST_ENTRY DriverListEntry, _In_opt_ PVOID Context); +typedef BOOLEAN (*PROCESS_MODULE_CALLBACK)(_In_ PPROCESS_MAP_MODULE_ENTRY Entry, + _In_opt_ PVOID Context); + NTSTATUS InitialiseDriverList(); @@ -52,16 +52,9 @@ OB_PREOP_CALLBACK_STATUS ObPreOpCallbackRoutine(_In_ PVOID RegistrationContext, _In_ POB_PRE_OPERATION_INFORMATION OperationInformation); -NTSTATUS -EnumerateProcessHandles(_In_ PPROCESS_LIST_ENTRY ProcessListEntry, - _In_opt_ PVOID Context); - NTSTATUS InitialiseThreadList(); -NTSTATUS -InitialiseProcessList(); - VOID ThreadCreateNotifyRoutine(_In_ HANDLE ProcessId, _In_ HANDLE ThreadId, @@ -76,28 +69,17 @@ VOID CleanupThreadListOnDriverUnload(); VOID -FindThreadListEntryByThreadAddress(_In_ PKTHREAD Thread, +FindThreadListEntryByThreadAddress(_In_ PKTHREAD Thread, _Out_ PTHREAD_LIST_ENTRY* Entry); -VOID -FindProcessListEntryByProcess(_In_ PKPROCESS Process, - _Out_ PPROCESS_LIST_ENTRY* Entry); - VOID EnumerateThreadListWithCallbackRoutine( _In_ THREADLIST_CALLBACK_ROUTINE CallbackRoutine, _In_opt_ PVOID Context); -VOID -EnumerateProcessListWithCallbackRoutine( - _In_ PROCESSLIST_CALLBACK_ROUTINE CallbackRoutine, _In_opt_ PVOID Context); - VOID FindDriverEntryByBaseAddress(_In_ PVOID ImageBase, _Out_ PDRIVER_LIST_ENTRY* Entry); -VOID -CleanupProcessListOnDriverUnload(); - VOID CleanupDriverListOnDriverUnload(); @@ -138,4 +120,25 @@ VOID DriverListEntryToExtendedModuleInfo(_In_ PDRIVER_LIST_ENTRY Entry, _Out_ PRTL_MODULE_EXTENDED_INFO Extended); +NTSTATUS +InitialiseProcessHashmap(); + +NTSTATUS +EnumerateProcessHandles(_In_ PPROCESS_LIST_ENTRY Entry, _In_opt_ PVOID Context); + +VOID +EnumerateAndPrintProcessHashmap(); + +VOID +CleanupProcessHashmap(); + +VOID +EnumerateProcessModuleList(_In_ HANDLE ProcessId, + _In_ PROCESS_MODULE_CALLBACK Callback, + _In_opt_ PVOID Context); + +VOID +FindOurUserModeModuleEntry(_In_ PROCESS_MODULE_CALLBACK Callback, + _In_opt_ PVOID Context); + #endif diff --git a/driver/common.h b/driver/common.h index 813be24..6705943 100644 --- a/driver/common.h +++ b/driver/common.h @@ -45,10 +45,16 @@ "donna-ac : [VERBOSE] : " fmt "\n", \ ##__VA_ARGS__) +#define HEX_DUMP(fmt, ...) \ + DbgPrintEx(DPFLTR_DEFAULT_ID, \ + LOG_VERBOSE_LEVEL, \ + fmt, \ + ##__VA_ARGS__) + #define STATIC static #define INLINE inline -#define MAX_MODULE_PATH 256 +#define MAX_MODULE_PATH 260 #define CONVERT_RELATIVE_ADDRESS(Cast, Base, Rel) \ ((Cast)((DWORD_PTR)(Base) + (DWORD_PTR)(Rel))) @@ -66,14 +72,6 @@ typedef struct _THREAD_LIST_HEAD { } THREAD_LIST_HEAD, *PTHREAD_LIST_HEAD; -typedef struct _PROCESS_LIST_HEAD { - SINGLE_LIST_ENTRY start; - volatile BOOLEAN active; - KGUARDED_MUTEX lock; - LOOKASIDE_LIST_EX lookaside_list; - -} PROCESS_LIST_HEAD, *PPROCESS_LIST_HEAD; - typedef struct _DRIVER_LIST_HEAD { SINGLE_LIST_ENTRY start; volatile ULONG count; @@ -97,11 +95,24 @@ typedef struct _THREAD_LIST_ENTRY { } THREAD_LIST_ENTRY, *PTHREAD_LIST_ENTRY; -typedef struct _PROCESS_LIST_ENTRY { - SINGLE_LIST_ENTRY list; - PKPROCESS process; - PKPROCESS parent; +typedef struct _PROCESS_MODULE_MAP_CONTEXT { + LOOKASIDE_LIST_EX pool; +} PROCESS_MODULE_MAP_CONTEXT, *PPROCESS_MODULE_MAP_CONTEXT; +typedef struct _PROCESS_MAP_MODULE_ENTRY { + LIST_ENTRY entry; + UINT64 base; + UINT32 size; + CHAR path[MAX_MODULE_PATH]; +} PROCESS_MAP_MODULE_ENTRY, *PPROCESS_MAP_MODULE_ENTRY; + +typedef struct _PROCESS_LIST_ENTRY { + /* IMPORTANT THIS IS FIRST!*/ + HANDLE process_id; + PEPROCESS process; + PEPROCESS parent; + LIST_ENTRY module_list; + volatile UINT32 list_count; } PROCESS_LIST_ENTRY, *PPROCESS_LIST_ENTRY; /* @@ -237,16 +248,28 @@ typedef struct _HEARTBEAT_CONFIGURATION { } HEARTBEAT_CONFIGURATION, *PHEARTBEAT_CONFIGURATION; +#define SHA_256_HASH_LENGTH 32 + +/* Contains information on our user mode module. */ +typedef struct _MODULE_INFORMATION { + PVOID base_address; + UINT32 size; + CHAR path[MAX_MODULE_PATH]; + CHAR module_hash[SHA_256_HASH_LENGTH]; + +} MODULE_INFORMATION, *PMODULE_INFORMATION; + typedef struct _SESSION_INITIATION_PACKET { - UINT32 cookie; - PVOID process_id; - UCHAR aes_key[AES_256_KEY_SIZE]; - UCHAR aes_iv[AES_256_IV_SIZE]; + UINT32 cookie; + PVOID process_id; + UCHAR aes_key[AES_256_KEY_SIZE]; + UCHAR aes_iv[AES_256_IV_SIZE]; + MODULE_INFORMATION module_info; } SESSION_INITIATION_PACKET, *PSESSION_INITIATION_PACKET; typedef struct _ACTIVE_SESSION { - BOOLEAN is_session_active; + volatile BOOLEAN is_session_active; PVOID um_handle; PVOID km_handle; PEPROCESS process; @@ -269,6 +292,8 @@ typedef struct _ACTIVE_SESSION { UINT32 heartbeat_count; }; + MODULE_INFORMATION module; + HEARTBEAT_CONFIGURATION heartbeat_config; KGUARDED_MUTEX lock; @@ -306,9 +331,13 @@ typedef struct _ACTIVE_SESSION { #define POOL_TAG_LIST_ITEM 'tsil' #define POOL_TAG_THREAD_LIST 'list' #define POOL_TAG_PROCESS_LIST 'plis' +#define POOL_TAG_USER_MODULE_LIST 'resu' +#define POOL_TAG_USER_MODULE_NODE 'edon' #define POOL_TAG_DRIVER_LIST 'drvl' #define POOL_TAG_IRP_QUEUE 'irpp' #define POOL_TAG_TIMER 'time' +#define POOL_TAG_MODULE_LIST 'elom' +#define POOL_TAG_HASHMAP 'hsah' #define IA32_APERF_MSR 0x000000E8 diff --git a/driver/crypt.c b/driver/crypt.c index 378bf53..85209e4 100644 --- a/driver/crypt.c +++ b/driver/crypt.c @@ -242,7 +242,7 @@ CryptInitialiseSessionCryptObjects() UINT32 data_copied = 0; PACTIVE_SESSION session = GetActiveSession(); PBCRYPT_KEY_DATA_BLOB_HEADER blob = NULL; - BCRYPT_ALG_HANDLE* handle = GetCryptAlgHandle(); + BCRYPT_ALG_HANDLE* handle = GetCryptHandle_AES(); blob = CryptBuildBlobForKeyImport(session); @@ -301,7 +301,7 @@ NTSTATUS CryptInitialiseProvider() { NTSTATUS status = STATUS_UNSUCCESSFUL; - BCRYPT_ALG_HANDLE* handle = GetCryptAlgHandle(); + BCRYPT_ALG_HANDLE* handle = GetCryptHandle_AES(); status = BCryptOpenAlgorithmProvider( handle, BCRYPT_AES_ALGORITHM, NULL, BCRYPT_PROV_DISPATCH); @@ -315,7 +315,7 @@ CryptInitialiseProvider() VOID CryptCloseProvider() { - BCRYPT_ALG_HANDLE* handle = GetCryptAlgHandle(); + BCRYPT_ALG_HANDLE* handle = GetCryptHandle_AES(); BCryptCloseAlgorithmProvider(*handle, 0); } @@ -448,6 +448,8 @@ TpmGetPtpInterfaceType(_In_ PVOID Register, return status; } + + NTSTATUS TpmExtractEndorsementKey() { diff --git a/driver/driver.c b/driver/driver.c index c347dc3..56739f0 100644 --- a/driver/driver.c +++ b/driver/driver.c @@ -95,14 +95,15 @@ typedef struct _DRIVER_CONFIG { /* terrible name..lol what is tis timer for ?? */ TIMER_OBJECT timer; - ACTIVE_SESSION session_information; - THREAD_LIST_HEAD thread_list; - DRIVER_LIST_HEAD driver_list; - PROCESS_LIST_HEAD process_list; - SHARED_MAPPING mapping; - BOOLEAN has_driver_loaded; + ACTIVE_SESSION session_information; + THREAD_LIST_HEAD thread_list; + DRIVER_LIST_HEAD driver_list; + RTL_HASHMAP process_hashmap; + SHARED_MAPPING mapping; + BOOLEAN has_driver_loaded; - BCRYPT_ALG_HANDLE alg_handle; + BCRYPT_ALG_HANDLE aes_hash; + BCRYPT_ALG_HANDLE sha256_hash; } DRIVER_CONFIG, *PDRIVER_CONFIG; @@ -123,9 +124,21 @@ PDRIVER_CONFIG g_DriverConfig = NULL; #define POOL_TAG_CONFIG 'conf' BCRYPT_ALG_HANDLE* -GetCryptAlgHandle() +GetCryptHandle_Sha256() { - return &g_DriverConfig->alg_handle; + return &g_DriverConfig->sha256_hash; +} + +PRTL_HASHMAP +GetProcessHashmap() +{ + return &g_DriverConfig->process_hashmap; +} + +BCRYPT_ALG_HANDLE* +GetCryptHandle_AES() +{ + return &g_DriverConfig->aes_hash; } BOOLEAN @@ -271,13 +284,6 @@ GetDriverList() return &g_DriverConfig->driver_list; } -PPROCESS_LIST_HEAD -GetProcessList() -{ - PAGED_CODE(); - return &g_DriverConfig->process_list; -} - /* * The question is, What happens if we attempt to register our callbacks after * we unregister them but before we free the pool? Hm.. No Good. @@ -350,7 +356,7 @@ VOID DrvUnloadFreeProcessList() { PAGED_CODE(); - CleanupProcessListOnDriverUnload(); + CleanupProcessHashmap(); } STATIC @@ -361,6 +367,14 @@ DrvUnloadFreeModuleValidationContext() CleanupValidationContextOnUnload(&g_DriverConfig->sys_val_context); } +STATIC +VOID +CloseHashingAlgorithmProvider() +{ + BCRYPT_ALG_HANDLE* handle = GetCryptHandle_Sha256(); + BCryptCloseAlgorithmProvider(*handle, 0); +} + STATIC VOID DriverUnload(_In_ PDRIVER_OBJECT DriverObject) @@ -391,6 +405,7 @@ DriverUnload(_In_ PDRIVER_OBJECT DriverObject) DrvUnloadFreeDriverList(); CryptCloseProvider(); + CloseHashingAlgorithmProvider(); DrvUnloadFreeConfigStrings(); DrvUnloadDeleteSymbolicLink(); @@ -470,7 +485,7 @@ DrvLoadSetupDriverLists() return status; } - status = InitialiseProcessList(); + status = InitialiseProcessHashmap(); if (!NT_SUCCESS(status)) { DEBUG_ERROR("InitialiseProcessList failed with status %x", status); @@ -826,6 +841,22 @@ DrvLoadInitialiseDriverConfig(_In_ PDRIVER_OBJECT DriverObject, return status; } +STATIC +NTSTATUS +InitialiseHashingAlgorithmProvider() +{ + NTSTATUS status = STATUS_UNSUCCESSFUL; + BCRYPT_ALG_HANDLE* handle = GetCryptHandle_Sha256(); + + status = BCryptOpenAlgorithmProvider( + handle, BCRYPT_SHA256_ALGORITHM, NULL, BCRYPT_PROV_DISPATCH); + + if (!NT_SUCCESS(status)) + DEBUG_ERROR("BCryptOpenAlgorithmProvider: %x", status); + + return status; +} + NTSTATUS DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegistryPath) { @@ -905,10 +936,23 @@ DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegistryPath) return status; } + status = InitialiseHashingAlgorithmProvider(); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("InitialiseHashingAlgorithmProvider failed with status %x", + status); + DrvUnloadFreeConfigStrings(); + DrvUnloadFreeTimerObject(); + DrvUnloadDeleteSymbolicLink(); + ImpIoDeleteDevice(DriverObject->DeviceObject); + return status; + } + status = DrvLoadSetupDriverLists(); if (!NT_SUCCESS(status)) { DEBUG_ERROR("DrvLoadSetupDriverLists failed with status %x", status); + CloseHashingAlgorithmProvider(); DrvUnloadFreeConfigStrings(); DrvUnloadFreeTimerObject(); DrvUnloadDeleteSymbolicLink(); diff --git a/driver/driver.h b/driver/driver.h index c568959..284878a 100644 --- a/driver/driver.h +++ b/driver/driver.h @@ -9,9 +9,13 @@ #include "modules.h" #include "integrity.h" #include "callbacks.h" +#include "map.h" BCRYPT_ALG_HANDLE* -GetCryptAlgHandle(); +GetCryptHandle_AES(); + +BCRYPT_ALG_HANDLE* +GetCryptHandle_Sha256(); NTSTATUS QueryActiveApcContextsForCompletion(); @@ -52,9 +56,6 @@ GetThreadList(); PDRIVER_LIST_HEAD GetDriverList(); -PPROCESS_LIST_HEAD -GetProcessList(); - PUINT64 GetApcContextArray(); @@ -82,4 +83,10 @@ IsNmiInProgress(); BOOLEAN HasDriverLoaded(); +PRTL_HASHMAP +GetProcessHashmap(); + +VOID +CleanupProcessTree(); + #endif \ No newline at end of file diff --git a/driver/driver.vcxproj b/driver/driver.vcxproj index a351622..86892ff 100644 --- a/driver/driver.vcxproj +++ b/driver/driver.vcxproj @@ -253,6 +253,7 @@ + @@ -275,6 +276,7 @@ + diff --git a/driver/driver.vcxproj.filters b/driver/driver.vcxproj.filters index 1205d2c..600e9d8 100644 --- a/driver/driver.vcxproj.filters +++ b/driver/driver.vcxproj.filters @@ -75,6 +75,9 @@ Source Files + + Source Files + @@ -146,6 +149,9 @@ Header Files + + Header Files + diff --git a/driver/hw.c b/driver/hw.c index 31d120d..087ed59 100644 --- a/driver/hw.c +++ b/driver/hw.c @@ -310,13 +310,13 @@ PciDeviceQueryCallback(_In_ PDEVICE_OBJECT DeviceObject, _In_opt_ PVOID Context) DEBUG_VERBOSE("Flagged DeviceID found. Device: %llx, DeviceId: %lx", (UINT64)DeviceObject, header.DeviceID); + ReportBlacklistedPcieDevice(DeviceObject, &header); } else { DEBUG_VERBOSE("Device: %llx, DeviceID: %lx, VendorID: %lx", DeviceObject, header.DeviceID, header.VendorID); - ReportBlacklistedPcieDevice(DeviceObject, &header); } return status; diff --git a/driver/integrity.c b/driver/integrity.c index 9835df7..045bb6f 100644 --- a/driver/integrity.c +++ b/driver/integrity.c @@ -257,6 +257,9 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer, MM_COPY_ADDRESS address = {0}; INTEGRITY_CHECK_HEADER header = {0}; + //DEBUG_VERBOSE("Storing x regions -> x86 module: %lx", (UINT32)IsModulex86); + //DEBUG_VERBOSE("MmIsAddressValid: %lx", MmIsAddressValid(ModuleBase)); + if (!ModuleBase || !ModuleSize) return STATUS_INVALID_PARAMETER; @@ -278,6 +281,10 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer, if (*Buffer == NULL) return STATUS_MEMORY_NOT_ALLOCATED; + /* For context, when we are hashing x86 modules, MmIsAddressValid will + * return FALSE. Yet we still need protection for when an invalid address is + * passed for a non-x86 based image.*/ + /* * The IMAGE_DOS_HEADER.e_lfanew stores the offset of the * IMAGE_NT_HEADER from the base of the image. @@ -434,7 +441,7 @@ ComputeHashOfBuffer(_In_ PVOID Buffer, PAGED_CODE(); NTSTATUS status = STATUS_UNSUCCESSFUL; - BCRYPT_ALG_HANDLE algo_handle = NULL; + BCRYPT_ALG_HANDLE* algo_handle = GetCryptHandle_Sha256(); BCRYPT_HASH_HANDLE hash_handle = NULL; ULONG bytes_copied = 0; ULONG resulting_hash_size = 0; @@ -445,21 +452,12 @@ ComputeHashOfBuffer(_In_ PVOID Buffer, *HashResult = NULL; *HashResultSize = 0; - status = BCryptOpenAlgorithmProvider( - &algo_handle, BCRYPT_SHA256_ALGORITHM, NULL, BCRYPT_PROV_DISPATCH); - - if (!NT_SUCCESS(status)) { - DEBUG_ERROR("BCryptOpenAlogrithmProvider failed with status %x", - status); - goto end; - } - /* * Request the size of the hash object buffer, this is different then * the buffer that will store the resulting hash, instead this will be * used to store the hash object used to create the hash. */ - status = BCryptGetProperty(algo_handle, + status = BCryptGetProperty(*algo_handle, BCRYPT_OBJECT_LENGTH, (PCHAR)&hash_object_size, sizeof(ULONG), @@ -483,7 +481,7 @@ ComputeHashOfBuffer(_In_ PVOID Buffer, * This call gets the size of the resulting hash, which we will use to * allocate the resulting hash buffer. */ - status = BCryptGetProperty(algo_handle, + status = BCryptGetProperty(*algo_handle, BCRYPT_HASH_LENGTH, (PCHAR)&resulting_hash_size, sizeof(ULONG), @@ -507,7 +505,7 @@ ComputeHashOfBuffer(_In_ PVOID Buffer, * Here we create our hash object and store it in the hash_object * buffer. */ - status = BCryptCreateHash(algo_handle, + status = BCryptCreateHash(*algo_handle, &hash_handle, hash_object, hash_object_size, @@ -549,9 +547,6 @@ ComputeHashOfBuffer(_In_ PVOID Buffer, end: - if (algo_handle) - BCryptCloseAlgorithmProvider(algo_handle, NULL); - if (hash_handle) BCryptDestroyHash(hash_handle); @@ -1027,6 +1022,66 @@ end: return status; } +NTSTATUS +HashUserModule(_In_ PPROCESS_MAP_MODULE_ENTRY Entry, + _Out_ PVOID OutBuffer, + _In_ UINT32 OutBufferSize) +{ + PAGED_CODE(); + + NTSTATUS status = STATUS_UNSUCCESSFUL; + KAPC_STATE apc_state = {0}; + PVAL_INTEGRITY_HEADER memory_buffer = NULL; + PVOID memory_hash = NULL; + ULONG memory_hash_size = 0; + SIZE_T bytes_written = 0; + PACTIVE_SESSION session = GetActiveSession(); + + /* + * Attach because the offsets given are from the process' context. + */ + ImpKeStackAttachProcess(session->process, &apc_state); + + status = StoreModuleExecutableRegionsInBuffer( + &memory_buffer, Entry->base, Entry->size, &bytes_written, FALSE); + + ImpKeUnstackDetachProcess(&apc_state); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR( + "StoreModuleExecutableRegionsInBuffer failed with status %x", + status); + goto end; + } + + status = ComputeHashOfBuffer(memory_buffer->section_base, + memory_buffer->section_header.SizeOfRawData, + &memory_hash, + &memory_hash_size); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("ComputeHashOfBuffer failed with status %x", status); + goto end; + } + + if (OutBufferSize > memory_hash_size) { + status = STATUS_BUFFER_TOO_SMALL; + goto end; + } + + RtlCopyMemory(OutBuffer, memory_hash, memory_hash_size); + +end: + + if (memory_buffer) + ImpExFreePoolWithTag(memory_buffer, POOL_TAG_INTEGRITY); + + if (memory_hash) + ImpExFreePoolWithTag(memory_hash, POOL_TAG_INTEGRITY); + + return status; +} + FORCEINLINE STATIC PCHAR @@ -1258,6 +1313,9 @@ GetAverageReadTimeAtRoutine(_In_ PVOID RoutineAddress, if (!RoutineAddress || !AverageTime) return STATUS_UNSUCCESSFUL; + if (!MmIsAddressValid(RoutineAddress)) + return STATUS_INVALID_ADDRESS; + *AverageTime = MeasureReads(RoutineAddress, EPT_CHECK_NUM_ITERATIONS); return *AverageTime == 0 ? STATUS_UNSUCCESSFUL : STATUS_SUCCESS; @@ -1304,8 +1362,11 @@ WCHAR PROTECTED_FUNCTIONS[EPT_PROTECTED_FUNCTIONS_COUNT] * call thereafter fails. So will be storing the routine addresses in arrays * since they dont change once the kernel is loaded. */ -UINT64 CONTROL_FUNCTION_ADDRESSES[EPT_CONTROL_FUNCTIONS_COUNT] = {0}; -UINT64 PROTECTED_FUNCTION_ADDRESSES[EPT_PROTECTED_FUNCTIONS_COUNT] = {0}; +#pragma section("NonPagedPool", read, write) +__declspec(allocate("NonPagedPool")) UINT64 + CONTROL_FUNCTION_ADDRESSES[EPT_CONTROL_FUNCTIONS_COUNT] = {0}; +__declspec(allocate("NonPagedPool")) UINT64 + PROTECTED_FUNCTION_ADDRESSES[EPT_PROTECTED_FUNCTIONS_COUNT] = {0}; STATIC NTSTATUS @@ -1407,7 +1468,7 @@ DetectEptHooksInKeyFunctions() } VOID -FindWinLogonProcess(_In_ PPROCESS_LIST_ENTRY Entry, _In_opt_ PVOID Context) +FindWinLogonProcess(_In_ PPROCESS_LIST_ENTRY Node, _In_opt_ PVOID Context) { LPCSTR process_name = NULL; PEPROCESS* process = (PEPROCESS*)Context; @@ -1415,10 +1476,10 @@ FindWinLogonProcess(_In_ PPROCESS_LIST_ENTRY Entry, _In_opt_ PVOID Context) if (!Context) return; - process_name = ImpPsGetProcessImageFileName(Entry->process); + process_name = ImpPsGetProcessImageFileName(Node->process); if (!strcmp(process_name, "winlogon.exe")) - *process = Entry->process; + *process = Node->process; } STATIC @@ -1431,7 +1492,7 @@ StoreModuleExecutableRegionsx86(_In_ PRTL_MODULE_EXTENDED_INFO Module, PEPROCESS process = NULL; KAPC_STATE apc_state = {0}; - EnumerateProcessListWithCallbackRoutine(FindWinLogonProcess, &process); + RtlEnumerateHashmap(GetProcessHashmap(), FindWinLogonProcess, &process); if (!process) return STATUS_NOT_FOUND; @@ -2224,33 +2285,6 @@ WaitForHeartbeatCompletion(_In_ PHEARTBEAT_CONFIGURATION Configuration) YieldProcessor(); } -STATIC -VOID -HeartbeatWorkItem(_In_ PDEVICE_OBJECT DeviceObject, _In_opt_ PVOID Context) -{ - UNREFERENCED_PARAMETER(DeviceObject); - - if (!ARGUMENT_PRESENT(Context)) - return; - - NTSTATUS status = STATUS_UNSUCCESSFUL; - PHEARTBEAT_CONFIGURATION config = (PHEARTBEAT_CONFIGURATION)Context; - - /* Ensure we wait until our heartbeats DPC has terminated. */ - KeFlushQueuedDpcs(); - FreeHeartbeatObjects(config); - - status = AllocateHeartbeatObjects(config); - - if (!NT_SUCCESS(status)) { - DEBUG_ERROR("AllocateHeartbeatObjects %x", status); - return; - } - - InitialiseHeartbeatObjects(config); - SetHeartbeatInactive(config); -} - FORCEINLINE STATIC VOID @@ -2291,22 +2325,17 @@ BuildHeartbeatPacket(_In_ UINT32 PacketSize) STATIC VOID -HeartbeatDpcRoutine(_In_ PKDPC Dpc, - _In_opt_ PVOID DeferredContext, - _In_opt_ PVOID SystemArgument1, - _In_opt_ PVOID SystemArgument2) +HeartbeatWorkItem(_In_ PDEVICE_OBJECT DeviceObject, _In_opt_ PVOID Context) { - UNREFERENCED_PARAMETER(Dpc); - UNREFERENCED_PARAMETER(SystemArgument1); - UNREFERENCED_PARAMETER(SystemArgument2); + UNREFERENCED_PARAMETER(DeviceObject); - if (!ARGUMENT_PRESENT(DeferredContext)) + if (!ARGUMENT_PRESENT(Context)) return; - NTSTATUS status = STATUS_UNSUCCESSFUL; - PHEARTBEAT_CONFIGURATION config = (PHEARTBEAT_CONFIGURATION)DeferredContext; - PHEARTBEAT_PACKET packet = NULL; UINT32 packet_size = 0; + NTSTATUS status = STATUS_UNSUCCESSFUL; + PHEARTBEAT_PACKET packet = NULL; + PHEARTBEAT_CONFIGURATION config = (PHEARTBEAT_CONFIGURATION)Context; DEBUG_VERBOSE("Heartbeat timer alerted. Generating heartbeat packet."); @@ -2321,14 +2350,45 @@ HeartbeatDpcRoutine(_In_ PKDPC Dpc, if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); ImpExFreePoolWithTag(packet, POOL_TAG_HEARTBEAT); - goto end; + goto queue_next; } IrpQueueSchedulePacket(packet, packet_size); IncrementHeartbeatCounter(config); } -end: +queue_next: + /* Ensure we wait until our heartbeats DPC has terminated. */ + KeFlushQueuedDpcs(); + FreeHeartbeatObjects(config); + + status = AllocateHeartbeatObjects(config); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("AllocateHeartbeatObjects %x", status); + return; + } + + InitialiseHeartbeatObjects(config); + SetHeartbeatInactive(config); +} + +STATIC +VOID +HeartbeatDpcRoutine(_In_ PKDPC Dpc, + _In_opt_ PVOID DeferredContext, + _In_opt_ PVOID SystemArgument1, + _In_opt_ PVOID SystemArgument2) +{ + UNREFERENCED_PARAMETER(Dpc); + UNREFERENCED_PARAMETER(SystemArgument1); + UNREFERENCED_PARAMETER(SystemArgument2); + + if (!ARGUMENT_PRESENT(DeferredContext)) + return; + + PHEARTBEAT_CONFIGURATION config = (PHEARTBEAT_CONFIGURATION)DeferredContext; + IoQueueWorkItem( config->work_item, HeartbeatWorkItem, NormalWorkQueue, config); } diff --git a/driver/integrity.h b/driver/integrity.h index c36b32a..d811072 100644 --- a/driver/integrity.h +++ b/driver/integrity.h @@ -119,7 +119,7 @@ DeferredModuleHashingCallback(_In_ PDEVICE_OBJECT DeviceObject, _In_opt_ PVOID Context); VOID -FindWinLogonProcess(_In_ PPROCESS_LIST_ENTRY Entry, _In_opt_ PVOID Context); +FindWinLogonProcess(_In_ PPROCESS_LIST_ENTRY Node, _In_opt_ PVOID Context); NTSTATUS InitialiseHeartbeatConfiguration( @@ -128,4 +128,9 @@ InitialiseHeartbeatConfiguration( VOID FreeHeartbeatConfiguration(_Inout_ PHEARTBEAT_CONFIGURATION Configuration); +NTSTATUS +HashUserModule(_In_ PPROCESS_MAP_MODULE_ENTRY Entry, + _Out_ PVOID OutBuffer, + _In_ UINT32 OutBufferSize); + #endif diff --git a/driver/io.c b/driver/io.c index 86ad5f5..b3b7df1 100644 --- a/driver/io.c +++ b/driver/io.c @@ -12,6 +12,7 @@ #include "list.h" #include "session.h" #include "hw.h" +#include "map.h" STATIC NTSTATUS @@ -446,19 +447,20 @@ SharedMappingWorkRoutine(_In_ PDEVICE_OBJECT DeviceObject, /* can maybe implement this better so we can extract a status * value */ - EnumerateProcessListWithCallbackRoutine(EnumerateProcessHandles, NULL); + RtlEnumerateHashmap(GetProcessHashmap(), EnumerateProcessHandles, NULL); break; case ssScanForUnlinkedProcesses: - DEBUG_INFO( - "SHARED_STATE_OPERATION_ID: ScanForUnlinkedProcesses Received"); + // DEBUG_INFO( + // "SHARED_STATE_OPERATION_ID: ScanForUnlinkedProcesses Received"); - status = FindUnlinkedProcesses(); + // status = FindUnlinkedProcesses(); - if (!NT_SUCCESS(status)) - DEBUG_ERROR("FindUnlinkedProcesses failed with status %x", status); + // if (!NT_SUCCESS(status)) + // DEBUG_ERROR("FindUnlinkedProcesses failed with status %x", + // status); break; @@ -896,7 +898,7 @@ DeviceControl(_In_ PDEVICE_OBJECT DeviceObject, _Inout_ PIRP Irp) /* can maybe implement this better so we can extract a status * value */ - EnumerateProcessListWithCallbackRoutine(EnumerateProcessHandles, NULL); + RtlEnumerateHashmap(GetProcessHashmap(), EnumerateProcessHandles, NULL); break; @@ -961,12 +963,13 @@ DeviceControl(_In_ PDEVICE_OBJECT DeviceObject, _Inout_ PIRP Irp) case IOCTL_SCAN_FOR_UNLINKED_PROCESS: - DEBUG_INFO("IOCTL_SCAN_FOR_UNLINKED_PROCESS Received"); + // DEBUG_INFO("IOCTL_SCAN_FOR_UNLINKED_PROCESS Received"); - status = FindUnlinkedProcesses(); + // status = FindUnlinkedProcesses(); - if (!NT_SUCCESS(status)) - DEBUG_ERROR("FindUnlinkedProcesses failed with status %x", status); + // if (!NT_SUCCESS(status)) + // DEBUG_ERROR("FindUnlinkedProcesses failed with status %x", + // status); break; @@ -1124,11 +1127,20 @@ DeviceControl(_In_ PDEVICE_OBJECT DeviceObject, _Inout_ PIRP Irp) DEBUG_INFO("IOCTL_VALIDATE_PCI_DEVICES Received"); - status = ValidatePciDevices(); + status = ImpPsCreateSystemThread(&handle, + PROCESS_ALL_ACCESS, + NULL, + NULL, + NULL, + ValidatePciDevices, + NULL); - if (!NT_SUCCESS(status)) - DEBUG_ERROR("ValidatePciDevices failed with status %x", status); + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("PsCreateSystemThread failed with status %x", status); + goto end; + } + ImpZwClose(handle); break; case IOCTL_VALIDATE_WIN32K_TABLES: @@ -1179,7 +1191,6 @@ DeviceCreate(_In_ PDEVICE_OBJECT DeviceObject, _Inout_ PIRP Irp) PAGED_CODE(); UNREFERENCED_PARAMETER(DeviceObject); DEBUG_INFO("Handle to driver opened."); - // NTSTATUS status = ValidatePciDevices(); // if (!NT_SUCCESS(status)) diff --git a/driver/map.c b/driver/map.c new file mode 100644 index 0000000..88f5aba --- /dev/null +++ b/driver/map.c @@ -0,0 +1,241 @@ +#include "map.h" + +NTSTATUS +RtlCreateHashmap(_In_ UINT32 BucketCount, + _In_ UINT32 EntryObjectSize, + _In_ HASH_FUNCTION HashFunction, + _In_ COMPARE_FUNCTION CompareFunction, + _In_ PVOID Context, + _Out_ PRTL_HASHMAP Hashmap) +{ + NTSTATUS status = STATUS_UNSUCCESSFUL; + UINT32 entry_size = sizeof(RTL_HASHMAP_ENTRY) + EntryObjectSize; + PRTL_HASHMAP_ENTRY entry = NULL; + + Hashmap->buckets = ExAllocatePool2( + POOL_FLAG_NON_PAGED, BucketCount * entry_size, POOL_TAG_HASHMAP); + + if (!Hashmap->buckets) + return STATUS_INSUFFICIENT_RESOURCES; + + for (UINT32 index = 0; index < BucketCount; index++) { + entry = &Hashmap->buckets[index]; + entry->in_use = FALSE; + InitializeListHead(&entry->entry); + } + + KeInitializeGuardedMutex(&Hashmap->lock); + + Hashmap->bucket_count = BucketCount; + Hashmap->hash_function = HashFunction; + Hashmap->compare_function = CompareFunction; + Hashmap->object_size = EntryObjectSize; + Hashmap->active = TRUE; + Hashmap->context = Context; + + return STATUS_SUCCESS; +} + +FORCEINLINE +STATIC +PRTL_HASHMAP_ENTRY +RtlFindUnusedHashmapEntry(_In_ PRTL_HASHMAP_ENTRY Head) +{ + PRTL_HASHMAP_ENTRY entry = Head; + + while (entry) { + if (entry->in_use == FALSE) + return entry; + + entry = CONTAINING_RECORD(entry->entry.Flink, RTL_HASHMAP_ENTRY, entry); + } + + return NULL; +} + +FORCEINLINE +STATIC +PRTL_HASHMAP_ENTRY +RtlAllocateBucketListEntry(_In_ PRTL_HASHMAP Hashmap) +{ + PRTL_HASHMAP_ENTRY entry = + ExAllocatePool2(POOL_FLAG_NON_PAGED, + Hashmap->object_size + sizeof(RTL_HASHMAP_ENTRY), + POOL_TAG_HASHMAP); + + if (!entry) + return NULL; + + entry->in_use = TRUE; + return entry; +} + +FORCEINLINE +STATIC +BOOLEAN +RtlIsIndexInHashmapRange(_In_ PRTL_HASHMAP Hashmap, _In_ UINT32 Index) +{ + return Index < Hashmap->bucket_count ? TRUE : FALSE; +} + +/* assumes map lock is held */ +PVOID +RtlInsertEntryHashmap(_In_ PRTL_HASHMAP Hashmap, _In_ UINT64 Key) +{ + UINT32 index = 0; + PLIST_ENTRY list_head = NULL; + PLIST_ENTRY list_entry = NULL; + PRTL_HASHMAP_ENTRY entry = NULL; + PRTL_HASHMAP_ENTRY new_entry = NULL; + + index = Hashmap->hash_function(Key); + + if (!RtlIsIndexInHashmapRange(Hashmap, index)) { + DEBUG_ERROR("Key is not in range of buckets"); + return NULL; + } + + list_head = &(&Hashmap->buckets[index])->entry; + list_entry = list_head->Flink; + + while (list_entry != list_head) { + entry = CONTAINING_RECORD(list_entry, RTL_HASHMAP_ENTRY, entry); + + if (entry->in_use == FALSE) { + entry->in_use = TRUE; + return entry->object; + } + + list_entry = list_entry->Flink; + } + + new_entry = RtlAllocateBucketListEntry(Hashmap); + + if (!new_entry) { + DEBUG_ERROR("Failed to allocate new entry"); + return NULL; + } + + InsertHeadList(list_head, &new_entry->entry); + return new_entry->object; +} + +/* Returns a pointer to the start of the entries caller defined data. i.e + * &PRTL_HASHMAP_ENTRY->Object + * + * Also assumes lock is held. + */ +PVOID +RtlLookupEntryHashmap(_In_ PRTL_HASHMAP Hashmap, + _In_ UINT64 Key, + _In_ PVOID Compare) +{ + UINT32 index = 0; + PRTL_HASHMAP_ENTRY entry = NULL; + + index = Hashmap->hash_function(Key); + + if (!RtlIsIndexInHashmapRange(Hashmap, index)) { + DEBUG_ERROR("Key is not in range of buckets"); + return NULL; + } + + entry = &Hashmap->buckets[index]; + + while (entry) { + if (entry->in_use == FALSE) + goto increment; + + if (Hashmap->compare_function(entry->object, Compare)) + return entry->object; + + increment: + entry = CONTAINING_RECORD(entry->entry.Flink, RTL_HASHMAP_ENTRY, entry); + } + + DEBUG_ERROR("Unable to find entry in hashmap."); + return NULL; +} + +/* Assumes lock is held */ +BOOLEAN +RtlDeleteEntryHashmap(_In_ PRTL_HASHMAP Hashmap, + _In_ UINT64 Key, + _In_ PVOID Compare) +{ + UINT32 index = 0; + PRTL_HASHMAP_ENTRY entry = NULL; + PRTL_HASHMAP_ENTRY next = NULL; + + index = Hashmap->hash_function(Key); + + if (!RtlIsIndexInHashmapRange(Hashmap, index)) { + DEBUG_ERROR("Key is not in range of buckets"); + return FALSE; + } + + entry = &Hashmap->buckets[index]; + + while (entry) { + if (entry->in_use == FALSE) { + next = + CONTAINING_RECORD(entry->entry.Flink, RTL_HASHMAP_ENTRY, entry); + + if (next == &Hashmap->buckets[index]) + break; + + entry = next; + continue; + } + + if (Hashmap->compare_function(entry->object, Compare)) { + if (entry == &Hashmap->buckets[index]) { + entry->in_use = FALSE; + } + else { + RemoveEntryList(&entry->entry); + ExFreePoolWithTag(entry, POOL_TAG_HASHMAP); + } + + return TRUE; + } + + next = CONTAINING_RECORD(entry->entry.Flink, RTL_HASHMAP_ENTRY, entry); + + if (next == &Hashmap->buckets[index]) + break; + + entry = next; + } + + return FALSE; +} + +VOID +RtlEnumerateHashmap(_In_ PRTL_HASHMAP Hashmap, + _In_ ENUMERATE_HASHMAP EnumerationCallback, + _In_opt_ PVOID Context) +{ + PRTL_HASHMAP_ENTRY entry = NULL; + + for (UINT32 index = 0; index < Hashmap->bucket_count; index++) { + PLIST_ENTRY list_head = &Hashmap->buckets[index]; + PLIST_ENTRY list_entry = list_head->Flink; + + while (list_entry != list_head) { + entry = CONTAINING_RECORD(list_entry, RTL_HASHMAP_ENTRY, entry); + + if (entry->in_use == TRUE) { + EnumerationCallback(entry->object, Context); + } + + list_entry = list_entry->Flink; + } + } +} + +VOID +RtlDeleteHashmap(_In_ PRTL_HASHMAP Hashmap) +{ + ExFreePoolWithTag(Hashmap->buckets, POOL_TAG_HASHMAP); +} \ No newline at end of file diff --git a/driver/map.h b/driver/map.h new file mode 100644 index 0000000..b207be5 --- /dev/null +++ b/driver/map.h @@ -0,0 +1,78 @@ +#ifndef MAP_H +#define MAP_H + +#include "common.h" + +typedef UINT32 (*HASH_FUNCTION)(_In_ UINT64 Key); + +/* Struct1 being the node being compared to the value in Struct 2*/ +typedef BOOLEAN (*COMPARE_FUNCTION)(_In_ PVOID Struct1, _In_ PVOID Struct2); + +/* To improve efficiency, each entry contains a common header + * RTL_HASHMAP_ENTRY*, reducing the need to store a seperate pointer to the + * entrys data. */ +typedef struct _RTL_HASHMAP_ENTRY { + LIST_ENTRY entry; + UINT32 in_use; + CHAR object[]; +} RTL_HASHMAP_ENTRY, *PRTL_HASHMAP_ENTRY; + +typedef VOID (*ENUMERATE_HASHMAP)(_In_ PRTL_HASHMAP_ENTRY Entry, + _In_opt_ PVOID Context); + +typedef struct _RTL_HASHMAP { + /* Array of RTL_HASHMAP_ENTRIES with length = bucket_count */ + PRTL_HASHMAP_ENTRY buckets; + + /* Number of buckets, ideally a prime number */ + UINT32 bucket_count; + + /* Size of each custom object existing after the RTL_HASHMAP_ENTRY */ + UINT32 object_size; + + /* Pointer to caller-designated callback routines */ + HASH_FUNCTION hash_function; + COMPARE_FUNCTION compare_function; + + KGUARDED_MUTEX lock; + + /* in the future bucket entries will use this */ + LOOKASIDE_LIST_EX pool; + + /* user allocated context */ + PVOID context; + volatile UINT32 active; + +} RTL_HASHMAP, *PRTL_HASHMAP; + +/* Hashmap is caller allocated */ +NTSTATUS +RtlCreateHashmap(_In_ UINT32 BucketCount, + _In_ UINT32 EntryObjectSize, + _In_ HASH_FUNCTION HashFunction, + _In_ COMPARE_FUNCTION CompareFunction, + _In_ PVOID Context, + _Out_ PRTL_HASHMAP Hashmap); + +PVOID +RtlInsertEntryHashmap(_In_ PRTL_HASHMAP Hashmap, _In_ UINT64 Key); + +PVOID +RtlLookupEntryHashmap(_In_ PRTL_HASHMAP Hashmap, + _In_ UINT64 Key, + _In_ PVOID Compare); + +BOOLEAN +RtlDeleteEntryHashmap(_In_ PRTL_HASHMAP Hashmap, + _In_ UINT64 Key, + _In_ PVOID Compare); + +VOID +RtlEnumerateHashmap(_In_ PRTL_HASHMAP Hashmap, + _In_ ENUMERATE_HASHMAP EnumerationCallback, + _In_opt_ PVOID Context); + +VOID +RtlDeleteHashmap(_In_ PRTL_HASHMAP Hashmap); + +#endif \ No newline at end of file diff --git a/driver/modules.c b/driver/modules.c index 1c399e1..98c48ec 100644 --- a/driver/modules.c +++ b/driver/modules.c @@ -1377,15 +1377,12 @@ end: /* todo: walk the chain of pointers to prevent jmp chaining */ STATIC -NTSTATUS +VOID ValidateTableDispatchRoutines(_In_ PVOID* Base, _In_ UINT32 Entries, _In_ PSYSTEM_MODULES Modules, _Out_ PVOID* Routine) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - BOOLEAN flag = FALSE; - for (UINT32 index = 0; index < Entries; index++) { if (!Base[index]) continue; @@ -1393,8 +1390,6 @@ ValidateTableDispatchRoutines(_In_ PVOID* Base, if (IsInstructionPointerInInvalidRegion(Base[index], Modules)) *Routine = Base[index]; } - - return status; } /* @@ -1447,24 +1442,14 @@ ValidateHalPrivateDispatchTable(_Out_ PVOID* Routine, base = (UINT64)table + sizeof(UINT64); count = GetHalPrivateDispatchTableRoutineCount(&os_info); - status = ValidateTableDispatchRoutines(base, count, Modules, Routine); - - if (!NT_SUCCESS(status)) { - DEBUG_ERROR("ValidateTableDispatchRoutines failed with status %x", - status); - return status; - } - + ValidateTableDispatchRoutines(base, count, Modules, Routine); return status; } STATIC -NTSTATUS +VOID ValidateHalDispatchTable(_Out_ PVOID* Routine, _In_ PSYSTEM_MODULES Modules) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - BOOLEAN flag = FALSE; - *Routine = NULL; DEBUG_VERBOSE("Validating HalDispatchTable."); @@ -1488,7 +1473,7 @@ ValidateHalDispatchTable(_Out_ PVOID* Routine, _In_ PSYSTEM_MODULES Modules) goto end; } - if (IsInstructionPointerInInvalidRegion(HalQueryBusSlots, Modules)) { + if (IsInstructionPointerInInvalidRegion(HalQueryBusSlots, Modules)) { *Routine = HalQueryBusSlots; goto end; } @@ -1573,7 +1558,7 @@ ValidateHalDispatchTable(_Out_ PVOID* Routine, _In_ PSYSTEM_MODULES Modules) } end: - return status; + return; } STATIC @@ -1627,12 +1612,7 @@ ValidateHalDispatchTables() return status; } - status = ValidateHalDispatchTable(&routine1, &modules); - - if (!NT_SUCCESS(status)) { - DEBUG_ERROR("ValidateHalDispatchTable failed with status %x", status); - goto end; - } + ValidateHalDispatchTable(&routine1, &modules); if (routine1) ReportDataTableInvalidRoutine(HalDispatch, routine1); @@ -1977,7 +1957,7 @@ ValidateWin32kBase_gDxgInterface() goto end; } - EnumerateProcessListWithCallbackRoutine(FindWinLogonProcess, &winlogon); + RtlEnumerateHashmap(GetProcessHashmap(), FindWinLogonProcess, &winlogon); if (!winlogon) { status = STATUS_UNSUCCESSFUL; diff --git a/driver/pe.c b/driver/pe.c index 46022ba..d183533 100644 --- a/driver/pe.c +++ b/driver/pe.c @@ -1,5 +1,19 @@ #include "pe.h" +PNT_HEADER_64 +PeGetNtHeaderSafe(_In_ PVOID Image) +{ + PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)Image; + + if (!MmIsAddressValid(Image)) + return NULL; + + if (dos->e_magic != IMAGE_DOS_SIGNATURE) + return NULL; + + return CONVERT_RELATIVE_ADDRESS(PNT_HEADER_64, Image, dos->e_lfanew); +} + PNT_HEADER_64 PeGetNtHeader(_In_ PVOID Image) { @@ -23,6 +37,21 @@ PeGetExportDataDirectory(_In_ PVOID Image) .DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT]; } +PIMAGE_DATA_DIRECTORY +PeGetExportDataDirectorySafe(_In_ PVOID Image) +{ + PNT_HEADER_64 nt = PeGetNtHeader(Image); + + if (!MmIsAddressValid(Image)) + return NULL; + + if (IMAGE_DIRECTORY_ENTRY_EXPORT >= nt->OptionalHeader.NumberOfRvaAndSizes) + return NULL; + + return (PIMAGE_DATA_DIRECTORY)&nt->OptionalHeader + .DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT]; +} + PIMAGE_EXPORT_DIRECTORY PeGetExportDirectory(_In_ PVOID Image, _In_ PIMAGE_DATA_DIRECTORY ExportDataDirectory) @@ -34,12 +63,35 @@ PeGetExportDirectory(_In_ PVOID Image, PIMAGE_EXPORT_DIRECTORY, Image, ExportDataDirectory->VirtualAddress); } +PIMAGE_EXPORT_DIRECTORY +PeGetExportDirectorySafe(_In_ PVOID Image, + _In_ PIMAGE_DATA_DIRECTORY ExportDataDirectory) +{ + if (!MmIsAddressValid(Image)) + return NULL; + + if (!ExportDataDirectory->VirtualAddress || !ExportDataDirectory->Size) + return NULL; + + return CONVERT_RELATIVE_ADDRESS( + PIMAGE_EXPORT_DIRECTORY, Image, ExportDataDirectory->VirtualAddress); +} + UINT32 GetSectionCount(_In_ PNT_HEADER_64 Header) { return Header->FileHeader.NumberOfSections; } +UINT32 +GetSectionCountSafe(_In_ PNT_HEADER_64 Header) +{ + if (!MmIsAddressValid(Header)) + return NULL; + + return Header->FileHeader.NumberOfSections; +} + PVOID PeFindExportByName(_In_ PVOID Image, _In_ PCHAR Name) { diff --git a/driver/pe.h b/driver/pe.h index 756ad8d..52274ac 100644 --- a/driver/pe.h +++ b/driver/pe.h @@ -22,4 +22,17 @@ PeGetExportDirectory(_In_ PVOID Image, UINT32 GetSectionCount(_In_ PNT_HEADER_64 Header); +PIMAGE_EXPORT_DIRECTORY +PeGetExportDirectorySafe(_In_ PVOID Image, + _In_ PIMAGE_DATA_DIRECTORY ExportDataDirectory); + +PIMAGE_DATA_DIRECTORY +PeGetExportDataDirectorySafe(_In_ PVOID Image); + +PNT_HEADER_64 +PeGetNtHeaderSafe(_In_ PVOID Image); + +UINT32 +GetSectionCountSafe(_In_ PNT_HEADER_64 Header); + #endif \ No newline at end of file diff --git a/driver/pool.c b/driver/pool.c index 4dc4624..695321e 100644 --- a/driver/pool.c +++ b/driver/pool.c @@ -83,13 +83,12 @@ WalkKernelPageTables(_In_ PPROCESS_SCAN_CONTEXT Context); STATIC VOID -IncrementProcessCounter(_In_ PPROCESS_LIST_ENTRY ProcessListEntry, - _Inout_opt_ PVOID Context); +IncrementProcessCounter(_In_ PPROCESS_LIST_ENTRY Node, _In_opt_ PVOID Context); STATIC VOID -CheckIfProcessAllocationIsInProcessList( - _In_ PPROCESS_LIST_ENTRY ProcessListEntry, _Inout_opt_ PVOID Context); +CheckIfProcessAllocationIsInProcessList(_In_ PPROCESS_LIST_ENTRY Node, + _In_opt_ PVOID Context); #ifdef ALLOC_PRAGMA # pragma alloc_text(PAGE, GetGlobalDebuggerData) @@ -629,12 +628,11 @@ WalkKernelPageTables(_In_ PPROCESS_SCAN_CONTEXT Context) STATIC VOID -IncrementProcessCounter(_In_ PPROCESS_LIST_ENTRY ProcessListEntry, - _Inout_opt_ PVOID Context) +IncrementProcessCounter(_In_ PPROCESS_LIST_ENTRY Node, _In_opt_ PVOID Context) { PAGED_CODE(); - UNREFERENCED_PARAMETER(ProcessListEntry); + UNREFERENCED_PARAMETER(Node); PPROCESS_SCAN_CONTEXT context = (PPROCESS_SCAN_CONTEXT)Context; @@ -646,8 +644,8 @@ IncrementProcessCounter(_In_ PPROCESS_LIST_ENTRY ProcessListEntry, STATIC VOID -CheckIfProcessAllocationIsInProcessList( - _In_ PPROCESS_LIST_ENTRY ProcessListEntry, _Inout_opt_ PVOID Context) +CheckIfProcessAllocationIsInProcessList(_In_ PPROCESS_LIST_ENTRY Node, + _In_opt_ PVOID Context) { PAGED_CODE(); @@ -660,9 +658,9 @@ CheckIfProcessAllocationIsInProcessList( for (INT i = 0; i < context->process_count; i++) { allocation_address = (PUINT64)context->process_buffer; - if ((UINT64)ProcessListEntry->process >= + if ((UINT64)Node->process >= allocation_address[i] - PROCESS_OBJECT_ALLOCATION_MARGIN && - (UINT64)ProcessListEntry->process <= + (UINT64)Node->process <= allocation_address[i] + PROCESS_OBJECT_ALLOCATION_MARGIN) { RtlZeroMemory((UINT64)context->process_buffer + i * sizeof(UINT64), sizeof(UINT64)); @@ -686,7 +684,7 @@ FindUnlinkedProcesses() UINT32 packet_size = CryptRequestRequiredBufferLength( sizeof(INVALID_PROCESS_ALLOCATION_REPORT)); - EnumerateProcessListWithCallbackRoutine(IncrementProcessCounter, &context); + RtlEnumerateHashmap(GetProcessHashmap(), IncrementProcessCounter, &context); if (context.process_count == 0) { DEBUG_ERROR("IncrementProcessCounter failed with no status."); @@ -703,8 +701,8 @@ FindUnlinkedProcesses() WalkKernelPageTables(&context); - EnumerateProcessListWithCallbackRoutine( - CheckIfProcessAllocationIsInProcessList, &context); + RtlEnumerateHashmap( + GetProcessHashmap(), CheckIfProcessAllocationIsInProcessList, &context); allocation_address = (PUINT64)context.process_buffer; diff --git a/driver/session.c b/driver/session.c index 9b883a0..100c9df 100644 --- a/driver/session.c +++ b/driver/session.c @@ -2,6 +2,7 @@ #include "imports.h" #include "crypt.h" +#include "util.h" NTSTATUS SessionInitialiseStructure() @@ -83,6 +84,34 @@ SessionTerminate() KeReleaseGuardedMutex(&session->lock); } +/* Return type for this doesnt matter */ +STATIC +BOOLEAN +HashOurUserModuleOnEntryCallback(_In_ PPROCESS_MAP_MODULE_ENTRY Entry, + _In_opt_ PVOID Context) +{ + NTSTATUS status = STATUS_UNSUCCESSFUL; + PACTIVE_SESSION session = (PACTIVE_SESSION)Context; + + if (!ARGUMENT_PRESENT(Context)) + return FALSE; + + status = HashUserModule(Entry, + session->module.module_hash, + sizeof(session->module.module_hash)); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("HashUserModule: %lx", status); + return FALSE; + } + + DEBUG_VERBOSE("User module hashed!"); + DumpBufferToKernelDebugger(session->module.module_hash, + sizeof(session->module.module_hash)); + + return TRUE; +} + NTSTATUS SessionInitialise(_In_ PIRP Irp) { @@ -94,7 +123,8 @@ SessionInitialise(_In_ PIRP Irp) DEBUG_VERBOSE("Initialising new session."); - status = ValidateIrpInputBuffer(Irp, sizeof(SESSION_INITIATION_PACKET)); + status = ValidateIrpInputBuffer( + Irp, sizeof(SESSION_INITIATION_PACKET) - SHA_256_HASH_LENGTH); if (!NT_SUCCESS(status)) { DEBUG_ERROR("ValidateIrpInputBuffer failed with status %x", status); @@ -123,6 +153,16 @@ SessionInitialise(_In_ PIRP Irp) RtlCopyMemory(session->aes_key, initiation->aes_key, AES_256_KEY_SIZE); RtlCopyMemory(session->iv, initiation->aes_iv, AES_256_IV_SIZE); + session->module.base_address = initiation->module_info.base_address; + session->module.size = initiation->module_info.size; + + RtlCopyMemory( + session->module.path, initiation->module_info.path, MAX_MODULE_PATH); + + DEBUG_VERBOSE("Module base: %llx", session->module.base_address); + DEBUG_VERBOSE("Module size: %lx ", session->module.size); + DEBUG_VERBOSE("Module path: %s", session->module.path); + status = CryptInitialiseSessionCryptObjects(); if (!NT_SUCCESS(status)) { @@ -137,6 +177,8 @@ SessionInitialise(_In_ PIRP Irp) goto end; } + FindOurUserModeModuleEntry(HashOurUserModuleOnEntryCallback, session); + end: KeReleaseGuardedMutex(&session->lock); return status; diff --git a/driver/util.c b/driver/util.c index cb41f53..2a4f8cf 100644 --- a/driver/util.c +++ b/driver/util.c @@ -40,4 +40,68 @@ MapAndReadPhysical(_In_ UINT64 PhysicalAddress, MmUnmapIoSpace(va, ReadLength); return STATUS_SUCCESS; +} + +NTSTATUS +UnicodeToCharBufString(_In_ PUNICODE_STRING UnicodeString, + _Out_ PVOID OutBuffer, + _In_ UINT32 OutBufferSize) +{ + ANSI_STRING string = {0}; + NTSTATUS status = STATUS_UNSUCCESSFUL; + + status = RtlUnicodeStringToAnsiString(&string, UnicodeString, TRUE); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("RtlUnicodeStringToAnsiString: %x", status); + return status; + } + + if (string.Length > OutBufferSize) { + RtlFreeAnsiString(&string); + return STATUS_BUFFER_TOO_SMALL; + } + + RtlCopyMemory(OutBuffer, string.Buffer, string.Length); + RtlFreeAnsiString(&string); + + return STATUS_SUCCESS; +} + +#define BYTES_PER_LINE 16 + +VOID +DumpBufferToKernelDebugger(_In_ PCHAR Buffer, _In_ UINT32 BufferLength) +{ + UINT32 i = 0; + UINT32 j = 0; + + for (i = 0; i < BufferLength; i += BYTES_PER_LINE) { + HEX_DUMP("%08x ", i); + + for (j = 0; j < BYTES_PER_LINE; ++j) { + if (i + j < BufferLength) { + HEX_DUMP("%02x ", (unsigned char)Buffer[i + j]); + } + else { + HEX_DUMP(" "); + } + } + + HEX_DUMP(" "); + + for (j = 0; j < BYTES_PER_LINE; ++j) { + if (i + j < BufferLength) { + char c = Buffer[i + j]; + if (c >= 32 && c <= 126) { + HEX_DUMP("%c", c); + } + else { + HEX_DUMP("."); + } + } + } + + HEX_DUMP("\n"); + } } \ No newline at end of file diff --git a/driver/util.h b/driver/util.h index 2436221..ef9a42a 100644 --- a/driver/util.h +++ b/driver/util.h @@ -12,4 +12,12 @@ MapAndReadPhysical(_In_ UINT64 PhysicalAddress, _Out_ PVOID OutputBuffer, _In_ UINT32 OutputBufferLength); +NTSTATUS +UnicodeToCharBufString(_In_ PUNICODE_STRING UnicodeString, + _Out_ PVOID OutBuffer, + _In_ UINT32 OutBufferSize); + +VOID +DumpBufferToKernelDebugger(_In_ PCHAR Buffer, _In_ UINT32 BufferLength); + #endif \ No newline at end of file diff --git a/module/dispatcher/dispatcher.cpp b/module/dispatcher/dispatcher.cpp index a3d37cd..e5b9a2c 100644 --- a/module/dispatcher/dispatcher.cpp +++ b/module/dispatcher/dispatcher.cpp @@ -1,16 +1,19 @@ #include "dispatcher.h" #include "../client/message_queue.h" -#include "../helper.h" #include "../crypt/crypt.h" +#include "../helper.h" #include #include dispatcher::dispatcher::dispatcher(LPCWSTR driver_name, - client::message_queue &message_queue) + client::message_queue &message_queue, + module::module_information *module_info) : thread_pool(DISPATCHER_THREAD_COUNT), - k_interface(driver_name, message_queue) {} + k_interface(driver_name, message_queue, module_info) { + this->module_info = module_info; +} void dispatcher::dispatcher::request_session_pk() { #ifdef NO_SERVER @@ -56,7 +59,7 @@ void dispatcher::dispatcher::run() { this->run_io_port_thread(); thread_pool.queue_job([this]() { k_interface.run_completion_port(); }); while (true) { - LOG_INFO("issueing kernel job!"); + LOG_INFO("issueing kernel job!"); this->issue_kernel_job(); helper::sleep_thread(DISPATCH_LOOP_SLEEP_TIME); } diff --git a/module/dispatcher/dispatcher.h b/module/dispatcher/dispatcher.h index d39a7c6..9d61e18 100644 --- a/module/dispatcher/dispatcher.h +++ b/module/dispatcher/dispatcher.h @@ -3,8 +3,11 @@ #include "threadpool.h" #include "timer.h" + #include "../kernel_interface/kernel_interface.h" +#include "../module.h" + namespace dispatcher { constexpr int DISPATCH_LOOP_SLEEP_TIME = 30; @@ -18,6 +21,7 @@ class dispatcher { timer timers; thread_pool thread_pool; kernel_interface::kernel_interface k_interface; + module::module_information *module_info; void issue_kernel_job(); void write_shared_mapping_operation(); @@ -27,7 +31,7 @@ class dispatcher { void request_session_pk(); public: - dispatcher(LPCWSTR driver_name, client::message_queue &queue); + dispatcher(LPCWSTR driver_name, client::message_queue &queue, module::module_information* module_info); void run(); }; } // namespace dispatcher \ No newline at end of file diff --git a/module/kernel_interface/kernel_interface.cpp b/module/kernel_interface/kernel_interface.cpp index ba505dd..5bbb421 100644 --- a/module/kernel_interface/kernel_interface.cpp +++ b/module/kernel_interface/kernel_interface.cpp @@ -3,8 +3,8 @@ #include #include "../common.h" -#include "../helper.h" #include "../crypt/crypt.h" +#include "../helper.h" #include #include @@ -95,9 +95,11 @@ void *kernel_interface::kernel_interface::get_buffer_from_event_object( } kernel_interface::kernel_interface::kernel_interface( - LPCWSTR driver_name, client::message_queue &queue) + LPCWSTR driver_name, client::message_queue &queue, + module::module_information *module_info) : message_queue(queue) { this->driver_name = driver_name; + this->module_info = module_info; this->port = INVALID_HANDLE_VALUE; this->driver_handle = CreateFileW( driver_name, GENERIC_WRITE | GENERIC_READ | GENERIC_EXECUTE, 0, 0, @@ -146,6 +148,8 @@ void kernel_interface::kernel_interface::notify_driver_on_process_launch() { packet.session_cookie = 123; memcpy(packet.aes_key, crypt::get_test_key(), 32); memcpy(packet.aes_iv, crypt::get_test_iv(), 16); + memcpy(&packet.module_info, (void*)this->module_info, + sizeof(module::module_information)); generic_driver_call_input(ioctl_code::NotifyDriverOnProcessLaunch, &packet, sizeof(session_initiation_packet), &bytes_returned); } diff --git a/module/kernel_interface/kernel_interface.h b/module/kernel_interface/kernel_interface.h index a8f0c67..01ab98d 100644 --- a/module/kernel_interface/kernel_interface.h +++ b/module/kernel_interface/kernel_interface.h @@ -4,6 +4,8 @@ #include "../client/message_queue.h" +#include "../module.h" + namespace kernel_interface { static constexpr int EVENT_COUNT = 5; @@ -130,17 +132,17 @@ struct process_module_validation_report { }; struct system_module_integrity_check_report { - report_header header; - uint64_t image_base; - uint32_t image_size; - char path_name[0x100]; + report_header header; + uint64_t image_base; + uint32_t image_size; + char path_name[0x100]; }; struct driver_self_integrity_check_report { - report_header header; - uint64_t image_base; - uint32_t image_size; - char path_name[0x100]; + report_header header; + uint64_t image_base; + uint32_t image_size; + char path_name[0x100]; }; struct heartbeat_packet { @@ -152,10 +154,10 @@ struct heartbeat_packet { }; struct blacklisted_pcie_device_report { - report_header header; - uint64_t device_object; - uint16_t device_id; - uint16_t vendor_id; + report_header header; + uint64_t device_object; + uint16_t device_id; + uint16_t vendor_id; }; enum apc_operation { operation_stackwalk = 0x1 }; @@ -224,6 +226,7 @@ class kernel_interface { void *process_id; unsigned char aes_key[32]; unsigned char aes_iv[16]; + struct module::module_information module_info; }; struct hv_detection_packet { @@ -247,6 +250,7 @@ class kernel_interface { HANDLE port; std::mutex lock; std::vector events; + module::module_information* module_info; struct shared_data { unsigned __int32 status; @@ -278,7 +282,8 @@ class kernel_interface { void generic_driver_call_apc(apc_operation operation); public: - kernel_interface(LPCWSTR driver_name, client::message_queue &queue); + kernel_interface(LPCWSTR driver_name, client::message_queue &queue, + module::module_information *module_info); ~kernel_interface(); void run_completion_port(); diff --git a/module/module.cpp b/module/module.cpp index 4d54631..ceaa966 100644 --- a/module/module.cpp +++ b/module/module.cpp @@ -6,6 +6,38 @@ #include "dispatcher/dispatcher.h" #include "crypt/crypt.h" +#include + +bool module::get_module_information(module_information *out) { + BOOL ret = FALSE; + HMODULE module = {0}; + MODULEINFO info = {0}; + + ret = GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | + GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, + (LPCSTR)get_module_information, &module); + + if (!ret) + return false; + + ret = GetModuleInformation(GetCurrentProcess(), module, (LPMODULEINFO)&info, + sizeof(info)); + + if (!ret) + return false; + + if (!GetModuleFileNameA(module, out->path, MAX_PATH)) + return false; + + out->base_address = info.lpBaseOfDll; + out->size = info.SizeOfImage; + + LOG_INFO("base: %llx", out->base_address); + LOG_INFO("size: %lx", out->size); + LOG_INFO("path: %s", out->path); + + return true; +} void module::run(HINSTANCE hinstDLL) { AllocConsole(); @@ -16,8 +48,18 @@ void module::run(HINSTANCE hinstDLL) { LPTSTR pipe_name = (LPTSTR)L"\\\\.\\pipe\\DonnaACPipe"; LPCWSTR driver_name = L"\\\\.\\DonnaAC"; + module::module_information info = {0}; + if (!module::get_module_information(&info)) { + LOG_ERROR("get_module_information: %x", GetLastError()); + fclose(stdout); + fclose(stdin); + FreeConsole(); + FreeLibraryAndExitThread(hinstDLL, 0); + return; + } + client::message_queue queue(pipe_name); - dispatcher::dispatcher dispatch(driver_name, queue); + dispatcher::dispatcher dispatch(driver_name, queue, &info); dispatch.run(); fclose(stdout); diff --git a/module/module.h b/module/module.h index 1c2419a..81dd0ae 100644 --- a/module/module.h +++ b/module/module.h @@ -7,4 +7,12 @@ namespace module { void run(HINSTANCE hinstDLL); void terminate(); + +struct module_information { + void *base_address; + uint32_t size; + char path[MAX_PATH]; +}; + +bool get_module_information(module_information *info); } // namespace module \ No newline at end of file