c to c++ #include <ntifs.h>
#include <ntddk.h>

// ==========================================================
// Definitions and Macros
// ==========================================================

#define TARGET_PROCESS_NAME L"notepad.exe" // Replace with your target process name
#define MY_WNF_STATE_NAME 0x41C64E6DA3BC1074ULL // Example WNF State Name (WNF_AUDC_CAPTURE)

// Windows Build Numbers
#define WINDOWS_1803 17134
#define WINDOWS_1809 17763
#define WINDOWS_1903 18362
#define WINDOWS_1909 18363
#define WINDOWS_2004 19041
#define WINDOWS_20H2 19042
#define WINDOWS_21H1 19043
#define WINDOWS_22H2 19044

// Page size definitions
#define PAGE_OFFSET_SIZE 12
#define PAGE_SIZE 0x1000
#define PMASK (~0xFFFULL)

// Custom type definitions
typedef unsigned __int64 QWORD;

// ==========================================================
// Function Pointer Types for WNF (Undocumented APIs)
// ==========================================================

typedef NTSTATUS(EXSUBSCRIBEWNFSTATECHANGE)(
PVOID Subscription,
ULONG64 StateName,
ULONG ChangeStamp,
PVOID Callback,
PVOID CallbackContext,
PVOID TypeId,
ULONG SerializationGroup,
ULONG Unknown
);

typedef NTSTATUS(EXPUBLISHWNFSTATEDATA)(
ULONG64 StateName,
PVOID Buffer,
ULONG Length,
PVOID TypeId,
PVOID ExplicitScope,
ULONG MatchingChangeStamp
);

typedef NTSTATUS(EXUNSUBSCRIBEWNFSTATECHANGE)(
PVOID Subscription
);

// ==========================================================
// Global Function Pointers for WNF
// ==========================================================

EXSUBSCRIBEWNFSTATECHANGE ExSubscribeWnfStateChange = NULL;
EXPUBLISHWNFSTATEDATA ExPublishWnfStateData = NULL;
EXUNSUBSCRIBEWNFSTATECHANGE* ExUnsubscribeWnfStateChange = NULL;

ULONG64 MY_WNF_STATE_NAME_VALUE = MY_WNF_STATE_NAME; // Ensure consistency

PVOID g_WnfSubscription = NULL; // Global WNF subscription handle

// ==========================================================
// Structures
// ==========================================================

#pragma pack(push, 1)
typedef struct _INSTRUCTIONS {
BYTE close; // Flag to close subscription
BYTE read; // Flag to perform read operation
BYTE write; // Flag to perform write operation
BYTE reqBase; // Flag to request base address
UINT64 processId; // Target process ID
UINT64 address; // Target virtual address
UINT64 size; // Size of read/write
PVOID bufferAddress; // Buffer address in user-mode
PVOID output; // Output buffer (for read operations)
UINT64 baseAddress; // Base address (for base address requests)
CHAR moduleName[256]; // Module name (if needed)
BYTE data[1024]; // Data buffer for read/write
} INSTRUCTIONS, PINSTRUCTIONS;
#pragma pack(pop)

// ==========================================================
// Function Declarations
// ==========================================================

// Physical Memory Access Functions
NTSTATUS ReadPhysicalAddress(PVOID TargetAddress, PVOID lpBuffer, SIZE_T Size, SIZE_T BytesRead);
NTSTATUS WritePhysicalAddress(PVOID TargetAddress, PVOID lpBuffer, SIZE_T Size, SIZE_T* BytesWritten);
DWORD GetUserDirectoryTableBaseOffset();
ULONG_PTR GetProcessCr3(PEPROCESS pProcess);
uint64_t TranslateLinearAddress(uint64_t directoryTableBase, uint64_t virtualAddress);

// Memory Operations
NTSTATUS ReadVirtualMemory(ULONG_PTR directoryTableBase, ULONG_PTR virtualAddress, PVOID buffer, SIZE_T size, SIZE_T* bytesRead);
NTSTATUS WriteVirtualMemory(ULONG_PTR directoryTableBase, ULONG_PTR virtualAddress, PVOID buffer, SIZE_T size, SIZE_T* bytesWritten);

// Process Memory Operations
NTSTATUS ReadProcessMemoryInternal(HANDLE pid, PVOID Address, PVOID Buffer, SIZE_T Size, SIZE_T* BytesRead);
NTSTATUS WriteProcessMemoryInternal(HANDLE pid, PVOID Address, PVOID Buffer, SIZE_T Size, SIZE_T* BytesWritten);

// WNF Communication
NTSTATUS OnWnfCallback(
PVOID Subscription,
ULONG64 StateName,
ULONG ChangeStamp,
PVOID TypeId,
PVOID CallbackContext,
PVOID Buffer,
ULONG BufferSize
);
NTSTATUS SubscribeToWNF();

// Base64 Encoding/Decoding
UCHAR Base64DecodeChar(char c);
NTSTATUS DecodeBase64(const CHAR* base64Input, UCHAR* outputBuffer, SIZE_T bufferSize, SIZE_T* decodedSize);
NTSTATUS EncodeBase64(const UCHAR* input, SIZE_T inputSize, CHAR* outputBuffer, SIZE_T bufferSize, SIZE_T* encodedSize);

// Process Notification
VOID ProcessNotifyCallback(
__in HANDLE ParentId,
__in HANDLE ProcessId,
__in BOOLEAN Create
);

// Driver Unload
VOID DriverUnloadRoutine(PDRIVER_OBJECT DriverObject);

// Driver Entry
NTSTATUS DriverEntryRoutine(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath);

// Helper Functions
PVOID GetProcessBaseAddress(HANDLE pid);

// ==========================================================
// Function Definitions
// ==========================================================

// Get the offset for UserDirectoryTableBase based on Windows version
DWORD GetUserDirectoryTableBaseOffset()
{
RTL_OSVERSIONINFOW ver = { 0 };
RtlGetVersion(&ver);
switch (ver.dwBuildNumber)
{
case WINDOWS_1803:
case WINDOWS_1809:
return 0x0278;
case WINDOWS_1903:
case WINDOWS_1909:
return 0x0280;
case WINDOWS_2004:
case WINDOWS_20H2:
case WINDOWS_21H1:
case WINDOWS_22H2:
default:
return 0x0388;
}
}

