#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  // Replace with your unique value

// 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 19569
#define WINDOWS_21H1 20180

// Page size definitions
#define PAGE_OFFSET_SIZE 12
#define PMASK (0xfull << 8) & 0xFFFFFFFFfull

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

typedef NTSTATUS (PExSubscribeWnfStateChange)(
    PWNF_STATE_NAME StateName,
    WNF_CHANGE_STAMP ChangeStamp,
    WNF_SUBSCRIPTION_FLAG SubscriptionFlag,
    PVOID Callback,
    PVOID CallbackContext,
    PCWNF_TYPE_ID TypeId,
    PVOID DeliveryDescriptor,
    ULONG DeliveryDescriptorSize
);

typedef NTSTATUS (PExPublishWnfStateData)(
    PCWNF_STATE_NAME StateName,
    PVOID Buffer,
    ULONG Length,
    PCWNF_TYPE_ID TypeId,
    PVOID ExplicitScope,
    ULONG MatchingChangeStamp
);

// Global function pointers
PExPublishWnfStateData ExPublishWnfStateDataFunc = NULL;

// ==========================================================
// Data Structures
// ==========================================================

typedef struct _INSTRUCTIONS {
    BOOLEAN close;
    BOOLEAN read;
    BOOLEAN reqBase;
    PVOID bufferAddress;
    UINT_PTR address;
    ULONGLONG size;
    PVOID output;
    ULONG64 baseAddress;
    const char moduleName;
} INSTRUCTIONS, PINSTRUCTIONS;

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

NTKERNELAPI
PVOID
PsGetProcessSectionBaseAddress(
    __in PEPROCESS Process
);

PVOID GetProcessBaseAddress(HANDLE pid);
DWORD GetUserDirectoryTableBaseOffset();
ULONG_PTR GetProcessCr3(PEPROCESS pProcess);
ULONG_PTR GetKernelDirBase();
NTSTATUS ReadVirtual(uint64_t dirbase, uint64_t address, uint8_t buffer, SIZE_T size, SIZE_T read);
NTSTATUS WriteVirtual(uint64_t dirbase, uint64_t address, uint8_t* buffer, SIZE_T size, SIZE_T* written);
NTSTATUS ReadPhysicalAddress(PVOID TargetAddress, PVOID lpBuffer, SIZE_T Size, SIZE_T* BytesRead);
NTSTATUS WritePhysicalAddress(PVOID TargetAddress, PVOID lpBuffer, SIZE_T Size, SIZE_T* BytesWritten);
uint64_t TranslateLinearAddress(uint64_t directoryTableBase, uint64_t virtualAddress);
NTSTATUS ReadProcessMemory(int pid, PVOID Address, PVOID AllocatedBuffer, SIZE_T size, SIZE_T* read);
NTSTATUS WriteProcessMemory(int pid, PVOID Address, PVOID AllocatedBuffer, SIZE_T size, SIZE_T* written);
VOID ProcessNotifyCallback(
    __in HANDLE ParentId,
    __in HANDLE ProcessId,
    __in BOOLEAN Create
);
VOID DriverUnload(PDRIVER_OBJECT DriverObject);

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

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

    NTSTATUS NtRet = PsLookupProcessByProcessId(pid, &pProcess);
    if (NtRet != STATUS_SUCCESS) return NULL;

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

// Get the offset for the 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:
    case WINDOWS_2004:
    case WINDOWS_20H2:
    case WINDOWS_21H1:
    default:
        return 0x0388;
    }
}

// Retrieve the CR3 (Page Directory Base) of a process
ULONG_PTR GetProcessCr3(PEPROCESS pProcess)
{
    PUCHAR process = (PUCHAR)pProcess;
#ifdef _WIN64
    ULONG_PTR process_dirbase = (PULONG_PTR)(process + 0x28); // dirbase for x64
#else
    ULONG_PTR process_dirbase = (PULONG_PTR)(process + 0x18); // dirbase for x86
#endif
    if (process_dirbase == 0)
    {
        DWORD UserDirOffset = GetUserDirectoryTableBaseOffset();
        ULONG_PTR process_userdirbase = (PULONG_PTR)(process + UserDirOffset);
        return process_userdirbase;
    }
    return process_dirbase;
}

