Merge feature/m6-extract-pe: PE verification and complete extraction

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "xex2tractor"
-version = "0.5.0"
+version = "0.6.0"
 edition = "2024"
 description = "A tool for extracting and inspecting Xbox 360 XEX2 executable files"
 license = "MIT"

README.md

@@ -4,11 +4,15 @@ A tool for extracting and inspecting Xbox 360 XEX2 executable files, written in
 
 ## Usage
 
+### Inspect
+
+Display XEX2 file information (headers, security info, etc.):
+
 ```sh
-xex2tractor <file.xex>
+xex2tractor inspect <file.xex>
 ```
 
-### Example Output
+#### Example Output
 
 ```
 === XEX2 Header ===
@@ -35,16 +39,9 @@ Header Count:     15
 [FILE_FORMAT_INFO]
   Encryption:  Normal (AES-128-CBC)
   Compression: Normal (LZX)
-
-[STATIC_LIBRARIES] (12 libraries)
-  XAPILIB  2.0.3215.0 (Unknown(64))
-  D3D9     2.0.3215.1 (Unknown(64))
+  Window Size: 0x8000 (32 KB)
   ...
 
-[IMPORT_LIBRARIES] (2 libraries)
-  xam.xex v2.0.4552.0 (min v2.0.4552.0) - 104 imports
-  xboxkrnl.exe v2.0.4552.0 (min v2.0.4552.0) - 294 imports
-
 === Security Info ===
 Header Size:          0x00000F34 (3892 bytes)
 Image Size:           0x00920000 (9568256 bytes)
@@ -54,12 +51,32 @@ Load Address:         0x82000000
 Region:               0xFFFFFFFF [ALL REGIONS]
 Allowed Media Types:  0x00000004 [DVD_CD]
 ...
-
-Page Descriptors (146 entries, 64KB pages):
-  #0    Unknown(0) 19   pages  ( 1245184 bytes)  offset +0x00000000  SHA1: B136058FBBAD...
-  ...
 ```
 