// Read physical memory
NTSTATUS ReadPhysicalAddress(PVOID TargetAddress, PVOID lpBuffer, SIZE_T Size, SIZE_T* BytesRead)
{
if (!TargetAddress || !lpBuffer || !BytesRead)
return STATUS_INVALID_PARAMETER;

MM_COPY_ADDRESS AddrToRead = { 0 };
AddrToRead.PhysicalAddress.QuadPart = (ULONGLONG)TargetAddress;

NTSTATUS status = MmCopyMemory(lpBuffer, AddrToRead, Size, MM_COPY_MEMORY_PHYSICAL, BytesRead);
return status;
}

// Write physical memory
NTSTATUS WritePhysicalAddress(PVOID TargetAddress, PVOID lpBuffer, SIZE_T Size, SIZE_T* BytesWritten)
{
if (!TargetAddress || !lpBuffer || !BytesWritten)
return STATUS_INVALID_PARAMETER;

PHYSICAL_ADDRESS AddrToWrite;
AddrToWrite.QuadPart = (ULONGLONG)TargetAddress;

PVOID pmapped_mem = MmMapIoSpaceEx(AddrToWrite, Size, PAGE_READWRITE);
if (!pmapped_mem)
return STATUS_UNSUCCESSFUL;

RtlCopyMemory(pmapped_mem, lpBuffer, Size);
BytesWritten = Size;
MmUnmapIoSpace(pmapped_mem, Size);
return STATUS_SUCCESS;
}

// Retrieve the CR3 (Page Directory Base) of a process
ULONG_PTR GetProcessCr3(PEPROCESS pProcess)
{
PUCHAR process = (PUCHAR)pProcess;
#ifdef _WIN64
ULONG_PTR process_dirbase_ptr = (ULONG_PTR*)(process + 0x28); // dirbase for x64
#else
ULONG_PTR* process_dirbase_ptr = (ULONG_PTR*)(process + 0x18); // dirbase for x86
#endif
ULONG_PTR process_dirbase = process_dirbase_ptr;

if (process_dirbase == 0)
{
DWORD UserDirOffset = GetUserDirectoryTableBaseOffset();
ULONG_PTR process_userdirbase = (ULONG_PTR)(process + UserDirOffset);
return process_userdirbase;
}
return process_dirbase;
}

// Translate a virtual address to a physical address using the provided directory table base
uint64_t TranslateLinearAddress(uint64_t directoryTableBase, uint64_t virtualAddress)
{
directoryTableBase &= PMASK; // Align directoryTableBase

uint64_t pageOffset = virtualAddress & 0xFFFULL;
uint64_t pteIndex = (virtualAddress >> 12) & 0x1FFULL;
uint64_t ptIndex = (virtualAddress >> 21) & 0x1FFULL;
uint64_t pdIndex = (virtualAddress >> 30) & 0x1FFULL;
uint64_t pdpIndex = (virtualAddress >> 39) & 0x1FFULL;

SIZE_T readsize = 0;
uint64_t pdpe = 0;

// Read PDPE
NTSTATUS status = ReadPhysicalAddress((PVOID)(directoryTableBase + 8 * pdpIndex), &pdpe, sizeof(pdpe), &readsize);
if (!NT_SUCCESS(status) || !(pdpe & 1))
return 0;

// Read PDE
uint64_t pde = 0;
status = ReadPhysicalAddress((PVOID)((pdpe & PMASK) + 8 * pdIndex), &pde, sizeof(pde), &readsize);
if (!NT_SUCCESS(status) || !(pde & 1))
return 0;

// 1GB large page
if (pde & (1 << 7)) // PS bit
return (pde & 0xFFFFFC0000000ULL) + (virtualAddress & 0x3FFFFFFFULL);

// Read PTE
uint64_t pteAddr = 0;
status = ReadPhysicalAddress((PVOID)((pde & PMASK) + 8 * ptIndex), &pteAddr, sizeof(pteAddr), &readsize);
if (!NT_SUCCESS(status) || !(pteAddr & 1))
return 0;

// 2MB large page
if (pteAddr & (1 << 7)) // PS bit
return (pteAddr & PMASK) + (virtualAddress & 0x1FFFFFULL);

// Read Page Table Entry
uint64_t pageTableEntry = 0;
status = ReadPhysicalAddress((PVOID)((pteAddr & PMASK) + 8 * pteIndex), &pageTableEntry, sizeof(pageTableEntry), &readsize);
if (!NT_SUCCESS(status) || !(pageTableEntry & 1))
return 0;

uint64_t physAddress = (pageTableEntry & PMASK) + pageOffset;
return physAddress;
}

// Read virtual memory using physical address translation
NTSTATUS ReadVirtualMemory(ULONG_PTR directoryTableBase, ULONG_PTR virtualAddress, PVOID buffer, SIZE_T size, SIZE_T bytesRead)
{
if (!buffer || !bytesRead)
return STATUS_INVALID_PARAMETER;

SIZE_T totalBytesRead = 0;
SIZE_T bytesToRead = size;
NTSTATUS status = STATUS_SUCCESS;

while (bytesToRead > 0)
{
uint64_t physAddress = TranslateLinearAddress(directoryTableBase, virtualAddress);
if (!physAddress)
{
status = STATUS_UNSUCCESSFUL;
break;
}

ULONG64 readSize = min(PAGE_SIZE - (physAddress & 0xFFFULL), bytesToRead);
SIZE_T localBytesRead = 0;
status = ReadPhysicalAddress((PVOID)physAddress, (PVOID)((ULONG_PTR)buffer + totalBytesRead), readSize, &localBytesRead);
if (!NT_SUCCESS(status) || localBytesRead == 0)
break;

bytesToRead -= localBytesRead;
virtualAddress += localBytesRead;
totalBytesRead += localBytesRead;
}

if (bytesRead)
bytesRead = totalBytesRead;
return status;
}

