in one code block only update my usermode to work with km 
um
using System;
using NtApiDotNet;
using System.Text;
using System.Threading;
using System.Runtime.InteropServices;

namespace WnfSubscriber
{
    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Ansi, Pack = 1)]
    public struct INSTRUCTIONS
    {
        public byte close;           // 1 byte
        public byte read;            // 1 byte
        public byte write;           // 1 byte
        public byte reqBase;         // 1 byte
        public ulong processId;      // 8 bytes (matches UINT_PTR in C code)
        public ulong address;        // 8 bytes
        public ulong size;           // 8 bytes
        public UIntPtr bufferAddress; // 8 bytes (unused, can be set to zero)
        public UIntPtr output;        // 8 bytes (unused, can be set to zero)
        public ulong baseAddress;    // 8 bytes

        [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 256)]
        public string moduleName;

        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 1024)]
        public byte[] data; // Data to write or read
    }


    public class KernelCom : IDisposable
    {
        private const ulong MY_WNF_STATE_NAME = 0x41C64E6DA3BC1074;
        private readonly WnfClient _client;
        private bool _initialized = false;

        public KernelCom()
        {
            _client = new WnfClient();
        }

        public void Initialize()
        {
            if (_initialized)
                return;

            _client.SubscribeState(new WnfStateName(MY_WNF_STATE_NAME), OnWnfStateChange);
            Console.WriteLine("[+] Subscribed to WNF state.");
            _initialized = true;
        }

        public void SendMessage(string message)
        {
            byte[] dataToSend = Encoding.ASCII.GetBytes(message);
            NtWnf.PublishStateData(new WnfStateName(MY_WNF_STATE_NAME), dataToSend);
            Console.WriteLine("[+] Sent message to kernel-mode: " + message);
        }

        public void SendInstructions(INSTRUCTIONS instructions)
        {
            int size = Marshal.SizeOf(typeof(INSTRUCTIONS));
            byte[] dataToSend = new byte[size];

            IntPtr ptr = Marshal.AllocHGlobal(size);
            try
            {
                Marshal.StructureToPtr(instructions, ptr, false);
                Marshal.Copy(ptr, dataToSend, 0, size);
                NtWnf.PublishStateData(new WnfStateName(MY_WNF_STATE_NAME), dataToSend);
                Console.WriteLine("[+] Sent instructions to kernel-mode.");
            }
            catch (Exception ex)
            {
               Console.WriteLine($"[!] Failed to send instructions: {ex.Message}");
            }
            finally
            {
                Marshal.FreeHGlobal(ptr);
            }
        }

        // New Method: ReadMemory
        public void ReadMemory(uint processId, ulong address, ulong size)
        {
            INSTRUCTIONS instructions = new INSTRUCTIONS
            {
                close = 0,
                read = 1,
                write = 0,
                reqBase = 0,
                processId = processId,
                address = address,
                size = size,
                bufferAddress = UIntPtr.Zero,
                output = UIntPtr.Zero,
                baseAddress = 0,
                moduleName = string.Empty,
                data = new byte[1024]
            };

            SendInstructions(instructions);
            Console.WriteLine($"[+] Requested to read {size} bytes from process {processId} at address 0x{address:X}.");
        }

        // New Method: WriteMemory
        public void WriteMemory(uint processId, ulong address, byte[] data)
        {
            if (data.Length > 1024) // Adjust the limit as per your requirements and kernel driver capability
            {
                Console.WriteLine("[!] Data size exceeds the maximum allowed (1024 bytes).");
                return;
            }
        
            INSTRUCTIONS instructions = new INSTRUCTIONS
            {
                close = 0,
                read = 0,
                write = 1,
                reqBase = 0,
                processId = processId,
                address = address,
                size = (ulong)data.Length,
                bufferAddress = UIntPtr.Zero,
                output = UIntPtr.Zero,
                baseAddress = 0,
                moduleName = string.Empty,
                data = new byte[1024]
            };

            Array.Copy(data, instructions.data, data.Length);
        
            SendInstructions(instructions);
            Console.WriteLine($"[+] Requested to write {data.Length} bytes to process {processId} at address 0x{address:X}.");
        }

        private void OnWnfStateChange(WnfStateChange change)
        {
            Console.WriteLine("[] WNF State Changed!");

            if (change.Data == null || change.Data.Length == 0)
            {
                Console.WriteLine("[!] Received empty data.");
                return;
            }

            // Attempt to parse the data as INSTRUCTIONS struct
            try
            {
                INSTRUCTIONS response = ByteArrayToStructure<INSTRUCTIONS>(change.Data);

                if (response.read == 1)
                {
                    Console.WriteLine($"[] Read Memory Response from process {response.processId}:");

                    if (response.size == 0)
                    {
                        Console.WriteLine("[!] No data received for read operation.");
                        return;
                    }

                    // Extract the data from the response
                    byte[] readData = new byte[response.size];
                    Array.Copy(response.data, readData, (int)response.size); // Ensure casting to int is safe

                    Console.WriteLine("[] Data: " + BitConverter.ToString(readData));
                }
                else if (response.write == 1)
                {
                    Console.WriteLine("[] Write Memory operation completed successfully.");
                }
                else if (response.close == 1)
                {
                    Console.WriteLine("[] Received instruction to close the driver.");
                    // Perform any cleanup if necessary
                }
                else if (response.reqBase == 1)
                {
                    Console.WriteLine($"[] Base Address for module '{response.moduleName}': 0x{response.baseAddress:X}");
                }
                else
                {
                    string receivedMessage = Encoding.ASCII.GetString(change.Data).TrimEnd('\0');
                    Console.WriteLine("[] Received message from kernel-mode: " + receivedMessage);

                    // Handle other messages or results
                }
            }
            catch (Exception)
            {
                // If parsing as INSTRUCTIONS struct fails, treat it as a simple string message
                string receivedMessage = Encoding.ASCII.GetString(change.Data).TrimEnd('\0');
                Console.WriteLine("[] Received message from kernel-mode: " + receivedMessage);
            }
        }

        private T ByteArrayToStructure<T>(byte[] bytes) where T : struct
        {
            int size = Marshal.SizeOf(typeof(T));
            if (bytes.Length < size)
                throw new ArgumentException("Byte array is too small for the specified structure.");

            IntPtr ptr = Marshal.AllocHGlobal(size);
            try
            {
                Marshal.Copy(bytes, 0, ptr, size);
                return Marshal.PtrToStructure<T>(ptr);
            }
            finally
            {
                Marshal.FreeHGlobal(ptr);
            }
        }

        public void CloseDriver()
        {
            INSTRUCTIONS instructions = new INSTRUCTIONS
            {
                close = 1,
                read = 0,
                write = 0,
                reqBase = 0,
                processId = 0,
                address = 0,
                size = 0,
                bufferAddress = UIntPtr.Zero,
                output = UIntPtr.Zero,
                baseAddress = 0,
                moduleName = string.Empty,
                data = new byte[1024]
            };

            SendInstructions(instructions);
        }

        public void GetModuleBaseAddress(string moduleName)
        {
            INSTRUCTIONS instructions = new INSTRUCTIONS
            {
                close = 0,
                read = 0,
                write = 0,
                reqBase = 1,
                processId = 0,
                address = 0,
                size = 0,
                bufferAddress = UIntPtr.Zero,
                output = UIntPtr.Zero,
                baseAddress = 0,
                moduleName = moduleName,
                data = new byte[1024]
            };

            SendInstructions(instructions);
        }

        public void Dispose()
        {
            _client.Dispose();
        }
    }

    class Program
    {
        static void Main(string[] args)
        {
            using (KernelCom kernelCom = new KernelCom())
            {
                // Initialize WNF subscription
                kernelCom.Initialize();

                Console.WriteLine("[>] Subscribed to WNF state. Sending 'Hi' to kernel-mode.");

                // Send 'Hi' message to kernel-mode
                kernelCom.SendMessage("Hi");

                Console.WriteLine("[>] Waiting for kernel-mode responses...");

                // Start the interactive menu in a separate thread to keep listening for WNF notifications
                Thread menuThread = new Thread(() => InteractiveMenu(kernelCom));
                menuThread.Start();

                // Keep the main thread alive to receive WNF notifications
                menuThread.Join();
            }
        }

        static void InteractiveMenu(KernelCom kernelCom)
        {
            while (true)
            {
                Console.WriteLine("\nSelect an option:");
                Console.WriteLine("1. Read Memory");
                Console.WriteLine("2. Write Memory");
                Console.WriteLine("3. Get Module Base Address");
                Console.WriteLine("4. Close Driver");
                Console.WriteLine("5. Exit");

                Console.Write("Choice: ");
                string choice = Console.ReadLine();

                switch (choice)
                {
                    case "1":
                        PerformReadMemory(kernelCom);
                        break;

                    case "2":
                        PerformWriteMemory(kernelCom);
                        break;

                    case "3":
                        PerformGetModuleBaseAddress(kernelCom);
                        break;

                    case "4":
                        kernelCom.CloseDriver();
                        Console.WriteLine("[>] Driver close instruction sent.");
                        break;

                    case "5":
                        Console.WriteLine("[>] Exiting application.");
                        return;

                    default:
                        Console.WriteLine("[!] Invalid choice. Please select a valid option.");
                        break;
                }

                // Allow some time for the kernel driver to respond
                Thread.Sleep(500);
            }
        }

        static void PerformReadMemory(KernelCom kernelCom)
        {
            Console.Write("Enter Process ID: ");
            if (uint.TryParse(Console.ReadLine(), out uint readPid))
            {
                Console.Write("Enter Address (hex, e.g., 7FF6C2D40000): ");
                if (ulong.TryParse(Console.ReadLine(), System.Globalization.NumberStyles.HexNumber, null, out ulong readAddress))
                {
                    Console.Write("Enter Size (in bytes): ");
                    if (ulong.TryParse(Console.ReadLine(), out ulong readSize))
                    {
                        kernelCom.ReadMemory(readPid, readAddress, readSize);
                    }
                    else
                    {
                        Console.WriteLine("[!] Invalid size. Please enter a numeric value.");
                    }
                }
                else
                {
                    Console.WriteLine("[!] Invalid address format. Please enter a hexadecimal value.");
                }
            }
            else
            {
                Console.WriteLine("[!] Invalid Process ID. Please enter a numeric value.");
            }
        }

        static void PerformWriteMemory(KernelCom kernelCom)
        {
            Console.Write("Enter Process ID: ");
            if (uint.TryParse(Console.ReadLine(), out uint writePid))
            {
                Console.Write("Enter Address (hex, e.g., 7FF6C2D40000): ");
                if (ulong.TryParse(Console.ReadLine(), System.Globalization.NumberStyles.HexNumber, null, out ulong writeAddress))
                {
                    Console.Write("Enter Data to Write (hex, e.g., DEADBEEF): ");
                    string dataInput = Console.ReadLine();

                    try
                    {
                        byte[] dataBytes = StringToByteArray(dataInput);
                        kernelCom.WriteMemory(writePid, writeAddress, dataBytes);
                    }
                    catch
                    {
                        Console.WriteLine("[!] Invalid data format. Please enter hexadecimal digits only.");
                    }
                }
                else
                {
                    Console.WriteLine("[!] Invalid address format. Please enter a hexadecimal value.");
                }
            }
            else
            {
                Console.WriteLine("[!] Invalid Process ID. Please enter a numeric value.");
            }
        }

        static void PerformGetModuleBaseAddress(KernelCom kernelCom)
        {
            Console.Write("Enter Process ID: ");
            if (uint.TryParse(Console.ReadLine(), out uint pid))
            {
                Console.Write("Enter Module Name (e.g., notepad.exe): ");
                string moduleName = Console.ReadLine();

                if (!string.IsNullOrWhiteSpace(moduleName))
                {
                    // Update the instructions to include the process ID
                    INSTRUCTIONS instructions = new INSTRUCTIONS
                    {
                        close = 0,
                        read = 0,
                        write = 0,
                        reqBase = 1,
                        processId = pid,
                        address = 0,
                        size = 0,
                        bufferAddress = UIntPtr.Zero,
                        output = UIntPtr.Zero,
                        baseAddress = 0,
                        moduleName = moduleName,
                        data = new byte[1024]
                    };

                    kernelCom.SendInstructions(instructions);
                }
                else
                {
                    Console.WriteLine("[!] Module name cannot be empty.");
                }
            }
            else
            {
                Console.WriteLine("[!] Invalid Process ID. Please enter a numeric value.");
            }
        }

        // Helper method to convert hex string to byte array
        static byte[] StringToByteArray(string hex)
        {
            hex = hex.Replace("0x", "").Replace(" ", "").Replace("-", "");
            if (hex.Length % 2 != 0)
                throw new FormatException("Invalid hex string length.");

            byte[] bytes = new byte[hex.Length / 2];
            for (int i = 0; i < bytes.Length; i++)
            {
                string byteValue = hex.Substring(i * 2, 2);
                bytes[i] = Convert.ToByte(byteValue, 16);
            }
            return bytes;
        }
    }
}
km
include <ntifs.h>
#include <ntddk.h>

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

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

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

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

