xll-utils 0.1.0

//! PE/COFF file parsing utilities.
//!
//! Provides functions for opening and parsing PE files, extracting
//! architecture information, export directories, and individual exports.

use std::path::Path;

use object::pe::{ImageExportDirectory, IMAGE_FILE_DLL};
use object::read::pe::{ExportTarget, ImageNtHeaders, ImageOptionalHeader};
use object::LittleEndian as LE;

use crate::error::{Error, Result};
use crate::types::{Architecture, ExportDirectory, ExportInfo, PeFile};

/// Open and parse a PE/COFF file from disk.
///
/// # Examples
///
/// ```no_run
/// use xll_utils::pe::parse_pe_file;
///
/// let pe = parse_pe_file("my_xll.xll").unwrap();
/// println!("{} exports", pe.exports().len());
/// ```
pub fn parse_pe_file(path: impl AsRef<Path>) -> Result<PeFile> {
    let path = path.as_ref();
    let data = std::fs::read(path).map_err(|e| Error::FileOpen {
        path: path.to_path_buf(),
        source: e,
    })?;
    let mut pe = parse_pe_bytes(&data)?;
    pe.path = path.to_path_buf();
    Ok(pe)
}

/// Parse a PE/COFF file from an in-memory byte slice.
pub fn parse_pe_bytes(data: &[u8]) -> Result<PeFile> {
    // Try PE32+ (64-bit) first since it's the most common for modern XLLs,
    // then fall back to PE32 (32-bit).
    if let Ok(pe) = object::read::pe::PeFile64::parse(data) {
        return build_pe_file(&pe, data);
    }
    if let Ok(pe) = object::read::pe::PeFile32::parse(data) {
        return build_pe_file(&pe, data);
    }
    Err(Error::InvalidFileFormat)
}

/// Extract all PE information from a parsed PE file (generic over 32/64-bit).
fn build_pe_file<'data, Pe>(
    pe: &object::read::pe::PeFile<'data, Pe, &'data [u8]>,
    _data: &'data [u8],
) -> Result<PeFile>
where
    Pe: ImageNtHeaders,
{
    let nt = pe.nt_headers();
    let fh = nt.file_header();

    let machine = fh.machine.get(LE);
    let architecture = Architecture::from_machine_type(machine);

    let characteristics = fh.characteristics.get(LE);
    let is_dll = (characteristics & IMAGE_FILE_DLL) != 0;

    let opt = nt.optional_header();
    let image_base = opt.image_base();
    let entry_point_rva = opt.address_of_entry_point();

    // Parse exports.
    let (export_directory, exports) = match pe.export_table() {
        Ok(Some(et)) => parse_exports(&et)?,
        Ok(None) => (None, Vec::new()),
        Err(e) => return Err(Error::PeParse(format!("failed to read export table: {e}"))),
    };

    Ok(PeFile {
        path: std::path::PathBuf::new(),
        architecture,
        is_dll,
        image_base,
        entry_point_rva,
        export_directory,
        exports,
    })
}

/// Extract the export directory metadata and individual exports from an `ExportTable`.
fn parse_exports(
    et: &object::read::pe::ExportTable<'_>,
) -> Result<(Option<ExportDirectory>, Vec<ExportInfo>)> {
    let dir: &ImageExportDirectory = et.directory();

    let export_directory = ExportDirectory {
        characteristics: dir.characteristics.get(LE),
        timestamp: dir.time_date_stamp.get(LE),
        major_version: dir.major_version.get(LE),
        minor_version: dir.minor_version.get(LE),
        name_rva: dir.name.get(LE),
        ordinal_base: dir.base.get(LE),
        address_table_entries: dir.number_of_functions.get(LE),
        number_of_name_pointers: dir.number_of_names.get(LE),
        address_table_rva: dir.address_of_functions.get(LE),
        name_pointer_rva: dir.address_of_names.get(LE),
        ordinal_table_rva: dir.address_of_name_ordinals.get(LE),
    };

    // Use the high-level exports() iterator.
    let raw_exports = et
        .exports()
        .map_err(|e| Error::PeParse(format!("failed to iterate exports: {e}")))?;

    let mut exports = Vec::with_capacity(raw_exports.len());
    for exp in &raw_exports {
        let name = exp
            .name
            .map(|n| {
                String::from_utf8(n.to_vec())
                    .map_err(|e| Error::PeParse(format!("invalid UTF-8 in export name: {e}")))
            })
            .transpose()?;

        let ordinal = exp.ordinal;

        let (is_forwarded, forward_to, relative_address) = match exp.target {
            ExportTarget::Address(rva) => (false, None, Some(rva)),
            ExportTarget::ForwardByName(dll, fname) => {
                let dll_s = String::from_utf8_lossy(dll);
                let fname_s = String::from_utf8_lossy(fname);
                (true, Some(format!("{dll_s}.{fname_s}")), None)
            }
            ExportTarget::ForwardByOrdinal(dll, ord) => {
                let dll_s = String::from_utf8_lossy(dll);
                (true, Some(format!("{dll_s}.#{ord}")), None)
            }
        };

        exports.push(ExportInfo {
            name,
            ordinal,
            is_forwarded,
            forward_to,
            relative_address,
        });
    }

    Ok((Some(export_directory), exports))
}

impl PeFile {
    /// Parse a PE file from the given path.
    ///
    /// ```no_run
    /// use xll_utils::PeFile;
    ///
    /// let pe = PeFile::parse("my_xll.xll").unwrap();
    /// assert!(pe.is_dll());
    /// ```
    pub fn parse(path: impl AsRef<Path>) -> Result<Self> {
        parse_pe_file(path.as_ref())
    }

    /// Parse a PE file from an in-memory byte slice.
    pub fn from_bytes(data: &[u8]) -> Result<Self> {
        parse_pe_bytes(data)
    }

    /// Get the export directory, if present.
    pub fn export_directory(&self) -> Option<&ExportDirectory> {
        self.export_directory.as_ref()
    }

    /// Get all exports.
    pub fn exports(&self) -> &[ExportInfo] {
        &self.exports
    }

    /// Get the CPU architecture.
    pub fn architecture(&self) -> Architecture {
        self.architecture
    }

    /// Check if this is a DLL file.
    pub fn is_dll(&self) -> bool {
        self.is_dll
    }

    /// Find an export by name.
    pub fn find_export(&self, name: &str) -> Option<&ExportInfo> {
        self.exports.iter().find(|e| e.name.as_deref() == Some(name))
    }

    /// Find an export by ordinal.
    pub fn find_export_by_ordinal(&self, ordinal: u32) -> Option<&ExportInfo> {
        self.exports.iter().find(|e| e.ordinal == ordinal)
    }

    /// Get export names as a sorted vector.
    pub fn export_names(&self) -> Vec<String> {
        let mut names: Vec<_> = self
            .exports
            .iter()
            .filter_map(|e| e.name.clone())
            .collect();
        names.sort();
        names
    }
}