// Write virtual memory using physical address translation
NTSTATUS WriteVirtualMemory(ULONG_PTR directoryTableBase, ULONG_PTR virtualAddress, PVOID buffer, SIZE_T size, SIZE_T bytesWritten)
{
if (!buffer || !bytesWritten)
return STATUS_INVALID_PARAMETER;

SIZE_T totalBytesWritten = 0;
SIZE_T bytesToWrite = size;
NTSTATUS status = STATUS_SUCCESS;

while (bytesToWrite > 0)
{
uint64_t physAddress = TranslateLinearAddress(directoryTableBase, virtualAddress);
if (!physAddress)
{
status = STATUS_UNSUCCESSFUL;
break;
}

ULONG64 writeSize = min(PAGE_SIZE - (physAddress & 0xFFFULL), bytesToWrite);
SIZE_T localBytesWritten = 0;
status = WritePhysicalAddress((PVOID)physAddress, (PVOID)((ULONG_PTR)buffer + totalBytesWritten), writeSize, &localBytesWritten);
if (!NT_SUCCESS(status) || localBytesWritten == 0)
break;

bytesToWrite -= localBytesWritten;
virtualAddress += localBytesWritten;
totalBytesWritten += localBytesWritten;
}

if (bytesWritten)
bytesWritten = totalBytesWritten;
return status;
}

// Internal ReadProcessMemory function
NTSTATUS ReadProcessMemoryInternal(HANDLE pid, PVOID Address, PVOID Buffer, SIZE_T Size, SIZE_T BytesRead)
{
if (pid == NULL || Address == NULL || Buffer == NULL || BytesRead == NULL)
return STATUS_INVALID_PARAMETER;

PEPROCESS pProcess = NULL;
NTSTATUS status = PsLookupProcessByProcessId(pid, &pProcess);
if (!NT_SUCCESS(status))
return status;

ULONG_PTR dirBase = GetProcessCr3(pProcess);
ObDereferenceObject(pProcess);

return ReadVirtualMemory(dirBase, (ULONG_PTR)Address, Buffer, Size, BytesRead);
}

// Internal WriteProcessMemory function
NTSTATUS WriteProcessMemoryInternal(HANDLE pid, PVOID Address, PVOID Buffer, SIZE_T Size, SIZE_T* BytesWritten)
{
if (pid == NULL || Address == NULL || Buffer == NULL || BytesWritten == NULL)
return STATUS_INVALID_PARAMETER;

PEPROCESS pProcess = NULL;
NTSTATUS status = PsLookupProcessByProcessId(pid, &pProcess);
if (!NT_SUCCESS(status))
return status;

ULONG_PTR dirBase = GetProcessCr3(pProcess);
ObDereferenceObject(pProcess);

return WriteVirtualMemory(dirBase, (ULONG_PTR)Address, Buffer, Size, BytesWritten);
}

// Base64 Decode a single character
UCHAR Base64DecodeChar(char c)
{
if (c >= 'A' && c <= 'Z') return c - 'A';
if (c >= 'a' && c <= 'z') return c - 'a' + 26;
if (c >= '0' && c <= '9') return c - '0' + 52;
if (c == '+') return 62;
if (c == '/') return 63;
return 0xFF; // Invalid character
}

// Decode Base64 string to binary
NTSTATUS DecodeBase64(const CHAR* base64Input, UCHAR* outputBuffer, SIZE_T bufferSize, SIZE_T* decodedSize)
{
if (!base64Input || !outputBuffer || !decodedSize)
return STATUS_INVALID_PARAMETER;

SIZE_T inputLen = strlen(base64Input);
if (inputLen % 4 != 0)
return STATUS_INVALID_PARAMETER;

SIZE_T outputLen = (inputLen / 4) * 3;
if (outputLen > bufferSize)
return STATUS_BUFFER_TOO_SMALL;

decodedSize = 0;

for (SIZE_T i = 0; i < inputLen; i += 4)
{
UCHAR sextet_a = Base64DecodeChar(base64Input[i]);
UCHAR sextet_b = Base64DecodeChar(base64Input[i + 1]);
UCHAR sextet_c = Base64DecodeChar(base64Input[i + 2]);
UCHAR sextet_d = Base64DecodeChar(base64Input[i + 3]);

if (sextet_a == 0xFF || sextet_b == 0xFF)
return STATUS_INVALID_PARAMETER;

UCHAR triple[3] = { 0 };

triple[0] = (sextet_a << 2) | (sextet_b >> 4);
if (base64Input[i + 2] != '=')
{
if (sextet_c == 0xFF)
return STATUS_INVALID_PARAMETER;
triple[1] = ((sextet_b & 0xF) << 4) | (sextet_c >> 2);
}
if (base64Input[i + 3] != '=')
{
if (sextet_d == 0xFF)
return STATUS_INVALID_PARAMETER;
triple[2] = ((sextet_c & 0x3) << 6) | sextet_d;
}

// Copy the decoded bytes to the output buffer
outputBuffer[decodedSize] = triple[0];
(decodedSize)++;

if (base64Input[i + 2] != '=')
{
outputBuffer[decodedSize] = triple[1];
(decodedSize)++;
}

if (base64Input[i + 3] != '=')
{
outputBuffer[decodedSize] = triple[2];
(decodedSize)++;
}
}

return STATUS_SUCCESS;
}

// Base64 Encode binary data to string
NTSTATUS EncodeBase64(const UCHAR input, SIZE_T inputSize, CHAR outputBuffer, SIZE_T bufferSize, SIZE_T encodedSize)
{
if (!input || !outputBuffer || !encodedSize)
return STATUS_INVALID_PARAMETER;

const char base64_table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
SIZE_T outputNeeded = ((inputSize + 2) / 3) * 4;
if (outputNeeded + 1 > bufferSize) // +1 for null terminator
return STATUS_BUFFER_TOO_SMALL;

encodedSize = 0;

for (SIZE_T i = 0; i < inputSize; i += 3)
{
UCHAR a = input[i];
UCHAR b = (i + 1 < inputSize) ? input[i + 1] : 0;
UCHAR c = (i + 2 < inputSize) ? input[i + 2] : 0;

outputBuffer[encodedSize] = base64_table[(a >> 2) & 0x3F];
(encodedSize)++;
outputBuffer[encodedSize] = base64_table[((a & 0x3) << 4) | ((b & 0xF0) >> 4)];
(encodedSize)++;
outputBuffer[encodedSize] = (i + 1 < inputSize) ? base64_table[((b & 0xF) << 2) | ((c & 0xC0) >> 6)] : '=';
(encodedSize)++;
outputBuffer[encodedSize] = (i + 2 < inputSize) ? base64_table[c & 0x3F] : '=';
(encodedSize)++;
}

// Null-terminate the output
if (encodedSize < bufferSize)
{
outputBuffer[encodedSize] = '\0';
}
else
{
outputBuffer[bufferSize - 1] = '\0'; // Ensure null termination
}

return STATUS_SUCCESS;
}

