/*
 * scanner.c
 *
 * PE pattern scanner: maps a PE as a SEC_IMAGE section, scans for a byte
 * pattern with wildcard mask, and decodes RIP-relative displacements to
 * recover unexported kernel symbol addresses.
 *
 * Usage: scanner.exe  (hardcoded to scan ntoskrnl.exe and CI.dll)
 *
 * Compile (MSVC, x64):
 *   cl /W4 /O2 scanner.c
 */

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <winternl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* RTL_PROCESS_MODULES -- not in public SDK headers */
typedef struct _RTL_PROCESS_MODULE_INFORMATION {
    HANDLE  Section;
    PVOID   MappedBase;
    PVOID   ImageBase;
    ULONG   ImageSize;
    ULONG   Flags;
    USHORT  LoadOrderIndex;
    USHORT  InitOrderIndex;
    USHORT  LoadCount;
    USHORT  OffsetToFileName;
    CHAR    FullPathName[256];
} RTL_PROCESS_MODULE_INFORMATION;

typedef struct _RTL_PROCESS_MODULES {
    ULONG                          NumberOfModules;
    RTL_PROCESS_MODULE_INFORMATION Modules[1];
} RTL_PROCESS_MODULES, *PRTL_PROCESS_MODULES;

typedef NTSTATUS (NTAPI *PFN_NtQSI)(ULONG, PVOID, ULONG, PULONG);

/* ------------------------------------------------------------------ */
/* PE mapping                                                           */
/* ------------------------------------------------------------------ */

static PBYTE map_pe(const wchar_t *path) {
    HANDLE fh = CreateFileW(path, GENERIC_READ, FILE_SHARE_READ,
                            NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (fh == INVALID_HANDLE_VALUE) {
        fprintf(stderr, "[-] CreateFileW failed: %lu\n", GetLastError());
        return NULL;
    }
    HANDLE fm = CreateFileMappingW(fh, NULL, PAGE_READONLY | SEC_IMAGE, 0, 0, NULL);
    CloseHandle(fh);
    if (!fm) {
        fprintf(stderr, "[-] CreateFileMappingW (SEC_IMAGE) failed: %lu\n", GetLastError());
        return NULL;
    }
    PBYTE view = (PBYTE)MapViewOfFile(fm, FILE_MAP_READ, 0, 0, 0);
    CloseHandle(fm);
    return view;
}

/* ------------------------------------------------------------------ */
/* Pattern scanner                                                      */
/* ------------------------------------------------------------------ */

static BOOL data_compare(const BYTE *data, const BYTE *pattern, const char *mask) {
    for (; *mask; ++mask, ++data, ++pattern) {
        if (*mask == 'x' && *data != *pattern)
            return FALSE;
    }
    return TRUE;
}

static ULONG_PTR pattern_scan(PBYTE base, const BYTE *pattern, const char *mask) {
    PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)base;
    if (dos->e_magic != IMAGE_DOS_SIGNATURE) return 0;

    PIMAGE_NT_HEADERS nt = (PIMAGE_NT_HEADERS)(base + dos->e_lfanew);
    PIMAGE_SECTION_HEADER sec = IMAGE_FIRST_SECTION(nt);

    for (WORD i = 0; i < nt->FileHeader.NumberOfSections; i++, sec++) {
        if (!(sec->Characteristics & IMAGE_SCN_MEM_EXECUTE)) continue;
        PBYTE start = base + sec->VirtualAddress;
        DWORD size  = sec->Misc.VirtualSize ? sec->Misc.VirtualSize : sec->SizeOfRawData;
        for (DWORD j = 0; j < size; j++) {
            if (data_compare(start + j, pattern, mask))
                return (ULONG_PTR)(start + j);
        }
    }
    return 0;
}

/* ------------------------------------------------------------------ */
/* Kernel module base                                                   */
/* ------------------------------------------------------------------ */

