#include <ntifs.h>
#include <ntddk.h>

// Include your provided functions here

NTKERNELAPI
PVOID
PsGetProcessSectionBaseAddress(
    __in PEPROCESS Process
);

PVOID GetProcessBaseAddress(int pid)
{
    PEPROCESS pProcess = NULL;
    if (pid == 0) return NULL;

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

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

// ... (Include the rest of your functions here)

#define TARGET_PROCESS_NAME L"notepad.exe" // Replace with your target process name

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

            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((int)(ULONG_PTR)ProcessId);
                        if (baseAddress)
                        {
                            DbgPrintEx(0, 0, "[Driver] Base address of %ws: %p\n", processName, baseAddress);

                            // 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(0, 0, "[Driver] Successfully read %zu bytes from %ws\n", bytesRead, processName);
                            }
                            else
                            {
                                DbgPrintEx(0, 0, "[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(0, 0, "[Driver] Successfully wrote %zu bytes to %ws\n", bytesWritten, processName);
                            // }
                            // else
                            // {
                            //     DbgPrintEx(0, 0, "[Driver] Failed to write process memory: 0x%X\n", ntStatus);
                            // }
                        }
                        else
                        {
                            DbgPrintEx(0, 0, "[Driver] Failed to get base address of %ws\n", processName);
                        }
                    }
                }
                else
                {
                    DbgPrintEx(0, 0, "[Driver] Failed to parse process name\n");
                }
                // Free the UNICODE_STRING allocated by SeLocateProcessImageName
                ExFreePool(uProcessImageName.Buffer);
            }
            else
            {
                DbgPrintEx(0, 0, "[Driver] SeLocateProcessImageName failed with status: 0x%X\n", status);
            }
            ObDereferenceObject(pProcess);
        }
        else
        {
            DbgPrintEx(0, 0, "[Driver] PsLookupProcessByProcessId failed with status: 0x%X\n", status);
        }
    }
}

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

    PsSetCreateProcessNotifyRoutine(ProcessNotifyCallback, TRUE);
    DbgPrintEx(0, 0, "[Driver] Driver unloaded\n");
}

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

    NTSTATUS status = PsSetCreateProcessNotifyRoutine(ProcessNotifyCallback, FALSE);
    if (!NT_SUCCESS(status))
    {
        DbgPrintEx(0, 0, "[Driver] PsSetCreateProcessNotifyRoutine failed: 0x%X\n", status);
        return status;
    }
    else
    {
        DbgPrintEx(0, 0, "[Driver] Process notify routine set successfully\n");
    }

    DriverObject->DriverUnload = DriverUnload;

    return STATUS_SUCCESS;
}

#include <ntifs.h>
#include <ntddk.h>

// Define a unique WNF State Name for your communication (this is just an example)
#define MY_WNF_STATE_NAME 0x41C64E6DA3BC1074ULL  // Replace with your unique value

// Function pointer types for WNF functions (since they are undocumented)
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 variables for function pointers
PExPublishWnfStateData ExPublishWnfStateData = NULL;

