#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