// Retrieve the base address of a process given its PID
PVOID GetProcessBaseAddress(HANDLE pid)
{
if (pid == NULL)
return NULL;

PEPROCESS pProcess = NULL;
NTSTATUS status = PsLookupProcessByProcessId(pid, &pProcess);
if (!NT_SUCCESS(status))
return NULL;

PVOID Base = PsGetProcessSectionBaseAddress(pProcess);
ObDereferenceObject(pProcess);
return Base;
}

// Read process memory
NTSTATUS ReadProcessMemoryInternal(HANDLE pid, PVOID Address, PVOID Buffer, SIZE_T Size, SIZE_T BytesRead)
{
if (pid == NULL || Address == NULL || Buffer == NULL || BytesRead == NULL)
return STATUS_INVALID_PARAMETER;

return ReadProcessMemoryInternal(pid, Address, Buffer, Size, BytesRead);
}

// Write process memory
NTSTATUS WriteProcessMemoryInternal(HANDLE pid, PVOID Address, PVOID Buffer, SIZE_T Size, SIZE_T BytesWritten)
{
if (pid == NULL || Address == NULL || Buffer == NULL || BytesWritten == NULL)
return STATUS_INVALID_PARAMETER;

return WriteProcessMemoryInternal(pid, Address, Buffer, Size, BytesWritten);
}

// ==========================================================
// WNF Callback Function
// ==========================================================

NTSTATUS OnWnfCallback(
PVOID Subscription,
ULONG64 StateName,
ULONG ChangeStamp,
PVOID TypeId,
PVOID CallbackContext,
PVOID Buffer,
ULONG BufferSize
)
{
UNREFERENCED_PARAMETER(Subscription);
UNREFERENCED_PARAMETER(StateName);
UNREFERENCED_PARAMETER(ChangeStamp);
UNREFERENCED_PARAMETER(TypeId);
UNREFERENCED_PARAMETER(CallbackContext);

NTSTATUS status = STATUS_SUCCESS;

if (Buffer && BufferSize >= sizeof(INSTRUCTIONS))
{
PINSTRUCTIONS instructions = (PINSTRUCTIONS)Buffer;

if (instructions->read)
{
// Handle read request
SIZE_T bytesRead = 0;
BYTE readBuffer[1024] = { 0 };

status = ReadProcessMemoryInternal(
(HANDLE)(ULONGLONG)instructions->processId,
(PVOID)(ULONGLONG)instructions->address,
readBuffer,
(SIZE_T)instructions->size,
&bytesRead
);

if (NT_SUCCESS(status))
{
// Encode the read data to Base64
SIZE_T base64Size = ((bytesRead + 2) / 3) * 4;
if (base64Size + 1 > sizeof(instructions->data))
{
DbgPrint("[Driver] Read operation failed: Encoded data exceeds buffer size.\n");
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Read operation failed: Encoded data exceeds buffer size.",
(ULONG)strlen("Read operation failed: Encoded data exceeds buffer size.") + 1,
NULL,
NULL,
0
);
return status;
}

CHAR encodedData[4096] = { 0 };
SIZE_T encodedLength = 0;
status = EncodeBase64(readBuffer, bytesRead, encodedData, sizeof(encodedData), &encodedLength);
if (!NT_SUCCESS(status))
{
DbgPrint("[Driver] EncodeBase64 failed with status: 0x%X\n", status);
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Read operation failed: Base64 encode failed.",
(ULONG)strlen("Read operation failed: Base64 encode failed.") + 1,
NULL,
NULL,
0
);
return status;
}

// Prepare instructions to send back
RtlZeroMemory(instructions->data, sizeof(instructions->data));
RtlCopyMemory(instructions->data, encodedData, min(encodedLength, sizeof(instructions->data) - 1));
instructions->size = bytesRead;

// Publish the data back via WNF
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
instructions,
sizeof(INSTRUCTIONS),
NULL,
NULL,
0
);

if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Read operation completed successfully. %zu bytes read, %zu bytes encoded.\n", bytesRead, encodedLength);
}
else
{
DbgPrint("[Driver] ExPublishWnfStateData failed with status: 0x%X\n", status);
}
}
else
{
// Handle read error
DbgPrint("[Driver] ReadProcessMemoryInternal failed with status: 0x%X\n", status);
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Read operation failed.",
(ULONG)strlen("Read operation failed.") + 1,
NULL,
NULL,
0
);
}
}
else if (instructions->write)
{
// Handle write request
SIZE_T bytesWritten = 0;

// Decode Base64 data
UCHAR decodedData[512] = { 0 };
SIZE_T decodedSize = 0;
status = DecodeBase64((CHAR)instructions->data, decodedData, sizeof(decodedData), &decodedSize);
if (!NT_SUCCESS(status))
{
DbgPrint("[Driver] DecodeBase64 failed with status: 0x%X\n", status);
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Write operation failed: Base64 decode failed.",
(ULONG)strlen("Write operation failed: Base64 decode failed.") + 1,
NULL,
NULL,
0
);
return status;
}

status = WriteProcessMemoryInternal(
(HANDLE)(ULONGLONG)instructions->processId,
(PVOID)(ULONGLONG)instructions->address,
decodedData,
decodedSize,
&bytesWritten
);

if (NT_SUCCESS(status) && bytesWritten == decodedSize)
{
// Prepare instruction to send back
instructions->size = bytesWritten;

// Publish the success message back via WNF
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
instructions,
sizeof(INSTRUCTIONS),
NULL,
NULL,
0
);

