libfreemkv/tests/crypto_tests.rs

//! Comprehensive roundtrip tests for CSS and AACS cryptographic implementations.
//!
//! These tests prove the cryptographic algorithms work end-to-end.
//! Tests that require access to private internals are placed as unit tests
//! inside the respective source files (css/lfsr.rs, css/crack.rs, aacs/handshake.rs).
//!
//! This file tests public API items accessible from integration tests.

use libfreemkv::aacs;
use libfreemkv::css;

// ── CSS Public API Tests ────────────────────────────────────────────────────

/// Test: css_descramble_sector modifies encrypted region and preserves header.
#[test]
fn css_descramble_sector_roundtrip_via_public_api() {
    let state = css::CssState {
        title_key: [0x42, 0x13, 0x37, 0xBE, 0xEF],
        crack_span: None,
    };

    let mut sector = vec![0xAAu8; 2048];
    // Real scrambled DVD sectors are MPEG-2 PS packs, and the public API now
    // requires the pack start code as well as the flag bits — byte 0x14 means
    // something else entirely in an IFO or UDF sector.
    sector[0x00..0x04].copy_from_slice(&[0x00, 0x00, 0x01, 0xBA]);
    sector[4] = 0x44; // '01': a 13818-1 pack
    sector[0x0D] = 0xF8; // pack_stuffing_length 0
    sector[0x14] = 0x30;
    sector[0x54..0x59].copy_from_slice(&[0xDE, 0xAD, 0xBE, 0xEF, 0x42]);
    let original = sector.clone();

    css::descramble_sector(&state, &mut sector);
    assert_eq!(sector[0x14] & 0x30, 0x00, "flag not cleared");
    // Header preserved (except flag byte)
    for i in 0..128 {
        if i == 0x14 {
            continue;
        }
        assert_eq!(sector[i], original[i], "header byte {} changed", i);
    }
    // Encrypted region modified
    assert_ne!(&sector[128..256], &original[128..256]);
}

/// Test: has_scramble_flag_bits detects the raw flag bits correctly.
#[test]
fn css_is_scrambled_detection() {
    let mut sector = vec![0u8; 2048];
    assert!(
        !css::has_scramble_flag_bits(&sector),
        "empty sector should not be scrambled"
    );

    sector[0x14] = 0x10; // bit 4 set
    assert!(
        css::has_scramble_flag_bits(&sector),
        "bit 4 set should be detected"
    );

    sector[0x14] = 0x20; // bit 5 set
    assert!(
        css::has_scramble_flag_bits(&sector),
        "bit 5 set should be detected"
    );

    sector[0x14] = 0x30; // both bits set
    assert!(
        css::has_scramble_flag_bits(&sector),
        "both bits set should be detected"
    );

    sector[0x14] = 0xCF; // bits 4-5 clear, other bits set
    assert!(
        !css::has_scramble_flag_bits(&sector),
        "bits 4-5 clear should not be scrambled"
    );
}

// ── AACS Public API Tests ───────────────────────────────────────────────────

