// peb-based runtime api resolution. walks InMemoryOrderModuleList via gs:[0x60].
// ror13 hash on module + export names, no GetProcAddress or GetModuleHandle needed.
const std = @import("std");
const windows = std.os.windows;
const pe = @import("pe.zig");

pub const LIST_ENTRY = extern struct {
    Flink: *LIST_ENTRY,
    Blink: *LIST_ENTRY,
};

pub const PEB_LDR_DATA = extern struct {
    Length: windows.ULONG,
    Initialized: windows.BYTE,
    Padding: [3]windows.BYTE,
    SsHandle: windows.LPVOID,
    InLoadOrderModuleList: LIST_ENTRY,
    InMemoryOrderModuleList: LIST_ENTRY,
    InInitializationOrderModuleList: LIST_ENTRY,
};

pub const LDR_DATA_TABLE_ENTRY = extern struct {
    InLoadOrderLinks: LIST_ENTRY,
    InMemoryOrderLinks: LIST_ENTRY,
    InInitializationOrderLinks: LIST_ENTRY,
    DllBase: windows.LPVOID,
    EntryPoint: windows.LPVOID,
    SizeOfImage: windows.ULONG,
    FullDllName: windows.UNICODE_STRING,
    BaseDllName: windows.UNICODE_STRING,
};

pub const PEB = extern struct {
    Reserved1: [2]windows.BYTE,
    BeingDebugged: windows.BYTE,
    Reserved2: [1]windows.BYTE,
    Reserved3: [2]windows.LPVOID,
    Ldr: *PEB_LDR_DATA,
};

pub fn hash_ror13(input: []const u8) u32 {
    var hash: u32 = 0;
    for (input) |c| {
        var c_upper = c;
        if (c_upper >= 'a' and c_upper <= 'z') c_upper -= 32;
        hash = (hash >> 13) | (hash << 19);
        hash +%= c_upper;
    }
    return hash;
}

pub fn get_peb() *PEB {
    const ptr: usize = asm volatile ("mov %%gs:0x60, %[ptr]"
        : [ptr] "=r" (-> usize),
    );
    return @as(*PEB, @ptrFromInt(ptr));
}

pub fn get_module_by_hash(hash: u32) ?*anyopaque {
    const peb_ptr = get_peb();
    const ldr = peb_ptr.Ldr;
    const head: *LIST_ENTRY = @ptrCast(&ldr.InMemoryOrderModuleList);
    var entry = head.Flink;

    while (entry != head) : (entry = entry.Flink) {
        const le: *LDR_DATA_TABLE_ENTRY = @fieldParentPtr("InMemoryOrderLinks", entry);

        if (le.BaseDllName.Length == 0) continue;

        const buf = le.BaseDllName.Buffer orelse continue;
        const len = le.BaseDllName.Length / 2;
        var mod_hash: u32 = 0;
        for (0..len) |i| {
            var c = buf[i];
            if (c >= 'a' and c <= 'z') c -= 32;
            mod_hash = (mod_hash >> 13) | (mod_hash << 19);
            mod_hash +%= @as(u32, c);
        }

        if (mod_hash == hash) return le.DllBase;
    }

    return null;
}

fn resolve_forward_string(base_bytes: [*]u8, forward_rva: u32) ?*anyopaque {
    const forward_str = @as([*:0]u8, @ptrCast(@alignCast(
        base_bytes + @as(usize, @intCast(forward_rva)),
    )));
    const str = std.mem.sliceTo(forward_str, 0);
    const dot_pos = std.mem.indexOfScalar(u8, str, '.') orelse return null;
    const dll_hash = hash_ror13(str[0..dot_pos]);
    const func_hash = hash_ror13(str[dot_pos + 1 ..]);
    const dll_base = get_module_by_hash(dll_hash) orelse return null;
    return get_func_by_hash(dll_base, func_hash);
}

