xex2tractor/tests/integration.rs

use xex2tractor::crypto;
use xex2tractor::extract;
use xex2tractor::header::{ModuleFlags, XEX2_MAGIC};
use xex2tractor::optional::{CompressionInfo, CompressionType, EncryptionType, SystemFlags};
use xex2tractor::security::{ImageFlags, MediaFlags, RegionFlags};

fn sample_data() -> Vec<u8> {
    let path = format!("{}/tests/data/default.xex", env!("CARGO_MANIFEST_DIR"));
    std::fs::read(&path).expect("sample file should exist at tests/data/default.xex")
}

// ── Header tests ──────────────────────────────────────────────────────────────

#[test]
fn test_full_parse() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();

    assert_eq!(xex.header.magic, XEX2_MAGIC);
    assert_eq!(xex.header.module_flags, ModuleFlags(0x00000001));
    assert_eq!(xex.header.header_size, 0x00003000);
    assert_eq!(xex.header.reserved, 0x00000000);
    assert_eq!(xex.header.security_offset, 0x00000090);
    assert_eq!(xex.header.header_count, 15);
}

#[test]
fn test_parse_empty_file() {
    let data = vec![];
    assert!(xex2tractor::parse(&data).is_err());
}

#[test]
fn test_parse_invalid_magic() {
    let mut data = sample_data();
    data[0] = 0x00;
    assert!(xex2tractor::parse(&data).is_err());
}

// ── Optional header tests ─────────────────────────────────────────────────────

#[test]
fn test_optional_headers_all_present() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();
    let opt = &xex.optional_headers;

    // All 15 entries should be parsed
    assert_eq!(opt.entries.len(), 15);

    // Verify presence of all expected headers
    assert!(opt.entry_point.is_some());
    assert!(opt.image_base_address.is_some());
    assert!(opt.default_stack_size.is_some());
    assert!(opt.system_flags.is_some());
    assert!(opt.execution_info.is_some());
    assert!(opt.file_format_info.is_some());
    assert!(opt.checksum_timestamp.is_some());
    assert!(opt.original_pe_name.is_some());
    assert!(opt.tls_info.is_some());
    assert!(opt.static_libraries.is_some());
    assert!(opt.import_libraries.is_some());
    assert!(opt.resource_info.is_some());
    assert!(opt.game_ratings.is_some());
    assert!(opt.lan_key.is_some());
    assert!(opt.xbox360_logo_size.is_some());
}

#[test]
fn test_optional_inline_values() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();
    let opt = &xex.optional_headers;

    assert_eq!(opt.entry_point.unwrap(), 0x824AB748);
    assert_eq!(opt.image_base_address.unwrap(), 0x82000000);
    assert_eq!(opt.default_stack_size.unwrap(), 0x00080000);
    assert_eq!(opt.system_flags.unwrap(), SystemFlags(0x00000400));
}

#[test]
fn test_optional_execution_info() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();
    let exec = xex.optional_headers.execution_info.as_ref().unwrap();

    assert_eq!(exec.title_id, 0x535107D4);
    assert_eq!(exec.media_id, 0x2D2E2EEB);
}

#[test]
fn test_optional_file_format() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();
    let fmt = xex.optional_headers.file_format_info.as_ref().unwrap();

    assert_eq!(fmt.encryption_type, EncryptionType::Normal);
    assert_eq!(fmt.compression_type, CompressionType::Normal);
}

#[test]
fn test_optional_static_libraries() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();
    let libs = xex.optional_headers.static_libraries.as_ref().unwrap();

    assert_eq!(libs.len(), 12);
    // Verify first and a few known libraries
    assert_eq!(libs[0].name, "XAPILIB");
    assert_eq!(libs[1].name, "D3D9");
    assert_eq!(libs[3].name, "XBOXKRNL");
}

#[test]
fn test_optional_import_libraries() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();
    let imports = xex.optional_headers.import_libraries.as_ref().unwrap();

    assert_eq!(imports.string_table.len(), 2);
    assert_eq!(imports.string_table[0], "xam.xex");
    assert_eq!(imports.string_table[1], "xboxkrnl.exe");
    assert!(!imports.libraries.is_empty());
}

// ── Security info tests ───────────────────────────────────────────────────────

#[test]
fn test_security_info_parsed() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();
    let sec = &xex.security_info;

    assert_eq!(sec.header_size, 0x00000F34);
    assert_eq!(sec.image_size, 0x00920000);
    assert_eq!(sec.unk_108, 0x00000174);
    assert_eq!(sec.load_address, 0x82000000);
    assert_eq!(sec.import_table_count, 2);
    assert_eq!(sec.export_table, 0);
}