DbgPrint("[Driver] Write operation completed successfully. %zu bytes written to process %llu at address 0x%llX.\n",
bytesWritten, instructions->processId, instructions->address);
}
else
{
// Handle write error
DbgPrint("[Driver] WriteProcessMemoryInternal failed with status: 0x%X and bytes written: %zu\n", status, bytesWritten);
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Write operation failed.",
(ULONG)strlen("Write operation failed.") + 1,
NULL,
NULL,
0
);
}
}
else if (instructions->close)
{
// Handle close instruction
// Currently, no specific cleanup is required
DbgPrint("[Driver] Received close instruction.\n");
}
else if (instructions->reqBase)
{
// Handle base address request
PEPROCESS pProcess = NULL;
NTSTATUS statusLookup = PsLookupProcessByProcessId((HANDLE)(ULONGLONG)instructions->processId, &pProcess);
if (NT_SUCCESS(statusLookup))
{
PVOID baseAddress = PsGetProcessSectionBaseAddress(pProcess);
instructions->baseAddress = (UINT64)baseAddress;
ObDereferenceObject(pProcess);

// Publish the base address back via WNF
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
instructions,
sizeof(INSTRUCTIONS),
NULL,
NULL,
0
);

DbgPrint("[Driver] Base address of process %llu: %p\n", instructions->processId, baseAddress);
}
else
{
// Handle error
DbgPrint("[Driver] PsLookupProcessByProcessId failed with status: 0x%X\n", statusLookup);
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Base address request failed: PsLookupProcessByProcessId failed.",
(ULONG)strlen("Base address request failed: PsLookupProcessByProcessId failed.") + 1,
NULL,
NULL,
0
);
}
}
// Additional instruction handling can be added here
}
else if (Buffer && BufferSize > 0)
{
// Handle simple messages like "Hi"
CHAR receivedMessage = (CHAR)ExAllocatePoolWithTag(NonPagedPoolNx, BufferSize + 1, 'msgK');
if (!receivedMessage)
{
DbgPrint("[Driver] Failed to allocate memory for received message.\n");
return STATUS_INSUFFICIENT_RESOURCES;
}

RtlZeroMemory(receivedMessage, BufferSize + 1);
RtlCopyMemory(receivedMessage, Buffer, min(BufferSize, sizeof(receivedMessage) - 1));

DbgPrint("[Driver] Kernel received message: %s\n", receivedMessage);

if (strcmp(receivedMessage, "Hi") == 0)
{
// Send handshake response
const CHAR responseMessage = "Hello from kernel mode";
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
(PVOID)responseMessage,
(ULONG)strlen(responseMessage) + 1,
NULL,
NULL,
0
);

if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Handshake response sent successfully.\n");
}
else
{
DbgPrint("[Driver] Failed to send handshake response with status: 0x%X\n", status);
}
}

ExFreePoolWithTag(receivedMessage, 'msgK');
}
else
{
// Handle unexpected or unsupported data
DbgPrint("[Driver] Received invalid or unsupported data.\n");
}

return status;
}

// ==========================================================
// WNF Subscription Function
// ==========================================================

// Subscribe to WNF state changes
NTSTATUS SubscribeToWNF()
{
if (!ExSubscribeWnfStateChange || !ExPublishWnfStateData || !ExUnsubscribeWnfStateChange)
{
DbgPrint("[Driver] WNF function pointers not initialized.\n");
return STATUS_UNSUCCESSFUL;
}

NTSTATUS status = ExSubscribeWnfStateChange(
g_WnfSubscription,
MY_WNF_STATE_NAME_VALUE,
0, // ChangeStamp (0 for all changes)
(PVOID)OnWnfCallback,
NULL, // CallbackContext
NULL, // TypeId
0, // SerializationGroup
0 // Unknown
);

if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Successfully subscribed to WNF state changes.\n");
}
else
{
DbgPrint("[Driver] ExSubscribeWnfStateChange failed with status: 0x%X\n", status);
}

return status;
}

// ==========================================================
// Process Notification Callback
// ==========================================================

VOID ProcessNotifyCallback(
__in HANDLE ParentId,
__in HANDLE ProcessId,
__in BOOLEAN Create
)
{
UNREFERENCED_PARAMETER(ParentId);

if (Create)
{
PEPROCESS pProcess = NULL;
NTSTATUS status = PsLookupProcessByProcessId(ProcessId, &pProcess);
if (NT_SUCCESS(status))
{
// Get process image file name
UNICODE_STRING uProcessImageName = { 0 };
status = SeLocateProcessImageName(pProcess, &uProcessImageName);
if (NT_SUCCESS(status))
{
// Extract file name from full path
PWSTR processName = wcsrchr(uProcessImageName.Buffer, L''); // Corrected backslash
if (processName)
{
processName++; // Move past the backslash

// Check if this is the target process
if (_wcsicmp(processName, TARGET_PROCESS_NAME) == 0)
{
// Perform memory operations here (if needed)
PVOID baseAddress = GetProcessBaseAddress(ProcessId);
if (baseAddress)
{
DbgPrint("[Driver] Target process '%ws' created with base address: %p\n", processName, baseAddress);

// Example: Read some memory from the process
BYTE buffer[256] = { 0 };
SIZE_T bytesRead = 0;
NTSTATUS ntStatus = ReadProcessMemoryInternal(
ProcessId,
baseAddress,
buffer,
sizeof(buffer),
&bytesRead
);

if (NT_SUCCESS(ntStatus))
{
DbgPrint("[Driver] Successfully read %zu bytes from '%ws'.\n", bytesRead, processName);
}
else
{
DbgPrint("[Driver] Failed to read memory of '%ws' with status: 0x%X\n", processName, ntStatus);
}

// Example: Write to process memory (commented out for safety)
/
BYTE writeData[256] = { / Your data here * / };
SIZE_T bytesWritten = 0;
ntStatus = WriteProcessMemoryInternal(
ProcessId,
baseAddress,
writeData,
sizeof(writeData),
&bytesWritten
);

if (NT_SUCCESS(ntStatus))
{
DbgPrint("[Driver] Successfully wrote %zu bytes to '%ws'.\n", bytesWritten, processName);
}
else
{
DbgPrint("[Driver] Failed to write memory of '%ws' with status: 0x%X\n", processName, ntStatus);
}
/
}
else
{
DbgPrint("[Driver] Failed to retrieve base address of process ID: %llu\n", ProcessId);
}
}
}
else
{
DbgPrint("[Driver] Failed to parse process name for PID: %llu\n", ProcessId);
}

// Free the allocated image name buffer
ExFreePoolWithTag(uProcessImageName.Buffer, 'imgN');
}
else
{
DbgPrint("[Driver] SeLocateProcessImageName failed with status: 0x%X for PID: %llu\n", status, ProcessId);
}

ObDereferenceObject(pProcess);
}
else
{
DbgPrint("[Driver] PsLookupProcessByProcessId failed with status: 0x%X for PID: %llu\n", status, ProcessId);
}
}
}

// ==========================================================
// Driver Unload Routine
// ==========================================================

