dyncvoke-manualmap 0.1.0

//! Manual PE mapping: relocations, IAT rewrite, section permissions, TLS.
//!
//! [`read_and_map_module`] loads a file, maps it, and returns
//! `(PeMetadata, base)`. [`manually_map_module`] does the same from a
//! buffer. Set `clean_dos_header` to wipe MZ / DOS stub IOCs.
//!
//! ```ignore
//! let (_pe, base) = manualmap::read_and_map_module(
//!     r"C:\Windows\System32\ntdll.dll",
//!     true,
//!     false,
//! ).unwrap();
//! ```

// dyncvoke - Manual PE Mapping Module
// @5mukx

#![cfg(windows)]
#![cfg_attr(docsrs, feature(doc_auto_cfg))]

use data::lc;
use data::LARGE_INTEGER;
use std::ffi::c_void;
use std::mem::size_of;
use std::path::Path;
use std::{fs, ptr};

use data::{section_ascii_name, RuntimeFunction, OBJECT_ATTRIBUTES, UNICODE_STRING};
use windows_sys::Win32::System::SystemServices::{IMAGE_BASE_RELOCATION, IMAGE_IMPORT_DESCRIPTOR};
use windows_sys::Win32::System::WindowsProgramming::{IMAGE_THUNK_DATA32, IMAGE_THUNK_DATA64};
use windows_sys::Win32::{
    Foundation::HANDLE,
    System::{
        Diagnostics::Debug::{IMAGE_OPTIONAL_HEADER32, IMAGE_SECTION_HEADER},
        IO::IO_STATUS_BLOCK,
    },
};

use data::{
    ImageFileHeader, ImageOptionalHeader64, PeManualMap, PeMetadata, FILE_EXECUTE,
    FILE_NON_DIRECTORY_FILE, FILE_READ_ATTRIBUTES, FILE_READ_DATA, FILE_SHARE_DELETE,
    FILE_SHARE_READ, FILE_SYNCHRONOUS_IO_NONALERT, MEM_COMMIT, MEM_RESERVE, PAGE_EXECUTE,
    PAGE_EXECUTE_READ, PAGE_EXECUTE_READWRITE, PAGE_READONLY, PAGE_READWRITE, PVOID,
    SECTION_ALL_ACCESS, SECTION_MEM_EXECUTE, SECTION_MEM_READ, SECTION_MEM_WRITE, SEC_IMAGE,
    SYNCHRONIZE,
};

/// Manually maps a PE from disk to the memory of the current process.
///
/// If the clean_headers parameters is set to true, the mapped pe's dos header will be removed during the
/// mapping process. Otherwise, the dos header will be kept untouched.
///
/// The third parameter determines whether TLS callbacks are executed (true) or not (false).
///
/// It will return either a pair (PeMetadata,usize) containing the mapped PE
/// metadata and its base address or a String with a descriptive error message.
///
/// # Examples
///
/// ```ignore
/// let ntdll = manualmap::read_and_map_module(r"c:\windows\system32\ntdll.dll", true, false);
///
/// match ntdll {
///     Ok(x) => if x.1 != 0 {println!("The base address of ntdll.dll is 0x{:X}.", x.1);},
///     Err(e) => println!("{}", e),
/// }
/// ```
pub fn read_and_map_module(
    filepath: &str,
    clean_dos_header: bool,
    run_callbacks: bool,
) -> Result<(PeMetadata, usize), String> {
    let file_content = fs::read(filepath).expect(&lc!("[x] Error opening the specified file."));
    let file_content_ptr = file_content.as_ptr() as *mut _;

    let result = manually_map_module(file_content_ptr, clean_dos_header, run_callbacks)?;

    unsafe {
        for i in 0..file_content.len() {
            *(file_content_ptr.add(i)) = 0u8;
        }

        Ok(result)
    }
}

/// Manually maps a PE into the current process.
///
/// If the clean_headers parameters is set to true, the mapped pe's dos header will be removed during the
/// mapping process. Otherwise, the dos header will be kept untouched.
///
/// The third parameter determines whether TLS callbacks are executed (true) or not (false).
///
/// It will return either a pair (PeMetadata,usize) containing the mapped PE
/// metadata and its base address or a String with a descriptive error message.
///
/// # Examples
///
/// ```ignore
/// use std::fs;
///
/// let file_content = fs::read("c:\\windows\\system32\\ntdll.dll").expect("[x] Error opening the specified file.");
/// let file_content_ptr = file_content.as_ptr();
/// let result = manualmap::manually_map_module(file_content_ptr, true, true);
/// ```
pub fn manually_map_module(
    file_ptr: *const u8,
    clean_dos_headers: bool,
    run_callbacks: bool,
) -> Result<(PeMetadata, usize), String> {
    let pe_info = get_pe_metadata(file_ptr, false)?;
    if (pe_info.is_32_bit && (size_of::<usize>() == 8))
        || (!pe_info.is_32_bit && (size_of::<usize>() == 4))
    {
        return Err(lc!(
            "[x] The module architecture does not match the process architecture."
        ));
    }

    let dwsize;
    if pe_info.is_32_bit {
        dwsize = pe_info.opt_header_32.SizeOfImage as usize;
    } else {
        dwsize = pe_info.opt_header_64.size_of_image as usize;
    }

    unsafe {
        let handle = dyncvoke_core::INVALID_HANDLE_VALUE;
        let mut a = usize::default();
        let base_address: *mut PVOID = &mut a as *mut usize as *mut PVOID;
        let zero_bits = 0 as usize;
        let size: *mut usize = std::mem::transmute(&dwsize);

        let ret = dyncvoke_core::nt_allocate_virtual_memory(
            handle,
            base_address,
            zero_bits,
            size,
            MEM_COMMIT | MEM_RESERVE,
            PAGE_READWRITE,
        );

        if ret != 0 {
            return Err(lc!("[x] Error allocating memory."));
        }

        let image_ptr = *base_address;

        map_module_to_memory(file_ptr, image_ptr, &pe_info)?;

        relocate_module(&pe_info, image_ptr);

        rewrite_module_iat(&pe_info, image_ptr)?;

        if clean_dos_headers {
            clean_dos_header(image_ptr);
        }

        set_module_section_permissions(&pe_info, image_ptr)?;

        add_runtime_table(&pe_info, image_ptr);

        if run_callbacks {
            run_tls_callbacks(&pe_info, image_ptr);
        }

        Ok((pe_info, image_ptr as usize))
    }
}