// Example function to get the kernel's CR3 (for demonstration; typically not used)
ULONG_PTR GetKernelDirBase()
{
    // Note: Accessing the kernel's CR3 is generally unsafe and not recommended
    // This is just for illustrative purposes
    return __readcr3();
}

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

    uint64_t pageOffset = virtualAddress & 0xFFF;
    uint64_t pte = (virtualAddress >> 12) & 0x1FF;
    uint64_t pt = (virtualAddress >> 21) & 0x1FF;
    uint64_t pd = (virtualAddress >> 30) & 0x1FF;
    uint64_t pdp = (virtualAddress >> 39) & 0x1FF;

    SIZE_T readsize = 0;
    uint64_t pdpe = 0;
    if (!NT_SUCCESS(ReadPhysicalAddress((PVOID)(directoryTableBase + 8 * pdp), &pdpe, sizeof(pdpe), &readsize)) || (pdpe & 1))
        return 0;

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

    // 1GB large page
    if (pde & 0x80)
        return (pde & 0x000FFFFFFFF000ULL) + (virtualAddress & 0x3FFFFFFFUL);

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

    // 2MB large page
    if (pteAddr & 0x80)
        return (pteAddr & PMASK) + (virtualAddress & 0x1FFFFFUL);

    uint64_t physPage = 0;
    if (!NT_SUCCESS(ReadPhysicalAddress((PVOID)(pteAddr & PMASK) + 8 * pte, &physPage, sizeof(physPage), &readsize)) || !(physPage & 1))
        return 0;

    physPage &= PMASK;
    return physPage + pageOffset;
}