// Custom type definitions
typedef unsigned __int64 QWORD;

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

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

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

typedef NTSTATUS(EXUNSUBSCRIBEWNFSTATECHANGE)(
PVOID Subscription
);

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

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

ULONG64 MY_WNF_STATE_NAME_VALUE = MY_WNF_STATE_NAME; // Ensure consistency

PVOID g_WnfSubscription = NULL; // Global WNF subscription handle

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

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

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

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

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

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

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

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

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

// Driver Unload
VOID DriverUnloadRoutine(PDRIVER_OBJECT DriverObject);

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

// Helper Functions
PVOID GetProcessBaseAddress(HANDLE pid);

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

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

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

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

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

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

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

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

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

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

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

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

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

    SIZE_T readsize = 0;
    uint64_t pdpe = 0;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    decodedSize = 0;

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

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

        UCHAR triple[3] = { 0 };

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

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

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

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

    return STATUS_SUCCESS;
}

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

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

    encodedSize = 0;

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

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

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

    return STATUS_SUCCESS;
}

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

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

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

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

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

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

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

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

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

    NTSTATUS status = STATUS_SUCCESS;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return status;
}

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

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

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

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

    return status;
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    NTSTATUS status;

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

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

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

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

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

        return status;
    }

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

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

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

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

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

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

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

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

    return status;
}

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

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

    NTSTATUS status = STATUS_SUCCESS;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return status;
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    NTSTATUS status;

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

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

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

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

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

        return status;
    }

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

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

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

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