VOID DriverUnloadRoutine(PDRIVER_OBJECT DriverObject)
{
UNREFERENCED_PARAMETER(DriverObject);

// Unsubscribe from WNF
if (g_WnfSubscription && ExUnsubscribeWnfStateChange)
{
NTSTATUS status = ExUnsubscribeWnfStateChange(g_WnfSubscription);
if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Successfully unsubscribed from WNF state changes.\n");
}
else
{
DbgPrint("[Driver] ExUnsubscribeWnfStateChange failed with status: 0x%X\n", status);
}
g_WnfSubscription = NULL;
}

// Remove the process notify routine
NTSTATUS status = PsSetCreateProcessNotifyRoutine(ProcessNotifyCallback, TRUE);
if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Successfully removed process notify routine.\n");
}
else
{
DbgPrint("[Driver] PsSetCreateProcessNotifyRoutine removal failed with status: 0x%X\n", status);
}

DbgPrint("[Driver] Driver unloaded.\n");
}

// ==========================================================
// Driver Entry Function
// ==========================================================

NTSTATUS DriverEntryRoutine(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
{
UNREFERENCED_PARAMETER(RegistryPath);

NTSTATUS status;

// Resolve WNF functions
UNICODE_STRING routineName;
RtlInitUnicodeString(&routineName, L"ExSubscribeWnfStateChange");
ExSubscribeWnfStateChange = (EXSUBSCRIBEWNFSTATECHANGE)MmGetSystemRoutineAddress(&routineName);
if (!ExSubscribeWnfStateChange)
{
DbgPrint("[Driver] ExSubscribeWnfStateChange not found.\n");
return STATUS_UNSUCCESSFUL;
}

RtlInitUnicodeString(&routineName, L"ExPublishWnfStateData");
ExPublishWnfStateData = (EXPUBLISHWNFSTATEDATA)MmGetSystemRoutineAddress(&routineName);
if (!ExPublishWnfStateData)
{
DbgPrint("[Driver] ExPublishWnfStateData not found.\n");
return STATUS_UNSUCCESSFUL;
}

RtlInitUnicodeString(&routineName, L"ExUnsubscribeWnfStateChange");
ExUnsubscribeWnfStateChange = (EXUNSUBSCRIBEWNFSTATECHANGE)MmGetSystemRoutineAddress(&routineName);
if (!ExUnsubscribeWnfStateChange)
{
DbgPrint("[Driver] ExUnsubscribeWnfStateChange not found.\n");
return STATUS_UNSUCCESSFUL;
}

// Subscribe to WNF state changes
status = SubscribeToWNF();
if (!NT_SUCCESS(status))
{
DbgPrint("[Driver] Failed to subscribe to WNF state changes with status: 0x%X\n", status);
return status;
}

// Register the process creation notify routine
status = PsSetCreateProcessNotifyRoutine(ProcessNotifyCallback, FALSE);
if (!NT_SUCCESS(status))
{
DbgPrint("[Driver] PsSetCreateProcessNotifyRoutine failed with status: 0x%X\n", status);
// Unsubscribe from WNF before exiting
if (g_WnfSubscription && ExUnsubscribeWnfStateChange)
ExUnsubscribeWnfStateChange(g_WnfSubscription);

return status;
}

// Set the unload routine
DriverObject->DriverUnload = DriverUnloadRoutine;

// Send handshake message to user-mode
const CHAR handshakeMessage = "Kernel driver loaded";
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
(PVOID)handshakeMessage,
(ULONG)strlen(handshakeMessage) + 1,
NULL,
NULL,
0
);

if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Handshake message published successfully.\n");
}
else
{
DbgPrint("[Driver] Failed to publish handshake message with status: 0x%X\n", status);
}

DbgPrint("[Driver] Driver loaded successfully.\n");
return STATUS_SUCCESS;
}

// ==========================================================
// WNF Subscription Function Implementation
// ==========================================================

NTSTATUS SubscribeToWNF()
{
if (!ExSubscribeWnfStateChange || !ExPublishWnfStateData || !ExUnsubscribeWnfStateChange)
{
DbgPrint("[Driver] WNF function pointers not initialized.\n");
return STATUS_UNSUCCESSFUL;
}

// Subscribe to WNF state changes
NTSTATUS status = ExSubscribeWnfStateChange(
g_WnfSubscription,
MY_WNF_STATE_NAME_VALUE,
0, // ChangeStamp (0 for all changes)
(PVOID)OnWnfCallback,
NULL, // CallbackContext
NULL, // TypeId
0, // SerializationGroup
0 // Unknown
);

if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Successfully subscribed to WNF state changes.\n");
}
else
{
DbgPrint("[Driver] ExSubscribeWnfStateChange failed with status: 0x%X\n", status);
}

return status;
}

// ==========================================================
// WNF Callback Function Implementation
// ==========================================================

