#include <ntifs.h>
// #include <wdm.h> // Removed to prevent redefinition
#include <stdint.h> // For fixed-width integer types like uint64_t
#include <ntddk.h>


// ==========================================================
// Definitions and Macros
// ==========================================================
typedef unsigned char BYTE;
typedef unsigned long DWORD;

#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
using QWORD = unsigned __int64;

// ==========================================================
// 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 = nullptr;
EXPUBLISHWNFSTATEDATA* ExPublishWnfStateData = nullptr;
EXUNSUBSCRIBEWNFSTATECHANGE* ExUnsubscribeWnfStateChange = nullptr;

ULONG64 MY_WNF_STATE_NAME_VALUE = MY_WNF_STATE_NAME; // Ensure consistency

PVOID g_WnfSubscription = nullptr; // Global WNF subscription handle

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

#pragma pack(push, 1)
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)
CHAR data[1024]; // Data buffer for read/write
};
#pragma pack(pop)

typedef INSTRUCTIONS* PINSTRUCTIONS;

// ==========================================================
// 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 WriteProcessMemoryInternal(HANDLE pid, PVOID Address, PVOID Buffer, SIZE_T Size, SIZE_T* BytesWritten);
NTSTATUS ReadProcessMemoryInternal(HANDLE pid, PVOID Address, PVOID Buffer, SIZE_T Size, SIZE_T* BytesRead);

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

extern "C" PVOID PsGetProcessSectionBaseAddress(PEPROCESS Process);

// Driver Unload
VOID DriverUnloadRoutine(PDRIVER_OBJECT DriverObject);

// Driver Entry
NTSTATUS DriverEntry(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 = reinterpret_cast<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 = reinterpret_cast<ULONGLONG>(TargetAddress);

PVOID pmapped_mem = MmMapIoSpace(AddrToWrite, Size, MmNonCached);
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)
{
ULONG_PTR process_dirbase_ptr = 0;

#ifdef _WIN64
process_dirbase_ptr = (ULONG_PTR)((PUCHAR)pProcess + 0x28); // Directory base for x64
#else
process_dirbase_ptr = (ULONG_PTR)((PUCHAR)pProcess + 0x18); // Directory base for x86
#endif

ULONG_PTR process_dirbase = process_dirbase_ptr;

if (process_dirbase == 0)
{
DWORD UserDirOffset = GetUserDirectoryTableBaseOffset();
ULONG_PTR process_userdirbase = (ULONG_PTR)((PUCHAR)pProcess + 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(reinterpret_cast<PVOID>(directoryTableBase + 8 * pdpIndex), &pdpe, sizeof(pdpe), &readsize);
if (!NT_SUCCESS(status) || !(pdpe & 1))
return 0;

// Read PDE
uint64_t pde = 0;
status = ReadPhysicalAddress(reinterpret_cast<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(reinterpret_cast<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(reinterpret_cast<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;
}

SIZE_T readSize = min(PAGE_SIZE - (physAddress & 0xFFFULL), bytesToRead);
SIZE_T localBytesRead = 0;
status = ReadPhysicalAddress(reinterpret_cast<PVOID>(physAddress), reinterpret_cast<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;
}

SIZE_T writeSize = min(PAGE_SIZE - (physAddress & 0xFFFULL), bytesToWrite);
SIZE_T localBytesWritten = 0;
status = WritePhysicalAddress(reinterpret_cast<PVOID>(physAddress), reinterpret_cast<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 (UCHAR)(c - 'A');
if (c >= 'a' && c <= 'z') return (UCHAR)(c - 'a' + 26);
if (c >= '0' && c <= '9') return (UCHAR)(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];

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

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

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

// 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 nullptr;

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

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

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

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

NTSTATUS status = STATUS_SUCCESS;

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

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

status = ReadProcessMemoryInternal(
(HANDLE)(ULONG_PTR)(instructions->processId),
(PVOID)(ULONG_PTR)(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),
nullptr,
nullptr,
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),
nullptr,
nullptr,
0
);
return status;
}

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

// Publish the data back via WNF
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
instructions,
sizeof(INSTRUCTIONS),
nullptr,
nullptr,
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),
nullptr,
nullptr,
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(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),
nullptr,
nullptr,
0
);
return status;
}

status = WriteProcessMemoryInternal(
(HANDLE)(ULONG_PTR)(instructions->processId),
(PVOID)(ULONG_PTR)(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),
nullptr,
nullptr,
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),
nullptr,
nullptr,
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 = nullptr;
NTSTATUS statusLookup = PsLookupProcessByProcessId((HANDLE)(ULONG_PTR)(instructions->processId), &pProcess);
if (NT_SUCCESS(statusLookup))
{
PVOID PsGetProcessSectionBaseAddress(PEPROCESS Process);
PVOID baseAddress = PsGetProcessSectionBaseAddress(pProcess);
instructions->baseAddress = (UINT64)(ULONG_PTR)baseAddress;
ObDereferenceObject(pProcess);

// Publish the base address back via WNF
status = ExPublishWnfStateData(
MY_WNF_STATE_NAME_VALUE,
instructions,
sizeof(INSTRUCTIONS),
nullptr,
nullptr,
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),
nullptr,
nullptr,
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
PCHAR receivedMessage = (PCHAR)ExAllocatePool2(POOL_FLAG_NON_PAGED, BufferSize + 1, 'msgK');
{
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, (BufferSize < (sizeof(CHAR) * (BufferSize))) ? BufferSize : sizeof(CHAR) * (BufferSize));

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((char)responseMessage) + 1),
nullptr,
nullptr,
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;
}

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

// Subscribe to WNF state changes
NTSTATUS status = ExSubscribeWnfStateChange(
&g_WnfSubscription, // Address of subscription handle
MY_WNF_STATE_NAME_VALUE,
0, // ChangeStamp (0 for all changes)
OnWnfCallback,
nullptr, // CallbackContext
nullptr, // 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 Implementation
// ==========================================================

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

if (Create)
{
PEPROCESS pProcess = nullptr;
NTSTATUS status = PsLookupProcessByProcessId(ProcessId, &pProcess);
if (NT_SUCCESS(status))
{
// Get process image file name
UNICODE_STRING uProcessImageName = { 0 };
PUNICODE_STRING puProcessImageName = &uProcessImageName;
status = SeLocateProcessImageName(pProcess, &puProcessImageName);
if (NT_SUCCESS(status))
{
if (uProcessImageName.Buffer != NULL) {
// 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", (ULONG_PTR)ProcessId);
}
}
}
else
{
DbgPrint("[Driver] Failed to parse process name for PID: %llu\n", (ULONG_PTR)ProcessId);
}
}
// Free the allocated image name buffer
if (uProcessImageName.Buffer != NULL)
{
ExFreePoolWithTag(uProcessImageName.Buffer, 'imgN');
}
}
else
{
DbgPrint("[Driver] SeLocateProcessImageName failed with status: 0x%X for PID: %llu\n", status, (ULONG_PTR)ProcessId);
}

ObDereferenceObject(pProcess);
}
else
{
DbgPrint("[Driver] PsLookupProcessByProcessId failed with status: 0x%X for PID: %llu\n", status, (ULONG_PTR)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 = nullptr;
}

// 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 Implementation
// ==========================================================

NTSTATUS DriverEntry(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;
}



// 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)(sizeof(handshakeMessage)), // Include null terminator
nullptr,
nullptr,
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;
}

make the driver compatable with kdmapper remember im using visual studio 2022 so it wont let me use old functions