static ULONG_PTR get_kernel_module_base(const char *module_name) {
    PFN_NtQSI NtQSI = (PFN_NtQSI)GetProcAddress(
        GetModuleHandleW(L"ntdll.dll"), "NtQuerySystemInformation");
    if (!NtQSI) return 0;

    ULONG    size = 0;
    PVOID    buf  = NULL;
    NTSTATUS st;

    NtQSI(11, NULL, 0, &size);
    if (!size) size = 0x10000;

    for (;;) {
        if (buf) VirtualFree(buf, 0, MEM_RELEASE);
        buf = VirtualAlloc(NULL, size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
        if (!buf) return 0;
        st = NtQSI(11, buf, size, &size);
        if (st != (NTSTATUS)0xC0000004L) break;
        size += 0x1000;
    }

    if (st != 0) { VirtualFree(buf, 0, MEM_RELEASE); return 0; }

    PRTL_PROCESS_MODULES mods = (PRTL_PROCESS_MODULES)buf;
    ULONG_PTR result = 0;
    for (ULONG i = 0; i < mods->NumberOfModules; i++) {
        const char *fname = (const char *)mods->Modules[i].FullPathName
                          + mods->Modules[i].OffsetToFileName;
        if (!_stricmp(fname, module_name)) {
            result = (ULONG_PTR)mods->Modules[i].ImageBase;
            break;
        }
    }
    VirtualFree(buf, 0, MEM_RELEASE);
    return result;
}

/* ------------------------------------------------------------------ */
/* Decode RIP-relative MOV RAX, [RIP+disp32]                           */
/* Opcode: 48 8B 05 dd dd dd dd  (7 bytes)                             */
/* ------------------------------------------------------------------ */

static ULONG_PTR decode_rip_relative(ULONG_PTR match_va, ULONG_PTR mapping_base,
                                     ULONG_PTR kernel_base) {
    INT32     disp         = *(INT32 *)(match_va + 3);
    ULONG_PTR next_va      = match_va + 7;
    ULONG_PTR ref_map      = next_va + (INT64)disp;
    return (ref_map - mapping_base) + kernel_base;
}

/* ------------------------------------------------------------------ */
/* main                                                                 */
/* ------------------------------------------------------------------ */

int wmain(void) {
    /* MiGetPteAddress: shr rcx,9 ; mov rax,imm64 */
    BYTE        pte_pat[]  = { 0x48, 0xC1, 0xE9, 0x09, 0x48, 0xB8 };
    const char *pte_mask   = "xxxxxx";

    /* CiValidateImageHeader in CI.dll (Win10 22H2 / Win11 23H2):
     * mov [rsp+20h],rbx ; push rbp ; push rsi ; push rdi
     * NOTE: verify this pattern produces exactly one match before use. */
    BYTE        ci_pat[]   = { 0x48, 0x89, 0x5C, 0x24, 0x20, 0x55, 0x56, 0x57 };
    const char *ci_mask    = "xxxxxxxx";

    ULONG_PTR ntos_base = get_kernel_module_base("ntoskrnl.exe");
    ULONG_PTR ci_base   = get_kernel_module_base("CI.dll");

    if (!ntos_base) { fprintf(stderr, "[-] could not find ntoskrnl.exe\n"); return 1; }
    printf("[*] ntoskrnl.exe kernel base: 0x%016llX\n", (unsigned long long)ntos_base);
    if (ci_base)
        printf("[*] CI.dll kernel base:       0x%016llX\n", (unsigned long long)ci_base);

    wchar_t ntos_path[MAX_PATH];
    GetSystemDirectoryW(ntos_path, MAX_PATH);
    wcscat_s(ntos_path, MAX_PATH, L"\\ntoskrnl.exe");

    PBYTE ntos_map = map_pe(ntos_path);
    if (!ntos_map) { fprintf(stderr, "[-] failed to map ntoskrnl.exe\n"); return 1; }

    ULONG_PTR match = pattern_scan(ntos_map, pte_pat, pte_mask);
    if (match) {
        ULONG_PTR pte_base = *(ULONG_PTR *)(match + 6);
        printf("[+] MiGetPteAddress offset: 0x%llX\n",
               (unsigned long long)(match - (ULONG_PTR)ntos_map));
        printf("[+] PTE_BASE (MOV RAX,imm64): 0x%016llX\n",
               (unsigned long long)pte_base);
        printf("[+] MiGetPteAddress kernel VA: 0x%016llX\n",
               (unsigned long long)((match - (ULONG_PTR)ntos_map) + ntos_base));
    } else {
        printf("[-] MiGetPteAddress pattern not found\n");
    }

    UnmapViewOfFile(ntos_map);

    if (ci_base) {
        wchar_t ci_path[MAX_PATH];
        GetSystemDirectoryW(ci_path, MAX_PATH);
        wcscat_s(ci_path, MAX_PATH, L"\\CI.dll");

        PBYTE ci_map = map_pe(ci_path);
        if (ci_map) {
            ULONG_PTR ci_match = pattern_scan(ci_map, ci_pat, ci_mask);
            if (ci_match) {
                /* 0x23 bytes before the matched instruction in Win10 22H2 / Win11 23H2.
                 * Build-specific: re-verify on other builds before use. */
                ULONG_PTR ci_fn_offset = (ci_match - (ULONG_PTR)ci_map) - 0x23;
                ULONG_PTR ci_fn_kva    = ci_fn_offset + ci_base;
                printf("[+] CiValidateImageHeader kernel VA: 0x%016llX\n",
                       (unsigned long long)ci_fn_kva);
            } else {
                printf("[-] CiValidateImageHeader pattern not found\n");
            }
            UnmapViewOfFile(ci_map);
        }
    }

    return 0;
}