pub fn get_func_by_hash(base: *anyopaque, func_hash: u32) ?*anyopaque {
    const base_bytes: [*]u8 = @ptrCast(base);
    const dos: *align(1) const pe.IMAGE_DOS_HEADER = @ptrCast(@alignCast(base_bytes));
    if (dos.e_magic != 0x5A4D) return null;

    const nt: *align(1) const pe.IMAGE_NT_HEADERS64 = @ptrCast(@alignCast(
        base_bytes + @as(usize, @intCast(dos.e_lfanew)),
    ));
    if (nt.Signature != 0x00004550) return null;

    const export_dir = nt.OptionalHeader.DataDirectory[0];
    if (export_dir.VirtualAddress == 0 or export_dir.Size == 0) return null;

    const exp: *align(1) const pe.IMAGE_EXPORT_DIRECTORY = @ptrCast(@alignCast(
        base_bytes + @as(usize, @intCast(export_dir.VirtualAddress)),
    ));

    if (exp.NumberOfNames == 0 or exp.AddressOfFunctions == 0 or
        exp.AddressOfNames == 0 or exp.AddressOfNameOrdinals == 0) return null;

    const names: [*]align(1) const u32 = @ptrCast(@alignCast(
        base_bytes + @as(usize, @intCast(exp.AddressOfNames)),
    ));
    const funcs: [*]align(1) const u32 = @ptrCast(@alignCast(
        base_bytes + @as(usize, @intCast(exp.AddressOfFunctions)),
    ));
    const ords: [*]align(1) const u16 = @ptrCast(@alignCast(
        base_bytes + @as(usize, @intCast(exp.AddressOfNameOrdinals)),
    ));

    var i: u32 = 0;
    while (i < exp.NumberOfNames) : (i += 1) {
        if (names[i] == 0) continue;

        const name_ptr: [*:0]const u8 = @ptrCast(@alignCast(
            base_bytes + @as(usize, @intCast(names[i])),
        ));
        const name = std.mem.sliceTo(name_ptr, 0);

        if (hash_ror13(name) == func_hash) {
            const ordinal = ords[i];
            if (ordinal >= exp.NumberOfFunctions) continue;
            const rva = funcs[ordinal];
            if (rva >= export_dir.VirtualAddress and rva < export_dir.VirtualAddress + export_dir.Size) {
                return resolve_forward_string(base_bytes, rva);
            }
            return @as(*anyopaque, @ptrCast(base_bytes + @as(usize, @intCast(rva))));
        }
    }

    return null;
}

var cached_kernel32: ?*anyopaque = null;
var cached_ntdll: ?*anyopaque = null;

pub fn resolve_api(module_hash: u32, func_hash: u32) ?*anyopaque {
    if (module_hash != 0) {
        if (get_module_by_hash(module_hash)) |base| {
            if (get_func_by_hash(base, func_hash)) |f| return f;
        }
    }

    if (cached_kernel32 == null) cached_kernel32 = get_module_by_hash(hash_ror13("kernel32.dll"));
    if (cached_kernel32) |base| {
        if (get_func_by_hash(base, func_hash)) |f| return f;
    }

    if (cached_ntdll == null) cached_ntdll = get_module_by_hash(hash_ror13("ntdll.dll"));
    if (cached_ntdll) |base| {
        if (get_func_by_hash(base, func_hash)) |f| return f;
    }

    const peb_ptr = get_peb();
    const ldr = peb_ptr.Ldr;
    const head: *LIST_ENTRY = @ptrCast(&ldr.InMemoryOrderModuleList);
    var entry = head.Flink;

    while (entry != head) : (entry = entry.Flink) {
        const le: *LDR_DATA_TABLE_ENTRY = @fieldParentPtr("InMemoryOrderLinks", entry);

        if (cached_kernel32) |ck| if (le.DllBase == ck) continue;
        if (cached_ntdll) |cn| if (le.DllBase == cn) continue;

        if (get_func_by_hash(le.DllBase, func_hash)) |f| return f;
    }

    return null;
}