#include <windows.h>
#include <winternl.h>
#include <array>
#include <iostream>

static void* NtOpenProcessAddress{ GetProcAddress(GetModuleHandleA("ntdll.dll"), "ZwOpenProcess") };

LONG WINAPI ExceptionHandler(EXCEPTION_POINTERS* const exceptionInfo)
{
    if (exceptionInfo->ExceptionRecord->ExceptionCode == STATUS_GUARD_PAGE_VIOLATION)
    {
        if (exceptionInfo->ExceptionRecord->ExceptionAddress == NtOpenProcessAddress)
        {
            printf("NtOpenProcess called \n");
        }

        // Set single step flag so that memory breakpoints are re-enabled
        // on the next instruction execution.
        exceptionInfo->ContextRecord->EFlags |= 0x100;

        return EXCEPTION_CONTINUE_EXECUTION;
    }

    if (exceptionInfo->ExceptionRecord->ExceptionCode == EXCEPTION_SINGLE_STEP) {

        // Re-enable memory breakpoint since a different address might
        // have caused the guard page violation.
        std::array<unsigned char, 11> NtProtectVirtualMemoryBytes =
        {
            0x4C, 0x8B, 0xD1,               /*mov r10, rcx*/
            0xB8, 0x50, 0x00, 0x00, 0x00,   /*mov eax, 0x50*/
            0x0F, 0x05,                     /*syscall*/
            0xC3                            /*ret*/
        };

        auto* NtProtectVirtualMemoryStub{ VirtualAlloc(nullptr,
            NtProtectVirtualMemoryBytes.size(),
            MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE) };

        memcpy(NtProtectVirtualMemoryStub, NtProtectVirtualMemoryBytes.data(), NtProtectVirtualMemoryBytes.size());

        using NtProtectVirtualMemoryFnc = NTSTATUS(NTAPI*)(HANDLE, PVOID*, SIZE_T*, ULONG, PULONG);
        auto NtProtectVirtualMemory{
            reinterpret_cast<NtProtectVirtualMemoryFnc>(NtProtectVirtualMemoryStub) };

        DWORD oldPermissions{};
        SIZE_T allocSize = 4096;
        auto result{ NtProtectVirtualMemory(GetCurrentProcess(), &NtOpenProcessAddress, &allocSize, PAGE_WRITECOPY | PAGE_GUARD, &oldPermissions) };

        return EXCEPTION_CONTINUE_EXECUTION;
    }

    return EXCEPTION_CONTINUE_SEARCH;
}

BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:
        AllocConsole();

        // Add a custom exception handler
        AddVectoredExceptionHandler(true, ExceptionHandler);

        // Set Breakpoint
        std::array<unsigned char, 11> NtProtectVirtualMemoryBytes =
        {
            0x4C, 0x8B, 0xD1,               /*mov r10, rcx*/
            0xB8, 0x50, 0x00, 0x00, 0x00,   /*mov eax, 0x50*/
            0x0F, 0x05,                     /*syscall*/
            0xC3                            /*ret*/
        };

        auto* NtProtectVirtualMemoryStub{ VirtualAlloc(nullptr,
            NtProtectVirtualMemoryBytes.size(),
            MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE) };

        memcpy(NtProtectVirtualMemoryStub, NtProtectVirtualMemoryBytes.data(),
            NtProtectVirtualMemoryBytes.size());

        using NtProtectVirtualMemoryFnc = NTSTATUS(NTAPI*)(HANDLE, PVOID*, SIZE_T*, ULONG, PULONG);
        auto NtProtectVirtualMemory{
            reinterpret_cast<NtProtectVirtualMemoryFnc>(NtProtectVirtualMemoryStub) };

        DWORD oldPermissions{};
        SIZE_T allocSize = 4096;
        auto result{ NtProtectVirtualMemory(GetCurrentProcess(), &NtOpenProcessAddress, &allocSize,
            PAGE_WRITECOPY | PAGE_GUARD, &oldPermissions) };

        printf("NtProtectVirtualMemory NTSTATUS: %016I64x\n", result);

        return TRUE;
    }

    return TRUE;
}