packet encryption!

driver/apc.c

@@ -7,7 +7,7 @@ VOID
 GetApcContextByIndex(_Out_ PVOID* Context, _In_ INT Index)
 {
     AcquireDriverConfigLock();
-    *Context = GetApcContextArray()[Index];
+    *Context = (PVOID)GetApcContextArray()[Index];
     ReleaseDriverConfigLock();
 }
 
@@ -39,14 +39,14 @@ unlock:
  * we hold the lock during the entire decrement and free process.
  */
 BOOLEAN
-FreeApcContextStructure(_Out_ PAPC_CONTEXT_HEADER Context)
+FreeApcContextStructure(_Inout_ PAPC_CONTEXT_HEADER Context)
 {
     DEBUG_VERBOSE("All APCs executed, freeing context structure");
 
     for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
         PUINT64 entry = GetApcContextArray();
 
-        if (entry[index] != Context)
+        if (entry[index] != (UINT64)Context)
             continue;
 
         if (Context->count > 0)
@@ -121,21 +121,21 @@ FreeApcAndDecrementApcCount(_Inout_ PRKAPC Apc, _In_ LONG ContextId)
 NTSTATUS
 QueryActiveApcContextsForCompletion()
 {
-  AcquireDriverConfigLock();
+    AcquireDriverConfigLock();
 
     for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
         PAPC_CONTEXT_HEADER entry = NULL;
         GetApcContextByIndex(&entry, index);
 
         if (!entry)
-          continue;
+            continue;
 
         if (entry->count > 0 || entry->allocation_in_progress == TRUE)
-          continue;
+            continue;
 
         switch (entry->context_id) {
         case APC_CONTEXT_ID_STACKWALK:
-            FreeApcStackwalkApcContextInformation(entry);
+            FreeApcStackwalkApcContextInformation((PAPC_STACKWALK_CONTEXT)entry);
             FreeApcContextStructure(entry);
             break;
         }
@@ -149,16 +149,15 @@ VOID
 InsertApcContext(_In_ PVOID Context)
 {
     if (IsDriverUnloading())
-        return STATUS_UNSUCCESSFUL;
+        return;
 
     AcquireDriverConfigLock();
-    PAPC_CONTEXT_HEADER header = Context;
 
     for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) {
         PUINT64 entry = GetApcContextArray();
 
         if (entry[index] == NULL) {
-            entry[index] = Context;
+            entry[index] = (UINT64)Context;
             goto end;
         }
     }
@@ -213,7 +212,7 @@ DrvUnloadFreeAllApcContextStructures()
 
         ImpExFreePoolWithTag(context, POOL_TAG_APC);
     }
-unlock:
+
     ReleaseDriverConfigLock();
     return TRUE;
 }

driver/apc.h

@@ -15,7 +15,7 @@ VOID
 GetApcContext(_Out_ PVOID* Context, _In_ LONG ContextIdentifier);
 
 BOOLEAN
-FreeApcContextStructure(_Out_ PAPC_CONTEXT_HEADER Context);
+FreeApcContextStructure(_Inout_ PAPC_CONTEXT_HEADER Context);
 
 VOID
 IncrementApcCount(_In_ LONG ContextId);

driver/callbacks.c

@@ -9,6 +9,7 @@
 #include "imports.h"
 #include "list.h"
 #include "session.h"
+#include "crypt.h"
 
 STATIC
 BOOLEAN