/// Returns a pair containing a pointer to the Exception data of an arbitrary module and the size of the  
/// corresponding PE section (.pdata). In case that it fails to retrieve this information, it returns
/// null values.
pub fn get_runtime_table(image_ptr: *mut c_void) -> (*mut data::RuntimeFunction, u32) {
    let mut size: u32 = 0;
    let module_metadata = get_pe_metadata(image_ptr as *const u8, false);
    if !module_metadata.is_ok() {
        return (ptr::null_mut(), size);
    }

    let metadata = module_metadata.unwrap();

    let mut runtime: *mut data::RuntimeFunction = ptr::null_mut();
    for section in &metadata.sections {
        if section_ascii_name(&section.Name) == b".pdata" {
            let base = image_ptr as usize;
            let addr = base + section.VirtualAddress as usize;
            runtime = std::ptr::with_exposed_provenance_mut::<RuntimeFunction>(addr);
            size = section.SizeOfRawData;
            break;
        }
    }

    return (runtime, size);
}

/// Retrieves PE headers information from the module base address.
///
/// It will return either a data::PeMetada struct containing the PE
/// metadata or a String with a descriptive error message.
///
/// # Examples
///
/// ```
/// use std::fs;
///
/// let file_content = fs::read("c:\\windows\\system32\\ntdll.dll").expect("[x] Error opening the specified file.");
/// let file_content_ptr = file_content.as_ptr();
/// let result = manualmap::get_pe_metadata(file_content_ptr, false);
/// ```
pub fn get_pe_metadata(module_ptr: *const u8, check_signature: bool) -> Result<PeMetadata, String> {
    if module_ptr.is_null() {
        return Err(lc!("[x] Null module pointer."));
    }
    let mut pe_metadata = PeMetadata::default();

    unsafe {
        let e_lfanew = *((module_ptr as usize + 0x3C) as *const u32);
        pe_metadata.pe = *((module_ptr as usize + e_lfanew as usize) as *const u32);

        if pe_metadata.pe != 0x4550 && check_signature {
            return Err(lc!("[x] Invalid PE signature."));
        }

        pe_metadata.image_file_header =
            *((module_ptr as usize + e_lfanew as usize + 0x4) as *mut ImageFileHeader);

        let opt_header: *const u16 = (module_ptr as usize + e_lfanew as usize + 0x18) as *const u16;
        let pe_arch = *(opt_header);

        if pe_arch == 0x010B {
            pe_metadata.is_32_bit = true;
            let opt_header_content: *const IMAGE_OPTIONAL_HEADER32 =
                std::mem::transmute(opt_header);
            pe_metadata.opt_header_32 = *opt_header_content;
        } else if pe_arch == 0x020B {
            pe_metadata.is_32_bit = false;
            let opt_header_content: *const ImageOptionalHeader64 = std::mem::transmute(opt_header);
            pe_metadata.opt_header_64 = *opt_header_content;
        } else {
            return Err(lc!("[x] Invalid magic value."));
        }

        let mut sections: Vec<IMAGE_SECTION_HEADER> = vec![];

        for i in 0..pe_metadata.image_file_header.number_of_sections {
            let section_ptr = (opt_header as usize
                + pe_metadata.image_file_header.size_of_optional_header as usize
                + (i * 0x28) as usize) as *const u8;
            let section_ptr: *const IMAGE_SECTION_HEADER = std::mem::transmute(section_ptr);
            sections.push(*section_ptr);
        }

        pe_metadata.sections = sections;

        Ok(pe_metadata)
    }
}