// Read a physical address into a buffer
NTSTATUS ReadPhysicalAddress(PVOID TargetAddress, PVOID lpBuffer, SIZE_T Size, SIZE_T BytesRead)
{
    if (TargetAddress == NULL || lpBuffer == NULL || BytesRead == NULL)
        return STATUS_INVALID_PARAMETER;

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

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

// Write a buffer to a physical address
NTSTATUS WritePhysicalAddress(PVOID TargetAddress, PVOID lpBuffer, SIZE_T Size, SIZE_T BytesWritten)
{
    if (!TargetAddress || !lpBuffer || !BytesWritten)
        return STATUS_INVALID_PARAMETER;

    PHYSICAL_ADDRESS AddrToWrite = { 0 };
    AddrToWrite.QuadPart = (ULONG_PTR)TargetAddress;

    // Limit MmMapIoSpaceEx to PAGE_SIZE (4096 bytes)
    if (Size > PAGE_SIZE)
        return STATUS_INVALID_PARAMETER;

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

// Read virtual memory from a process
NTSTATUS ReadVirtual(uint64_t dirbase, uint64_t address, uint8_t buffer, SIZE_T size, SIZE_T read)
{
    if (!buffer || !read)
        return STATUS_INVALID_PARAMETER;

    uint64_t paddress;
    NTSTATUS status = STATUS_SUCCESS;
    SIZE_T CurOffset = 0;
    SIZE_T TotalSize = size;

    while (TotalSize > 0)
    {
        paddress = TranslateLinearAddress(dirbase, address + CurOffset);
        if (!paddress)
            return STATUS_UNSUCCESSFUL;

        ULONG64 ReadSize = min(PAGE_SIZE - (paddress & 0xFFF), TotalSize);
        SIZE_T BytesRead = 0;

        status = ReadPhysicalAddress((PVOID)paddress, buffer + CurOffset, ReadSize, &BytesRead);
        if (!NT_SUCCESS(status) || BytesRead == 0)
            break;

        TotalSize -= BytesRead;
        CurOffset += BytesRead;
    }

    read = CurOffset;
    return status;
}

// Write virtual memory to a process
NTSTATUS WriteVirtual(uint64_t dirbase, uint64_t address, uint8_t buffer, SIZE_T size, SIZE_T* written)
{
    if (!buffer || !written)
        return STATUS_INVALID_PARAMETER;

    uint64_t paddress;
    NTSTATUS status = STATUS_SUCCESS;
    SIZE_T CurOffset = 0;
    SIZE_T TotalSize = size;

    while (TotalSize > 0)
    {
        paddress = TranslateLinearAddress(dirbase, address + CurOffset);
        if (!paddress)
            return STATUS_UNSUCCESSFUL;

        ULONG64 WriteSize = min(PAGE_SIZE - (paddress & 0xFFF), TotalSize);
        SIZE_T BytesWrittenLocal = 0;

        status = WritePhysicalAddress((PVOID)paddress, buffer + CurOffset, WriteSize, &BytesWrittenLocal);
        if (!NT_SUCCESS(status) || BytesWrittenLocal == 0)
            break;

        TotalSize -= BytesWrittenLocal;
        CurOffset += BytesWrittenLocal;
    }

    written = CurOffset;
    return status;
}

// Read the memory of a process by PID
NTSTATUS ReadProcessMemory(int pid, PVOID Address, PVOID AllocatedBuffer, SIZE_T size, SIZE_T read)
{
    if (pid == 0 || Address == NULL || AllocatedBuffer == NULL || read == NULL)
        return STATUS_INVALID_PARAMETER;

    PEPROCESS pProcess = NULL;
    NTSTATUS NtRet = PsLookupProcessByProcessId((HANDLE)(ULONG_PTR)pid, &pProcess);
    if (NtRet != STATUS_SUCCESS)
        return NtRet;

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

    NtRet = ReadVirtual(process_dirbase, (ULONG_PTR)Address, (uint8_t*)AllocatedBuffer, size, read);
    return NtRet;
}

// Write the memory of a process by PID
NTSTATUS WriteProcessMemory(int pid, PVOID Address, PVOID AllocatedBuffer, SIZE_T size, SIZE_T* written)
{
    if (pid == 0 || Address == NULL || AllocatedBuffer == NULL || written == NULL)
        return STATUS_INVALID_PARAMETER;

    PEPROCESS pProcess = NULL;
    NTSTATUS NtRet = PsLookupProcessByProcessId((HANDLE)(ULONG_PTR)pid, &pProcess);
    if (NtRet != STATUS_SUCCESS)
        return NtRet;

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

    NtRet = WriteVirtual(process_dirbase, (ULONG_PTR)Address, (uint8_t*)AllocatedBuffer, size, written);
    return NtRet;
}

// Callback function for process creation/termination
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
            WCHAR imageName[300] = { 0 };
            UNICODE_STRING uProcessImageName = { 0 };

            // Use RtlQueryInformationProcess if SeLocateProcessImageName is unavailable
            status = SeLocateProcessImageName(pProcess, &uProcessImageName);
            if (NT_SUCCESS(status))
            {
                // Extract file name from full path
                PWSTR processName = wcsrchr(uProcessImageName.Buffer, L'\');
                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

                        // Get base address
                        PVOID baseAddress = GetProcessBaseAddress(ProcessId);
                        if (baseAddress)
                        {
                            DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_INFO_LEVEL, "[Driver] Base address of %ws: %p\n", processName, baseAddress);

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

                            if (NT_SUCCESS(ntStatus))
                            {
                                // Do something with the read data
                                DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_INFO_LEVEL, "[Driver] Successfully read %zu bytes from %ws\n", bytesRead, processName);
                            }
                            else
                            {
                                DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Driver] Failed to read process memory: 0x%X\n", ntStatus);
                            }

                            // Example: Write to process memory (be cautious with actual addresses)
                            /
                            UCHAR writeData[256] = { / ... * / };
                            SIZE_T bytesWritten = 0;
                            ntStatus = WriteProcessMemory((int)(ULONG_PTR)ProcessId, baseAddress, writeData, sizeof(writeData), &bytesWritten);

                            if (NT_SUCCESS(ntStatus))
                            {
                                DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_INFO_LEVEL, "[Driver] Successfully wrote %zu bytes to %ws\n", bytesWritten, processName);
                            }
                            else
                            {
                                DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Driver] Failed to write process memory: 0x%X\n", ntStatus);
                            }
                            */
                        }
                        else
                        {
                            DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Driver] Failed to get base address of %ws\n", processName);
                        }
                    }
                }
                else
                {
                    DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Driver] Failed to parse process name\n");
                }

                // Free the UNICODE_STRING allocated by SeLocateProcessImageName
                ExFreePoolWithTag(uProcessImageName.Buffer, 'imgN');
            }
            else
            {
                DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Driver] SeLocateProcessImageName failed with status: 0x%X\n", status);
            }

            ObDereferenceObject(pProcess);
        }
        else
        {
            DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Driver] PsLookupProcessByProcessId failed with status: 0x%X\n", status);
        }
    }
}

