xex2tractor/src/imports.rs

/// Import record decoding and resolution for XEX2 PE images.
///
/// Reads import records from the extracted PE image, decodes their type and
/// ordinal, and resolves them against the Xbox 360 export database.
use std::fmt;

use crate::error::{Result, Xex2Error};
use crate::exports::{self, ExportInfo};
use crate::util::read_u32_be;
use crate::Xex2File;

/// The type of an import record.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ImportRecordType {
    /// A 4-byte variable slot (record_type 0x00).
    Variable,
    /// A 16-byte function thunk stub (record_type 0x01).
    Thunk,
}

impl fmt::Display for ImportRecordType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ImportRecordType::Variable => write!(f, "variable"),
            ImportRecordType::Thunk => write!(f, "thunk"),
        }
    }
}

/// A decoded import record from the PE image.
#[derive(Debug, Clone)]
pub struct ImportRecord {
    /// The original memory address from the XEX2 import table.
    pub address: u32,
    /// Offset within the PE image (address - load_address).
    pub pe_offset: usize,
    /// Whether this is a variable slot or function thunk.
    pub record_type: ImportRecordType,
    /// The ordinal number identifying the imported function/variable.
    pub ordinal: u16,
}

/// An import record resolved against the export database.
#[derive(Debug, Clone)]
pub struct ResolvedImport {
    /// The decoded import record.
    pub record: ImportRecord,
    /// The library filename (e.g. "xboxkrnl.exe").
    pub library: String,
    /// The export info if the ordinal was found in the database.
    pub export: Option<ExportInfo>,
}

/// Decodes all import records from the PE image and resolves them against
/// the export database.
///
/// Iterates over each import library's addresses, reads the u32 value at
/// the corresponding PE offset, extracts record_type and ordinal, then
/// looks up the ordinal in the export database.
pub fn decode_import_records(pe_image: &[u8], xex: &Xex2File) -> Result<Vec<ResolvedImport>> {
    let imports = xex
        .optional_headers
        .import_libraries
        .as_ref()
        .ok_or_else(|| Xex2Error::InvalidPeImage("no import libraries header".into()))?;

    let load_address = xex.security_info.load_address;
    let mut resolved = Vec::new();

    for lib in &imports.libraries {
        for &addr in &lib.import_addresses {
            let pe_offset = (addr.wrapping_sub(load_address)) as usize;

            let raw = read_u32_be(pe_image, pe_offset).map_err(|_| {
                Xex2Error::InvalidPeImage(format!(
                    "import address 0x{addr:08X} (offset 0x{pe_offset:08X}) out of bounds"
                ))
            })?;

            let record_type_byte = (raw >> 24) & 0xFF;
            let ordinal = (raw & 0xFFFF) as u16;

            let record_type = match record_type_byte {
                0x00 => ImportRecordType::Variable,
                0x01 => ImportRecordType::Thunk,
                _ => {
                    // Unknown record types — treat as variable (best effort)
                    ImportRecordType::Variable
                }
            };

            let record = ImportRecord {
                address: addr,
                pe_offset,
                record_type,
                ordinal,
            };

            let export = exports::lookup(&lib.name, ordinal).cloned();

            resolved.push(ResolvedImport {
                record,
                library: lib.name.clone(),
                export,
            });
        }
    }

    Ok(resolved)
}

/// Summary of import resolution results.
#[derive(Debug)]
pub struct ImportResolutionSummary {
    /// Total number of import records processed.
    pub total: usize,
    /// Number of variable slots written.
    pub variables_written: usize,
    /// Number of thunk stubs written.
    pub thunks_written: usize,
    /// Per-library breakdown: (library_name, count).
    pub per_library: Vec<(String, usize)>,
}

impl fmt::Display for ImportResolutionSummary {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Resolved {} imports ({} variables, {} thunks)",
            self.total, self.variables_written, self.thunks_written
        )?;
        for (lib, count) in &self.per_library {
            write!(f, "\n  {lib}: {count}")?;
        }
        Ok(())
    }
}

