#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 compatible with kdmapper remember im using visual studio 2022 so it wont let me use old functions stick to the use of windows notification Facility to communicate in the driver