// Test 6: build a synthetic aligned unit with TS sync bytes, encrypt via
// AACS, decrypt with decrypt_unit(), and verify the plaintext matches.
#[test]
fn aacs_decrypt_unit_roundtrip() {
    use aes::Aes128;
    use aes::cipher::{Array, BlockCipherEncrypt, KeyInit};

    let unit_key = [0xAAu8; 16];
    let aacs_iv: [u8; 16] = [
        0x0B, 0xA0, 0xF8, 0xDD, 0xFE, 0xA6, 0x1F, 0xB3, 0xD8, 0xDF, 0x9F, 0x56, 0x6A, 0x05, 0x0F,
        0x78,
    ];

    // Build plaintext unit with TS sync bytes every 192 bytes starting at offset 4
    let mut plain = vec![0u8; aacs::content::ALIGNED_UNIT_LEN];
    let mut offset = 4;
    while offset < aacs::content::ALIGNED_UNIT_LEN {
        plain[offset] = 0x47; // TS sync byte
        offset += 192;
    }
    // Flag the unit encrypted via the CPI bits (byte 0), the gate `decrypt_unit`
    // consults. Set before snapshotting `expected`: decrypt preserves the
    // plaintext header, so the round-tripped unit carries these bits too.
    plain[0] |= 0xC0;

    // Save original plaintext for comparison
    let expected = plain.clone();

    // Encrypt: replicate the AACS encryption algorithm (reverse of decrypt_unit)
    let header: [u8; 16] = plain[..16].try_into().unwrap();

    // Step 1: AES-ECB encrypt header with unit key
    let cipher_header = Aes128::new(&unit_key.into());
    let mut block: Array<u8, _> = header.into();
    cipher_header.encrypt_block(&mut block);
    let mut derived = [0u8; 16];
    derived.copy_from_slice(&block);

    // Step 2: XOR to get per-unit decryption key
    let mut encrypt_key = [0u8; 16];
    for i in 0..16 {
        encrypt_key[i] = derived[i] ^ header[i];
    }

    // Step 3: AES-CBC encrypt bytes 16..6144
    let cipher = Aes128::new(&encrypt_key.into());
    let mut prev = aacs_iv;
    let num_blocks = (aacs::content::ALIGNED_UNIT_LEN - 16) / 16;
    for i in 0..num_blocks {
        let off = 16 + i * 16;
        for j in 0..16 {
            plain[off + j] ^= prev[j];
        }
        let mut c_blk = [0u8; 16];
        c_blk.copy_from_slice(&plain[off..off + 16]);
        let mut blk: Array<u8, _> = c_blk.into();
        cipher.encrypt_block(&mut blk);
        plain[off..off + 16].copy_from_slice(&blk);
        prev.copy_from_slice(&plain[off..off + 16]);
    }

    // Verify it looks encrypted (body TS syncs scrambled)
    assert!(!aacs::content::is_clean(
        &plain,
        libfreemkv::disc::ContentFormat::BdTs
    ));

    // Now decrypt
    libfreemkv::test_util::decrypt_unit(&mut plain, &unit_key);
    assert!(
        aacs::content::is_clean(&plain, libfreemkv::disc::ContentFormat::BdTs),
        "decrypted unit should read as clear (TS syncs restored)"
    );

    // Verify TS sync bytes at expected positions (flag byte is cleared by decrypt)
    let mut sync_count = 0;
    let mut off = 4;
    while off < aacs::content::ALIGNED_UNIT_LEN {
        if plain[off] == 0x47 {
            sync_count += 1;
        }
        off += 192;
    }
    let expected_syncs = (aacs::content::ALIGNED_UNIT_LEN - 4) / 192 + 1;
    assert_eq!(
        sync_count, expected_syncs,
        "TS sync bytes not recovered: got {}, expected {}",
        sync_count, expected_syncs
    );

    // Decryption clears no flag, so the unit round-trips byte-for-byte.
    assert_eq!(
        &plain[..],
        &expected[..],
        "decrypted unit does not match original plaintext"
    );
}

/// Test 7: aacs_disc_hash_deterministic
///
/// compute disc_hash on the same data twice, verify identical results.
#[test]
fn aacs_disc_hash_deterministic() {
    let data1 = b"Unit_Key_RO.inf test data for deterministic hashing";
    let data2 = b"Different data should produce different hash";

    let hash1a = aacs::inf::disc_hash(data1);
    let hash1b = aacs::inf::disc_hash(data1);
    assert_eq!(hash1a, hash1b, "disc_hash not deterministic on same input");

    let hash2 = aacs::inf::disc_hash(data2);
    assert_ne!(
        hash1a, hash2,
        "different inputs should produce different hashes"
    );

    // Verify it is a 20-byte SHA-1 hash
    assert_eq!(hash1a.len(), 20);

    // Verify disc_hash_hex formatting
    let hex = aacs::inf::disc_hash_hex(&hash1a);
    assert!(hex.starts_with("0x"), "hex should start with 0x prefix");
    assert_eq!(
        hex.len(),
        42,
        "hex string should be 42 chars (0x + 40 hex digits)"
    );
}

// Test: verify that encrypting a unit key with AES-ECB and decrypting it
// with decrypt_unit_key recovers the original.
#[test]
fn aacs_decrypt_unit_key_roundtrip() {
    use aes::Aes128;
    use aes::cipher::{Array, BlockCipherEncrypt, KeyInit};

    let vuk = [
        0x11u8, 0x14, 0x36, 0x0B, 0x10, 0xEE, 0x6E, 0xAC, 0x78, 0xAA, 0x4A, 0xC0, 0xB7, 0x52, 0xEA,
        0xEB,
    ];
    let original_unit_key = [
        0x9E, 0x5D, 0x13, 0x10, 0x33, 0x74, 0x43, 0xE8, 0x11, 0xA5, 0x2E, 0xBB, 0xEA, 0xE0, 0x47,
        0x0F,
    ];

    // Encrypt: AES-ECB encrypt the unit key with VUK
    let cipher = Aes128::new(&vuk.into());
    let mut block: Array<u8, _> = original_unit_key.into();
    cipher.encrypt_block(&mut block);
    let mut encrypted_uk = [0u8; 16];
    encrypted_uk.copy_from_slice(&block);

    // Decrypt with the public API
    let decrypted = aacs::derive::decrypt_unit_key(&vuk, &encrypted_uk);
    assert_eq!(
        decrypted, original_unit_key,
        "decrypt_unit_key did not recover original unit key"
    );
}