/// Resolves imports in the PE image by writing Xenia-style values.
///
/// For variable imports (type 0): writes `0xD000BEEF | (ordinal & 0xFFF) << 16`.
/// For thunk imports (type 1): rewrites the first 8 bytes to valid PowerPC:
///   - word 0: `li r3, 0`       (0x38600000)
///   - word 1: `li r4, <ordinal>` (0x38800000 | ordinal)
///   - words 2-3 are left unchanged (already mtspr CTR, r11 + bctr)
pub fn resolve_imports(
    pe_image: &mut [u8],
    xex: &Xex2File,
) -> Result<ImportResolutionSummary> {
    let imports = xex
        .optional_headers
        .import_libraries
        .as_ref()
        .ok_or_else(|| Xex2Error::InvalidPeImage("no import libraries header".into()))?;

    let load_address = xex.security_info.load_address;
    let mut total = 0;
    let mut variables_written = 0;
    let mut thunks_written = 0;
    let mut per_library: Vec<(String, usize)> = Vec::new();

    for lib in &imports.libraries {
        let mut lib_count = 0;

        for &addr in &lib.import_addresses {
            let pe_offset = (addr.wrapping_sub(load_address)) as usize;

            let raw = read_u32_be(pe_image, pe_offset).map_err(|_| {
                Xex2Error::InvalidPeImage(format!(
                    "import address 0x{addr:08X} (offset 0x{pe_offset:08X}) out of bounds"
                ))
            })?;

            let record_type_byte = (raw >> 24) & 0xFF;
            let ordinal = (raw & 0xFFFF) as u16;

            match record_type_byte {
                0x00 => {
                    // Variable: write 0xD000BEEF | (ordinal & 0xFFF) << 16
                    let resolved_val: u32 =
                        0xD000BEEF | ((ordinal as u32 & 0xFFF) << 16);
                    pe_image[pe_offset..pe_offset + 4]
                        .copy_from_slice(&resolved_val.to_be_bytes());
                    variables_written += 1;
                }
                0x01 => {
                    // Thunk: rewrite first 8 bytes to valid PPC
                    if pe_offset + 16 > pe_image.len() {
                        return Err(Xex2Error::InvalidPeImage(format!(
                            "thunk at 0x{addr:08X} (offset 0x{pe_offset:08X}) extends past PE image"
                        )));
                    }
                    // li r3, 0
                    let word0: u32 = 0x38600000;
                    // li r4, <ordinal>
                    let word1: u32 = 0x38800000 | ordinal as u32;
                    pe_image[pe_offset..pe_offset + 4]
                        .copy_from_slice(&word0.to_be_bytes());
                    pe_image[pe_offset + 4..pe_offset + 8]
                        .copy_from_slice(&word1.to_be_bytes());
                    // words 2-3 (mtspr CTR, r11 + bctr) left unchanged
                    thunks_written += 1;
                }
                _ => {
                    // Unknown record type — skip
                }
            }

            lib_count += 1;
            total += 1;
        }

        per_library.push((lib.name.clone(), lib_count));
    }

    Ok(ImportResolutionSummary {
        total,
        variables_written,
        thunks_written,
        per_library,
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_import_record_type_display() {
        assert_eq!(ImportRecordType::Variable.to_string(), "variable");
        assert_eq!(ImportRecordType::Thunk.to_string(), "thunk");
    }

    #[test]
    fn test_resolution_summary_display() {
        let summary = ImportResolutionSummary {
            total: 10,
            variables_written: 5,
            thunks_written: 5,
            per_library: vec![("xboxkrnl.exe".into(), 10)],
        };
        let s = summary.to_string();
        assert!(s.contains("10 imports"));
        assert!(s.contains("5 variables"));
        assert!(s.contains("5 thunks"));
        assert!(s.contains("xboxkrnl.exe: 10"));
    }
}