NTSTATUS OnWnfCallback(
PVOID Subscription,
ULONG64 StateName,
ULONG ChangeStamp,
PVOID TypeId,
PVOID CallbackContext,
PVOID Buffer,
ULONG BufferSize
)
{
UNREFERENCED_PARAMETER(Subscription);
UNREFERENCED_PARAMETER(StateName);
UNREFERENCED_PARAMETER(ChangeStamp);
UNREFERENCED_PARAMETER(TypeId);
UNREFERENCED_PARAMETER(CallbackContext);

NTSTATUS status = STATUS_SUCCESS;

if (Buffer && BufferSize >= sizeof(INSTRUCTIONS))
{
PINSTRUCTIONS instructions = (PINSTRUCTIONS)Buffer;

if (instructions->read)
{
// Handle read request
SIZE_T bytesRead = 0;
BYTE readBuffer[1024] = { 0 };

status = ReadProcessMemoryInternal(
(HANDLE)(ULONGLONG)instructions->processId,
(PVOID)(ULONGLONG)instructions->address,
readBuffer,
(SIZE_T)instructions->size,
&bytesRead
);

if (NT_SUCCESS(status))
{
// Encode the read data to Base64
SIZE_T base64Size = ((bytesRead + 2) / 3) * 4;
if (base64Size + 1 > sizeof(instructions->data))
{
DbgPrint("[Driver] Read operation failed: Encoded data exceeds buffer size.\n");
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Read operation failed: Encoded data exceeds buffer size.",
(ULONG)strlen("Read operation failed: Encoded data exceeds buffer size.") + 1,
NULL,
NULL,
0
);
return status;
}

CHAR encodedData[4096] = { 0 };
SIZE_T encodedLength = 0;
status = EncodeBase64(readBuffer, bytesRead, encodedData, sizeof(encodedData), &encodedLength);
if (!NT_SUCCESS(status))
{
DbgPrint("[Driver] EncodeBase64 failed with status: 0x%X\n", status);
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Read operation failed: Base64 encode failed.",
(ULONG)strlen("Read operation failed: Base64 encode failed.") + 1,
NULL,
NULL,
0
);
return status;
}

// Prepare instructions to send back
RtlZeroMemory(instructions->data, sizeof(instructions->data));
RtlCopyMemory(instructions->data, encodedData, min(encodedLength, sizeof(instructions->data) - 1));
instructions->size = bytesRead;

// Publish the data back via WNF
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
instructions,
sizeof(INSTRUCTIONS),
NULL,
NULL,
0
);

if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Read operation completed successfully. %zu bytes read, %zu bytes encoded.\n", bytesRead, encodedLength);
}
else
{
DbgPrint("[Driver] ExPublishWnfStateData failed with status: 0x%X\n", status);
}
}
else
{
// Handle read error
DbgPrint("[Driver] ReadProcessMemoryInternal failed with status: 0x%X\n", status);
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Read operation failed.",
(ULONG)strlen("Read operation failed.") + 1,
NULL,
NULL,
0
);
}
}
else if (instructions->write)
{
// Handle write request
SIZE_T bytesWritten = 0;

// Decode Base64 data
UCHAR decodedData[512] = { 0 };
SIZE_T decodedSize = 0;
status = DecodeBase64((CHAR)instructions->data, decodedData, sizeof(decodedData), &decodedSize);
if (!NT_SUCCESS(status))
{
DbgPrint("[Driver] DecodeBase64 failed with status: 0x%X\n", status);
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Write operation failed: Base64 decode failed.",
(ULONG)strlen("Write operation failed: Base64 decode failed.") + 1,
NULL,
NULL,
0
);
return status;
}

status = WriteProcessMemoryInternal(
(HANDLE)(ULONGLONG)instructions->processId,
(PVOID)(ULONGLONG)instructions->address,
decodedData,
decodedSize,
&bytesWritten
);

if (NT_SUCCESS(status) && bytesWritten == decodedSize)
{
// Prepare instruction to send back
instructions->size = bytesWritten;

// Publish the success message back via WNF
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
instructions,
sizeof(INSTRUCTIONS),
NULL,
NULL,
0
);

DbgPrint("[Driver] Write operation completed successfully. %zu bytes written to process %llu at address 0x%llX.\n",
bytesWritten, instructions->processId, instructions->address);
}
else
{
// Handle write error
DbgPrint("[Driver] WriteProcessMemoryInternal failed with status: 0x%X and bytes written: %zu\n", status, bytesWritten);
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Write operation failed.",
(ULONG)strlen("Write operation failed.") + 1,
NULL,
NULL,
0
);
}
}
else if (instructions->close)
{
// Handle close instruction
// Currently, no specific cleanup is required
DbgPrint("[Driver] Received close instruction.\n");
// Additional cleanup logic can be implemented here if needed
}
else if (instructions->reqBase)
{
// Handle base address request
PEPROCESS pProcess = NULL;
NTSTATUS statusLookup = PsLookupProcessByProcessId((HANDLE)(ULONGLONG)instructions->processId, &pProcess);
if (NT_SUCCESS(statusLookup))
{
PVOID baseAddress = PsGetProcessSectionBaseAddress(pProcess);
instructions->baseAddress = (UINT64)baseAddress;
ObDereferenceObject(pProcess);

// Publish the base address back via WNF
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
instructions,
sizeof(INSTRUCTIONS),
NULL,
NULL,
0
);

DbgPrint("[Driver] Base address of process %llu: %p\n", instructions->processId, baseAddress);
}
else
{
// Handle error
DbgPrint("[Driver] PsLookupProcessByProcessId failed with status: 0x%X\n", statusLookup);
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
"Base address request failed: PsLookupProcessByProcessId failed.",
(ULONG)strlen("Base address request failed: PsLookupProcessByProcessId failed.") + 1,
NULL,
NULL,
0
);
}
}
// Additional instruction handling can be added here
}
else if (Buffer && BufferSize > 0)
{
// Handle simple messages like "Hi"
// Allocate memory for received message, ensuring no buffer overflow
CHAR receivedMessage = (CHAR)ExAllocatePoolWithTag(NonPagedPoolNx, BufferSize + 1, 'msgK'); // 'msgK' tag
if (!receivedMessage)
{
DbgPrint("[Driver] Failed to allocate memory for received message.\n");
return STATUS_INSUFFICIENT_RESOURCES;
}

// Zero-initialize the allocated memory
RtlZeroMemory(receivedMessage, BufferSize + 1);

// Copy the received data, ensuring not to exceed the allocated buffer
RtlCopyMemory(receivedMessage, Buffer, min(BufferSize, (ULONG)(sizeof(receivedMessage) - 1)));

DbgPrint("[Driver] Kernel received message: %s\n", receivedMessage);

// Check if the message is 'Hi' and respond accordingly
if (strcmp(receivedMessage, "Hi") == 0)
{
// Send handshake response
const CHAR responseMessage = "Hello from kernel mode";
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
(PVOID)responseMessage,
(ULONG)strlen(responseMessage) + 1, // Include null terminator
NULL,
NULL,
0
);

if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Handshake response sent successfully.\n");
}
else
{
DbgPrint("[Driver] Failed to send handshake response with status: 0x%X\n", status);
}
}

// Free the allocated memory for the received message
ExFreePoolWithTag(receivedMessage, 'msgK');
}
else
{
// Handle unexpected or unsupported data
DbgPrint("[Driver] Received invalid or unsupported data.\n");
}

return status;
}

// ==========================================================
// Process Notification Callback Implementation
// ==========================================================