/// Test: `is_clean` distinguishes clean vs scrambled units via the TS proof floor.
#[test]
fn aacs_is_clean_detection() {
    // A clear unit: TS sync (0x47) intact at every 192-byte packet → not
    // scrambled. (Flag bits play no role.)
    let mut clear = vec![0u8; aacs::content::ALIGNED_UNIT_LEN];
    let mut off = 4;
    while off < aacs::content::ALIGNED_UNIT_LEN {
        clear[off] = 0x47;
        off += 192;
    }
    assert!(
        aacs::content::is_clean(&clear, libfreemkv::disc::ContentFormat::BdTs),
        "clear unit (syncs intact) must not be scrambled"
    );

    // Flag bits do NOT make a synced unit look scrambled.
    let mut flagged = clear.clone();
    flagged[0] = 0xC0; // copy-control bits
    flagged[7] = 0xC0; // TSC bits
    assert!(
        aacs::content::is_clean(&flagged, libfreemkv::disc::ContentFormat::BdTs),
        "flag bits must not be read as encryption"
    );

    // A scrambled body (syncs destroyed) → scrambled.
    let scrambled = vec![0x99u8; aacs::content::ALIGNED_UNIT_LEN];
    assert!(
        !aacs::content::is_clean(&scrambled, libfreemkv::disc::ContentFormat::BdTs),
        "unit with no intact TS syncs must read as scrambled"
    );

    // Too short
    let short = vec![0xFFu8; 100];
    assert!(
        aacs::content::is_clean(&short, libfreemkv::disc::ContentFormat::BdTs),
        "short buffer should not be detected"
    );
}

// Test: decrypt_unit is PURE (applies the key unconditionally). The "leave
// a clear unit untouched" policy lives at the caller's gate
// aacs_unit_encrypted: a CPI-clear unit reports not-encrypted.
#[test]
fn aacs_clear_unit_reports_not_encrypted() {
    let mut unit = vec![0x42u8; aacs::content::ALIGNED_UNIT_LEN];
    // Intact TS syncs every 192 bytes → not scrambled.
    let mut off = 4;
    while off < aacs::content::ALIGNED_UNIT_LEN {
        unit[off] = 0x47;
        off += 192;
    }
    // CPI bits (byte 0) CLEAR → the authoritative gate reads this as plaintext.
    unit[0] &= 0x3F;

    assert!(aacs::content::is_clean(
        &unit,
        libfreemkv::disc::ContentFormat::BdTs
    ));
    assert!(
        !aacs::content::aacs_unit_encrypted(&unit, libfreemkv::disc::ContentFormat::BdTs),
        "CPI-clear unit reports not-encrypted; the caller never decrypts it"
    );
}

// ── AACS cross-validation with independent AES implementation ──────────────

/// Independent AES-128-ECB encrypt (uses `aes` crate directly, NOT our library).
fn ref_aes_ecb_encrypt(key: &[u8; 16], data: &[u8; 16]) -> [u8; 16] {
    use aes::Aes128;
    use aes::cipher::{Array, BlockCipherEncrypt, KeyInit};
    let cipher = Aes128::new(&(*key).into());
    let mut block: Array<u8, _> = (*data).into();
    cipher.encrypt_block(&mut block);
    let mut out = [0u8; 16];
    out.copy_from_slice(&block);
    out
}

/// Independent AES-128-CBC encrypt (uses `aes` crate directly, NOT our library).
fn ref_aes_cbc_encrypt(key: &[u8; 16], iv: &[u8; 16], data: &mut [u8]) {
    use aes::Aes128;
    use aes::cipher::{Array, BlockCipherEncrypt, KeyInit};
    let cipher = Aes128::new(&(*key).into());
    let mut prev = *iv;
    let num_blocks = data.len() / 16;
    for i in 0..num_blocks {
        let off = i * 16;
        for j in 0..16 {
            data[off + j] ^= prev[j];
        }
        let mut c_block = [0u8; 16];
        c_block.copy_from_slice(&data[off..off + 16]);
        let mut block: Array<u8, _> = c_block.into();
        cipher.encrypt_block(&mut block);
        data[off..off + 16].copy_from_slice(&block);
        prev.copy_from_slice(&data[off..off + 16]);
    }
}