+### Extract
+
+Extract the decrypted and decompressed PE image from a XEX2 file:
+
+```sh
+xex2tractor extract <file.xex> [output.exe]
+```
+
+If no output path is given, defaults to the input filename with `.exe` extension.
+
+#### Example
+
+```sh
+$ xex2tractor extract default.xex default.exe
+Encryption:  Normal (AES-128-CBC)
+Compression: Normal (LZX)
+Extracted PE image (9568256 bytes) -> default.exe
+```
+
+Supports:
+- AES-128-CBC decryption (retail, devkit, and XEX1 master keys)
+- No compression, basic (zero-fill), and normal (LZX) decompression
+- PE header verification (MZ signature, PE signature, POWERPCBE machine type)
+
 ## Building
 
 ```sh

src/error.rs

@@ -22,6 +22,8 @@ pub enum Xex2Error {
     DecryptionFailed,
     /// Decompression failed.
     DecompressionFailed(String),
+    /// The extracted PE image is invalid.
+    InvalidPeImage(String),
 }
 
 impl fmt::Display for Xex2Error {
@@ -49,6 +51,7 @@ impl fmt::Display for Xex2Error {
                 write!(f, "decryption failed: no master key produced valid output")
             }
             Xex2Error::DecompressionFailed(msg) => write!(f, "decompression failed: {msg}"),
+            Xex2Error::InvalidPeImage(msg) => write!(f, "invalid PE image: {msg}"),
         }
     }
 }

src/extract.rs

@@ -1,15 +1,24 @@
-/// PE image extraction pipeline: decrypt → decompress → raw PE bytes.
+/// PE image extraction pipeline: decrypt → decompress → verify → raw PE bytes.
 use crate::crypto;
 use crate::decompress;
 use crate::error::{Result, Xex2Error};
-use crate::optional::{CompressionInfo, CompressionType, EncryptionType};
+use crate::optional::{CompressionInfo, EncryptionType};
 use crate::Xex2File;
 
+/// Expected MZ (DOS) signature at the start of a PE image.
+const MZ_SIGNATURE: u16 = 0x5A4D;
+
+/// Expected PE signature: "PE\0\0" = 0x00004550.
+const PE_SIGNATURE: u32 = 0x00004550;
+
+/// IMAGE_FILE_MACHINE_POWERPCBE — Xbox 360 PowerPC big-endian.
+const MACHINE_POWERPCBE: u16 = 0x01F2;
+
 /// Extracts the PE image from a parsed XEX2 file.
 ///
 /// Reads the encrypted/compressed payload from `data` (the full XEX2 file),
-/// decrypts it if needed, decompresses it based on the file format info, and
-/// returns the raw PE image bytes.
+/// decrypts it if needed, decompresses it based on the file format info,
+/// verifies the PE headers, and returns the raw PE image bytes.
 pub fn extract_pe_image(data: &[u8], xex: &Xex2File) -> Result<Vec<u8>> {
     let fmt = xex
         .optional_headers
@@ -33,7 +42,25 @@ pub fn extract_pe_image(data: &[u8], xex: &Xex2File) -> Result<Vec<u8>> {
 
     // Step 1: Decrypt if needed
     if fmt.encryption_type == EncryptionType::Normal {
-        let session_key = crypto::derive_session_key(&xex.security_info.aes_key);
+        let session_key = crypto::derive_session_key_with_validation(
+            &xex.security_info.aes_key,
+            |key| {
+                // Quick validation: decrypt the first 16 bytes and check for patterns
+                // that indicate a valid decryption (non-random data).
+                // Full MZ validation happens after decompression.
+                let mut test_block = payload[..16.min(payload.len())].to_vec();
+                crypto::decrypt_in_place(key, &mut test_block);
+                // For uncompressed data, check MZ signature directly
+                if matches!(fmt.compression_info, CompressionInfo::None) {
+                    test_block.len() >= 2 && test_block[0] == 0x4D && test_block[1] == 0x5A
+                } else {
+                    // For compressed data, any successful decryption might be valid.
+                    // Accept the key if the decrypted data isn't all zeros or all 0xFF.
+                    !test_block.iter().all(|&b| b == 0)
+                        && !test_block.iter().all(|&b| b == 0xFF)
+                }
+            },
+        )?;
         crypto::decrypt_in_place(&session_key, &mut payload);
     }
 
@@ -57,12 +84,128 @@ pub fn extract_pe_image(data: &[u8], xex: &Xex2File) -> Result<Vec<u8>> {
         }
     };
 
-    // Sanity check: for unencrypted files, treat Unknown compression type through Delta path
-    if fmt.compression_type == CompressionType::Delta {
-        return Err(Xex2Error::DecompressionFailed(
-            "delta compression is not supported".into(),
-        ));
-    }
+    // Step 3: Verify PE headers
+    verify_pe_image(&pe_image)?;
 
     Ok(pe_image)
 }
+
+/// Verifies that the extracted data is a valid PE image.
+///
+/// Checks:
+/// - MZ (DOS) signature at offset 0
+/// - PE signature at e_lfanew offset
+/// - Machine type is POWERPCBE (0x01F2)
+pub fn verify_pe_image(pe_data: &[u8]) -> Result<()> {
+    if pe_data.len() < 0x40 {
+        return Err(Xex2Error::InvalidPeImage(
+            "PE image too small for DOS header".into(),
+        ));
+    }
+
+    // Check MZ signature (little-endian per PE spec)
+    let mz = u16::from_le_bytes([pe_data[0], pe_data[1]]);
+    if mz != MZ_SIGNATURE {
+        return Err(Xex2Error::InvalidPeImage(format!(
+            "invalid MZ signature: 0x{mz:04X} (expected 0x{MZ_SIGNATURE:04X})"
+        )));
+    }
+
+    // Read e_lfanew (little-endian per PE spec)
+    let e_lfanew = u32::from_le_bytes([
+        pe_data[0x3C],
+        pe_data[0x3D],
+        pe_data[0x3E],
+        pe_data[0x3F],
+    ]) as usize;
+
+    if e_lfanew + 6 > pe_data.len() {
+        return Err(Xex2Error::InvalidPeImage(format!(
+            "e_lfanew (0x{e_lfanew:X}) points past end of image"
+        )));
+    }
+
+    // Check PE signature
+    let pe_sig = u32::from_le_bytes([
+        pe_data[e_lfanew],
+        pe_data[e_lfanew + 1],
+        pe_data[e_lfanew + 2],
+        pe_data[e_lfanew + 3],
+    ]);
+    if pe_sig != PE_SIGNATURE {
+        return Err(Xex2Error::InvalidPeImage(format!(
+            "invalid PE signature: 0x{pe_sig:08X} (expected 0x{PE_SIGNATURE:08X})"
+        )));
+    }
+
+    // Check machine type (little-endian per PE spec)
+    let machine = u16::from_le_bytes([pe_data[e_lfanew + 4], pe_data[e_lfanew + 5]]);
+    if machine != MACHINE_POWERPCBE {
+        return Err(Xex2Error::InvalidPeImage(format!(
+            "unexpected machine type: 0x{machine:04X} (expected 0x{MACHINE_POWERPCBE:04X} POWERPCBE)"
+        )));
+    }
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_verify_pe_image_too_small() {
+        let data = vec![0x4D, 0x5A]; // MZ but too short
+        assert!(verify_pe_image(&data).is_err());
+    }
+
+    #[test]
+    fn test_verify_pe_image_bad_mz() {
+        let mut data = vec![0u8; 0x100];
+        data[0] = 0x00; // Not MZ
+        assert!(verify_pe_image(&data).is_err());
+    }
+
+    #[test]
+    fn test_verify_pe_image_bad_pe_sig() {
+        let mut data = vec![0u8; 0x200];
+        data[0] = 0x4D;
+        data[1] = 0x5A; // MZ
+        // e_lfanew = 0x80 (little-endian)
+        data[0x3C] = 0x80;
+        // No PE signature at 0x80
+        assert!(verify_pe_image(&data).is_err());
+    }
+
+    #[test]
+    fn test_verify_pe_image_valid() {
+        let mut data = vec![0u8; 0x200];
+        data[0] = 0x4D;
+        data[1] = 0x5A; // MZ
+        data[0x3C] = 0x80; // e_lfanew = 0x80
+        // PE signature at 0x80
+        data[0x80] = b'P';
+        data[0x81] = b'E';
+        data[0x82] = 0;
+        data[0x83] = 0;
+        // Machine = 0x01F2 (little-endian: F2 01)
+        data[0x84] = 0xF2;
+        data[0x85] = 0x01;
+        assert!(verify_pe_image(&data).is_ok());
+    }
+
+    #[test]
+    fn test_verify_pe_image_wrong_machine() {
+        let mut data = vec![0u8; 0x200];
+        data[0] = 0x4D;
+        data[1] = 0x5A;
+        data[0x3C] = 0x80;
+        data[0x80] = b'P';
+        data[0x81] = b'E';
+        data[0x82] = 0;
+        data[0x83] = 0;
+        data[0x84] = 0x4C; // x86 machine type
+        data[0x85] = 0x01;
+        assert!(verify_pe_image(&data).is_err());
+    }
+}

src/main.rs

@@ -50,6 +50,12 @@ fn cmd_extract(path: &PathBuf, output: Option<PathBuf>) {
     let data = read_file(path);
     let xex = parse_xex(&data);
 
+    // Display encryption/compression info
+    if let Some(ref fmt) = xex.optional_headers.file_format_info {
+        println!("Encryption:  {}", fmt.encryption_type);
+        println!("Compression: {}", fmt.compression_type);
+    }
+
     let pe_image = match xex2tractor::extract::extract_pe_image(&data, &xex) {
         Ok(img) => img,
         Err(e) => {
@@ -58,15 +64,16 @@ fn cmd_extract(path: &PathBuf, output: Option<PathBuf>) {
         }
     };
 
-    // TODO(M6): verify PE headers before writing
     if let Err(e) = std::fs::write(&output_path, &pe_image) {
         eprintln!("Error writing {}: {e}", output_path.display());
         process::exit(1);
     }
 
-    println!("Extracted PE image to {}", output_path.display());
-    println!("  Input:  {} ({} bytes)", path.display(), data.len());
-    println!("  Output: {} ({} bytes)", output_path.display(), pe_image.len());
+    println!(
+        "Extracted PE image ({} bytes) -> {}",
+        pe_image.len(),
+        output_path.display()
+    );
 }
 
 fn read_file(path: &PathBuf) -> Vec<u8> {

tests/integration.rs

@@ -409,3 +409,34 @@ fn test_cli_extract_default_output_path() {
     // Clean up
     let _ = std::fs::remove_file(&expected_output);
 }
+
+#[test]
+fn test_extract_pe_verification_runs() {
+    let data = sample_data();
+    let xex = xex2tractor::parse(&data).unwrap();
+
+    let pe_image = extract::extract_pe_image(&data, &xex).unwrap();
+
+    // verify_pe_image should succeed on the extracted image
+    assert!(extract::verify_pe_image(&pe_image).is_ok());
+}
+
+#[test]
+fn test_cli_extract_shows_format_info() {
+    let path = format!("{}/tests/data/default.xex", env!("CARGO_MANIFEST_DIR"));
+    let output_path = format!("{}/target/test_extract_info.exe", env!("CARGO_MANIFEST_DIR"));
+    let _ = std::fs::remove_file(&output_path);
+
+    let output = std::process::Command::new(env!("CARGO_BIN_EXE_xex2tractor"))
+        .args(["extract", &path, &output_path])
+        .output()
+        .expect("failed to run xex2tractor");
+
+    assert!(output.status.success());
+    let stdout = String::from_utf8_lossy(&output.stdout);
+    assert!(stdout.contains("Encryption:"));
+    assert!(stdout.contains("Compression:"));
+    assert!(stdout.contains("Extracted PE image"));
+
+    let _ = std::fs::remove_file(&output_path);
+}