VOID ProcessNotifyCallback(
__in HANDLE ParentId,
__in HANDLE ProcessId,
_in BOOLEAN Create
)
{
UNREFERENCED_PARAMETER(ParentId);

if (Create)
{
PEPROCESS pProcess = NULL;
NTSTATUS status = PsLookupProcessByProcessId(ProcessId, &pProcess);
if (NT_SUCCESS(status))
{
// Get process image file name
UNICODE_STRING uProcessImageName = { 0 };
status = SeLocateProcessImageName(pProcess, &uProcessImageName);
if (NT_SUCCESS(status))
{
// Extract file name from full path
PWSTR processName = wcsrchr(uProcessImageName.Buffer, L''); // Corrected backslash
if (processName)
{
processName++; // Move past the backslash

// Check if this is the target process
if (wcsicmp(processName, TARGET_PROCESS_NAME) == 0)
{
// Perform memory operations here (if needed)
PVOID baseAddress = GetProcessBaseAddress(ProcessId);
if (baseAddress)
{
DbgPrint("[Driver] Target process '%ws' created with base address: %p\n", processName, baseAddress);

// Example: Read some memory from the process
BYTE buffer[256] = { 0 };
SIZE_T bytesRead = 0;
NTSTATUS ntStatus = ReadProcessMemoryInternal(
(HANDLE)ProcessId,
baseAddress,
buffer,
sizeof(buffer),
&bytesRead
);

if (NT_SUCCESS(ntStatus))
{
DbgPrint("[Driver] Successfully read %zu bytes from '%ws'.\n", bytesRead, processName);
}
else
{
DbgPrint("[Driver] Failed to read memory of '%ws' with status: 0x%X\n", processName, ntStatus);
}

// Example: Write to process memory (commented out for safety)
/
BYTE writeData[256] = { / Your data here * / };
SIZE_T bytesWritten = 0;
ntStatus = WriteProcessMemoryInternal(
(HANDLE)ProcessId,
baseAddress,
writeData,
sizeof(writeData),
&bytesWritten
);

if (NT_SUCCESS(ntStatus))
{
DbgPrint("[Driver] Successfully wrote %zu bytes to '%ws'.\n", bytesWritten, processName);
}
else
{
DbgPrint("[Driver] Failed to write memory of '%ws' with status: 0x%X\n", processName, ntStatus);
}
/
}
else
{
DbgPrint("[Driver] Failed to retrieve base address of process ID: %llu\n", ProcessId);
}
}
}
else
{
DbgPrint("[Driver] Failed to parse process name for PID: %llu\n", ProcessId);
}

// Free the allocated image name buffer
ExFreePoolWithTag(uProcessImageName.Buffer, 'imgN');
}
else
{
DbgPrint("[Driver] SeLocateProcessImageName failed with status: 0x%X for PID: %llu\n", status, ProcessId);
}

ObDereferenceObject(pProcess);
}
else
{
DbgPrint("[Driver] PsLookupProcessByProcessId failed with status: 0x%X for PID: %llu\n", status, ProcessId);
}
}
}

// ==========================================================
// Driver Unload Routine Implementation
// ==========================================================

VOID DriverUnloadRoutine(PDRIVER_OBJECT DriverObject)
{
UNREFERENCED_PARAMETER(DriverObject);

// Unsubscribe from WNF
if (g_WnfSubscription && ExUnsubscribeWnfStateChange)
{
NTSTATUS status = ExUnsubscribeWnfStateChange(g_WnfSubscription);
if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Successfully unsubscribed from WNF state changes.\n");
}
else
{
DbgPrint("[Driver] ExUnsubscribeWnfStateChange failed with status: 0x%X\n", status);
}
g_WnfSubscription = NULL;
}

// Remove the process notify routine
NTSTATUS status = PsSetCreateProcessNotifyRoutine(ProcessNotifyCallback, TRUE);
if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Successfully removed process notify routine.\n");
}
else
{
DbgPrint("[Driver] PsSetCreateProcessNotifyRoutine removal failed with status: 0x%X\n", status);
}

DbgPrint("[Driver] Driver unloaded.\n");
}

// ==========================================================
// DriverEntry Function Implementation
// ==========================================================

NTSTATUS DriverEntryRoutine(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
{
UNREFERENCED_PARAMETER(RegistryPath);

NTSTATUS status;

// Resolve WNF functions
UNICODE_STRING routineName;
RtlInitUnicodeString(&routineName, L"ExSubscribeWnfStateChange");
ExSubscribeWnfStateChange = (EXSUBSCRIBEWNFSTATECHANGE)MmGetSystemRoutineAddress(&routineName);
if (!ExSubscribeWnfStateChange)
{
DbgPrint("[Driver] ExSubscribeWnfStateChange not found.\n");
return STATUS_UNSUCCESSFUL;
}

RtlInitUnicodeString(&routineName, L"ExPublishWnfStateData");
ExPublishWnfStateData = (EXPUBLISHWNFSTATEDATA*)MmGetSystemRoutineAddress(&routineName);
if (!ExPublishWnfStateData)
{
DbgPrint("[Driver] ExPublishWnfStateData not found.\n");
return STATUS_UNSUCCESSFUL;
}

RtlInitUnicodeString(&routineName, L"ExUnsubscribeWnfStateChange");
ExUnsubscribeWnfStateChange = (EXUNSUBSCRIBEWNFSTATECHANGE*)MmGetSystemRoutineAddress(&routineName);
if (!ExUnsubscribeWnfStateChange)
{
DbgPrint("[Driver] ExUnsubscribeWnfStateChange not found.\n");
return STATUS_UNSUCCESSFUL;
}

// Subscribe to WNF state changes
status = SubscribeToWNF();
if (!NT_SUCCESS(status))
{
DbgPrint("[Driver] Failed to subscribe to WNF state changes with status: 0x%X\n", status);
return status;
}

// Register the process creation notify routine
status = PsSetCreateProcessNotifyRoutine(ProcessNotifyCallback, FALSE);
if (!NT_SUCCESS(status))
{
DbgPrint("[Driver] PsSetCreateProcessNotifyRoutine failed with status: 0x%X\n", status);
// Unsubscribe from WNF before exiting
if (g_WnfSubscription && ExUnsubscribeWnfStateChange)
ExUnsubscribeWnfStateChange(g_WnfSubscription);

return status;
}

// Set the unload routine
DriverObject->DriverUnload = DriverUnloadRoutine;

// Send handshake message to user-mode
const CHAR* handshakeMessage = "Kernel driver loaded";
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
(PVOID)handshakeMessage,
(ULONG)strlen(handshakeMessage) + 1,
NULL,
NULL,
0
);

if (NT_SUCCESS(status))
{
DbgPrint("[Driver] Handshake message published successfully.\n");
}
else
{
DbgPrint("[Driver] Failed to publish handshake message with status: 0x%X\n", status);
}

DbgPrint("[Driver] Driver loaded successfully.\n");
return STATUS_SUCCESS;
}