// Driver unload routine
VOID DriverUnload(PDRIVER_OBJECT DriverObject)
{
    UNREFERENCED_PARAMETER(DriverObject);

    // Remove the process notify routine
    NTSTATUS status = PsSetCreateProcessNotifyRoutine(ProcessNotifyCallback, TRUE);
    if (!NT_SUCCESS(status))
    {
        DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Driver] Failed to remove process notify routine: 0x%X\n", status);
    }
    else
    {
        DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_INFO_LEVEL, "[Driver] Process notify routine removed successfully\n");
    }

    DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_INFO_LEVEL, "[Driver] Driver unloaded\n");
}

// ==========================================================
// WNF Functions and DriverEntry Implementation
// ==========================================================

// Function to initialize WNF communication
NTSTATUS InitializeWNF()
{
    // Dynamically resolve the ExPublishWnfStateData function
    UNICODE_STRING functionName;
    RtlInitUnicodeString(&functionName, L"ExPublishWnfStateData");
    ExPublishWnfStateDataFunc = (PExPublishWnfStateData)MmGetSystemRoutineAddress(&functionName);

    if (ExPublishWnfStateDataFunc == NULL)
    {
        DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Driver] Failed to resolve ExPublishWnfStateData\n");
        return STATUS_UNSUCCESSFUL;
    }

    // Prepare your data
    INSTRUCTIONS instructions = { 0 };
    instructions.reqBase = TRUE;
    instructions.baseAddress = 0x12345678;  // Example base address
    instructions.output = (PVOID)instructions.baseAddress;

    // Publish the data to the WNF state
    NTSTATUS status = ExPublishWnfStateDataFunc(
        (PCWNF_STATE_NAME)&MY_WNF_STATE_NAME,
        &instructions,
        sizeof(INSTRUCTIONS),
        NULL,
        NULL,
        0
    );

    if (!NT_SUCCESS(status))
    {
        DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Driver] ExPublishWnfStateData failed: 0x%08X\n", status);
        return status;
    }

    DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_INFO_LEVEL, "[Driver] Published WNF state data successfully\n");
    return STATUS_SUCCESS;
}

// Main DriverEntry function
NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
{
    UNREFERENCED_PARAMETER(RegistryPath);

    NTSTATUS status;

    // Register the process creation notify routine
    status = PsSetCreateProcessNotifyRoutine(ProcessNotifyCallback, FALSE);
    if (!NT_SUCCESS(status))
    {
        DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Driver] PsSetCreateProcessNotifyRoutine failed: 0x%X\n", status);
        return status;
    }
    else
    {
        DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_INFO_LEVEL, "[Driver] Process notify routine set successfully\n");
    }

    // Initialize WNF communication
    status = InitializeWNF();
    if (!NT_SUCCESS(status))
    {
        // If WNF initialization fails, unload the driver
        PsSetCreateProcessNotifyRoutine(ProcessNotifyCallback, TRUE);
        DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Driver] WNF initialization failed. Driver will unload.\n");
        return status;
    }

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

    DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_INFO_LEVEL, "[Driver] Driver loaded successfully\n");
    return STATUS_SUCCESS;
}

make a user mode for this