/// Maps a module to a valid memory space in the current process.
///
/// The parameters required are a vector with the module content, the base address where the module should be
/// mapped and the module's metadata.
pub fn map_module_to_memory(
    module_ptr: *const u8,
    image_ptr: *mut c_void,
    pe_info: &PeMetadata,
) -> Result<(), String> {
    if (pe_info.is_32_bit && (size_of::<usize>() == 8))
        || (!pe_info.is_32_bit && (size_of::<usize>() == 4))
    {
        return Err(lc!(
            "[x] The module architecture does not match the process architecture."
        ));
    }

    let nsize;
    if pe_info.is_32_bit {
        nsize = pe_info.opt_header_32.SizeOfHeaders as usize;
    } else {
        nsize = pe_info.opt_header_64.size_of_headers as usize;
    }

    unsafe {
        let handle = dyncvoke_core::INVALID_HANDLE_VALUE;
        let base_address: *mut c_void = std::mem::transmute(image_ptr);
        let buffer: *mut c_void = std::mem::transmute(module_ptr);
        let mut written: usize = 0;
        let bytes_written: *mut usize = &mut written;
        let ret = dyncvoke_core::nt_write_virtual_memory(
            handle,
            base_address,
            buffer,
            nsize,
            bytes_written,
        );

        if ret != 0 {
            return Err(lc!("[x] Error writing PE headers to the allocated memory."));
        }

        for section in &pe_info.sections {
            let section_base_ptr =
                (image_ptr as usize + section.VirtualAddress as usize) as *mut u8;
            let section_content_ptr =
                (module_ptr as usize + section.PointerToRawData as usize) as *mut u8;

            let base_address: *mut c_void = std::mem::transmute(section_base_ptr);
            let buffer: *mut c_void = std::mem::transmute(section_content_ptr);
            let nsize = section.SizeOfRawData as usize;
            let bytes_written: *mut usize = std::mem::transmute(&written);
            let ret = dyncvoke_core::nt_write_virtual_memory(
                handle,
                base_address,
                buffer,
                nsize,
                bytes_written,
            );

            if ret != 0 || *bytes_written != nsize {
                return Err(lc!(
                    "[x] Failed to write PE sections to the allocated memory."
                ));
            }
        }

        Ok(())
    }
}

/// Relocates a module in memory.
///
/// The parameters required are the module's metadata information and a
/// pointer to the base address where the module is mapped in memory.
pub fn relocate_module(pe_info: &PeMetadata, image_ptr: *mut c_void) {
    unsafe {
        let module_memory_base: *mut usize = std::mem::transmute(image_ptr);
        let image_data_directory;
        let image_delta: isize;
        if pe_info.is_32_bit {
            image_data_directory = pe_info.opt_header_32.DataDirectory[5]; // BaseRelocationTable
            image_delta = module_memory_base as isize - pe_info.opt_header_32.ImageBase as isize;
        } else {
            image_data_directory = pe_info.opt_header_64.datas_directory[5]; // BaseRelocationTable
            image_delta = module_memory_base as isize - pe_info.opt_header_64.image_base as isize;
        }

        if image_data_directory.VirtualAddress == 0 || image_data_directory.Size == 0 {
            return;
        }

        let reloc_table_start = module_memory_base as usize
            + image_data_directory.VirtualAddress as usize;
        let reloc_table_end = reloc_table_start + image_data_directory.Size as usize;

        let mut reloc_table_ptr = reloc_table_start as *mut i32;

        while (reloc_table_ptr as usize) + size_of::<IMAGE_BASE_RELOCATION>()
            <= reloc_table_end
        {
            let ibr: *mut IMAGE_BASE_RELOCATION = std::mem::transmute(reloc_table_ptr);
            let image_base_relocation = *ibr;

            if image_base_relocation.VirtualAddress == 0
                || (image_base_relocation.SizeOfBlock as usize)
                    < size_of::<IMAGE_BASE_RELOCATION>()
            {
                break;
            }

            let reloc_count: isize = (image_base_relocation.SizeOfBlock as isize
                - size_of::<IMAGE_BASE_RELOCATION>() as isize)
                / 2;

            for i in 0..reloc_count {
                let reloc_entry_ptr = (reloc_table_ptr as usize
                    + size_of::<IMAGE_BASE_RELOCATION>() as usize
                    + (i * 2) as usize) as *mut u16;
                let reloc_value = *reloc_entry_ptr;

                let reloc_type = reloc_value >> 12;
                let reloc_patch = reloc_value & 0xfff;

                if reloc_type != 0 {
                    let patch_addr = module_memory_base as usize
                        + image_base_relocation.VirtualAddress as usize
                        + reloc_patch as usize;

                    if reloc_type == 0x3 {
                        // IMAGE_REL_BASED_HIGHLOW: 32-bit; target must be 4-byte aligned
                        if patch_addr % 4 == 0 {
                            let patch_ptr = patch_addr as *mut i32;
                            let original_ptr = *patch_ptr;
                            let patch = original_ptr + image_delta as i32;
                            *patch_ptr = patch;
                        }
                    } else {
                        // IMAGE_REL_BASED_DIR64 (0xA) and others: 64-bit; target must be 8-byte aligned
                        // Skip misaligned entries to avoid misaligned pointer dereference (e.g. ntdll on some builds)
                        if patch_addr % 8 == 0 {
                            let patch_ptr = patch_addr as *mut isize;
                            let original_ptr = *patch_ptr;
                            let patch = original_ptr + image_delta as isize;
                            *patch_ptr = patch;
                        }
                    }
                }
            }

            reloc_table_ptr = (reloc_table_ptr as usize
                + image_base_relocation.SizeOfBlock as usize)
                as *mut i32;
        }
    }
}