#[test]
fn test_security_flags() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();
    let sec = &xex.security_info;

    assert_eq!(sec.image_flags, ImageFlags(0x00000008));
    assert_eq!(sec.region, RegionFlags(0xFFFFFFFF));
    assert_eq!(sec.allowed_media_types, MediaFlags(0x00000004));
}

#[test]
fn test_security_page_descriptors() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();
    let sec = &xex.security_info;

    assert_eq!(sec.page_descriptor_count, 146);
    assert_eq!(sec.page_descriptors.len(), 146);

    // First descriptor has page_count = 19
    assert_eq!(sec.page_descriptors[0].page_count, 19);

    // Page size should be 64KB (4KB flag is not set)
    assert_eq!(sec.image_flags.page_size(), 0x10000);

    // Each page descriptor should have a valid page_count
    for desc in &sec.page_descriptors {
        assert!(desc.page_count > 0, "page_count should be positive");
    }
}

#[test]
fn test_security_crypto_fields() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();
    let sec = &xex.security_info;

    // RSA signature starts with 2C94EBE6
    assert_eq!(sec.rsa_signature[0..4], [0x2C, 0x94, 0xEB, 0xE6]);

    // XGD2 media ID starts with 3351
    assert_eq!(sec.xgd2_media_id[0..2], [0x33, 0x51]);

    // AES key starts with EACB
    assert_eq!(sec.aes_key[0..2], [0xEA, 0xCB]);
}

// ── Compression info tests ────────────────────────────────────────────────────

#[test]
fn test_compression_info_normal() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();
    let fmt = xex.optional_headers.file_format_info.as_ref().unwrap();

    match &fmt.compression_info {
        CompressionInfo::Normal {
            window_size,
            first_block,
        } => {
            // Window size should be a power of 2
            assert!(window_size.is_power_of_two(), "window_size should be power of 2");
            assert!(*window_size > 0);
            // First block should have non-zero size
            assert!(first_block.block_size > 0);
            // Block hash should not be all zeros
            assert!(!first_block.block_hash.iter().all(|&b| b == 0));
        }
        other => panic!("expected Normal compression info, got {other:?}"),
    }
}

// ── Crypto tests ─────────────────────────────────────────────────────────────

#[test]
fn test_session_key_derivation() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();

    let session_key = crypto::derive_session_key(&xex.security_info.aes_key);
    // Session key should be non-zero
    assert!(!session_key.iter().all(|&b| b == 0));
    // Session key should differ from encrypted key
    assert_ne!(&session_key[..], &xex.security_info.aes_key[..]);
}

#[test]
fn test_payload_decryption() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();

    let session_key = crypto::derive_session_key(&xex.security_info.aes_key);

    // Decrypt the first 256 bytes of payload
    let payload_start = xex.header.header_size as usize;
    let mut payload_head = data[payload_start..payload_start + 256].to_vec();
    let original = payload_head.clone();

    crypto::decrypt_in_place(&session_key, &mut payload_head);

    // Decrypted data should differ from encrypted
    assert_ne!(payload_head, original);
}

// ── CLI tests ─────────────────────────────────────────────────────────────────

#[test]
fn test_cli_inspect_with_sample() {
    let path = format!("{}/tests/data/default.xex", env!("CARGO_MANIFEST_DIR"));
    let output = std::process::Command::new(env!("CARGO_BIN_EXE_xex2tractor"))
        .args(["inspect", &path])
        .output()
        .expect("failed to run xex2tractor");

    assert!(output.status.success(), "CLI inspect should exit successfully");
    let stdout = String::from_utf8_lossy(&output.stdout);

    // Header section
    assert!(stdout.contains("XEX2 Header"));
    assert!(stdout.contains("0x58455832"));
    assert!(stdout.contains("TITLE"));
    assert!(stdout.contains("Header Count:     15"));

    // Optional headers section
    assert!(stdout.contains("Optional Headers (15 entries)"));
    assert!(stdout.contains("[ENTRY_POINT] 0x824AB748"));
    assert!(stdout.contains("[IMAGE_BASE_ADDRESS] 0x82000000"));
    assert!(stdout.contains("EXECUTION_INFO"));
    assert!(stdout.contains("0x535107D4")); // title ID
    assert!(stdout.contains("FILE_FORMAT_INFO"));
    assert!(stdout.contains("Normal (AES-128-CBC)"));
    assert!(stdout.contains("Normal (LZX)"));
    assert!(stdout.contains("Window Size:")); // new compression info
    assert!(stdout.contains("First Block:")); // new compression info
    assert!(stdout.contains("STATIC_LIBRARIES"));
    assert!(stdout.contains("XAPILIB"));
    assert!(stdout.contains("IMPORT_LIBRARIES"));
    assert!(stdout.contains("xboxkrnl.exe"));
    assert!(stdout.contains("default.pe")); // original PE name
    assert!(stdout.contains("PAL50_INCOMPATIBLE")); // system flags

    // Security info section
    assert!(stdout.contains("Security Info"));
    assert!(stdout.contains("0x00000F34")); // header size
    assert!(stdout.contains("0x00920000")); // image size
    assert!(stdout.contains("XGD2_MEDIA_ONLY")); // image flags
    assert!(stdout.contains("ALL REGIONS")); // region
    assert!(stdout.contains("DVD_CD")); // media type
    assert!(stdout.contains("Page Descriptors")); // page descriptors section
}