/// The standard AACS IV, copied here independently so we are NOT importing
/// the library's constant — this IS the cross-validation reference value.
const CROSS_AACS_IV: [u8; 16] = [
    0x0B, 0xA0, 0xF8, 0xDD, 0xFE, 0xA6, 0x1F, 0xB3, 0xD8, 0xDF, 0x9F, 0x56, 0x6A, 0x05, 0x0F, 0x78,
];

/// Build a plaintext aligned unit with TS sync markers and recognisable
/// content, encrypt it using only the `aes` crate (independent of the
/// library), then decrypt with `decrypt_unit()` and verify the match.
#[test]
fn aacs_cross_validation_encrypt_then_decrypt() {
    let unit_key: [u8; 16] = [
        0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF, 0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32,
        0x10,
    ];

    let mut plaintext = vec![0u8; aacs::content::ALIGNED_UNIT_LEN];
    // TS sync bytes every 192 bytes starting at offset 4
    let mut off = 4;
    while off < aacs::content::ALIGNED_UNIT_LEN {
        plaintext[off] = 0x47;
        off += 192;
    }
    // Fill the rest with a recognisable pattern (prime modulus avoids artefacts).
    // Index-arithmetic — clearer with a counted loop than an enumerate chain.
    #[allow(clippy::needless_range_loop)]
    for i in 16..aacs::content::ALIGNED_UNIT_LEN {
        if plaintext[i] == 0 {
            plaintext[i] = (i % 251) as u8;
        }
    }
    // Flag the unit encrypted via the CPI bits (byte 0) — the gate `decrypt_unit`
    // consults. Set before snapshotting `expected`; decrypt preserves the header.
    plaintext[0] |= 0xC0;

    let expected = plaintext.clone();

    // -- Encrypt with independent implementation --
    let mut header = [0u8; 16];
    header.copy_from_slice(&plaintext[..16]);
    let derived = ref_aes_ecb_encrypt(&unit_key, &header);
    let mut dk = [0u8; 16];
    for i in 0..16 {
        dk[i] = derived[i] ^ header[i];
    }
    ref_aes_cbc_encrypt(
        &dk,
        &CROSS_AACS_IV,
        &mut plaintext[16..aacs::content::ALIGNED_UNIT_LEN],
    );

    // Sanity: ciphertext should differ
    assert_ne!(
        &plaintext[16..32],
        &expected[16..32],
        "encryption did not change ciphertext region"
    );

    // -- Decrypt with the library --
    libfreemkv::test_util::decrypt_unit(&mut plaintext, &unit_key);

    // Decryption clears no flag, so the unit round-trips byte-for-byte.
    assert_eq!(
        &plaintext[..],
        &expected[..],
        "decrypted unit does not match original plaintext"
    );
}

/// Same cross-validation with a different key and all-0xFF payload to
/// exercise different AES round-key schedules.
#[test]
fn aacs_cross_validation_alternate_key() {
    let unit_key: [u8; 16] = [
        0xDE, 0xAD, 0xBE, 0xEF, 0xCA, 0xFE, 0xBA, 0xBE, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
        0x08,
    ];

    let mut plaintext = vec![0xFFu8; aacs::content::ALIGNED_UNIT_LEN];
    let mut off = 4;
    while off < aacs::content::ALIGNED_UNIT_LEN {
        plaintext[off] = 0x47;
        off += 192;
    }
    // Byte 0 stays 0xFF (CPI bits set); decrypt_unit ignores the flag. The CBC-encrypted
    // body scrambles the packet syncs.
    let expected = plaintext.clone();

    let mut header = [0u8; 16];
    header.copy_from_slice(&plaintext[..16]);
    let derived = ref_aes_ecb_encrypt(&unit_key, &header);
    let mut dk = [0u8; 16];
    for i in 0..16 {
        dk[i] = derived[i] ^ header[i];
    }
    ref_aes_cbc_encrypt(
        &dk,
        &CROSS_AACS_IV,
        &mut plaintext[16..aacs::content::ALIGNED_UNIT_LEN],
    );

    libfreemkv::test_util::decrypt_unit(&mut plaintext, &unit_key);

    // Decryption clears no flag, so the unit round-trips byte-for-byte.
    assert_eq!(&plaintext[..], &expected[..]);
}

// (`decrypt_bus` is a crate-internal layer — its cross-validation lives in-crate
// in `aacs::content`'s unit tests, not here.)

// ── CSS roundtrip test vectors ─────────────────────────────────────────────