/// Rewrites the IAT of a manually mapped module.
///
/// The parameters required are the module's metadata information and a
/// pointer to the base address where the module is mapped in memory.
pub fn rewrite_module_iat(pe_info: &PeMetadata, image_ptr: *mut c_void) -> Result<(), String> {
    unsafe {
        let module_memory_base: *mut usize = std::mem::transmute(image_ptr);
        let image_data_directory;
        if pe_info.is_32_bit {
            image_data_directory = pe_info.opt_header_32.DataDirectory[1]; // ImportTable
        } else {
            image_data_directory = pe_info.opt_header_64.datas_directory[1]; // ImportTable
        }

        if image_data_directory.VirtualAddress == 0 {
            return Ok(());
        }

        let import_table_ptr = (module_memory_base as usize
            + image_data_directory.VirtualAddress as usize)
            as *mut usize;

        let api_set_dict = dyncvoke_core::get_api_mapping();

        let mut counter = 0;
        let mut image_import_descriptor_ptr = (import_table_ptr as usize
            + size_of::<IMAGE_IMPORT_DESCRIPTOR>() as usize * counter)
            as *mut IMAGE_IMPORT_DESCRIPTOR;
        let mut image_import_descriptor = *image_import_descriptor_ptr;

        while image_import_descriptor.Name != 0 {
            let mut dll_name = "".to_string();
            let mut c: char = ' ';
            let mut ptr =
                (module_memory_base as usize + image_import_descriptor.Name as usize) as *mut u8;
            while c != '\0' {
                c = *ptr as char;
                if c != '\0' {
                    dll_name.push(c);
                    ptr = ptr.add(1);
                }
            }

            if dll_name == "" {
                return Ok(());
            } else {
                let lookup_key = if dll_name.len() >= 6 {
                    format!("{}{}", &dll_name[..dll_name.len() - 6], ".dll")
                } else {
                    dll_name.clone()
                };

                if (dll_name.starts_with("api-") || dll_name.starts_with("ext-"))
                    && api_set_dict.contains_key(&lookup_key)
                {
                    let key = match api_set_dict.get(&lookup_key) {
                        Some(x) => x.to_string(),
                        None => "".to_string(),
                    };

                    if key.len() > 0 {
                        dll_name = key.to_string();
                    }
                }

                let mut module_handle = dyncvoke_core::get_module_base_address(&dll_name) as usize;

                if module_handle == 0 {
                    module_handle = dyncvoke_core::load_library_a(&dll_name) as usize;

                    if module_handle == 0 {
                        return Err(format!(
                            "{}{}",
                            lc!("[x] Unable to find the specified module: "),
                            dll_name
                        ));
                    }
                }

                if pe_info.is_32_bit {
                    let mut i: isize = 0;

                    loop {
                        let image_thunk_data = (module_memory_base as usize
                            + image_import_descriptor.Anonymous.OriginalFirstThunk as usize
                            + i as usize * size_of::<u32>() as usize)
                            as *mut IMAGE_THUNK_DATA32;
                        let image_thunk_data = *image_thunk_data;
                        let ft_itd = (module_memory_base as usize
                            + image_import_descriptor.FirstThunk as usize
                            + i as usize * size_of::<u32>() as usize)
                            as *mut i32;
                        if image_thunk_data.u1.AddressOfData == 0 {
                            break;
                        }

                        if image_thunk_data.u1.AddressOfData < 0x80000000 {
                            let mut imp_by_name_ptr = (module_memory_base as usize
                                + image_thunk_data.u1.AddressOfData as usize
                                + size_of::<u16>() as usize)
                                as *mut u8;
                            let mut import_name: String = "".to_string();
                            let mut c: char = ' ';
                            while c != '\0' {
                                c = *imp_by_name_ptr as char;
                                if c != '\0' {
                                    import_name.push(c);
                                }

                                imp_by_name_ptr = imp_by_name_ptr.add(1);
                            }

                            let func_ptr =
                                dyncvoke_core::get_function_address(module_handle, &import_name);
                            *ft_itd = func_ptr as i32;
                        } else {
                            let f_ordinal = (image_thunk_data.u1.AddressOfData & 0xFFFF) as u32;
                            let func_ptr = dyncvoke_core::get_function_address_by_ordinal(
                                module_handle,
                                f_ordinal,
                            );
                            let func_ptr = func_ptr as *mut i32;
                            *ft_itd = func_ptr as i32;
                        }

                        i = i + 1;
                    }
                } else {
                    let mut i: isize = 0;

                    loop {
                        let image_thunk_data = (module_memory_base as u64
                            + image_import_descriptor.Anonymous.OriginalFirstThunk as u64
                            + i as u64 * size_of::<u64>() as u64)
                            as *mut IMAGE_THUNK_DATA64;
                        let image_thunk_data = *image_thunk_data;
                        let ft_itd = (module_memory_base as u64
                            + image_import_descriptor.FirstThunk as u64
                            + i as u64 * size_of::<u64>() as u64)
                            as *mut isize;

                        if image_thunk_data.u1.AddressOfData == 0 {
                            break;
                        }

                        if image_thunk_data.u1.AddressOfData < 0x8000000000000000 {
                            let mut imp_by_name_ptr = (module_memory_base as u64
                                + image_thunk_data.u1.AddressOfData as u64
                                + size_of::<u16>() as u64)
                                as *mut u8;
                            let mut import_name: String = "".to_string();
                            let mut c: char = ' ';
                            while c != '\0' {
                                c = *imp_by_name_ptr as char;
                                if c != '\0' {
                                    import_name.push(c);
                                }

                                imp_by_name_ptr = imp_by_name_ptr.add(1);
                            }

                            let func_ptr =
                                dyncvoke_core::get_function_address(module_handle, &import_name)
                                    as *mut isize;
                            *ft_itd = func_ptr as isize;
                        } else {
                            let f_ordinal = (image_thunk_data.u1.AddressOfData & 0xFFFF) as u32;
                            let func_ptr = dyncvoke_core::get_function_address_by_ordinal(
                                module_handle,
                                f_ordinal,
                            );
                            *ft_itd = func_ptr as isize;
                        }

                        i = i + 1;
                    }
                }
            }

            counter = counter + 1;
            image_import_descriptor_ptr = (import_table_ptr as usize
                + size_of::<IMAGE_IMPORT_DESCRIPTOR>() as usize * counter)
                as *mut IMAGE_IMPORT_DESCRIPTOR;
            image_import_descriptor = *image_import_descriptor_ptr;
        }

        Ok(())
    }
}