@@ -124,7 +125,7 @@ EnumerateThreadListWithCallbackRoutine(
 
     while (entry) {
         CallbackRoutine(entry, Context);
-        entry = entry->list.Next;
+        entry = (PTHREAD_LIST_ENTRY)entry->list.Next;
     }
 
 unlock:
@@ -145,7 +146,7 @@ EnumerateProcessListWithCallbackRoutine(
 
     while (entry) {
         CallbackRoutine(entry, Context);
-        entry = entry->list.Next;
+        entry = (PPROCESS_LIST_ENTRY)entry->list.Next;
     }
 
 unlock:
@@ -166,7 +167,7 @@ EnumerateDriverListWithCallbackRoutine(
 
     while (entry) {
         CallbackRoutine(entry, Context);
-        entry = entry->list.Next;
+        entry = (PDRIVER_LIST_ENTRY)entry->list.Next;
     }
 
 unlock:
@@ -249,7 +250,6 @@ InitialiseDriverList()
 
     list->active = TRUE;
 
-end:
     if (modules.address)
         ImpExFreePoolWithTag(modules.address, SYSTEM_MODULES_POOL);
 
@@ -288,12 +288,13 @@ ImageLoadNotifyRoutineCallback(_In_opt_ PUNICODE_STRING FullImageName,
                                _In_ HANDLE              ProcessId,
                                _In_ PIMAGE_INFO         ImageInfo)
 {
-    NTSTATUS                 status             = STATUS_UNSUCCESSFUL;
+    UNREFERENCED_PARAMETER(ProcessId);
-    PDRIVER_LIST_ENTRY       entry              = NULL;
+
-    RTL_MODULE_EXTENDED_INFO module             = {0};
+    NTSTATUS                 status    = STATUS_UNSUCCESSFUL;
-    PDRIVER_LIST_HEAD        list               = GetDriverList();
+    PDRIVER_LIST_ENTRY       entry     = NULL;
-    ANSI_STRING              ansi_path          = {0};
+    RTL_MODULE_EXTENDED_INFO module    = {0};
-    UINT32                   ansi_string_length = 0;
+    PDRIVER_LIST_HEAD        list      = GetDriverList();
+    ANSI_STRING              ansi_path = {0};
 
     if (InterlockedExchange(&list->active, list->active) == FALSE)
         return;
@@ -486,8 +487,6 @@ ProcessCreateNotifyRoutine(_In_ HANDLE  ParentId,
 
     process_name = ImpPsGetProcessImageFileName(process);
 
-    DEBUG_INFO("process create notify: %s", process_name);
-
     if (Create) {
         entry = ExAllocateFromLookasideListEx(&list->lookaside_list);
 
@@ -585,6 +584,49 @@ ObPostOpCallbackRoutine(_In_ PVOID RegistrationContext,
     UNREFERENCED_PARAMETER(OperationInformation);
 }
 
+#define MAX_PROCESS_NAME_LENGTH             30
+#define PROCESS_HANDLE_OPEN_DOWNGRADE_COUNT 4
+
+CHAR PROCESS_HANDLE_OPEN_DOWNGRADE[PROCESS_HANDLE_OPEN_DOWNGRADE_COUNT]
+                                  [MAX_PROCESS_NAME_LENGTH] = {"lsass.exe",
+                                                               "csrss.exe",
+                                                               "WerFault.exe",
+                                                               "MsMpEng.exe"};
+
+#define PROCESS_HANDLE_OPEN_WHITELIST_COUNT 3
+
+CHAR PROCESS_HANDLE_OPEN_WHITELIST[PROCESS_HANDLE_OPEN_WHITELIST_COUNT]
+                                  [MAX_PROCESS_NAME_LENGTH] = {"Discord.exe",
+                                                               "svchost.exe",
+                                                               "explorer.exe"};
+
+STATIC
+BOOLEAN
+IsWhitelistedHandleOpenProcess(_In_ LPCSTR ProcessName)
+{
+    for (UINT32 index = 0; index < PROCESS_HANDLE_OPEN_WHITELIST_COUNT;
+         index++) {
+        if (!strcmp(ProcessName, PROCESS_HANDLE_OPEN_WHITELIST[index]))
+            return TRUE;
+    }
+
+    return FALSE;
+}
+
+STATIC
+BOOLEAN
+IsDowngradeHandleOpenProcess(_In_ LPCSTR ProcessName)
+{
+    DEBUG_INFO("proc name: %s", ProcessName);
+    for (UINT32 index = 0; index < PROCESS_HANDLE_OPEN_DOWNGRADE_COUNT;
+         index++) {
+        if (!strcmp(ProcessName, PROCESS_HANDLE_OPEN_DOWNGRADE[index]))
+            return TRUE;
+    }
+
+    return FALSE;
+}
+
 // https://www.sysnative.com/forums/threads/object-headers-handles-and-types.34987/
 #define GET_OBJECT_HEADER_FROM_HANDLE(x) ((x << 4) | 0xffff000000000000);
 
@@ -603,6 +645,7 @@ ObPreOpCallbackRoutine(_In_ PVOID                         RegistrationContext,
      * This callback routine is executed in the context of the thread that
      * is requesting to open said handle
      */
+    NTSTATUS  status                 = STATUS_UNSUCCESSFUL;
     PEPROCESS process_creator        = PsGetCurrentProcess();
     PEPROCESS protected_process      = NULL;
     PEPROCESS target_process         = (PEPROCESS)OperationInformation->Object;
@@ -612,6 +655,7 @@ ObPreOpCallbackRoutine(_In_ PVOID                         RegistrationContext,
     LPCSTR    target_process_name    = NULL;
     LPCSTR    protected_process_name = NULL;
     POB_CALLBACKS_CONFIG configuration = NULL;
+    UINT32               report_size   = 0;
 
     /*
      * This is to prevent the condition where the thread executing this
@@ -648,10 +692,7 @@ ObPreOpCallbackRoutine(_In_ PVOID                         RegistrationContext,
      * todo: perform stricter checks rather then the image name.
      * perhapds check some certificate or something.
      */
-    if (!strcmp(process_creator_name, "lsass.exe") ||
+    if (IsDowngradeHandleOpenProcess(process_creator_name) ||
-        !strcmp(process_creator_name, "csrss.exe") ||
-        !strcmp(process_creator_name, "WerFault.exe") ||
-        !strcmp(process_creator_name, "MsMpEng.exe") ||
         !strcmp(process_creator_name, target_process_name)) {
         /* We will downgrade these handles later */
         // DEBUG_LOG("Handles created by CSRSS, LSASS and
@@ -674,27 +715,19 @@ ObPreOpCallbackRoutine(_In_ PVOID                         RegistrationContext,
          * atleast.
          */
 
-        if (!strcmp(process_creator_name, "Discord.exe") ||
+        if (IsWhitelistedHandleOpenProcess(process_creator_name))
-            !strcmp(process_creator_name, "svchost.exe") ||
-            !strcmp(process_creator_name, "explorer.exe"))
             goto end;
 
-        POPEN_HANDLE_FAILURE_REPORT report =
+        report_size = CryptRequestRequiredBufferLength(
-            ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
+            sizeof(OPEN_HANDLE_FAILURE_REPORT));
-                               sizeof(OPEN_HANDLE_FAILURE_REPORT),
+
-                               REPORT_POOL_TAG);
+        POPEN_HANDLE_FAILURE_REPORT report = ImpExAllocatePool2(
+            POOL_FLAG_NON_PAGED, report_size, REPORT_POOL_TAG);
 
         if (!report)
             goto end;
 
-        INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+        INIT_REPORT_PACKET(report, REPORT_ILLEGAL_HANDLE_OPERATION, 0);
-        INIT_REPORT_HEADER(
-            &report->report_header, REPORT_ILLEGAL_HANDLE_OPERATION, 0);
-
-        DEBUG_INFO("packet type: %hx", report->header.packet_type);
-        DEBUG_INFO("report code: %lx", report->report_header.report_code);
-        DEBUG_INFO("report subcode: %lx",
-                   report->report_header.report_sub_type);
 
         report->is_kernel_handle = OperationInformation->KernelHandle;
         report->process_id       = process_creator_id;
@@ -706,7 +739,15 @@ ObPreOpCallbackRoutine(_In_ PVOID                         RegistrationContext,
                       process_creator_name,
                       HANDLE_REPORT_PROCESS_NAME_MAX_LENGTH);
 
-        IrpQueueCompletePacket(report, sizeof(OPEN_HANDLE_FAILURE_REPORT));
+        status = CryptEncryptBuffer(report, report_size);
+
+        if (!NT_SUCCESS(status)) {
+            DEBUG_ERROR("CryptEncryptBuffer: %x", status);
+            ExFreePoolWithTag(report, report_size);
+            goto end;
+        }
+
+        IrpQueueCompletePacket(report, report_size);
     }
 
 end:
@@ -727,7 +768,18 @@ ExUnlockHandleTableEntry(IN PHANDLE_TABLE       HandleTable,
     old_value = InterlockedOr((PLONG)&HandleTableEntry->VolatileLowValue, 1);
 
     /* Unblock any waiters */
+#pragma warning(push)
+#pragma warning(disable : C6387)
     ImpExfUnblockPushLock(&HandleTable->HandleContentionEvent, NULL);
+#pragma warning(pop)
+}
+
+FORCEINLINE
+STATIC
+ACCESS_MASK
+GetHandleAccessMask(_In_ PHANDLE_TABLE_ENTRY Entry)
+{
+    return (ACCESS_MASK)Entry->GrantedAccessBits;
 }
 
 static UNICODE_STRING OBJECT_TYPE_PROCESS = RTL_CONSTANT_STRING(L"Process");
@@ -742,6 +794,9 @@ EnumHandleCallback(_In_ PHANDLE_TABLE       HandleTable,
 {
     PAGED_CODE();
 
+    UNREFERENCED_PARAMETER(Context);
+
+    NTSTATUS     status                 = STATUS_UNSUCCESSFUL;
     PVOID        object                 = NULL;
     PVOID        object_header          = NULL;
     POBJECT_TYPE object_type            = NULL;
@@ -750,12 +805,12 @@ EnumHandleCallback(_In_ PHANDLE_TABLE       HandleTable,
     LPCSTR       process_name           = NULL;
     LPCSTR       protected_process_name = NULL;
     ACCESS_MASK  handle_access_mask     = 0;
+    UINT32       report_size            = 0;
 
     object_header = GET_OBJECT_HEADER_FROM_HANDLE(Entry->ObjectPointerBits);
 
     /* Object header is the first 30 bytes of the object */
-    object = (uintptr_t)object_header + OBJECT_HEADER_SIZE;
+    object      = (uintptr_t)object_header + OBJECT_HEADER_SIZE;
-
     object_type = ImpObGetObjectType(object);
 
     /* TODO: check for threads aswell */
@@ -778,7 +833,7 @@ EnumHandleCallback(_In_ PHANDLE_TABLE       HandleTable,
         "Handle references our protected process with access mask: %lx",
         (ACCESS_MASK)Entry->GrantedAccessBits);
 
-    handle_access_mask = (ACCESS_MASK)Entry->GrantedAccessBits;
+    handle_access_mask = GetHandleAccessMask(Entry);
 
     /* These permissions can be stripped from every process
      * including CSRSS and LSASS */
@@ -851,10 +906,11 @@ EnumHandleCallback(_In_ PHANDLE_TABLE       HandleTable,
         DEBUG_VERBOSE("Stripped PROCESS_VM_WRITE");
     }
 
+    report_size =
+        CryptRequestRequiredBufferLength(sizeof(OPEN_HANDLE_FAILURE_REPORT));
+
     POPEN_HANDLE_FAILURE_REPORT report =
-        ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
+        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, report_size, REPORT_POOL_TAG);
-                           sizeof(OPEN_HANDLE_FAILURE_REPORT),
-                           REPORT_POOL_TAG);
 
     if (!report)
         goto end;
@@ -867,9 +923,7 @@ EnumHandleCallback(_In_ PHANDLE_TABLE       HandleTable,
      * also don't think its worth creating another queue
      * specifically for open handle reports since they will be rare.
      */
-    INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+    INIT_REPORT_PACKET(report, REPORT_ILLEGAL_HANDLE_OPERATION, 0);
-    INIT_REPORT_HEADER(
-        &report->report_header, REPORT_ILLEGAL_HANDLE_OPERATION, 0);
 
     report->is_kernel_handle = Entry->Attributes & OBJ_KERNEL_HANDLE;
     report->process_id       = ImpPsGetProcessId(process);
@@ -880,12 +934,16 @@ EnumHandleCallback(_In_ PHANDLE_TABLE       HandleTable,
                   process_name,
                   HANDLE_REPORT_PROCESS_NAME_MAX_LENGTH);
 
-    if (!NT_SUCCESS(IrpQueueCompletePacket(
+    status = CryptEncryptBuffer(report, report_size);
-            report, sizeof(OPEN_HANDLE_FAILURE_REPORT)))) {
+
-        DEBUG_ERROR("IrpQueueCompleteIrp failed with no status.");
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+        ImpExFreePoolWithTag(report, report_size);
         goto end;
     }
 
+    IrpQueueCompletePacket(report, report_size);
+
 end:
     ExUnlockHandleTableEntry(HandleTable, Entry);
     return FALSE;
@@ -919,8 +977,7 @@ EnumerateProcessHandles(_In_ PPROCESS_LIST_ENTRY ProcessListEntry,
 #pragma warning(push)
 #pragma warning(suppress : 6387)
 
-    BOOLEAN result =
+    ImpExEnumHandleTable(handle_table, EnumHandleCallback, NULL, NULL);
-        ImpExEnumHandleTable(handle_table, EnumHandleCallback, NULL, NULL);
 
 #pragma warning(pop)
 
@@ -929,8 +986,7 @@ EnumerateProcessHandles(_In_ PPROCESS_LIST_ENTRY ProcessListEntry,
 
 #define REPEAT_TIME_10_SEC 10000
 
-ULONG value = 10;
+STATIC
-
 VOID
 TimerObjectWorkItemRoutine(_In_ PDEVICE_OBJECT DeviceObject,
                            _In_opt_ PVOID      Context)
@@ -939,6 +995,11 @@ TimerObjectWorkItemRoutine(_In_ PDEVICE_OBJECT DeviceObject,
     PTIMER_OBJECT     timer  = (PTIMER_OBJECT)Context;
     PDRIVER_LIST_HEAD list   = GetDriverList();
 
+    UNREFERENCED_PARAMETER(DeviceObject);
+
+    if (!ARGUMENT_PRESENT(Context))
+        return;
+
     if (!list->active)
         goto end;
 
@@ -960,17 +1021,22 @@ end:
 /*
  * This routine is executed every x seconds, and is run at IRQL = DISPATCH_LEVEL
  */
+STATIC
 VOID
 TimerObjectCallbackRoutine(_In_ PKDPC     Dpc,
                            _In_opt_ PVOID DeferredContext,
                            _In_opt_ PVOID SystemArgument1,
                            _In_opt_ PVOID SystemArgument2)
 {
-    PTIMER_OBJECT timer = (PTIMER_OBJECT)DeferredContext;
+    UNREFERENCED_PARAMETER(Dpc);
+    UNREFERENCED_PARAMETER(SystemArgument1);
+    UNREFERENCED_PARAMETER(SystemArgument2);
 
-    if (!HasDriverLoaded())
+    if (!HasDriverLoaded() || !ARGUMENT_PRESENT(DeferredContext))
         return;
 
+    PTIMER_OBJECT timer = (PTIMER_OBJECT)DeferredContext;
+
     /* we dont want to queue our work item if it hasnt executed */
     if (timer->state)
         return;
@@ -987,10 +1053,7 @@ TimerObjectCallbackRoutine(_In_ PKDPC     Dpc,
 NTSTATUS
 InitialiseTimerObject(_Out_ PTIMER_OBJECT Timer)
 {
-    LARGE_INTEGER due_time = {0};
+    LARGE_INTEGER due_time = {.QuadPart = -ABSOLUTE(SECONDS(5))};
-    LONG          period   = 0;
-
-    due_time.QuadPart = -ABSOLUTE(SECONDS(5));
 
     Timer->work_item = IoAllocateWorkItem(GetDriverDeviceObject());
 
@@ -1006,7 +1069,7 @@ InitialiseTimerObject(_Out_ PTIMER_OBJECT Timer)
 }
 
 VOID
-CleanupDriverTimerObjects(_Out_ PTIMER_OBJECT Timer)
+CleanupDriverTimerObjects(_Inout_ PTIMER_OBJECT Timer)
 {
     /* this routine blocks until all queued DPCs on all processors have
      * executed. */
@@ -1044,16 +1107,14 @@ RegisterProcessObCallbacks()
 {
     PAGED_CODE();
 
-    NTSTATUS        status = STATUS_UNSUCCESSFUL;
+    NTSTATUS                  status                 = STATUS_UNSUCCESSFUL;
-    PACTIVE_SESSION config = GetActiveSession();
+    PACTIVE_SESSION           config                 = GetActiveSession();
+    OB_CALLBACK_REGISTRATION  callback_registration  = {0};
+    OB_OPERATION_REGISTRATION operation_registration = {0};
 
     DEBUG_VERBOSE("Enabling ObRegisterCallbacks.");
     AcquireDriverConfigLock();
 
-    OB_CALLBACK_REGISTRATION          callback_registration  = {0};
-    OB_OPERATION_REGISTRATION         operation_registration = {0};
-    PCREATE_PROCESS_NOTIFY_ROUTINE_EX notify_routine         = {0};
-
     operation_registration.ObjectType = PsProcessType;
     operation_registration.Operations |= OB_OPERATION_HANDLE_CREATE;
     operation_registration.Operations |= OB_OPERATION_HANDLE_DUPLICATE;

driver/callbacks.h

@@ -77,11 +77,11 @@ CleanupThreadListOnDriverUnload();
 
 VOID
 FindThreadListEntryByThreadAddress(_In_ PKTHREAD               Thread,
-                                   _Inout_ PTHREAD_LIST_ENTRY* Entry);
+                                   _Out_ PTHREAD_LIST_ENTRY* Entry);
 
 VOID
 FindProcessListEntryByProcess(_In_ PKPROCESS               Process,
-                              _Inout_ PPROCESS_LIST_ENTRY* Entry);
+                              _Out_ PPROCESS_LIST_ENTRY* Entry);
 
 VOID
 EnumerateThreadListWithCallbackRoutine(
@@ -110,7 +110,7 @@ NTSTATUS
 InitialiseTimerObject(_Out_ PTIMER_OBJECT Timer);
 
 VOID
-CleanupDriverTimerObjects(_Out_ PTIMER_OBJECT Timer);
+CleanupDriverTimerObjects(_Inout_ PTIMER_OBJECT Timer);
 
 VOID
 UnregisterProcessCreateNotifyRoutine();

driver/common.h

@@ -6,6 +6,7 @@
 #include "io.h"
 
 #include "types/types.h"
+#include <bcrypt.h>
 
 /*
  * For numbers < 32, these are equivalent to 0ul << x.
@@ -80,10 +81,10 @@ typedef struct _DRIVER_LIST_HEAD {
     KGUARDED_MUTEX    lock;
 
     /* modules that need to be hashed later. */
-    PIO_WORKITEM  work_item;
+    PIO_WORKITEM     work_item;
-    LIST_ENTRY    deferred_list;
+    LIST_ENTRY       deferred_list;
     volatile BOOLEAN deferred_complete;
-    volatile LONG can_hash_x86;
+    volatile LONG    can_hash_x86;
 
 } DRIVER_LIST_HEAD, *PDRIVER_LIST_HEAD;
 
@@ -213,7 +214,8 @@ typedef struct _IRP_QUEUE_ENTRY {
  * the target process to protect is open / closed / changes etc.
  */
 
-#define AES_128_KEY_SIZE 16
+#define AES_256_KEY_SIZE 32
+#define AES_256_IV_SIZE  16
 
 typedef struct _HEARTBEAT_CONFIGURATION {
     volatile UINT32 counter;
@@ -233,6 +235,14 @@ typedef struct _HEARTBEAT_CONFIGURATION {
 
 } HEARTBEAT_CONFIGURATION, *PHEARTBEAT_CONFIGURATION;
 
+typedef struct _SESSION_INITIATION_PACKET {
+    UINT32 cookie;
+    PVOID  process_id;
+    UCHAR  aes_key[AES_256_KEY_SIZE];
+    UCHAR  aes_iv[AES_256_IV_SIZE];
+
+} SESSION_INITIATION_PACKET, *PSESSION_INITIATION_PACKET;
+
 typedef struct _ACTIVE_SESSION {
     BOOLEAN             is_session_active;
     PVOID               um_handle;
@@ -240,8 +250,16 @@ typedef struct _ACTIVE_SESSION {
     PEPROCESS           process;
     OB_CALLBACKS_CONFIG callback_configuration;
 
-    UINT32 session_cookie;
+    struct {
-    CHAR   session_aes_key[AES_128_KEY_SIZE];
+        UINT32            cookie;
+        UINT32            magic_number;
+        PUCHAR            aes_key[AES_256_KEY_SIZE];
+        PUCHAR            iv[AES_256_IV_SIZE];
+        BCRYPT_KEY_HANDLE key_handle;
+
+        PUCHAR key_object;
+        UINT32 key_object_length;
+    };
 
     struct SESSION_STATISTICS {
         UINT32 irps_received;
@@ -250,7 +268,7 @@ typedef struct _ACTIVE_SESSION {
     };
 
     HEARTBEAT_CONFIGURATION heartbeat_config;
-    KGUARDED_MUTEX          lock;
+    KSPIN_LOCK              lock;
 
 } ACTIVE_SESSION, *PACTIVE_SESSION;
 
@@ -259,6 +277,7 @@ typedef struct _ACTIVE_SESSION {
 #define INVALID_DRIVER_LIST_HEAD_POOL  'rwar'
 #define INVALID_DRIVER_LIST_ENTRY_POOL 'gaah'
 #define POOL_TAG_APC                   'apcc'
+#define POOL_TAG_CRYPT                 'tpcr'
 #define POOL_TAG_HW                    'hwhw'
 #define POOL_TAG_DPC                   'apcc'
 #define POOL_TAG_HEARTBEAT             'teab'

driver/crypt.c

@@ -1,8 +1,10 @@
 #include "crypt.h"
 
-#include <immintrin.h>
 #include "imports.h"
+#include "session.h"
+#include "driver.h"
 
+#include <immintrin.h>
 #include <bcrypt.h>
 
 #define XOR_KEY_1 0x1122334455667788
@@ -58,7 +60,7 @@ CryptDecryptImportBlock(_In_ PUINT64 Array, _In_ UINT32 BlockIndex)
     return _mm256_xor_si256(load_block, CryptGenerateSseXorKey());
 }
 
-STATIC
+FORCEINLINE
 INLINE
 VOID
 CryptFindContainingBlockForArrayIndex(_In_ UINT32   EntryIndex,
@@ -127,16 +129,188 @@ CryptDecryptImportsArrayEntry(_In_ PUINT64 Array,
     return pointer;
 }
 
-/*
+STATIC
- * simple for now.. just to get it working
+PBCRYPT_KEY_DATA_BLOB_HEADER
- */
+CryptBuildBlobForKeyImport(_In_ PACTIVE_SESSION Session)
-VOID
-CryptDecryptBufferWithCookie(_In_ PVOID  Buffer,
-                             _In_ UINT32 BufferSize,
-                             _In_ UINT32 Cookie)
 {
-    PCHAR buffer = (PCHAR)Buffer;
+    PBCRYPT_KEY_DATA_BLOB_HEADER blob =
-    for (UINT32 index = 0; index < BufferSize; index++) {
+        ExAllocatePool2(POOL_FLAG_NON_PAGED,
-        buffer[index] ^= Cookie;
+                        sizeof(BCRYPT_KEY_DATA_BLOB_HEADER) + AES_256_KEY_SIZE,
-    }
+                        POOL_TAG_CRYPT);
+
+    if (!blob)
+        return NULL;
+
+    blob->dwMagic   = BCRYPT_KEY_DATA_BLOB_MAGIC;
+    blob->dwVersion = BCRYPT_KEY_DATA_BLOB_VERSION1;
+    blob->cbKeyData = AES_256_KEY_SIZE;
+
+    RtlCopyMemory((UINT64)blob + sizeof(BCRYPT_KEY_DATA_BLOB_HEADER),
+                  Session->aes_key,
+                  AES_256_KEY_SIZE);
+
+    return blob;
 }
+
+#define AES_256_BLOCK_SIZE 16
+
+UINT32
+CryptRequestRequiredBufferLength(_In_ UINT32 BufferLength)
+{
+    // status = BCryptEncrypt(session->key_handle,
+    //                        lol,
+    //                        BufferLength,
+    //                        NULL,
+    //                        session->iv,
+    //                        sizeof(session->iv),
+    //                        NULL,
+    //                        0,
+    //                        RequiredLength,
+    //                        0);
+
+    // if (!NT_SUCCESS(status))
+    //     DEBUG_ERROR("CryptRequestRequiredBufferLength -> BCryptEncrypt: %x",
+    //                 status);
+
+    return (BufferLength + AES_256_BLOCK_SIZE - 1) / AES_256_BLOCK_SIZE *
+           AES_256_BLOCK_SIZE;
+}
+
+/* Encrypts in place! */
+NTSTATUS
+CryptEncryptBuffer(_In_ PVOID Buffer, _In_ UINT32 BufferLength)
+{
+    NTSTATUS        status                        = STATUS_UNSUCCESSFUL;
+    UINT32          data_copied                   = 0;
+    PACTIVE_SESSION session                       = GetActiveSession();
+    UCHAR           local_iv[sizeof(session->iv)] = {0};
+    UINT64          buffer                        = (UINT64)Buffer;
+    UINT32          length                        = BufferLength;
+
+    /* The IV is consumed during every encrypt / decrypt procedure, so to ensure
+     * we have access to the iv we need to create a local copy.*/
+    RtlCopyMemory(local_iv, session->iv, sizeof(session->iv));
+
+    /* We arent encrypting the first 16 bytes */
+    buffer = buffer + AES_256_BLOCK_SIZE;
+    length = length - AES_256_BLOCK_SIZE;
+
+    status = BCryptEncrypt(session->key_handle,
+                           buffer,
+                           length,
+                           NULL,
+                           local_iv,
+                           sizeof(local_iv),
+                           buffer,
+                           length,
+                           &data_copied,
+                           0);
+
+    if (!NT_SUCCESS(status))
+        DEBUG_ERROR("CryptEncryptBuffer -> BCryptEncrypt: %x", status);
+
+    return status;
+}
+
+/* Lock is held */
+VOID
+CryptCloseSessionCryptObjects()
+{
+    PACTIVE_SESSION session = GetActiveSession();
+
+    if (session->key_handle) {
+        BCryptDestroyKey(session->key_handle);
+        session->key_handle = NULL;
+    }
+
+    if (session->key_object) {
+        ExFreePoolWithTag(session->key_object, POOL_TAG_CRYPT);
+        session->key_object = NULL;
+    }
+
+    session->key_object_length = 0;
+}
+
+NTSTATUS
+CryptInitialiseSessionCryptObjects()
+{
+    NTSTATUS                     status      = STATUS_UNSUCCESSFUL;
+    UINT32                       data_copied = 0;
+    PACTIVE_SESSION              session     = GetActiveSession();
+    PBCRYPT_KEY_DATA_BLOB_HEADER blob        = NULL;
+    BCRYPT_ALG_HANDLE*           handle      = GetCryptAlgHandle();
+
+    blob = CryptBuildBlobForKeyImport(session);
+
+    if (!blob)
+        return STATUS_INSUFFICIENT_RESOURCES;
+
+    status = BCryptGetProperty(*handle,
+                               BCRYPT_OBJECT_LENGTH,
+                               &session->key_object_length,
+                               sizeof(UINT32),
+                               &data_copied,
+                               0);
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("BCryptGetProperty: %x", status);
+        goto end;
+    }
+
+    session->key_object = ExAllocatePool2(
+        POOL_FLAG_NON_PAGED, session->key_object_length, POOL_TAG_CRYPT);
+
+    if (!session->key_object) {
+        status = STATUS_INSUFFICIENT_RESOURCES;
+        goto end;
+    }
+
+    DEBUG_INFO("key object: %llx, key_object_length: %lx",
+               session->key_object,
+               session->key_object_length);
+
+    status =
+        BCryptImportKey(*handle,
+                        NULL,
+                        BCRYPT_KEY_DATA_BLOB,
+                        &session->key_handle,
+                        session->key_object,
+                        session->key_object_length,
+                        blob,
+                        sizeof(BCRYPT_KEY_DATA_BLOB_HEADER) + AES_256_KEY_SIZE,
+                        0);
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("BCryptImportKey: %x", status);
+        ExFreePoolWithTag(session->key_object, POOL_TAG_CRYPT);
+        goto end;
+    }
+
+end:
+    if (blob)
+        ExFreePoolWithTag(blob, POOL_TAG_CRYPT);
+
+    return status;
+}
+
+NTSTATUS
+CryptInitialiseProvider()
+{
+    NTSTATUS           status = STATUS_UNSUCCESSFUL;
+    BCRYPT_ALG_HANDLE* handle = GetCryptAlgHandle();
+
+    status = BCryptOpenAlgorithmProvider(
+        handle, BCRYPT_AES_ALGORITHM, NULL, BCRYPT_PROV_DISPATCH);
+
+    if (!NT_SUCCESS(status))
+        DEBUG_ERROR("BCryptOpenAlgorithmProvider: %x", status);
+
+    return status;
+}
+
+VOID
+CryptCloseProvider()
+{
+    BCRYPT_ALG_HANDLE* handle = GetCryptAlgHandle();
+    BCryptCloseAlgorithmProvider(*handle, 0);
+}

driver/crypt.h

@@ -11,9 +11,22 @@ CryptDecryptImportsArrayEntry(_In_ PUINT64 Array,
                               _In_ UINT32  Entries,
                               _In_ UINT32  EntryIndex);
 
+NTSTATUS
+CryptInitialiseProvider();
+
+UINT32
+CryptRequestRequiredBufferLength(_In_ UINT32 BufferLength);
+
+NTSTATUS
+CryptEncryptBuffer(_In_ PVOID Buffer, _In_ UINT32 BufferLength);
+
+NTSTATUS
+CryptInitialiseSessionCryptObjects();
+
 VOID
-CryptDecryptBufferWithCookie(_In_ PVOID  Buffer,
+CryptCloseSessionCryptObjects();
-                             _In_ UINT32 BufferSize,
+
-                             _In_ UINT32 Cookie);
+VOID
+CryptCloseProvider();
 
 #endif

driver/driver.c

@@ -93,14 +93,16 @@ typedef struct _DRIVER_CONFIG {
     IRP_QUEUE_HEAD         irp_queue;
 
     /* terrible name..lol what is tis timer for ?? */
-    TIMER_OBJECT           timer;
+    TIMER_OBJECT timer;
 
-    ACTIVE_SESSION         session_information;
+    ACTIVE_SESSION    session_information;
-    THREAD_LIST_HEAD       thread_list;
+    THREAD_LIST_HEAD  thread_list;
-    DRIVER_LIST_HEAD       driver_list;
+    DRIVER_LIST_HEAD  driver_list;
-    PROCESS_LIST_HEAD      process_list;
+    PROCESS_LIST_HEAD process_list;
-    SHARED_MAPPING         mapping;
+    SHARED_MAPPING    mapping;
-    BOOLEAN                has_driver_loaded;
+    BOOLEAN           has_driver_loaded;
+
+    BCRYPT_ALG_HANDLE alg_handle;
 
 } DRIVER_CONFIG, *PDRIVER_CONFIG;
 
@@ -120,6 +122,12 @@ PDRIVER_CONFIG g_DriverConfig = NULL;
 
 #define POOL_TAG_CONFIG 'conf'
 
+BCRYPT_ALG_HANDLE*
+GetCryptAlgHandle()
+{
+    return &g_DriverConfig->alg_handle;
+}
+
 BOOLEAN
 HasDriverLoaded()
 {
@@ -382,6 +390,8 @@ DriverUnload(_In_ PDRIVER_OBJECT DriverObject)
     DrvUnloadFreeProcessList();
     DrvUnloadFreeDriverList();
 
+    CryptCloseProvider();
+
     DrvUnloadFreeConfigStrings();
     DrvUnloadDeleteSymbolicLink();
     ImpIoDeleteDevice(DriverObject->DeviceObject);
@@ -863,7 +873,15 @@ DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegistryPath)
         return status;
     }
 
-    SessionInitialiseStructure();
+    status = SessionInitialiseStructure();
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("SessionInitialiseStructure failed with status %x", status);
+        DrvUnloadFreeConfigStrings();
+        DrvUnloadFreeTimerObject();
+        ImpIoDeleteDevice(DriverObject->DeviceObject);
+        return status;
+    }
 
     status = IoCreateSymbolicLink(g_DriverConfig->device_symbolic_link,
                                   g_DriverConfig->device_name);

driver/driver.h

@@ -10,6 +10,9 @@
 #include "integrity.h"
 #include "callbacks.h"
 
+BCRYPT_ALG_HANDLE*
+GetCryptAlgHandle();
+
 NTSTATUS
 QueryActiveApcContextsForCompletion();
 

driver/hw.c

@@ -64,7 +64,7 @@ STATIC
 NTSTATUS
 QueryPciDeviceConfigurationSpace(_In_ PDEVICE_OBJECT DeviceObject,
                                  _In_ UINT32         Offset,
-                                 _Out_ PVOID         Buffer,
+                                 _Out_opt_ PVOID     Buffer,
                                  _In_ UINT32         BufferLength)
 {
     NTSTATUS           status            = STATUS_UNSUCCESSFUL;
@@ -100,7 +100,7 @@ QueryPciDeviceConfigurationSpace(_In_ PDEVICE_OBJECT DeviceObject,
 
     status = IoCallDriver(DeviceObject, irp);
 
-    if (status = STATUS_PENDING) {
+    if (status == STATUS_PENDING) {
         KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
         status = io.Status;
     }
@@ -127,6 +127,7 @@ EnumerateDriverObjectDeviceObjects(_In_ PDRIVER_OBJECT    DriverObject,
     UINT32          buffer_size  = 0;
 
     *DeviceObjectArray = NULL;
+    *ArrayEntries      = 0;
 
     status = IoEnumerateDeviceObjectList(DriverObject, NULL, 0, &object_count);
 
@@ -200,7 +201,6 @@ EnumeratePciDeviceObjects(_In_ PCI_DEVICE_CALLBACK CallbackRoutine,
     PDEVICE_OBJECT* pci_device_objects = NULL;
     PDEVICE_OBJECT  current_device     = NULL;
     UINT32          pci_device_objects_count = 0;
-    USHORT          vendor_id                = 0;
 
     status = GetDriverObjectByDriverName(&pci, &pci_driver_object);
 
@@ -238,7 +238,6 @@ EnumeratePciDeviceObjects(_In_ PCI_DEVICE_CALLBACK CallbackRoutine,
         ObDereferenceObject(current_device);
     }
 
-end:
     if (pci_device_objects)
         ExFreePoolWithTag(pci_device_objects, POOL_TAG_HW);
 
@@ -260,6 +259,8 @@ STATIC
 NTSTATUS
 PciDeviceQueryCallback(_In_ PDEVICE_OBJECT DeviceObject, _In_opt_ PVOID Context)
 {
+    UNREFERENCED_PARAMETER(Context);
+
     NTSTATUS          status = STATUS_UNSUCCESSFUL;
     PCI_COMMON_HEADER header = {0};
 

driver/integrity.c

@@ -9,6 +9,7 @@
 #include "session.h"
 #include "util.h"
 #include "pe.h"
+#include "crypt.h"
 
 #include <bcrypt.h>
 #include <initguid.h>
@@ -174,6 +175,8 @@ GetModuleInformationByName(_Out_ PRTL_MODULE_EXTENDED_INFO ModuleInfo,
         return status;
     }
 
+    /* TODO: think this remains from testing, we only use this to find our
+     * driver anyway but should be fixed. */
     driver_info = FindSystemModuleByName(driver_name, &modules);
 
     if (!driver_info) {
@@ -244,7 +247,6 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID*  Buffer,
     PAGED_CODE();
 
     NTSTATUS               status                  = STATUS_UNSUCCESSFUL;
-    PIMAGE_DOS_HEADER      dos_header              = NULL;
     PNT_HEADER_64          nt_header               = NULL;
     PIMAGE_SECTION_HEADER  section                 = NULL;
     ULONG                  total_packet_size       = 0;
@@ -865,24 +867,32 @@ STATIC
 VOID
 ReportInvalidProcessModule(_In_ PPROCESS_MODULE_INFORMATION Module)
 {
+    NTSTATUS status      = STATUS_UNSUCCESSFUL;
+    UINT32   report_size = CryptRequestRequiredBufferLength(
+        sizeof(PROCESS_MODULE_VALIDATION_REPORT));
+
     PPROCESS_MODULE_VALIDATION_REPORT report =
-        ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
+        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, report_size, REPORT_POOL_TAG);
-                           sizeof(PROCESS_MODULE_VALIDATION_REPORT),
-                           REPORT_POOL_TAG);
 
     if (!report)
         return;
 
-    INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+    INIT_REPORT_PACKET(report, REPORT_INVALID_PROCESS_MODULE, 0);
-    INIT_REPORT_HEADER(
-        &report->report_header, REPORT_INVALID_PROCESS_MODULE, 0);
 
     report->image_base = Module->module_base;
     report->image_size = Module->module_size;
     RtlCopyMemory(
         report->module_path, Module->module_path, sizeof(report->module_path));
 
-    IrpQueueCompletePacket(report, sizeof(PROCESS_MODULE_VALIDATION_REPORT));
+    status = CryptEncryptBuffer(report, report_size);
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+        ImpExFreePoolWithTag(report, report_size);
+        return;
+    }
+
+    IrpQueueCompletePacket(report, report_size);
 }
 
 /*
@@ -961,7 +971,7 @@ ValidateProcessLoadedModule(_Inout_ PIRP Irp)
     }
 
     status = MapDiskImageIntoVirtualAddressSpace(
-        &section_handle, &section, &module_path, &section_size, 0);
+        &section_handle, &section, &module_path, &section_size);
 
     if (!NT_SUCCESS(status)) {
         DEBUG_ERROR("MapDiskImageIntoVirtualAddressSpace failed with status %x",
@@ -1450,8 +1460,12 @@ Enablex86Hashing(_In_ PDRIVER_LIST_HEAD Head)
 }
 
 VOID
-DeferredModuleHashingCallback()
+DeferredModuleHashingCallback(_In_ PDEVICE_OBJECT DeviceObject,
+                              _In_opt_ PVOID      Context)
 {
+    UNREFERENCED_PARAMETER(Context);
+    UNREFERENCED_PARAMETER(DeviceObject);
+
     NTSTATUS                 status        = STATUS_UNSUCCESSFUL;
     RTL_MODULE_EXTENDED_INFO module        = {0};
     PDRIVER_LIST_HEAD        driver_list   = GetDriverList();
@@ -1488,6 +1502,7 @@ DeferredModuleHashingCallback()
     }
 
 end:
+
     DEBUG_VERBOSE("All deferred modules hashed.");
     ImpIoFreeWorkItem(driver_list->work_item);
     driver_list->work_item = NULL;
@@ -1510,7 +1525,7 @@ HashModule(_In_ PRTL_MODULE_EXTENDED_INFO Module, _Out_ PVOID Hash)
 
     if (!ansi_string.Buffer) {
         DEBUG_ERROR("RtlInitAnsiString failed with status %x", status);
-        return;
+        return STATUS_UNSUCCESSFUL;
     }
 
     status = ImpRtlAnsiStringToUnicodeString(&path, &ansi_string, TRUE);
@@ -1518,7 +1533,7 @@ HashModule(_In_ PRTL_MODULE_EXTENDED_INFO Module, _Out_ PVOID Hash)
     if (!NT_SUCCESS(status)) {
         DEBUG_ERROR("RtlAnsiStringToUnicodeString failed with status %x",
                     status);
-        goto end;
+        return status;
     }
 
     /*
@@ -1622,6 +1637,7 @@ ValidateSystemModule(_In_ PRTL_MODULE_EXTENDED_INFO Module)
                       Module->FullPathName);
 
 end:
+
     if (hash)
         ExFreePoolWithTag(hash, POOL_TAG_INTEGRITY);
 }
@@ -1754,6 +1770,8 @@ VOID
 SystemModuleVerificationDispatchFunction(_In_ PDEVICE_OBJECT DeviceObject,
                                          _In_ PSYS_MODULE_VAL_CONTEXT Context)
 {
+    UNREFERENCED_PARAMETER(DeviceObject);
+
     IncrementActiveThreadCount(Context);
 
     UINT32 count = GetCurrentVerificationIndex(Context);
@@ -2068,7 +2086,7 @@ AllocateHeartbeatObjects(_Inout_ PHEARTBEAT_CONFIGURATION Configuration)
  * intervals. */
 STATIC
 LARGE_INTEGER
-GenerateHeartbeatDueTime(_In_ PHEARTBEAT_CONFIGURATION Configuration)
+GenerateHeartbeatDueTime()
 {
     LARGE_INTEGER ticks = {0};
     KeQueryTickCount(&ticks);
@@ -2089,9 +2107,8 @@ InitialiseHeartbeatObjects(_Inout_ PHEARTBEAT_CONFIGURATION Configuration)
 {
     KeInitializeDpc(Configuration->dpc, HeartbeatDpcRoutine, Configuration);
     KeInitializeTimer(Configuration->timer);
-    KeSetTimer(Configuration->timer,
+    KeSetTimer(
-               GenerateHeartbeatDueTime(Configuration),
+        Configuration->timer, GenerateHeartbeatDueTime(), Configuration->dpc);
-               Configuration->dpc);
 }
 
 FORCEINLINE
@@ -2158,16 +2175,17 @@ IncrementHeartbeatCounter(_In_ PHEARTBEAT_CONFIGURATION Configuration)
 FORCEINLINE
 STATIC
 PHEARTBEAT_PACKET
-BuildHeartbeatPacket(_In_ PHEARTBEAT_CONFIGURATION Configuration)
+BuildHeartbeatPacket(_In_ UINT32 PacketSize)
 {
-    PIRP_QUEUE_HEAD   queue  = GetIrpQueueHead();
+    PIRP_QUEUE_HEAD queue = GetIrpQueueHead();
-    PHEARTBEAT_PACKET packet = ImpExAllocatePool2(
+
-        POOL_FLAG_NON_PAGED, sizeof(HEARTBEAT_PACKET), POOL_TAG_HEARTBEAT);
+    PHEARTBEAT_PACKET packet =
+        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, PacketSize, POOL_TAG_HEARTBEAT);
 
     if (!packet)
         return NULL;
 
-    INIT_PACKET_HEADER(&packet->header, PACKET_TYPE_HEARTBEAT);
+    INIT_HEARTBEAT_PACKET(packet);
 
     /* This routine always runs at DPC level */
     KeAcquireSpinLockAtDpcLevel(&queue->lock);
@@ -2199,22 +2217,32 @@ HeartbeatDpcRoutine(_In_ PKDPC     Dpc,
     if (!ARGUMENT_PRESENT(DeferredContext))
         return;
 
+    NTSTATUS                 status = STATUS_UNSUCCESSFUL;
     PHEARTBEAT_CONFIGURATION config = (PHEARTBEAT_CONFIGURATION)DeferredContext;
     PHEARTBEAT_PACKET        packet = NULL;
+    UINT32                   packet_size = 0;
 
     DEBUG_VERBOSE("Heartbeat timer alerted. Generating heartbeat packet.");
 
     SetHeartbeatActive(config);
 
-    packet = BuildHeartbeatPacket(config);
+    packet_size = CryptRequestRequiredBufferLength(sizeof(HEARTBEAT_PACKET));
+    packet      = BuildHeartbeatPacket(packet_size);
 
     if (packet) {
-        IrpQueueCompletePacket(packet, sizeof(HEARTBEAT_PACKET));
+        status = CryptEncryptBuffer(packet, packet_size);
+
+        if (!NT_SUCCESS(status)) {
+            DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+            ImpExFreePoolWithTag(packet, POOL_TAG_HEARTBEAT);
+            goto end;
+        }
+
+        IrpQueueCompletePacket(packet, packet_size);
         IncrementHeartbeatCounter(config);
     }
 
 end:
-
     IoQueueWorkItem(
         config->work_item, HeartbeatWorkItem, NormalWorkQueue, config);
 }

driver/integrity.h

@@ -115,7 +115,8 @@ BOOLEAN
 ValidateOurDriversDispatchRoutines();
 
 VOID
-DeferredModuleHashingCallback();
+DeferredModuleHashingCallback(_In_ PDEVICE_OBJECT DeviceObject,
+                              _In_opt_ PVOID      Context);
 
 VOID
 FindWinLogonProcess(_In_ PPROCESS_LIST_ENTRY Entry, _In_opt_ PVOID Context);

driver/modules.c

@@ -8,6 +8,7 @@
 #include "apc.h"
 #include "thread.h"
 #include "pe.h"
+#include "crypt.h"
 
 #define WHITELISTED_MODULE_TAG 'whte'
 
@@ -340,18 +341,17 @@ STATIC
 VOID
 ReportInvalidDriverObject(_In_ PDRIVER_OBJECT Driver, _In_ UINT32 ReportSubType)
 {
-    PMODULE_VALIDATION_FAILURE report =
+    NTSTATUS status = STATUS_UNSUCCESSFUL;
-        ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
+    UINT32   packet_size =
-                           sizeof(MODULE_VALIDATION_FAILURE),
+        CryptRequestRequiredBufferLength(sizeof(MODULE_VALIDATION_FAILURE));
-                           POOL_TAG_INTEGRITY);
+
+    PMODULE_VALIDATION_FAILURE report = ImpExAllocatePool2(
+        POOL_FLAG_NON_PAGED, packet_size, POOL_TAG_INTEGRITY);
 
     if (!report)
         return;
 
-    INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+    INIT_REPORT_PACKET(report, REPORT_MODULE_VALIDATION_FAILURE, ReportSubType);
-    INIT_REPORT_HEADER(&report->report_header,
-                       REPORT_MODULE_VALIDATION_FAILURE,
-                       ReportSubType);
 
     report->driver_base_address = Driver->DriverStart;
     report->driver_size         = Driver->DriverSize;
@@ -363,7 +363,16 @@ ReportInvalidDriverObject(_In_ PDRIVER_OBJECT Driver, _In_ UINT32 ReportSubType)
 
     /* Continue regardless of result */
     ImpRtlUnicodeStringToAnsiString(&string, &Driver->DriverName, FALSE);
-    IrpQueueCompletePacket(report, sizeof(MODULE_VALIDATION_FAILURE));
+
+    status = CryptEncryptBuffer(report, packet_size);
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+        ImpExFreePoolWithTag(report, REPORT_POOL_TAG);
+        return;
+    }
+
+    IrpQueueCompletePacket(report, packet_size);
 }
 
 FORCEINLINE
@@ -562,20 +571,31 @@ STATIC
 VOID
 ReportNmiBlocking()
 {
-    PNMI_CALLBACK_FAILURE report = ImpExAllocatePool2(
+    NTSTATUS status = STATUS_UNSUCCESSFUL;
-        POOL_FLAG_NON_PAGED, sizeof(NMI_CALLBACK_FAILURE), REPORT_POOL_TAG);
+    UINT32   packet_size =
+        CryptRequestRequiredBufferLength(sizeof(NMI_CALLBACK_FAILURE));
+
+    PNMI_CALLBACK_FAILURE report =
+        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG);
 
     if (!report)
         return STATUS_INSUFFICIENT_RESOURCES;
 
-    INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+    INIT_REPORT_PACKET(report, REPORT_NMI_CALLBACK_FAILURE, 0);
-    INIT_REPORT_HEADER(&report->report_header, REPORT_NMI_CALLBACK_FAILURE, 0);
 
     report->kthread_address    = NULL;
     report->invalid_rip        = NULL;
     report->were_nmis_disabled = TRUE;
 
-    IrpQueueCompletePacket(report, sizeof(NMI_CALLBACK_FAILURE));
+    status = CryptEncryptBuffer(report, packet_size);
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+        ImpExFreePoolWithTag(report, REPORT_POOL_TAG);
+        return;
+    }
+
+    IrpQueueCompletePacket(report, packet_size);
 }
 
 STATIC
@@ -585,16 +605,17 @@ ReportMissingCidTableEntry(_In_ PNMI_CONTEXT Context)
     DEBUG_WARNING("Thread: %llx was not found in the pspcid table.",
                   Context->kthread);
 
+    NTSTATUS status = STATUS_UNSUCCESSFUL;
+    UINT32   packet_size =
+        CryptRequestRequiredBufferLength(sizeof(HIDDEN_SYSTEM_THREAD_REPORT));
+
     PHIDDEN_SYSTEM_THREAD_REPORT report =
-        ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
+        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG);
-                           sizeof(HIDDEN_SYSTEM_THREAD_REPORT),
-                           REPORT_POOL_TAG);
 
     if (!report)
         return;
 
-    INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+    INIT_REPORT_PACKET(report, REPORT_HIDDEN_SYSTEM_THREAD, 0);
-    INIT_REPORT_HEADER(&report->report_header, REPORT_HIDDEN_SYSTEM_THREAD, 0);
 
     report->found_in_kthreadlist = FALSE; // wip
     report->found_in_pspcidtable = FALSE;
@@ -602,29 +623,47 @@ ReportMissingCidTableEntry(_In_ PNMI_CONTEXT Context)
     report->thread_address       = Context->kthread;
 
     RtlCopyMemory(report->thread, Context->kthread, sizeof(report->thread));
-    IrpQueueCompletePacket(report, sizeof(HIDDEN_SYSTEM_THREAD_REPORT));
+
+    status = CryptEncryptBuffer(report, packet_size);
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+        ImpExFreePoolWithTag(report, REPORT_POOL_TAG);
+        return;
+    }
+
+    IrpQueueCompletePacket(report, packet_size);
 }
 
 STATIC
 VOID
 ReportInvalidRipFoundDuringNmi(_In_ PNMI_CONTEXT Context)
 {
+    NTSTATUS status = STATUS_UNSUCCESSFUL;
+    UINT32   packet_size =
+        CryptRequestRequiredBufferLength(sizeof(HIDDEN_SYSTEM_THREAD_REPORT));
+
     PNMI_CALLBACK_FAILURE report =
-        ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
+        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG);
-                           sizeof(HIDDEN_SYSTEM_THREAD_REPORT),
-                           REPORT_POOL_TAG);
 
     if (!report)
         return;
 
-    INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+    INIT_REPORT_PACKET(report, REPORT_NMI_CALLBACK_FAILURE, 0);
-    INIT_REPORT_HEADER(&report->report_header, REPORT_NMI_CALLBACK_FAILURE, 0);
 
     report->kthread_address    = Context->kthread;
     report->invalid_rip        = Context->interrupted_rip;
     report->were_nmis_disabled = FALSE;
 
-    IrpQueueCompletePacket(report, sizeof(HIDDEN_SYSTEM_THREAD_REPORT));
+    status = CryptEncryptBuffer(report, packet_size);
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+        ImpExFreePoolWithTag(report, REPORT_POOL_TAG);
+        return;
+    }
+
+    IrpQueueCompletePacket(report, packet_size);
 }
 
 /*
@@ -728,6 +767,9 @@ NmiCallback(_Inout_opt_ PVOID Context, _In_ BOOLEAN Handled)
     TASK_STATE_SEGMENT_64* tss           = NULL;
     PMACHINE_FRAME         machine_frame = NULL;
 
+    if (!ARGUMENT_PRESENT(Context))
+        return TRUE;
+
     /*
      * To find the IRETQ frame (MACHINE_FRAME) we need to find the top of
      * the NMI ISR stack. This is stored at TSS->Ist[3]. To find the TSS, we
@@ -890,20 +932,31 @@ STATIC
 VOID
 ReportApcStackwalkViolation(_In_ UINT64 Rip)
 {
-    PAPC_STACKWALK_REPORT report = ImpExAllocatePool2(
+    NTSTATUS status = STATUS_UNSUCCESSFUL;
-        POOL_FLAG_NON_PAGED, sizeof(APC_STACKWALK_REPORT), REPORT_POOL_TAG);
+    UINT32   packet_size =
+        CryptRequestRequiredBufferLength(sizeof(APC_STACKWALK_REPORT));
+
+    PAPC_STACKWALK_REPORT report =
+        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG);
 
     if (!report)
         return;
 
-    INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+    INIT_REPORT_PACKET(report, REPORT_APC_STACKWALK, 0);
-    INIT_REPORT_HEADER(&report->report_header, REPORT_APC_STACKWALK, 0);
 
     report->kthread_address = (UINT64)KeGetCurrentThread();
     report->invalid_rip     = Rip;
     // report->driver ?? todo!
 
-    IrpQueueCompletePacket(report, sizeof(APC_STACKWALK_REPORT));
+    status = CryptEncryptBuffer(report, packet_size);
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+        ImpExFreePoolWithTag(report, REPORT_POOL_TAG);
+        return;
+    }
+
+    IrpQueueCompletePacket(report, packet_size);
 }
 
 /*
@@ -1004,6 +1057,9 @@ ValidateThreadViaKernelApcCallback(_In_ PTHREAD_LIST_ENTRY ThreadListEntry,
     LPCSTR                 process_name  = NULL;
     PAPC_STACKWALK_CONTEXT context       = (PAPC_STACKWALK_CONTEXT)Context;
 
+    if (!ARGUMENT_PRESENT(Context))
+        return;
+
     process_name =
         ImpPsGetProcessImageFileName(ThreadListEntry->owning_process);
 
@@ -1165,6 +1221,12 @@ DpcStackwalkCallbackRoutine(_In_ PKDPC     Dpc,
                             _In_opt_ PVOID SystemArgument1,
                             _In_opt_ PVOID SystemArgument2)
 {
+    UNREFERENCED_PARAMETER(Dpc);
+    UNREFERENCED_PARAMETER(SystemArgument2);
+
+    if (!ARGUMENT_PRESENT(DeferredContext))
+        return;
+
     PDPC_CONTEXT context =
         &((PDPC_CONTEXT)DeferredContext)[KeGetCurrentProcessorNumber()];
 
@@ -1174,7 +1236,11 @@ DpcStackwalkCallbackRoutine(_In_ PKDPC     Dpc,
                                     &context->stack_frame,
                                     NULL);
     InterlockedExchange(&context->executed, TRUE);
+
+#pragma warning(push)
+#pragma warning(disable : C6387)
     ImpKeSignalCallDpcDone(SystemArgument1);
+#pragma warning(pop)
 
     DEBUG_VERBOSE("Executed DPC on core: %lx, with %lx frames captured.",
                   KeGetCurrentProcessorNumber(),
@@ -1198,14 +1264,17 @@ STATIC
 VOID
 ReportDpcStackwalkViolation(_In_ PDPC_CONTEXT Context, _In_ UINT64 Frame)
 {
-    PDPC_STACKWALK_REPORT report = ImpExAllocatePool2(
+    NTSTATUS status = STATUS_UNSUCCESSFUL;
-        POOL_FLAG_NON_PAGED, sizeof(DPC_STACKWALK_REPORT), REPORT_POOL_TAG);
+    UINT32   packet_size =
+        CryptRequestRequiredBufferLength(sizeof(DPC_STACKWALK_REPORT));
+
+    PDPC_STACKWALK_REPORT report =
+        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG);
 
     if (!report)
         return;
 
-    INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+    INIT_REPORT_PACKET(report, REPORT_DPC_STACKWALK, 0);
-    INIT_REPORT_HEADER(&report->report_header, REPORT_DPC_STACKWALK, 0);
 
     report->kthread_address = PsGetCurrentThread();
     report->invalid_rip     = Frame;
@@ -1215,7 +1284,15 @@ ReportDpcStackwalkViolation(_In_ PDPC_CONTEXT Context, _In_ UINT64 Frame)
     //               - 0x50,
     //               APC_STACKWALK_BUFFER_SIZE);
 
-    IrpQueueCompletePacket(report, sizeof(DPC_STACKWALK_REPORT));
+    status = CryptEncryptBuffer(report, packet_size);
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+        ImpExFreePoolWithTag(report, REPORT_POOL_TAG);
+        return;
+    }
+
+    IrpQueueCompletePacket(report, packet_size);
 }
 
 STATIC
@@ -1500,10 +1577,12 @@ STATIC
 VOID
 ReportDataTableInvalidRoutine(_In_ TABLE_ID TableId, _In_ UINT64 Address)
 {
+    NTSTATUS status = STATUS_UNSUCCESSFUL;
+    UINT32   packet_size =
+        CryptRequestRequiredBufferLength(sizeof(DATA_TABLE_ROUTINE_REPORT));
+
     PDATA_TABLE_ROUTINE_REPORT report =
-        ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
+        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG);
-                           sizeof(DATA_TABLE_ROUTINE_REPORT),
-                           REPORT_POOL_TAG);
 
     if (!report)
         return;
@@ -1512,17 +1591,22 @@ ReportDataTableInvalidRoutine(_In_ TABLE_ID TableId, _In_ UINT64 Address)
                   TableId,
                   Address);
 
-    INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+    INIT_REPORT_PACKET(report, REPORT_DATA_TABLE_ROUTINE, 0);
-    INIT_REPORT_HEADER(&report->report_header, REPORT_DATA_TABLE_ROUTINE, 0);
 
     report->address  = Address;
     report->table_id = TableId;
     report->index    = 0;
     RtlCopyMemory(report->routine, Address, DATA_TABLE_ROUTINE_BUF_SIZE);
 
-    if (!NT_SUCCESS(
+    status = CryptEncryptBuffer(report, packet_size);
-            IrpQueueCompletePacket(report, sizeof(DATA_TABLE_ROUTINE_REPORT))))
+
-        DEBUG_ERROR("IrpQueueCompleteIrp failed with no status.");
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+        ImpExFreePoolWithTag(report, REPORT_POOL_TAG);
+        return;
+    }
+
+    IrpQueueCompletePacket(report, packet_size);
 }
 
 NTSTATUS
@@ -1835,16 +1919,17 @@ VOID
 ReportWin32kBase_DxgInterfaceViolation(_In_ UINT32 TableIndex,
                                        _In_ UINT64 Address)
 {
+    NTSTATUS status = STATUS_UNSUCCESSFUL;
+    UINT32   packet_size =
+        CryptRequestRequiredBufferLength(sizeof(DATA_TABLE_ROUTINE_REPORT));
+
     PDATA_TABLE_ROUTINE_REPORT report =
-        ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
+        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG);
-                           sizeof(DATA_TABLE_ROUTINE_REPORT),
-                           REPORT_POOL_TAG);
 
     if (!report)
         return;
 
-    INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+    INIT_REPORT_PACKET(report, REPORT_DATA_TABLE_ROUTINE, 0);
-    INIT_REPORT_HEADER(&report->report_header, REPORT_DATA_TABLE_ROUTINE, 0);
 
     report->address  = Address;
     report->table_id = Win32kBase_gDxgInterface;
@@ -1852,7 +1937,15 @@ ReportWin32kBase_DxgInterfaceViolation(_In_ UINT32 TableIndex,
     // todo! report->routine = ??
     // todo: maybe get routine by name from index ?
 
-    IrpQueueCompletePacket(report, sizeof(DPC_STACKWALK_REPORT));
+    status = CryptEncryptBuffer(report, packet_size);
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+        ImpExFreePoolWithTag(report, REPORT_POOL_TAG);
+        return;
+    }
+
+    IrpQueueCompletePacket(report, packet_size);
 }
 
 STATIC

driver/modules.h

@@ -48,9 +48,6 @@ FindSystemModuleByName(_In_ LPCSTR          ModuleName,
 NTSTATUS
 HandleNmiIOCTL();
 
-BOOLEAN
-FreeApcContextStructure(_Inout_ PAPC_CONTEXT_HEADER Context);
-
 NTSTATUS
 ValidateThreadsViaKernelApc();
 

driver/pool.c

@@ -6,6 +6,7 @@
 #include "queue.h"
 #include "ia32.h"
 #include "imports.h"
+#include "crypt.h"
 
 #define PAGE_BASE_SIZE 0x1000
 #define POOL_TAG_SIZE  0x004
@@ -678,9 +679,12 @@ FindUnlinkedProcesses()
 {
     PAGED_CODE();
 
+    NTSTATUS                           status             = STATUS_UNSUCCESSFUL;
     PUINT64                            allocation_address = NULL;
     PROCESS_SCAN_CONTEXT               context            = {0};
     PINVALID_PROCESS_ALLOCATION_REPORT report             = NULL;
+    UINT32 packet_size = CryptRequestRequiredBufferLength(
+        sizeof(INVALID_PROCESS_ALLOCATION_REPORT));
 
     EnumerateProcessListWithCallbackRoutine(IncrementProcessCounter, &context);
 
@@ -722,25 +726,26 @@ FindUnlinkedProcesses()
             "Potentially found an unlinked process allocation at address: %llx",
             allocation);
 
-        report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED,
+        report = ImpExAllocatePool2(
-                                    sizeof(INVALID_PROCESS_ALLOCATION_REPORT),
+            POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG);
-                                    REPORT_POOL_TAG);
 
         if (!report)
             continue;
 
-        INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+        INIT_REPORT_PACKET(report, REPORT_INVALID_PROCESS_ALLOCATION, 0);
-        INIT_REPORT_HEADER(
-            &report->report_header, REPORT_INVALID_PROCESS_ALLOCATION, 0);
 
         RtlCopyMemory(
             report->process, allocation, REPORT_INVALID_PROCESS_BUFFER_SIZE);
 
-        if (!NT_SUCCESS(IrpQueueCompletePacket(
+        status = CryptEncryptBuffer(report, packet_size);
-                report, sizeof(INVALID_PROCESS_ALLOCATION_REPORT)))) {
+
-            DEBUG_ERROR("IrpQueueCompleteIrp failed with no status.");
+        if (!NT_SUCCESS(status)) {
+            DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+            ImpExFreePoolWithTag(report, REPORT_POOL_TAG);
             continue;
         }
+
+        IrpQueueCompletePacket(report, packet_size);
     }
 
 end:

driver/session.c

@@ -1,64 +1,61 @@
 #include "session.h"
 
 #include "imports.h"
+#include "crypt.h"
 
-/* for now, lets just xor the aes key with our cookie */
+NTSTATUS
-
-typedef struct _SESSION_INITIATION_PACKET {
-    UINT32 session_cookie;
-    CHAR   session_aes_key[AES_128_KEY_SIZE];
-    PVOID  protected_process_id;
-
-} SESSION_INITIATION_PACKET, *PSESSION_INITIATION_PACKET;
-
-VOID
 SessionInitialiseStructure()
 {
-    PAGED_CODE();
+    NTSTATUS        status  = STATUS_UNSUCCESSFUL;
-    ImpKeInitializeGuardedMutex(&GetActiveSession()->lock);
+    PACTIVE_SESSION session = GetActiveSession();
+
+    KeInitializeSpinLock(&session->lock);
+
+    status = CryptInitialiseProvider();
+
+    if (!NT_SUCCESS(status))
+        DEBUG_ERROR("CryptInitialiseProvider: %x", status);
+
+    return status;
 }
 
 VOID
 SessionInitialiseCallbackConfiguration()
 {
-    PAGED_CODE();
     InitialiseObCallbacksConfiguration(GetActiveSession());
 }
 
 VOID
 SessionIsActive(_Out_ PBOOLEAN Flag)
 {
-    PAGED_CODE();
+    KIRQL irql = KeAcquireSpinLockRaiseToDpc(&GetActiveSession()->lock);
-    ImpKeAcquireGuardedMutex(&GetActiveSession()->lock);
+    *Flag      = GetActiveSession()->is_session_active;
-    *Flag = GetActiveSession()->is_session_active;
+    KeReleaseSpinLock(&GetActiveSession()->lock, irql);
-    ImpKeReleaseGuardedMutex(&GetActiveSession()->lock);
 }
 
 VOID
 SessionGetProcess(_Out_ PEPROCESS* Process)
 {
-    PAGED_CODE();
+    KIRQL irql = KeAcquireSpinLockRaiseToDpc(&GetActiveSession()->lock);
-    ImpKeAcquireGuardedMutex(&GetActiveSession()->lock);
+    *Process   = GetActiveSession()->process;
-    *Process = GetActiveSession()->process;
+    KeReleaseSpinLock(&GetActiveSession()->lock, irql);
-    ImpKeReleaseGuardedMutex(&GetActiveSession()->lock);
 }
 
 VOID
 SessionGetProcessId(_Out_ PLONG ProcessId)
 {
-    PAGED_CODE();
+    KIRQL irql = KeAcquireSpinLockRaiseToDpc(&GetActiveSession()->lock);
-    ImpKeAcquireGuardedMutex(&GetActiveSession()->lock);
     *ProcessId = GetActiveSession()->km_handle;
-    ImpKeReleaseGuardedMutex(&GetActiveSession()->lock);
+    KeReleaseSpinLock(&GetActiveSession()->lock, irql);
 }
 
 VOID
 SessionGetCallbackConfiguration(
     _Out_ POB_CALLBACKS_CONFIG* CallbackConfiguration)
 {
-    ImpKeAcquireGuardedMutex(&GetActiveSession()->lock);
+    KIRQL irql = KeAcquireSpinLockRaiseToDpc(&GetActiveSession()->lock);
     *CallbackConfiguration = &GetActiveSession()->callback_configuration;
-    ImpKeReleaseGuardedMutex(&GetActiveSession()->lock);
+    KeReleaseSpinLock(&GetActiveSession()->lock, irql);
 }
 
 STATIC
@@ -71,29 +68,29 @@ SessionTerminateHeartbeat(_In_ PHEARTBEAT_CONFIGURATION Configuration)
 VOID
 SessionTerminate()
 {
-    PAGED_CODE();
     DEBUG_INFO("Termination active session.");
 
     PACTIVE_SESSION session = GetActiveSession();
+    KIRQL           irql    = {0};
 
-    ImpKeAcquireGuardedMutex(&session->lock);
+    KeAcquireSpinLock(&session->lock, &irql);
     session->km_handle         = NULL;
     session->um_handle         = NULL;
     session->process           = NULL;
     session->is_session_active = FALSE;
     SessionTerminateHeartbeat(&session->heartbeat_config);
-    ImpKeReleaseGuardedMutex(&session->lock);
+    CryptCloseSessionCryptObjects();
+    KeReleaseSpinLock(&GetActiveSession()->lock, irql);
 }
 
 NTSTATUS
 SessionInitialise(_In_ PIRP Irp)
 {
-    PAGED_CODE();
+    NTSTATUS                   status     = STATUS_UNSUCCESSFUL;
-
+    PEPROCESS                  process    = NULL;
-    NTSTATUS                   status      = STATUS_UNSUCCESSFUL;
+    PSESSION_INITIATION_PACKET initiation = NULL;
-    PEPROCESS                  process     = NULL;
+    PACTIVE_SESSION            session    = GetActiveSession();
-    PSESSION_INITIATION_PACKET information = NULL;
+    KIRQL                      irql       = {0};
-    PACTIVE_SESSION            session     = GetActiveSession();
 
     DEBUG_VERBOSE("Initialising new session.");
 
@@ -104,11 +101,11 @@ SessionInitialise(_In_ PIRP Irp)
         return status;
     }
 
-    information = (PSESSION_INITIATION_PACKET)Irp->AssociatedIrp.SystemBuffer;
+    initiation = (PSESSION_INITIATION_PACKET)Irp->AssociatedIrp.SystemBuffer;
 
-    ImpKeAcquireGuardedMutex(&session->lock);
+    KeAcquireSpinLock(&session->lock, &irql);
 
-    session->um_handle = information->protected_process_id;
+    session->um_handle = initiation->process_id;
 
     /* What if we pass an invalid handle here? not good. */
     status = ImpPsLookupProcessByProcessId(session->um_handle, &process);
@@ -121,11 +118,17 @@ SessionInitialise(_In_ PIRP Irp)
     session->km_handle         = ImpPsGetProcessId(process);
     session->process           = process;
     session->is_session_active = TRUE;
-    session->session_cookie    = information->session_cookie;
+    session->cookie            = initiation->cookie;
 
-    RtlCopyMemory(session->session_aes_key,
+    RtlCopyMemory(session->aes_key, initiation->aes_key, AES_256_KEY_SIZE);
-                  information->session_aes_key,
+    RtlCopyMemory(session->iv, initiation->aes_iv, AES_256_IV_SIZE);
-                  AES_128_KEY_SIZE);
+
+    status = CryptInitialiseSessionCryptObjects();
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptInitialiseSessionCryptObjects: %x", status);
+        goto end;
+    }
 
     status = InitialiseHeartbeatConfiguration(&session->heartbeat_config);
 
@@ -135,15 +138,13 @@ SessionInitialise(_In_ PIRP Irp)
     }
 
 end:
-    ImpKeReleaseGuardedMutex(&session->lock);
+    KeReleaseSpinLock(&GetActiveSession()->lock, irql);
     return status;
 }
 
 VOID
 SessionTerminateProcess()
 {
-    PAGED_CODE();
-
     NTSTATUS status     = STATUS_UNSUCCESSFUL;
     ULONG    process_id = 0;
 
@@ -174,23 +175,23 @@ SessionTerminateProcess()
 VOID
 SessionIncrementIrpsProcessedCount()
 {
-    ImpKeAcquireGuardedMutex(&GetActiveSession()->lock);
+    KIRQL irql = KeAcquireSpinLockRaiseToDpc(&GetActiveSession()->lock);
     GetActiveSession()->irps_received;
-    ImpKeReleaseGuardedMutex(&GetActiveSession()->lock);
+    KeReleaseSpinLock(&GetActiveSession()->lock, irql);
 }
 
 VOID
 SessionIncrementReportCount()
 {
-    ImpKeAcquireGuardedMutex(&GetActiveSession()->lock);
+    KIRQL irql = KeAcquireSpinLockRaiseToDpc(&GetActiveSession()->lock);
     GetActiveSession()->report_count++;
-    ImpKeReleaseGuardedMutex(&GetActiveSession()->lock);
+    KeReleaseSpinLock(&GetActiveSession()->lock, irql);
 }
 
 VOID
 SessionIncrementHeartbeatCount()
 {
-    ImpKeAcquireGuardedMutex(&GetActiveSession()->lock);
+    KIRQL irql = KeAcquireSpinLockRaiseToDpc(&GetActiveSession()->lock);
     GetActiveSession()->heartbeat_count++;
-    ImpKeReleaseGuardedMutex(&GetActiveSession()->lock);
+    KeReleaseSpinLock(&GetActiveSession()->lock, irql);
 }

driver/session.h

@@ -5,7 +5,7 @@
 
 #include "driver.h"
 
-VOID
+NTSTATUS
 SessionInitialiseStructure();
 
 VOID

driver/thread.c

@@ -8,6 +8,7 @@
 #include "queue.h"
 #include "session.h"
 #include "imports.h"
+#include "crypt.h"
 
 #ifdef ALLOC_PRAGMA
 #    pragma alloc_text(PAGE, DetectThreadsAttachedToProtectedProcess)
@@ -81,8 +82,11 @@ DetectAttachedThreadsProcessCallback(_In_ PTHREAD_LIST_ENTRY ThreadListEntry,
 {
     UNREFERENCED_PARAMETER(Context);
 
+    NTSTATUS    status            = STATUS_UNSUCCESSFUL;
     PKAPC_STATE apc_state         = NULL;
     PEPROCESS   protected_process = NULL;
+    UINT32      packet_size =
+        CryptRequestRequiredBufferLength(sizeof(ATTACH_PROCESS_REPORT));
 
     SessionGetProcess(&protected_process);
 
@@ -106,21 +110,26 @@ DetectAttachedThreadsProcessCallback(_In_ PTHREAD_LIST_ENTRY ThreadListEntry,
     DEBUG_WARNING("Thread is attached to our protected process: %llx",
                   (UINT64)ThreadListEntry->thread);
 
-    PATTACH_PROCESS_REPORT report = ImpExAllocatePool2(
+    PATTACH_PROCESS_REPORT report =
-        POOL_FLAG_NON_PAGED, sizeof(ATTACH_PROCESS_REPORT), REPORT_POOL_TAG);
+        ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG);
 
     if (!report)
         return;
 
-    INIT_PACKET_HEADER(&report->header, PACKET_TYPE_REPORT);
+    INIT_REPORT_PACKET(report, REPORT_ILLEGAL_ATTACH_PROCESS, 0);
-    INIT_REPORT_HEADER(
-        &report->report_header, REPORT_ILLEGAL_ATTACH_PROCESS, 0);
 
     report->thread_id      = ImpPsGetThreadId(ThreadListEntry->thread);
     report->thread_address = ThreadListEntry->thread;
 
-    if (!NT_SUCCESS(IrpQueueCompletePacket(report, sizeof(ATTACH_PROCESS_REPORT))))
+    status = CryptEncryptBuffer(report, packet_size);
-        DEBUG_ERROR("IrpQueueCompleteIrp failed with no status.");
+
+    if (!NT_SUCCESS(status)) {
+        DEBUG_ERROR("CryptEncryptBuffer: %lx", status);
+        ImpExFreePoolWithTag(report, REPORT_POOL_TAG);
+        return;
+    }
+
+    IrpQueueCompletePacket(report, packet_size);
 }
 
 VOID

driver/types/types.h

@@ -15,34 +15,54 @@
 #define REPORT_INVALID_PROCESS_MODULE     140
 
 #define REPORT_SUBTYPE_NO_BACKING_MODULE 0x0
-#define REPORT_SUBTYPE_INVALID_DISPATCH 0x1
+#define REPORT_SUBTYPE_INVALID_DISPATCH  0x1
 
 #define PACKET_TYPE_REPORT    0x0
 #define PACKET_TYPE_HEARTBEAT 0x1
 
-#define INIT_PACKET_HEADER(header, type) \
+#define PACKET_MAGIC_NUMBER 0x1337
-    {                                    \
+
-        (header)->packet_type = type;    \
+#define INIT_REPORT_PACKET(report, code, subcode)                          \
+    {                                                                      \
+        (report)->header.packet_header.packet_type  = PACKET_TYPE_REPORT;  \
+        (report)->header.packet_header.magic_number = PACKET_MAGIC_NUMBER; \
+        (report)->header.report_code                = code;                \
+        (report)->header.report_sub_type            = subcode;             \
     }
 
-#define INIT_REPORT_HEADER(report, code, subcode) \
+#define INIT_HEARTBEAT_PACKET(packet)                                        \
-    {                                             \
+    {                                                                        \
-        (report)->report_code     = code;         \
+        (packet)->header.packet_header.packet_type  = PACKET_TYPE_HEARTBEAT; \
-        (report)->report_sub_type = subcode;      \
+        (packet)->header.packet_header.magic_number = PACKET_MAGIC_NUMBER;   \
     }
 
 /* use a UINT16 rather then enum to explicitly state the size */
 typedef struct _PACKET_HEADER {
-    UINT16 packet_type;
+    UINT32 packet_type;
+    UINT32 magic_number;
 
 } PACKET_HEADER, *PPACKET_HEADER;
 
+/* unencrypted header structures, should always == AES block size i.e 16 */
 typedef struct _REPORT_PACKET_HEADER {
-    UINT32 report_code;
+    PACKET_HEADER packet_header;
-    UINT32 report_sub_type;
+    UINT32        report_code;
+    UINT32        report_sub_type;
 
 } REPORT_PACKET_HEADER, *PREPORT_PACKET_HEADER;
 
+typedef struct _HEARTBEAT_PACKET_HEADER {
+    PACKET_HEADER packet_header;
+    UINT32        unused[2];
+} HEARTBEAT_PACKET_HEADER, *PHEARTBEAT_PACKET_HEADER;
+
+#define AES_256_BLOCK_SIZE 16
+
+static_assert(sizeof(HEARTBEAT_PACKET_HEADER) == AES_256_BLOCK_SIZE,
+              "invalid heartbeat header size");
+static_assert(sizeof(REPORT_PACKET_HEADER) == AES_256_BLOCK_SIZE,
+              "invalid report header size");
+
 typedef enum _TABLE_ID {
     HalDispatch = 0,
     HalPrivateDispatch,
@@ -50,8 +70,7 @@ typedef enum _TABLE_ID {
 } TABLE_ID;
 
 typedef struct _HYPERVISOR_DETECTION_REPORT {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
     UINT8                aperf_msr_timing_check;
     UINT8                invd_emulation_check;
 
@@ -60,8 +79,7 @@ typedef struct _HYPERVISOR_DETECTION_REPORT {
 #define APC_STACKWALK_BUFFER_SIZE 500
 
 typedef struct _APC_STACKWALK_REPORT {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
     UINT64               kthread_address;
     UINT64               invalid_rip;
     CHAR                 driver[APC_STACKWALK_BUFFER_SIZE];
@@ -69,8 +87,7 @@ typedef struct _APC_STACKWALK_REPORT {
 } APC_STACKWALK_REPORT, *PAPC_STACKWALK_REPORT;
 
 typedef struct _DPC_STACKWALK_REPORT {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
     UINT64               kthread_address;
     UINT64               invalid_rip;
     CHAR                 driver[APC_STACKWALK_BUFFER_SIZE];
@@ -78,8 +95,7 @@ typedef struct _DPC_STACKWALK_REPORT {
 } DPC_STACKWALK_REPORT, *PDPC_STACKWALK_REPORT;
 
 typedef struct _MODULE_VALIDATION_FAILURE {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
     UINT64               driver_base_address;
     UINT64               driver_size;
     CHAR                 driver_name[128];
@@ -89,8 +105,7 @@ typedef struct _MODULE_VALIDATION_FAILURE {
 #define DATA_TABLE_ROUTINE_BUF_SIZE 256
 
 typedef struct _DATA_TABLE_ROUTINE_REPORT {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
     TABLE_ID             table_id;
     UINT64               address;
     UINT32               index;
@@ -99,8 +114,7 @@ typedef struct _DATA_TABLE_ROUTINE_REPORT {
 } DATA_TABLE_ROUTINE_REPORT, *PDATA_TABLE_ROUTINE_REPORT;
 
 typedef struct _NMI_CALLBACK_FAILURE {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
     UINT8                were_nmis_disabled;
     UINT64               kthread_address;
     UINT64               invalid_rip;
@@ -110,15 +124,13 @@ typedef struct _NMI_CALLBACK_FAILURE {
 #define REPORT_INVALID_PROCESS_BUFFER_SIZE 500
 
 typedef struct _INVALID_PROCESS_ALLOCATION_REPORT {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
     CHAR                 process[REPORT_INVALID_PROCESS_BUFFER_SIZE];
 
 } INVALID_PROCESS_ALLOCATION_REPORT, *PINVALID_PROCESS_ALLOCATION_REPORT;
 
 typedef struct _HIDDEN_SYSTEM_THREAD_REPORT {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
     UINT8                found_in_kthreadlist;
     UINT8                found_in_pspcidtable;
     UINT64               thread_address;
@@ -128,16 +140,14 @@ typedef struct _HIDDEN_SYSTEM_THREAD_REPORT {
 } HIDDEN_SYSTEM_THREAD_REPORT, *PHIDDEN_SYSTEM_THREAD_REPORT;
 
 typedef struct _ATTACH_PROCESS_REPORT {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
     UINT32               thread_id;
     UINT64               thread_address;
 
 } ATTACH_PROCESS_REPORT, *PATTACH_PROCESS_REPORT;
 
 typedef struct _KPRCB_THREAD_VALIDATION_CTX {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
     UINT64               thread;
     BOOLEAN              thread_found_in_pspcidtable;
     // BOOLEAN thread_found_in_kthreadlist;
@@ -148,9 +158,8 @@ typedef struct _KPRCB_THREAD_VALIDATION_CTX {
 #define HANDLE_REPORT_PROCESS_NAME_MAX_LENGTH 64
 
 typedef struct _OPEN_HANDLE_FAILURE_REPORT {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
+    UINT32               is_kernel_handle;
-    UINT8                is_kernel_handle;
     UINT32               process_id;
     UINT32               thread_id;
     UINT32               access;
@@ -161,8 +170,7 @@ typedef struct _OPEN_HANDLE_FAILURE_REPORT {
 #define MODULE_PATH_LEN 256
 
 typedef struct _PROCESS_MODULE_VALIDATION_REPORT {
-    PACKET_HEADER        header;
+    REPORT_PACKET_HEADER header;
-    REPORT_PACKET_HEADER report_header;
     UINT64               image_base;
     UINT32               image_size;
     WCHAR                module_path[MODULE_PATH_LEN];
@@ -170,11 +178,11 @@ typedef struct _PROCESS_MODULE_VALIDATION_REPORT {
 } PROCESS_MODULE_VALIDATION_REPORT, *PPROCESS_MODULE_VALIDATION_REPORT;
 
 typedef struct _HEARTBEAT_PACKET {
-    PACKET_HEADER header;
+    HEARTBEAT_PACKET_HEADER header;
-    UINT32        heartbeat_count;
+    UINT32                  heartbeat_count;
-    UINT32        total_reports_completed;
+    UINT32                  total_reports_completed;
-    UINT32        total_irps_completed;
+    UINT32                  total_irps_completed;
-    UINT32        total_heartbeats_completed;
+    UINT32                  total_heartbeats_completed;
 
 } HEARTBEAT_PACKET, *PHEARTBEAT_PACKET;
 

module/crypt/crypt.cpp

@@ -0,0 +1,128 @@
+#include "crypt.h"
+
+#include "../common.h"
+
+#include <bcrypt.h>
+#include <iomanip>
+#include <iostream>
+#include <stdio.h>
+#include <winternl.h>
+
+#pragma comment(lib, "bcrypt.lib")
+
+BCRYPT_ALG_HANDLE alg_handle = NULL;
+BCRYPT_KEY_HANDLE key_handle = NULL;
+
+namespace crypt {
+#define STATUS_UNSUCCESSFUL ((NTSTATUS)0xC0000001L)
+
+namespace globals {
+
+#define TEST_AES_KEY_LENGTH 0x32
+#define TEST_AES_IV_LENGTH 0x16
+
+const unsigned char TEST_KEY[] = {
+    0xAA, 0x50, 0xA7, 0x00, 0x79, 0xF1, 0x6C, 0x2D, 0x6B, 0xAD, 0xAC,
+    0x19, 0x18, 0x66, 0xFB, 0xEF, 0xCA, 0x9B, 0x6D, 0x3E, 0xA3, 0x7D,
+    0x2D, 0xF6, 0x10, 0x95, 0xB3, 0xB3, 0x8D, 0x34, 0x69, 0xF1};
+
+const unsigned char TEST_IV[] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05,
+                                 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
+                                 0x0C, 0x0D, 0x0E, 0x0F};
+
+PBCRYPT_KEY_DATA_BLOB_HEADER blob = nullptr;
+
+static PUCHAR key_object = NULL;
+static UINT32 key_object_length = 0;
+
+} // namespace globals
+
+boolean initialise_session_key() {
+  globals::blob = reinterpret_cast<PBCRYPT_KEY_DATA_BLOB_HEADER>(
+      malloc(sizeof(BCRYPT_KEY_DATA_BLOB_HEADER) + sizeof(globals::TEST_KEY)));
+
+  if (!globals::blob)
+    return false;
+
+  globals::blob->dwMagic = BCRYPT_KEY_DATA_BLOB_MAGIC;
+  globals::blob->dwVersion = BCRYPT_KEY_DATA_BLOB_VERSION1;
+  globals::blob->cbKeyData = sizeof(globals::TEST_KEY);
+  memcpy((void *)((UINT64)globals::blob + sizeof(BCRYPT_KEY_DATA_BLOB_HEADER)),
+         (void *)globals::TEST_KEY, sizeof(globals::TEST_KEY));
+
+  return true;
+}
+
+boolean initialise_provider() {
+  UINT32 data_copied = 0;
+  NTSTATUS status =
+      BCryptOpenAlgorithmProvider(&alg_handle, BCRYPT_AES_ALGORITHM, NULL, 0);
+
+  if (!NT_SUCCESS(status)) {
+    LOG_ERROR("BCryptOpenAlgorithmProvider: %x", status);
+    return false;
+  }
+
+  status = BCryptGetProperty(alg_handle, BCRYPT_OBJECT_LENGTH,
+                             (PUCHAR)&globals::key_object_length,
+                             sizeof(UINT32), (PULONG)&data_copied, 0);
+
+  if (!NT_SUCCESS(status)) {
+    LOG_ERROR("BCryptGetProperty: %x", status);
+    return false;
+  }
+
+  globals::key_object = (PUCHAR)malloc(globals::key_object_length);
+
+  if (!globals::key_object)
+    return false;
+
+  if (!initialise_session_key())
+    return false;
+
+  status = BCryptImportKey(
+      alg_handle, NULL, BCRYPT_KEY_DATA_BLOB, &key_handle, globals::key_object,
+      globals::key_object_length, (PUCHAR)globals::blob,
+      sizeof(BCRYPT_KEY_DATA_BLOB_HEADER) + sizeof(globals::TEST_KEY), 0);
+
+  if (!NT_SUCCESS(status)) {
+    LOG_ERROR("BCryptImportKey: %x", status);
+    return false;
+  }
+
+  return true;
+}
+
+boolean decrypt_packet(void *packet, uint32_t packet_length) {
+  ULONG data_copied = 0;
+  unsigned char local_iv[sizeof(globals::TEST_IV)] = {0};
+  memcpy((void *)local_iv, (void *)globals::TEST_IV, sizeof(globals::TEST_IV));
+
+  void* buffer = (void*)((UINT64)packet + 16);
+  uint32_t length = packet_length - 16;
+
+  NTSTATUS status = BCryptDecrypt(
+      key_handle, (PUCHAR)buffer, length, NULL, (PUCHAR)local_iv,
+      sizeof(globals::TEST_IV), (PUCHAR)buffer, length, &data_copied, 0);
+
+  if (!NT_SUCCESS(status)) {
+    LOG_ERROR("BCryptDecrypt: %x", status);
+    return false;
+  }
+
+  return true;
+}
+
+uint32_t get_padded_packet_size(uint32_t original_size) {
+  uint32_t remainder = original_size % 16;
+
+  if (remainder != 0) {
+    original_size += 16 - remainder;
+  }
+
+  return original_size;
+}
+
+const unsigned char *get_test_key() { return globals::TEST_KEY; }
+const unsigned char *get_test_iv() { return globals::TEST_IV; }
+} // namespace crypt

module/crypt/crypt.h

@@ -0,0 +1,12 @@
+#pragma once
+
+#include <cstdint>
+#include <windows.h>
+
+namespace crypt {
+const unsigned char *get_test_key();
+const unsigned char *get_test_iv();
+boolean initialise_provider();
+boolean decrypt_packet(void *packet, uint32_t packet_length);
+uint32_t get_padded_packet_size(uint32_t original_size);
+} // namespace crypt

module/dispatcher/dispatcher.cpp

@@ -2,6 +2,7 @@
 
 #include "../client/message_queue.h"
 #include "../helper.h"
+#include "../crypt/crypt.h"
 
 #include <bcrypt.h>
 #include <chrono>
@@ -48,6 +49,7 @@ void dispatcher::dispatcher::run_io_port_thread() {
 
 void dispatcher::dispatcher::run() {
   // helper::generate_rand_seed();
+  crypt::initialise_provider();
   std::srand(std::time(nullptr));
   this->init_timer_callbacks();
   this->run_timer_thread();

module/helper.cpp

@@ -3,6 +3,8 @@
 #include <chrono>
 #include <random>
 
+#include "crypt/crypt.h"
+
 void helper::generate_rand_seed() { srand(time(0)); }
 
 int helper::generate_rand_int(int max) { return std::rand() % max; }
@@ -49,127 +51,143 @@ kernel_interface::report_id helper::get_kernel_report_type(void *buffer) {
   }
 }
 
+void
+print_report_packet(void* buffer)
+{
+    kernel_interface::report_header* report_header =
+        (kernel_interface::report_header*)buffer;
+
+    LOG_INFO("report code: %lx", report_header->report_code);
+    LOG_INFO("report sub code: %lx", report_header->report_sub_type);
+
+    switch (report_header->report_code) {
+    case kernel_interface::report_id::report_nmi_callback_failure: {
+        kernel_interface::nmi_callback_failure* r1 =
+            reinterpret_cast<kernel_interface::nmi_callback_failure*>(buffer);
+        LOG_INFO("were_nmis_disabled: %lx", r1->were_nmis_disabled);
+        LOG_INFO("kthread_address: %llx", r1->kthread_address);
+        LOG_INFO("invalid_rip: %llx", r1->invalid_rip);
+        LOG_INFO("********************************");
+        break;
+    }
+    case kernel_interface::report_id::report_invalid_process_allocation: {
+        kernel_interface::invalid_process_allocation_report* r2 =
+            reinterpret_cast<
+                kernel_interface::invalid_process_allocation_report*>(buffer);
+        LOG_INFO("********************************");
+        break;
+    }
+    case kernel_interface::report_id::report_hidden_system_thread: {
+        kernel_interface::hidden_system_thread_report* r3 =
+            reinterpret_cast<kernel_interface::hidden_system_thread_report*>(
+                buffer);
+        LOG_INFO("found_in_kthreadlist: %lx", r3->found_in_kthreadlist);
+        LOG_INFO("found_in_pspcidtable: %lx", r3->found_in_pspcidtable);
+        LOG_INFO("thread_address: %llx", r3->thread_address);
+        LOG_INFO("thread_id: %lx", r3->thread_id);
+        LOG_INFO("********************************");
+        break;
+    }
+    case kernel_interface::report_id::report_illegal_attach_process: {
+        kernel_interface::attach_process_report* r4 =
+            reinterpret_cast<kernel_interface::attach_process_report*>(buffer);
+        LOG_INFO("report type: attach_process_report");
+        LOG_INFO("report code: %lx", r4->report_code);
+        LOG_INFO("thread_id: %lx", r4->thread_id);
+        LOG_INFO("thread_address: %llx", r4->thread_address);
+        LOG_INFO("********************************");
+        break;
+    }
+    case kernel_interface::report_id::report_illegal_handle_operation: {
+        kernel_interface::open_handle_failure_report* r5 =
+            reinterpret_cast<kernel_interface::open_handle_failure_report*>(
+                buffer);
+        LOG_INFO("is_kernel_handle: %lx", r5->is_kernel_handle);
+        LOG_INFO("process_id: %lx", r5->process_id);
+        LOG_INFO("thread_id: %lx", r5->thread_id);
+        LOG_INFO("access: %lx", r5->access);
+        LOG_INFO("process_name: %s", r5->process_name);
+        LOG_INFO("********************************");
+        break;
+    }
+    case kernel_interface::report_id::report_invalid_process_module: {
+        kernel_interface::process_module_validation_report* r6 =
+            reinterpret_cast<
+                kernel_interface::process_module_validation_report*>(buffer);
+        LOG_INFO("image_base: %llx", r6->image_base);
+        LOG_INFO("image_size: %u", r6->image_size);
+        LOG_INFO("module_path: %ls", r6->module_path);
+        LOG_INFO("********************************");
+        break;
+    }
+    case kernel_interface::report_id::report_apc_stackwalk: {
+        kernel_interface::apc_stackwalk_report* r7 =
+            reinterpret_cast<kernel_interface::apc_stackwalk_report*>(buffer);
+        LOG_INFO("kthread_address: %llx", r7->kthread_address);
+        LOG_INFO("invalid_rip: %llx", r7->invalid_rip);
+        LOG_INFO("********************************");
+        break;
+    }
+    case kernel_interface::report_id::report_dpc_stackwalk: {
+        kernel_interface::dpc_stackwalk_report* r8 =
+            reinterpret_cast<kernel_interface::dpc_stackwalk_report*>(buffer);
+        LOG_INFO("kthread_address: %llx", r8->kthread_address);
+        LOG_INFO("invalid_rip: %llx", r8->invalid_rip);
+        LOG_INFO("********************************");
+        break;
+    }
+    case kernel_interface::report_id::report_data_table_routine: {
+        kernel_interface::data_table_routine_report* r9 =
+            reinterpret_cast<kernel_interface::data_table_routine_report*>(
+                buffer);
+        LOG_INFO("id: %d", r9->id);
+        LOG_INFO("address: %llx", r9->address);
+        LOG_INFO("routine: %s", r9->routine);
+        LOG_INFO("********************************");
+        break;
+    }
+    case kernel_interface::report_id::report_module_validation_failure: {
+        kernel_interface::module_validation_failure* r10 =
+            reinterpret_cast<kernel_interface::module_validation_failure*>(
+                buffer);
+        LOG_INFO("driver_base_address: %llx", r10->driver_base_address);
+        LOG_INFO("driver_size: %llx", r10->driver_size);
+        LOG_INFO("driver_name: %s", r10->driver_name);
+        LOG_INFO("********************************");
+        break;
+    }
+    default: LOG_INFO("Invalid report type."); break;
+    }
+}
+
+void
+print_heartbeat_packet(void* buffer)
+{
+    kernel_interface::heartbeat_packet* hb =
+        reinterpret_cast<kernel_interface::heartbeat_packet*>(buffer);
+    LOG_INFO("Heartbeat Count: %lx", hb->heartbeat_count);
+    LOG_INFO("Total Reports Completed: %lx", hb->total_reports_completed);
+    LOG_INFO("Total IRPs Completed: %lx", hb->total_irps_completed);
+    LOG_INFO("Total Heartbeats Completed: %lx", hb->total_heartbeats_completed);
+    LOG_INFO("********************************");
+}
+
 void helper::print_kernel_report(void *buffer) {
+  uint32_t size = crypt::get_padded_packet_size(
+      sizeof(kernel_interface::open_handle_failure_report));
+  crypt::decrypt_packet(buffer, size);
+
   kernel_interface::packet_header *header =
       reinterpret_cast<kernel_interface::packet_header *>(buffer);
+
   LOG_INFO("packet type: %lx", header->packet_type);
 
-  if (header->packet_type == 1)
+  switch (header->packet_type)
   {
-      kernel_interface::heartbeat_packet* hb =
+  case 0: print_report_packet(buffer); break;
-          reinterpret_cast<kernel_interface::heartbeat_packet*>(buffer);
+  case 1: print_heartbeat_packet(buffer); break;
-      LOG_INFO("Heartbeat Count: %lx", hb->heartbeat_count);
-      LOG_INFO("Total Reports Completed: %lx", hb->total_reports_completed);
-      LOG_INFO("Total IRPs Completed: %lx", hb->total_irps_completed);
-      LOG_INFO("Total Heartbeats Completed: %lx",
-               hb->total_heartbeats_completed);
-      LOG_INFO("********************************");
-      return;
-  }
-
-  kernel_interface::report_header *report_header =
-      reinterpret_cast<kernel_interface::report_header *>(
-          (uint64_t)buffer + sizeof(kernel_interface::packet_header));
-  LOG_INFO("report code: %lx", report_header->report_code);
-  LOG_INFO("report sub code: %lx", report_header->report_sub_type);
-  switch (report_header->report_code) {
-  case kernel_interface::report_id::report_nmi_callback_failure: {
-    kernel_interface::nmi_callback_failure *r1 =
-        reinterpret_cast<kernel_interface::nmi_callback_failure *>(buffer);
-    LOG_INFO("were_nmis_disabled: %lx", r1->were_nmis_disabled);
-    LOG_INFO("kthread_address: %llx", r1->kthread_address);
-    LOG_INFO("invalid_rip: %llx", r1->invalid_rip);
-    LOG_INFO("********************************");
-    break;
-  }
-  case kernel_interface::report_id::report_invalid_process_allocation: {
-    kernel_interface::invalid_process_allocation_report *r2 =
-        reinterpret_cast<kernel_interface::invalid_process_allocation_report *>(
-            buffer);
-    LOG_INFO("********************************");
-    break;
-  }
-  case kernel_interface::report_id::report_hidden_system_thread: {
-    kernel_interface::hidden_system_thread_report *r3 =
-        reinterpret_cast<kernel_interface::hidden_system_thread_report *>(
-            buffer);
-    LOG_INFO("found_in_kthreadlist: %lx", r3->found_in_kthreadlist);
-    LOG_INFO("found_in_pspcidtable: %lx", r3->found_in_pspcidtable);
-    LOG_INFO("thread_address: %llx", r3->thread_address);
-    LOG_INFO("thread_id: %lx", r3->thread_id);
-    LOG_INFO("********************************");
-    break;
-  }
-  case kernel_interface::report_id::report_illegal_attach_process: {
-    kernel_interface::attach_process_report *r4 =
-        reinterpret_cast<kernel_interface::attach_process_report *>(buffer);
-    LOG_INFO("report type: attach_process_report");
-    LOG_INFO("report code: %lx", r4->report_code);
-    LOG_INFO("thread_id: %lx", r4->thread_id);
-    LOG_INFO("thread_address: %llx", r4->thread_address);
-    LOG_INFO("********************************");
-    break;
-  }
-  case kernel_interface::report_id::report_illegal_handle_operation: {
-    kernel_interface::open_handle_failure_report *r5 =
-        reinterpret_cast<kernel_interface::open_handle_failure_report *>(
-            buffer);
-    LOG_INFO("is_kernel_handle: %lx", r5->is_kernel_handle);
-    LOG_INFO("process_id: %lx", r5->process_id);
-    LOG_INFO("thread_id: %lx", r5->thread_id);
-    LOG_INFO("access: %lx", r5->access);
-    LOG_INFO("process_name: %s", r5->process_name);
-    LOG_INFO("********************************");
-    break;
-  }
-  case kernel_interface::report_id::report_invalid_process_module: {
-    kernel_interface::process_module_validation_report *r6 =
-        reinterpret_cast<kernel_interface::process_module_validation_report *>(
-            buffer);
-    LOG_INFO("image_base: %llx", r6->image_base);
-    LOG_INFO("image_size: %u", r6->image_size);
-    LOG_INFO("module_path: %ls", r6->module_path);
-    LOG_INFO("********************************");
-    break;
-  }
-  case kernel_interface::report_id::report_apc_stackwalk: {
-    kernel_interface::apc_stackwalk_report *r7 =
-        reinterpret_cast<kernel_interface::apc_stackwalk_report *>(buffer);
-    LOG_INFO("kthread_address: %llx", r7->kthread_address);
-    LOG_INFO("invalid_rip: %llx", r7->invalid_rip);
-    LOG_INFO("********************************");
-    break;
-  }
-  case kernel_interface::report_id::report_dpc_stackwalk: {
-    kernel_interface::dpc_stackwalk_report *r8 =
-        reinterpret_cast<kernel_interface::dpc_stackwalk_report *>(buffer);
-    LOG_INFO("kthread_address: %llx", r8->kthread_address);
-    LOG_INFO("invalid_rip: %llx", r8->invalid_rip);
-    LOG_INFO("********************************");
-    break;
-  }
-  case kernel_interface::report_id::report_data_table_routine: {
-    kernel_interface::data_table_routine_report *r9 =
-        reinterpret_cast<kernel_interface::data_table_routine_report *>(buffer);
-    LOG_INFO("id: %d", r9->id);
-    LOG_INFO("address: %llx", r9->address);
-    LOG_INFO("routine: %s", r9->routine);
-    LOG_INFO("********************************");
-    break;
-  }
-  case kernel_interface::report_id::report_module_validation_failure: {
-    kernel_interface::module_validation_failure *r10 =
-        reinterpret_cast<kernel_interface::module_validation_failure *>(buffer);
-    LOG_INFO("driver_base_address: %llx", r10->driver_base_address);
-    LOG_INFO("driver_size: %llx", r10->driver_size);
-    LOG_INFO("driver_name: %s", r10->driver_name);
-    LOG_INFO("********************************");
-    break;
-  }
-  default:
-    LOG_INFO("Invalid report type.");
-    break;
   }
+ 
 }
 
 unsigned __int64 helper::seconds_to_nanoseconds(int seconds) {

module/kernel_interface/kernel_interface.cpp

@@ -4,6 +4,7 @@
 
 #include "../common.h"
 #include "../helper.h"
+#include "../crypt/crypt.h"
 
 #include <TlHelp32.h>
 #include <winternl.h>
@@ -140,7 +141,10 @@ void kernel_interface::kernel_interface::generic_driver_call_apc(
 void kernel_interface::kernel_interface::notify_driver_on_process_launch() {
   unsigned long bytes_returned = 0;
   session_initiation_packet packet = {0};
-  packet.protected_process_id = reinterpret_cast<void *>(GetCurrentProcessId());
+  packet.process_id = reinterpret_cast<void *>(GetCurrentProcessId());
+  packet.session_cookie = 123;
+  memcpy(packet.aes_key, crypt::get_test_key(), 32);
+  memcpy(packet.aes_iv, crypt::get_test_iv(), 16);
   generic_driver_call_input(ioctl_code::NotifyDriverOnProcessLaunch, &packet,
                             sizeof(session_initiation_packet), &bytes_returned);
 }

module/kernel_interface/kernel_interface.h

@@ -25,15 +25,27 @@ enum report_id {
   report_invalid_process_module = 140
 };
 
+#define AES_256_BLOCK_SIZE 16
+
 struct packet_header {
-  uint16_t packet_type;
+  uint32_t packet_type;
+  uint32_t magic_number;
+};
+
+struct heartbeat_header {
+  packet_header header;
+  uint32_t unused[2];
 };
 
 struct report_header {
+  struct packet_header header;
   uint32_t report_code;
   uint32_t report_sub_type;
 };
 
+static_assert(sizeof(heartbeat_header) == AES_256_BLOCK_SIZE);
+static_assert(sizeof(report_header) == AES_256_BLOCK_SIZE);
+
 constexpr int APC_STACKWALK_BUFFER_SIZE = 500;
 constexpr int DATA_TABLE_ROUTINE_BUF_SIZE = 256;
 constexpr int REPORT_INVALID_PROCESS_BUFFER_SIZE = 500;
@@ -41,7 +53,6 @@ constexpr int HANDLE_REPORT_PROCESS_NAME_MAX_LENGTH = 64;
 constexpr int MODULE_PATH_LEN = 256;
 
 struct apc_stackwalk_report {
-  packet_header header;
   report_header report_header;
   uint64_t kthread_address;
   uint64_t invalid_rip;
@@ -49,7 +60,6 @@ struct apc_stackwalk_report {
 };
 
 struct dpc_stackwalk_report {
-  packet_header header;
   report_header report_header;
   uint64_t kthread_address;
   uint64_t invalid_rip;
@@ -57,7 +67,6 @@ struct dpc_stackwalk_report {
 };
 
 struct module_validation_failure {
-  packet_header header;
   report_header report_header;
   uint64_t driver_base_address;
   uint64_t driver_size;
@@ -67,7 +76,6 @@ struct module_validation_failure {
 enum table_id { hal_dispatch = 0, hal_private_dispatch };
 
 struct data_table_routine_report {
-  packet_header header;
   report_header report_header;
   table_id id;
   uint64_t address;
@@ -76,7 +84,6 @@ struct data_table_routine_report {
 };
 
 struct nmi_callback_failure {
-  packet_header header;
   report_header report_header;
   uint8_t were_nmis_disabled;
   uint64_t kthread_address;
@@ -84,13 +91,11 @@ struct nmi_callback_failure {
 };
 
 struct invalid_process_allocation_report {
-  packet_header header;
   report_header report_header;
   char process[REPORT_INVALID_PROCESS_BUFFER_SIZE];
 };
 
 struct hidden_system_thread_report {
-  packet_header header;
   report_header report_header;
   uint8_t found_in_kthreadlist;
   uint8_t found_in_pspcidtable;
@@ -106,9 +111,8 @@ struct attach_process_report {
 };
 
 struct open_handle_failure_report {
-  packet_header header;
   report_header report_header;
-  uint8_t is_kernel_handle;
+  uint32_t is_kernel_handle;
   uint32_t process_id;
   uint32_t thread_id;
   uint32_t access;
@@ -116,7 +120,6 @@ struct open_handle_failure_report {
 };
 
 struct process_module_validation_report {
-  packet_header header;
   report_header report_header;
   uint64_t image_base;
   uint32_t image_size;
@@ -124,11 +127,11 @@ struct process_module_validation_report {
 };
 
 struct heartbeat_packet {
-    packet_header header;
+  heartbeat_header header;
-    uint32_t        heartbeat_count;
+  uint32_t heartbeat_count;
-    uint32_t        total_reports_completed;
+  uint32_t total_reports_completed;
-    uint32_t        total_irps_completed;
+  uint32_t total_irps_completed;
-    uint32_t        total_heartbeats_completed;
+  uint32_t total_heartbeats_completed;
 };
 
 enum apc_operation { operation_stackwalk = 0x1 };
@@ -194,8 +197,9 @@ struct event_dispatcher {
 class kernel_interface {
   struct session_initiation_packet {
     unsigned __int32 session_cookie;
-    char session_aes_key[AES_128_KEY_SIZE];
+    void *process_id;
-    void *protected_process_id;
+    unsigned char aes_key[32];
+    unsigned char aes_iv[16];
   };
 
   struct hv_detection_packet {

module/module.cpp

@@ -5,9 +5,11 @@
 #include "client/message_queue.h"
 #include "dispatcher/dispatcher.h"
 
+#include "crypt/crypt.h"
+
 void module::run(HINSTANCE hinstDLL) {
   AllocConsole();
-  FILE *file;
+  FILE *file = NULL;
   freopen_s(&file, "CONOUT$", "w", stdout);
   freopen_s(&file, "CONIN$", "r", stdin);
 

module/module.vcxproj

@@ -209,6 +209,7 @@
     </Link>
   </ItemDefinitionGroup>
   <ItemGroup>
+    <ClCompile Include="crypt\crypt.cpp" />
     <ClCompile Include="dispatcher\timer.cpp" />
     <ClCompile Include="helper.cpp" />
     <ClCompile Include="imports.cpp" />
@@ -223,6 +224,7 @@
   <ItemGroup>
     <ClInclude Include="client\message_queue.h" />
     <ClInclude Include="client\pipe.h" />
+    <ClInclude Include="crypt\crypt.h" />
     <ClInclude Include="dispatcher\dispatcher.h" />
     <ClInclude Include="dispatcher\threadpool.h" />
     <ClInclude Include="common.h" />

module/module.vcxproj.filters

@@ -11,6 +11,7 @@
     <ClCompile Include="dispatcher\threadpool.cpp" />
     <ClCompile Include="main.cpp" />
     <ClCompile Include="kernel_interface\kernel_interface.cpp" />
+    <ClCompile Include="crypt\crypt.cpp" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="client\message_queue.h" />
@@ -23,5 +24,6 @@
     <ClInclude Include="imports.h" />
     <ClInclude Include="kernel_interface\kernel_interface.h" />
     <ClInclude Include="module.h" />
+    <ClInclude Include="crypt\crypt.h" />
   </ItemGroup>
 </Project>