/// CSS descramble modifies encrypted region deterministically.
#[test]
fn css_roundtrip_with_snapshot() {
    let title_key: [u8; 5] = [0x42, 0x13, 0x37, 0xBE, 0xEF];

    let mut sector = vec![0x00u8; 2048];
    sector[0x14] = 0x30;
    sector[0x54..0x59].copy_from_slice(&[0xDE, 0xAD, 0xBE, 0xEF, 0x42]);
    #[allow(clippy::needless_range_loop)]
    for i in 0x80..2048 {
        sector[i] = ((i * 7 + 3) & 0xFF) as u8;
    }
    let original = sector.clone();

    css::lfsr::descramble_sector(&title_key, &mut sector);
    assert_eq!(sector[0x14] & 0x30, 0x00, "flag not cleared");
    assert_ne!(&sector[0x80..0xA0], &original[0x80..0xA0]);

    // Determinism: same input → same output
    let mut sector2 = original.clone();
    css::lfsr::descramble_sector(&title_key, &mut sector2);
    assert_eq!(
        &sector[0x80..2048],
        &sector2[0x80..2048],
        "not deterministic"
    );
}

/// Multiple key/seed combinations produce different outputs.
#[test]
fn css_roundtrip_multiple_keys() {
    // The cipher seeds its LFSRs from `key XOR seed`, so cases need DISTINCT
    // key^seed values to produce distinct keystreams (key=0/seed=0 and
    // key=FF/seed=FF both give key^seed=0 and are excluded here).
    let cases: &[([u8; 5], [u8; 5])] = &[
        (
            [0x00, 0x00, 0x00, 0x00, 0x00],
            [0x00, 0x00, 0x00, 0x00, 0x00],
        ),
        (
            [0xFF, 0xFF, 0xFF, 0xFF, 0xFF],
            [0x00, 0x00, 0x00, 0x00, 0x00],
        ),
        (
            [0x01, 0x02, 0x03, 0x04, 0x05],
            [0xAA, 0xBB, 0xCC, 0xDD, 0xEE],
        ),
        (
            [0xAB, 0xCD, 0xEF, 0x01, 0x23],
            [0x12, 0x34, 0x56, 0x78, 0x9A],
        ),
    ];

    let mut results = Vec::new();
    for (idx, (key, seed)) in cases.iter().enumerate() {
        let mut sector = vec![0xAAu8; 2048];
        sector[0x14] = 0x30;
        sector[0x54..0x59].copy_from_slice(seed);

        css::lfsr::descramble_sector(key, &mut sector);
        assert_eq!(sector[0x14] & 0x30, 0x00, "case {}: flag not cleared", idx);
        results.push(sector[0x80..0xA0].to_vec());
    }
    // Different keys/seeds should produce different outputs
    for i in 0..results.len() {
        for j in (i + 1)..results.len() {
            assert_ne!(
                results[i], results[j],
                "cases {} and {} produced same output",
                i, j
            );
        }
    }
}

// ── CSS keyless recovery tests ─────────────────────────────────────────────

/// Test: aacs_parse_unit_key_ro with minimal valid data
#[test]
fn aacs_parse_unit_key_ro_minimal() {
    // Build a minimal Unit_Key_RO.inf structure
    // Header: first 4 bytes = BE32 offset to key storage area
    let uk_pos: u32 = 100;
    let mut data = vec![0u8; 200];

    // Key storage offset
    data[0..4].copy_from_slice(&uk_pos.to_be_bytes());
    // app_type
    data[16] = 1; // BD-ROM
    // num_bdmv_dir
    data[17] = 1;
    // flags
    data[18] = 0;

    // At uk_pos: num_unit_keys = 1
    let pos = uk_pos as usize;
    data[pos] = 0;
    data[pos + 1] = 1; // 1 key

    // At uk_pos + 48: first encrypted key (16 bytes)
    let key_pos = pos + 48;
    for i in 0..16 {
        data[key_pos + i] = (0xA0 + i) as u8;
    }

    let result = aacs::inf::parse_unit_key_ro(&data, aacs::mkb::AacsVersion::V10);
    assert!(
        result.is_some(),
        "parse_unit_key_ro should succeed on valid data"
    );

    let ukf = result.unwrap();
    assert_eq!(ukf.app_type, 1);
    assert_eq!(ukf.num_bdmv_dir, 1);
    assert_eq!(ukf.encrypted_keys.len(), 1);
    assert_eq!(ukf.disc_hash.len(), 20);

    // disc_hash should be deterministic
    let hash = aacs::inf::disc_hash(&data);
    assert_eq!(ukf.disc_hash, hash);
}