fn clean_dos_header(image_ptr: *mut c_void) {
    unsafe {
        let mut base_addr = image_ptr as *mut u8;
        let pe_header = image_ptr as isize + 0x3C;
        while (base_addr as isize) < pe_header {
            *base_addr = 0;
            base_addr = base_addr.add(1);
        }
        base_addr = base_addr.add(4);

        let e_lfanew = *((image_ptr as usize + 0x3C) as *const u32);
        let pe = image_ptr as isize + e_lfanew as isize;

        while (base_addr as isize) < pe {
            *base_addr = 0;
            base_addr = base_addr.add(1);
        }

        let pe = pe as *mut u16;
        *pe = 0;
    }
}

pub fn add_runtime_table(pe_info: &PeMetadata, image_ptr: *mut c_void) {
    unsafe {
        for section in &pe_info.sections {
            if section_ascii_name(&section.Name) == b".pdata" {
                let entry_count = (section.SizeOfRawData / 12) as i32; // 12 = size_of RUNTIME_FUNCTION

                let func: data::RtlAddFunctionTable;
                let _ret: Option<bool>;
                let k32 = dyncvoke_core::get_module_base_address(&lc!("kernel32.dll"));
                let function_table_addr: usize =
                    image_ptr as usize + section.VirtualAddress as usize;
                dyncvoke_core::dynamic_invoke!(
                    k32,
                    &lc!("RtlAddFunctionTable"),
                    func,
                    _ret,
                    function_table_addr,
                    entry_count,
                    image_ptr as usize
                );
            }
        }
    }
}

/// Sets correct module section permissions for a manually mapped module.
///
/// The parameters required are the module's metadata information and a
/// pointer to the base address where the module is mapped in memory.
pub fn set_module_section_permissions(
    pe_info: &PeMetadata,
    image_ptr: *mut c_void,
) -> Result<(), String> {
    unsafe {
        let base_of_code;

        if pe_info.is_32_bit {
            base_of_code = pe_info.opt_header_32.BaseOfCode as usize;
        } else {
            base_of_code = pe_info.opt_header_64.base_of_code as usize;
        }

        let handle = dyncvoke_core::INVALID_HANDLE_VALUE;
        let mut image_ptr = image_ptr;
        let base_address: *mut PVOID = &mut image_ptr as *mut *mut c_void as *mut PVOID;
        let mut size: usize = base_of_code;
        let mut old_protection: u32 = 0;
        let _ret = dyncvoke_core::nt_protect_virtual_memory(
            handle,
            base_address,
            &mut size,
            PAGE_READONLY,
            &mut old_protection,
        );

        for section in &pe_info.sections {
            let is_read = (section.Characteristics & SECTION_MEM_READ) != 0;
            let is_write = (section.Characteristics & SECTION_MEM_WRITE) != 0;
            let is_execute = (section.Characteristics & SECTION_MEM_EXECUTE) != 0;
            let new_protect: u32;

            if is_read & !is_write & !is_execute {
                new_protect = PAGE_READONLY;
            } else if is_read & is_write & !is_execute {
                new_protect = PAGE_READWRITE;
            } else if is_read & is_write & is_execute {
                new_protect = PAGE_EXECUTE_READWRITE;
            } else if is_read & !is_write & is_execute {
                new_protect = PAGE_EXECUTE_READ;
            } else if !is_read & !is_write & is_execute {
                new_protect = PAGE_EXECUTE;
            } else {
                return Err(lc!("[x] Unknown section permission."));
            }

            let mut address: *mut c_void =
                (image_ptr as usize + section.VirtualAddress as usize) as *mut c_void;
            let base_address: *mut PVOID = &mut address as *mut *mut c_void as *mut PVOID;
            size = section.Misc.VirtualSize as usize;
            let mut old_protection: u32 = 0;
            let ret = dyncvoke_core::nt_protect_virtual_memory(
                handle,
                base_address,
                &mut size,
                new_protect,
                &mut old_protection,
            );

            if ret != 0 {
                return Err(lc!("[x] Error changing section permission."));
            }
        }

        Ok(())
    }
}