// Your data structure
typedef struct _INSTRUCTIONS {
    BOOLEAN close;
    BOOLEAN read;
    BOOLEAN reqBase;
    void bufferAddress;
    UINT_PTR address;
    ULONGLONG size;
    void output;
    ULONG64 baseAddress;
    const char* moduleName;
} INSTRUCTIONS, PINSTRUCTIONS;

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

    // Since we cannot use driver objects, we won't set any major functions

    // Dynamically resolve the ExPublishWnfStateData function
    UNICODE_STRING functionName;
    RtlInitUnicodeString(&functionName, L"ExPublishWnfStateData");
    ExPublishWnfStateData = (PExPublishWnfStateData)MmGetSystemRoutineAddress(&functionName);

    if (ExPublishWnfStateData == NULL)
    {
        DbgPrintEx(0, 0, "[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 = (void)instructions.baseAddress;

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

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

    DbgPrintEx(0, 0, "[Driver] Published WNF state data successfully\n");

    // Since we don't use system threads or further driver operations, we can return
    return STATUS_SUCCESS;
}


complete the implemitation


NTKERNELAPI
PVOID
PsGetProcessSectionBaseAddress(
__in PEPROCESS Process
);

PVOID GetProcessBaseAddress(int pid)
{
PEPROCESS pProcess = NULL;
if (pid == 0) return STATUS_UNSUCCESSFUL;

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

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

//https://ntdiff.github.io/
#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

DWORD GetUserDirectoryTableBaseOffset()
{
RTL_OSVERSIONINFOW ver = { 0 };
RtlGetVersion(&ver);

switch (ver.dwBuildNumber)
{
case WINDOWS_1803:
return 0x0278;
break;
case WINDOWS_1809:
return 0x0278;
break;
case WINDOWS_1903:
return 0x0280;
break;
case WINDOWS_1909:
return 0x0280;
break;
case WINDOWS_2004:
return 0x0388;
break;
case WINDOWS_20H2:
return 0x0388;
break;
case WINDOWS_21H1:
return 0x0388;
break;
default:
return 0x0388;
}
}

//check normal dirbase if 0 then get from UserDirectoryTableBas
ULONG_PTR GetProcessCr3(PEPROCESS pProcess)
{
PUCHAR process = (PUCHAR)pProcess;
ULONG_PTR process_dirbase = *(PULONG_PTR)(process + 0x28); //dirbase x64, 32bit is 0x18
if (process_dirbase==0)
{
DWORD UserDirOffset = GetUserDirectoryTableBaseOffset();
ULONG_PTR process_userdirbase = *(PULONG_PTR)(process + UserDirOffset);
return process_userdirbase;
}
return process_dirbase;
}
ULONG_PTR GetKernelDirBase()
{
PUCHAR process = (PUCHAR)PsGetCurrentProcess();
ULONG_PTR cr3 = (PULONG_PTR)(process + 0x28); //dirbase x64, 32bit is 0x18
return cr3;
}

NTSTATUS ReadVirtual(uint64_t dirbase, uint64_t address, uint8_t buffer, SIZE_T size, SIZE_T read)
{
uint64_t paddress = TranslateLinearAddress(dirbase, address);
return ReadPhysicalAddress(paddress, buffer, size, read);
}

NTSTATUS WriteVirtual(uint64_t dirbase, uint64_t address, uint8_t buffer, SIZE_T size, SIZE_T *written)
{
uint64_t paddress = TranslateLinearAddress(dirbase, address);
return WritePhysicalAddress(paddress, buffer, size, written);
}

NTSTATUS ReadPhysicalAddress(PVOID TargetAddress, PVOID lpBuffer, SIZE_T Size, SIZE_T BytesRead)
{
MM_COPY_ADDRESS AddrToRead = { 0 };
AddrToRead.PhysicalAddress.QuadPart = TargetAddress;
return MmCopyMemory(lpBuffer, AddrToRead, Size, MM_COPY_MEMORY_PHYSICAL, BytesRead);
}

//MmMapIoSpaceEx limit is page 4096 byte
NTSTATUS WritePhysicalAddress(PVOID TargetAddress, PVOID lpBuffer, SIZE_T Size, SIZE_T BytesWritten)
{
if (!TargetAddress)
return STATUS_UNSUCCESSFUL;

PHYSICAL_ADDRESS AddrToWrite = { 0 };
AddrToWrite.QuadPart = TargetAddress;

PVOID pmapped_mem = MmMapIoSpaceEx(AddrToWrite, Size, PAGE_READWRITE);

if (!pmapped_mem)
return STATUS_UNSUCCESSFUL;

memcpy(pmapped_mem, lpBuffer, Size);

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

#define PAGE_OFFSET_SIZE 12
static const uint64_t PMASK = (0xfull << 8) & 0xfffffffffull;

uint64_t TranslateLinearAddress(uint64_t directoryTableBase, uint64_t virtualAddress) {
directoryTableBase &= ~0xf;

uint64_t pageOffset = virtualAddress & ~(0ul << PAGE_OFFSET_SIZE);
uint64_t pte = ((virtualAddress >> 12) & (0x1ffll));
uint64_t pt = ((virtualAddress >> 21) & (0x1ffll));
uint64_t pd = ((virtualAddress >> 30) & (0x1ffll));
uint64_t pdp = ((virtualAddress >> 39) & (0x1ffll));

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

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

/ 1GB large page, use pde's 12-34 bits /
if (pde & 0x80)
return (pde & (0ull << 42 >> 12)) + (virtualAddress & ~(0ull << 30));

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

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

virtualAddress = 0;
ReadPhysicalAddress((pteAddr & PMASK) + 8 * pte, &virtualAddress, sizeof(virtualAddress), &readsize);
virtualAddress &= PMASK;

if (!virtualAddress)
return 0;

return virtualAddress + pageOffset;
}


//
NTSTATUS ReadProcessMemory(int pid, PVOID Address, PVOID AllocatedBuffer, SIZE_T size, SIZE_T read)
{
PEPROCESS pProcess = NULL;
if (pid == 0) return STATUS_UNSUCCESSFUL;

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

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

SIZE_T CurOffset = 0;
SIZE_T TotalSize = size;
while (TotalSize)
{

uint64_t CurPhysAddr = TranslateLinearAddress(process_dirbase, (ULONG64)Address + CurOffset);
if (!CurPhysAddr) return STATUS_UNSUCCESSFUL;

ULONG64 ReadSize = min(PAGE_SIZE - (CurPhysAddr & 0xFFF), TotalSize);
SIZE_T BytesRead = 0;
NtRet = ReadPhysicalAddress(CurPhysAddr, (PVOID)((ULONG64)AllocatedBuffer + CurOffset), ReadSize, &BytesRead);
TotalSize -= BytesRead;
CurOffset += BytesRead;
if (NtRet != STATUS_SUCCESS) break;
if(BytesRead==0) break;
}

read = CurOffset;
return NtRet;
}

NTSTATUS WriteProcessMemory(int pid, PVOID Address, PVOID AllocatedBuffer, SIZE_T size, SIZE_T written)
{
PEPROCESS pProcess = NULL;
if (pid == 0) return STATUS_UNSUCCESSFUL;

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

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

SIZE_T CurOffset = 0;
SIZE_T TotalSize = size;
while (TotalSize)
{
uint64_t CurPhysAddr = TranslateLinearAddress(process_dirbase, (ULONG64)Address + CurOffset);
if (!CurPhysAddr) return STATUS_UNSUCCESSFUL;

ULONG64 WriteSize = min(PAGE_SIZE - (CurPhysAddr & 0xFFF), TotalSize);
SIZE_T BytesWritten = 0;
NtRet = WritePhysicalAddress(CurPhysAddr, (PVOID)((ULONG64)AllocatedBuffer + CurOffset), WriteSize, &BytesWritten);
TotalSize -= BytesWritten;
CurOffset += BytesWritten;
if (NtRet != STATUS_SUCCESS) break;
if(BytesWritten==0) break;
}

*written = CurOffset;
return NtRet;
}