#[test]
fn test_cli_no_args() {
    let output = std::process::Command::new(env!("CARGO_BIN_EXE_xex2tractor"))
        .output()
        .expect("failed to run xex2tractor");

    assert!(!output.status.success());
    let stderr = String::from_utf8_lossy(&output.stderr);
    assert!(stderr.contains("Usage"));
}

#[test]
fn test_cli_inspect_missing_file() {
    let output = std::process::Command::new(env!("CARGO_BIN_EXE_xex2tractor"))
        .args(["inspect", "/nonexistent/file.xex"])
        .output()
        .expect("failed to run xex2tractor");

    assert!(!output.status.success());
}

// ── Extraction tests ─────────────────────────────────────────────────────────

#[test]
fn test_extract_pe_image() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();

    let pe_image = extract::extract_pe_image(&data, &xex).unwrap();

    // Output size should match security_info.image_size
    assert_eq!(pe_image.len(), xex.security_info.image_size as usize);
}

#[test]
fn test_extract_pe_starts_with_mz() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();

    let pe_image = extract::extract_pe_image(&data, &xex).unwrap();

    // PE image must start with MZ signature (0x4D5A)
    assert_eq!(pe_image[0], 0x4D, "first byte should be 'M'");
    assert_eq!(pe_image[1], 0x5A, "second byte should be 'Z'");
}

#[test]
fn test_extract_pe_has_valid_pe_header() {
    let data = sample_data();
    let xex = xex2tractor::parse(&data).unwrap();

    let pe_image = extract::extract_pe_image(&data, &xex).unwrap();

    // Read e_lfanew from DOS header (offset 0x3C, little-endian per PE spec)
    let e_lfanew = u32::from_le_bytes([
        pe_image[0x3C],
        pe_image[0x3D],
        pe_image[0x3E],
        pe_image[0x3F],
    ]) as usize;

    // PE signature at e_lfanew: "PE\0\0"
    assert_eq!(&pe_image[e_lfanew..e_lfanew + 4], b"PE\0\0");

    // Machine type at e_lfanew + 4: 0x01F2 (IMAGE_FILE_MACHINE_POWERPCBE, little-endian)
    let machine = u16::from_le_bytes([pe_image[e_lfanew + 4], pe_image[e_lfanew + 5]]);
    assert_eq!(machine, 0x01F2, "machine should be POWERPCBE");
}

#[test]
fn test_cli_extract_writes_file() {
    let path = format!("{}/tests/data/default.xex", env!("CARGO_MANIFEST_DIR"));
    let output_path = format!("{}/target/test_extract_output.exe", env!("CARGO_MANIFEST_DIR"));

    // Clean up any previous test output
    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(), "CLI extract should succeed");

    // Verify output file exists and starts with MZ
    let extracted = std::fs::read(&output_path).expect("should be able to read extracted file");
    assert!(extracted.len() > 2);
    assert_eq!(extracted[0], 0x4D); // 'M'
    assert_eq!(extracted[1], 0x5A); // 'Z'

    // Clean up
    let _ = std::fs::remove_file(&output_path);
}

#[test]
fn test_cli_extract_default_output_path() {
    let path = format!("{}/tests/data/default.xex", env!("CARGO_MANIFEST_DIR"));
    let expected_output = format!("{}/tests/data/default.exe", env!("CARGO_MANIFEST_DIR"));

    // Clean up any previous test output
    let _ = std::fs::remove_file(&expected_output);

    let output = std::process::Command::new(env!("CARGO_BIN_EXE_xex2tractor"))
        .args(["extract", &path])
        .output()
        .expect("failed to run xex2tractor");

    assert!(output.status.success(), "CLI extract should succeed");

    // Verify default output path was used
    assert!(
        std::fs::metadata(&expected_output).is_ok(),
        "default output file should exist"
    );

    // 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);
}