/// Executes any registered TLS Callback function.
///
/// The parameters required are the module's metadata information and a
/// pointer to the base address where the module is mapped in memory.
pub fn run_tls_callbacks(pe_info: &PeMetadata, image_ptr: *mut c_void) {
    unsafe {
        let entry_point;
        if pe_info.is_32_bit {
            entry_point = image_ptr as isize + pe_info.opt_header_32.AddressOfEntryPoint as isize;
        } else {
            entry_point =
                image_ptr as isize + pe_info.opt_header_64.address_of_entry_point as isize;
        }

        let (tls_dir_rva, tls_cb_off) = if pe_info.is_32_bit {
            (
                pe_info.opt_header_32.NumberOfRvaAndSizes >= 10,
                pe_info.opt_header_32.DataDirectory[9].VirtualAddress as usize,
            )
        } else {
            (
                pe_info.opt_header_64.number_of_rva_and_sizes >= 10,
                pe_info.opt_header_64.datas_directory[9].VirtualAddress as usize,
            )
        };
        if tls_dir_rva && tls_cb_off != 0 {
            let address: *mut u8 = (image_ptr as usize + tls_cb_off) as *mut u8;
            // IMAGE_TLS_DIRECTORY64.AddressOfCallBacks is at offset 24;
            // IMAGE_TLS_DIRECTORY32.AddressOfCallBacks is at offset 12.
            // winnt.h IMAGE_TLS_DIRECTORY.
            let cb_off = if pe_info.is_32_bit { 12 } else { 24 };
            let address_of_tls_callback = address.add(cb_off) as *mut usize;
            let mut address_of_tls_callback_array: *mut usize =
                std::ptr::with_exposed_provenance_mut::<usize>(*address_of_tls_callback);

            while *address_of_tls_callback_array != 0 {
                let tls_callback: extern "system" fn(isize, u32, PVOID) =
                    std::mem::transmute(*address_of_tls_callback_array);
                tls_callback(entry_point, 1, ptr::null_mut());
                address_of_tls_callback_array = address_of_tls_callback_array.add(1);
            }
        }
    }
}

/// Map a module to a memory section.
///
/// The parameter required is the file path of the module that should be mapped.
pub fn map_to_section(module_path: &str) -> Result<(PeManualMap, HANDLE), String> {
    unsafe {
        if !Path::new(&module_path).is_file() {
            return Err(lc!("[x] Filepath not found."));
        }

        let module_path = format!("{}{}", "\\??\\", module_path);
        let mut module_path_utf16: Vec<u16> = module_path.encode_utf16().collect();
        module_path_utf16.push(0);

        let mut o_name: UNICODE_STRING = std::mem::zeroed();
        dyncvoke_core::rtl_init_unicode_string(&mut o_name, module_path_utf16.as_ptr());

        let mut object_attributes: OBJECT_ATTRIBUTES = std::mem::zeroed();
        object_attributes.Length = size_of::<OBJECT_ATTRIBUTES>() as u32;
        object_attributes.ObjectName = &mut o_name;
        object_attributes.Attributes = 0x40; // OBJ_CASE_INSENSITIVE

        let mut io: IO_STATUS_BLOCK = std::mem::zeroed();
        let mut hfile: HANDLE = std::ptr::null_mut();
        let r = dyncvoke_core::nt_open_file(
            &mut hfile,
            FILE_READ_DATA | FILE_EXECUTE | FILE_READ_ATTRIBUTES | SYNCHRONIZE,
            &mut object_attributes,
            &mut io,
            FILE_SHARE_READ | FILE_SHARE_DELETE,
            FILE_SYNCHRONOUS_IO_NONALERT | FILE_NON_DIRECTORY_FILE,
        );

        if r != 0 {
            return Err(lc!("[x] Error opening file."));
        }

        let mut max_size: LARGE_INTEGER = LARGE_INTEGER::default();
        let mut hsection: HANDLE = std::ptr::null_mut();
        let r = dyncvoke_core::nt_create_section(
            &mut hsection,
            SECTION_ALL_ACCESS,
            ptr::null_mut(),
            &mut max_size,
            PAGE_READONLY,
            SEC_IMAGE,
            hfile,
        );

        if r != 0 {
            return Err(lc!("[x] Error creating file section in memory."));
        }

        let mut offset: LARGE_INTEGER = LARGE_INTEGER::default();
        let mut b = usize::default();
        let base_address: *mut PVOID = &mut b as *mut usize as *mut PVOID;
        let mut v = usize::default();
        let r = dyncvoke_core::nt_map_view_of_section(
            hsection,
            dyncvoke_core::INVALID_HANDLE_VALUE,
            base_address,
            0,
            0,
            &mut offset,
            &mut v,
            0x2,
            0x0,
            PAGE_READWRITE,
        );

        if r != 0 {
            return Err(lc!("[x] Error mapping file section."));
        }

        let base_address: *const u8 = std::mem::transmute(*base_address);
        let sec_object: PeManualMap = PeManualMap {
            pe_info: get_pe_metadata(base_address, false).unwrap(),
            base_address: base_address as usize,
            decoy_module: module_path,
        };

        let _r = dyncvoke_core::close_handle(hfile);

        Ok((sec_object, hsection))
    }
}

pub fn map_to_allocated_memory(
    module_ptr: *const u8,
    image_ptr: *mut c_void,
    pe_info: &PeMetadata,
) -> Result<(), String> {
    map_module_to_memory(module_ptr, image_ptr, &pe_info)?;

    relocate_module(&pe_info, image_ptr);

    rewrite_module_iat(&pe_info, image_ptr)?;

    clean_dos_header(image_ptr);

    set_module_section_permissions(&pe_info, image_ptr)?;

    add_runtime_table(&pe_info, image_ptr);

    Ok(())
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use data::{PeMetadata, ImageOptionalHeader64, ImageFileHeader};
    use windows_sys::Win32::System::Diagnostics::Debug::IMAGE_SECTION_HEADER;
    use std::collections::BTreeMap;

    // Test PeMetadata creation
    #[test]
    fn test_pe_metadata_creation() {
        let pe = PeMetadata {
            pe: 0x4550, // "PE\0\0"
            is_32_bit: false,
            image_file_header: ImageFileHeader {
                machine: 0x8664, // AMD64
                number_of_sections: 3,
                time_data_stamp: 1234567890,
                pointer_to_symbol_table: 0,
                number_of_symbols: 0,
                size_of_optional_header: 0xF0,
                characteristics: 0x2002, // Executable, Large address aware
            },
            opt_header_32: unsafe { std::mem::zeroed() },
            opt_header_64: ImageOptionalHeader64 {
                magic: 0x20B, // PE32+
                major_linker_version: 14,
                minor_linker_version: 0,
                size_of_code: 0x1000,
                size_of_initialized_data: 0,
                size_of_unitialized_data: 0,
                address_of_entry_point: 0x1000,
                base_of_code: 0x1000,
                image_base: 0x140000000,
                section_alignment: 0x1000,
                file_alignment: 0x200,
                major_operating_system_version: 10,
                minor_operating_system_version: 0,
                major_image_version: 0,
                minor_image_version: 0,
                major_subsystem_version: 10,
                minor_subsystem_version: 0,
                win32_version_value: 0,
                size_of_image: 0x7000,
                size_of_headers: 0x400,
                checksum: 0,
                subsystem: 3, // WINDOWS_CUI
                dll_characteristics: 0x8160,
                size_of_stack_reserve: 0x40000,
                size_of_stack_commit: 0x1000,
                size_of_heap_reserve: 0x100000,
                size_of_heap_commit: 0x1000,
                loader_flags: 0,
                number_of_rva_and_sizes: 16,
                datas_directory: [unsafe { std::mem::zeroed() }; 16],
            },
            sections: Vec::new(),
        };

        assert_eq!(pe.pe, 0x4550);
        assert!(!pe.is_32_bit);
        assert_eq!(pe.image_file_header.machine, 0x8664);
        assert_eq!(pe.image_file_header.number_of_sections, 3);
    }

    // Test PeMetadata default
    #[test]
    fn test_pe_metadata_default() {
        let pe = PeMetadata::default();
        assert_eq!(pe.pe, 0);
        assert!(!pe.is_32_bit);
    }


    #[test]
    fn test_image_section_header_size() {
        assert_eq!(std::mem::size_of::<IMAGE_SECTION_HEADER>(), 40);
    }

    // Test PeManualMap creation
    #[test]
    fn test_pe_manual_map_creation() {
        let pmm = PeManualMap {
            decoy_module: "kernel32.dll".to_string(),
            base_address: 0x12340000,
            pe_info: PeMetadata::default(),
        };

        assert_eq!(pmm.decoy_module, "kernel32.dll");
        assert_eq!(pmm.base_address, 0x12340000);
    }

    // Test module parsing validation
    #[test]
    fn test_pe_header_validation() {
        // Test that we can validate a PE header
        let valid_pe: [u8; 2] = [0x50, 0x45]; // "PE"
        assert_eq!(valid_pe[0], 0x50);
        assert_eq!(valid_pe[1], 0x45);

        let invalid_pe: [u8; 2] = [0x00, 0x00];
        assert_ne!(invalid_pe[0], 0x50);
    }

    // Test DOS header magic
    #[test]
    fn test_dos_header_magic() {
        let dos_magic: [u8; 2] = [0x5A, 0x4D]; // "MZ"
        assert_eq!(dos_magic[0], 0x5A); // 'M'
        assert_eq!(dos_magic[1], 0x4D); // 'Z'
    }

    // Test section characteristics
    #[test]
    fn test_section_characteristics() {
        // CODE section
        const IMAGE_SCN_CNT_CODE: u32 = 0x20000000;
        const IMAGE_SCN_MEM_EXECUTE: u32 = 0x20000000;
        const IMAGE_SCN_MEM_READ: u32 = 0x40000000;
        const IMAGE_SCN_MEM_WRITE: u32 = 0x80000000;

        let code_section = IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ;
        assert!(code_section & IMAGE_SCN_CNT_CODE != 0);
        assert!(code_section & IMAGE_SCN_MEM_EXECUTE != 0);
        assert!(code_section & IMAGE_SCN_MEM_READ != 0);

        let data_section = IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE;
        assert!(data_section & IMAGE_SCN_MEM_READ != 0);
        assert!(data_section & IMAGE_SCN_MEM_WRITE != 0);
        assert!(data_section & IMAGE_SCN_MEM_EXECUTE == 0);
    }

    // Test optional header magic values
    #[test]
    fn test_optional_header_magic() {
        const PE32_MAGIC: u16 = 0x10B;
        const PE32_PLUS_MAGIC: u16 = 0x20B;

        assert_eq!(PE32_MAGIC, 0x10B);
        assert_eq!(PE32_PLUS_MAGIC, 0x20B);
    }

    // Test subsystem values
    #[test]
    fn test_subsystem_values() {
        const IMAGE_SUBSYSTEM_UNKNOWN: u16 = 0;
        const IMAGE_SUBSYSTEM_NATIVE: u16 = 1;
        const IMAGE_SUBSYSTEM_WINDOWS_GUI: u16 = 2;
        const IMAGE_SUBSYSTEM_WINDOWS_CUI: u16 = 3;

        assert_eq!(IMAGE_SUBSYSTEM_UNKNOWN, 0);
        assert_eq!(IMAGE_SUBSYSTEM_NATIVE, 1);
        assert_eq!(IMAGE_SUBSYSTEM_WINDOWS_GUI, 2);
        assert_eq!(IMAGE_SUBSYSTEM_WINDOWS_CUI, 3);
    }

    // Test DLL characteristics
    #[test]
    fn test_dll_characteristics() {
        const IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE: u16 = 0x0040;
        const IMAGE_DLL_CHARACTERISTICS_NX_COMPAT: u16 = 0x0100;
        const IMAGE_DLL_CHARACTERISTICS_NO_ISOLATION: u16 = 0x0200;
        const IMAGE_DLL_CHARACTERISTICS_NO_SEH: u16 = 0x0400;

        let chars = IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE | IMAGE_DLL_CHARACTERISTICS_NX_COMPAT;
        assert!(chars & IMAGE_DLL_CHARACTERISTICS_DYNAMIC_BASE != 0);
        assert!(chars & IMAGE_DLL_CHARACTERISTICS_NX_COMPAT != 0);
        assert!(chars & IMAGE_DLL_CHARACTERISTICS_NO_ISOLATION == 0);
        assert!(chars & IMAGE_DLL_CHARACTERISTICS_NO_SEH == 0);
    }

    // Test machine types
    #[test]
    fn test_machine_types() {
        const IMAGE_FILE_MACHINE_I386: u16 = 0x014C;
        const IMAGE_FILE_MACHINE_AMD64: u16 = 0x8664;
        const IMAGE_FILE_MACHINE_ARM: u16 = 0x01C0;
        const IMAGE_FILE_MACHINE_ARM64: u16 = 0xAA64;

        assert_eq!(IMAGE_FILE_MACHINE_I386, 0x014C);
        assert_eq!(IMAGE_FILE_MACHINE_AMD64, 0x8664);
        assert_eq!(IMAGE_FILE_MACHINE_ARM, 0x01C0);
        assert_eq!(IMAGE_FILE_MACHINE_ARM64, 0xAA64);
    }

    // Test data directories
    #[test]
    fn test_data_directories() {
        const IMAGE_DIRECTORY_ENTRY_EXPORT: usize = 0;
        const IMAGE_DIRECTORY_ENTRY_IMPORT: usize = 1;
        const IMAGE_DIRECTORY_ENTRY_RESOURCE: usize = 2;
        const IMAGE_DIRECTORY_ENTRY_EXCEPTION: usize = 3;
        const IMAGE_DIRECTORY_ENTRY_TLS: usize = 9;

        assert_eq!(IMAGE_DIRECTORY_ENTRY_EXPORT, 0);
        assert_eq!(IMAGE_DIRECTORY_ENTRY_IMPORT, 1);
        assert_eq!(IMAGE_DIRECTORY_ENTRY_RESOURCE, 2);
        assert_eq!(IMAGE_DIRECTORY_ENTRY_EXCEPTION, 3);
        assert_eq!(IMAGE_DIRECTORY_ENTRY_TLS, 9);
    }

    // Test runtime function structure
    #[test]
    fn test_runtime_function_structure() {
        use data::RuntimeFunction;

        let rf = RuntimeFunction {
            begin_addr: 0x1000,
            end_addr: 0x2000,
            unwind_addr: 0x1500,
        };

        assert_eq!(rf.begin_addr, 0x1000);
        assert_eq!(rf.end_addr, 0x2000);
        assert_eq!(rf.unwind_addr, 0x1500);
    }

    // Test module size calculations
    #[test]
    fn test_module_size_calculations() {
        let section_alignment = 0x1000u32;
        let file_alignment = 0x200u32;
        let raw_size = 0x1800u32;

        // Calculate aligned size
        let aligned_size = ((raw_size + section_alignment - 1) / section_alignment) * section_alignment;
        assert_eq!(aligned_size, 0x2000);

        // File size alignment
        let file_aligned = ((raw_size + file_alignment - 1) / file_alignment) * file_alignment;
        assert_eq!(file_aligned, 0x1800);
    }

    // Test RVA to file offset calculation
    #[test]
    fn test_rva_to_offset() {
        // Simulate section info
        struct SectionInfo {
            virtual_address: u32,
            pointer_to_raw_data: u32,
            virtual_size: u32,
        }

        let sections = vec![
            SectionInfo { virtual_address: 0x1000, pointer_to_raw_data: 0x400, virtual_size: 0x1000 },
            SectionInfo { virtual_address: 0x2000, pointer_to_raw_data: 0x1400, virtual_size: 0x800 },
        ];

        // Test RVA in first section
        // RVA 0x1500 is in first section (0x1000-0x2000)
        // offset = 0x400 + (0x1500 - 0x1000) = 0x400 + 0x500 = 0x900
        let rva = 0x1500u32;
        let offset = sections.iter().find(|s| rva >= s.virtual_address && rva < s.virtual_address + s.virtual_size)
            .map(|s| s.pointer_to_raw_data + (rva - s.virtual_address));

        assert_eq!(offset, Some(0x900));
    }

    // Test import lookup
    #[test]
    fn test_import_lookup() {
        let mut imports: BTreeMap<String, usize> = BTreeMap::new();
        imports.insert("kernel32.dll".to_string(), 0x1000);
        imports.insert("ntdll.dll".to_string(), 0x2000);
        imports.insert("user32.dll".to_string(), 0x3000);

        assert_eq!(imports.len(), 3);
        assert_eq!(imports.get("kernel32.dll"), Some(&0x1000));
        assert!(imports.get("nonexistent.dll").is_none());
    }
}