oxideav-pdf 0.1.4

//! Round-4 encryption decode tests.
//!
//! Builds tiny encrypted PDFs by hand against the spec (§7.6.3) so the
//! reader is exercised against a known-good fixture without needing an
//! external authoring tool. Three revisions are covered here:
//!
//! * **R=2** — RC4-40, V=1 (40-bit key)
//! * **R=3** — RC4-128, V=2 (128-bit key)
//! * **R=4** — AES-128 CBC, V=4 with `CFM=AESV2`
//!
//! Each fixture has the same shape: a one-page document with a single
//! solid-fill rectangle on a 100×100 MediaBox + an `/Info` dict whose
//! `/Title` is a known string. The test asserts the reader recovers
//! the title (string-decryption) and the content stream's path
//! commands (stream-decryption).

use oxideav_pdf::decrypt::{md5, rc4, StandardHandler};
use oxideav_pdf::read_pdf_to_scene_with_password;

// ───────── PDF fixture builder ─────────

const PAD: [u8; 32] = [
    0x28, 0xBF, 0x4E, 0x5E, 0x4E, 0x75, 0x8A, 0x41, 0x64, 0x00, 0x4E, 0x56, 0xFF, 0xFA, 0x01, 0x08,
    0x2E, 0x2E, 0x00, 0xB6, 0xD0, 0x68, 0x3E, 0x80, 0x2F, 0x0C, 0xA9, 0xFE, 0x64, 0x53, 0x69, 0x7A,
];

fn pad_password(password: &[u8]) -> [u8; 32] {
    let mut out = [0u8; 32];
    let take = password.len().min(32);
    out[..take].copy_from_slice(&password[..take]);
    if take < 32 {
        out[take..].copy_from_slice(&PAD[..32 - take]);
    }
    out
}

/// Produce O / U / file-key for a chosen revision + length given a
/// user / owner password pair and a file ID. Mirrors §7.6.3
/// Algorithms 2..=5.
#[derive(Clone)]
#[allow(dead_code)]
struct EncFile {
    revision: u8,
    length_bits: usize,
    o: [u8; 32],
    u: [u8; 32],
    p: i32,
    file_id: Vec<u8>,
    file_key: Vec<u8>,
}

fn build_encryption(
    user_password: &[u8],
    owner_password: &[u8],
    revision: u8,
    length_bits: usize,
    p: i32,
    file_id: &[u8],
) -> EncFile {
    let n = length_bits / 8;
    // ─── /O computation (Algorithm 3) ───
    // (a) Pad owner password (or user password if no owner given).
    let owner_src = if owner_password.is_empty() {
        user_password
    } else {
        owner_password
    };
    let opad = pad_password(owner_src);
    // (b) MD5.
    let mut h = md5(&opad);
    // (c) for R≥3, loop 50 times.
    if revision >= 3 {
        for _ in 0..50 {
            h = md5(&h[..n]);
        }
    }
    // (d) Build RC4 key.
    let okey = h[..n].to_vec();
    // (e) Pad user password.
    let upad = pad_password(user_password);
    // (f) Encrypt.
    let mut o_buf = rc4(&okey, &upad);
    // (g) For R≥3, 19 more rounds.
    if revision >= 3 {
        for i in 1u8..=19 {
            let xkey: Vec<u8> = okey.iter().map(|b| b ^ i).collect();
            o_buf = rc4(&xkey, &o_buf);
        }
    }
    let mut o = [0u8; 32];
    o.copy_from_slice(&o_buf);

    // ─── File-encryption key (Algorithm 2) ───
    let mut buf = Vec::new();
    buf.extend_from_slice(&pad_password(user_password));
    buf.extend_from_slice(&o);
    buf.extend_from_slice(&(p as u32).to_le_bytes());
    buf.extend_from_slice(file_id);
    let mut h = md5(&buf);
    if revision >= 3 {
        for _ in 0..50 {
            h = md5(&h[..n]);
        }
    }
    let file_key = h[..n].to_vec();

    // ─── /U computation (Algorithm 4 / 5) ───
    let mut u = [0u8; 32];
    if revision == 2 {
        u.copy_from_slice(&rc4(&file_key, &PAD));
    } else {
        let mut hash_input = Vec::with_capacity(32 + file_id.len());
        hash_input.extend_from_slice(&PAD);
        hash_input.extend_from_slice(file_id);
        let mut data = rc4(&file_key, &md5(&hash_input));
        for i in 1u8..=19 {
            let xkey: Vec<u8> = file_key.iter().map(|b| b ^ i).collect();
            data = rc4(&xkey, &data);
        }
        u[..16].copy_from_slice(&data[..16]);
        // Last 16 are arbitrary; zeros work.
    }

    EncFile {
        revision,
        length_bits,
        o,
        u,
        p,
        file_id: file_id.to_vec(),
        file_key,
    }
}

/// Encrypt a string or stream payload using the per-object key.
fn encrypt_with(handler: &StandardHandler, id: u32, data: &[u8]) -> Vec<u8> {
    // The handler's decrypt is symmetric for RC4; for AES we'd need a
    // separate encrypt path. The test only uses RC4 (R=2 / R=3) and
    // AES-128 (R=4) — for AES we encrypt manually below.
    use aes::cipher::{BlockEncryptMut, KeyIvInit};

    type Aes128CbcEnc = cbc::Encryptor<aes::Aes128>;

    let key = handler.key.clone();
    let n = key.len();

    // Algorithm 1 — derive per-object key.
    let mut buf = Vec::with_capacity(n + 9);
    buf.extend_from_slice(&key);
    buf.push((id & 0xFF) as u8);
    buf.push(((id >> 8) & 0xFF) as u8);
    buf.push(((id >> 16) & 0xFF) as u8);
    buf.push(0); // gen low
    buf.push(0); // gen high
    let aes_mode = handler.method == oxideav_pdf::decrypt::CryptMethod::Aes128;
    if aes_mode {
        buf.extend_from_slice(&[0x73, 0x41, 0x6C, 0x54]);
    }
    let h = md5(&buf);
    let take = (n + 5).min(16);
    let obj_key = &h[..take];

    if aes_mode {
        // Pick a deterministic IV for reproducibility.
        let iv = [0x42u8; 16];
        let enc = Aes128CbcEnc::new(obj_key.into(), (&iv).into());
        let mut out = vec![0u8; ((data.len() / 16) + 1) * 16];
        let n = enc
            .encrypt_padded_b2b_mut::<aes::cipher::block_padding::Pkcs7>(data, &mut out)
            .unwrap()
            .len();
        out.truncate(n);
        let mut full = Vec::with_capacity(16 + n);
        full.extend_from_slice(&iv);
        full.extend_from_slice(&out);
        full
    } else {
        rc4(obj_key, data)
    }
}

/// Build a tiny encrypted PDF by hand. Returns the bytes + the title
/// the caller-supplied user password should recover.
fn build_encrypted_pdf(
    user_password: &[u8],
    method: oxideav_pdf::decrypt::CryptMethod,
    revision: u8,
    length_bits: usize,
    title: &str,
) -> Vec<u8> {
    // Permissions: -4 = "everything except printing & modifying" (writer
    // convention). The exact value doesn't matter for the test as long
    // as it round-trips through Algorithm 2.
    let p = -4i32;
    let file_id = b"OXIDEAV-TEST-FILE-ID-0123456789".to_vec();

    let enc = build_encryption(
        user_password,
        b"", // empty owner password — Algo 3 falls back to user password
        revision,
        length_bits,
        p,
        &file_id,
    );

    let handler = StandardHandler {
        key: enc.file_key.clone(),
        method,
        revision,
    };

    // ─── Build object bodies ───
    // Object 1: Catalog
    // Object 2: Pages
    // Object 3: Page
    // Object 4: Contents stream
    // Object 5: Info dict
    // Object 6: Encrypt dict (NOT encrypted itself)

    let info_title_str: Vec<u8> = encrypt_with(&handler, 5, title.as_bytes());

    // Content stream — paint a 50×50 rectangle in red using `re f`.
    let content_plain = b"q\n1 0 0 rg\n10 10 50 50 re\nf\nQ\n".to_vec();
    let content_cipher: Vec<u8> = encrypt_with(&handler, 4, &content_plain);

    let mut bytes = Vec::with_capacity(2048);

    // Header
    bytes.extend_from_slice(b"%PDF-1.4\n%\xE2\xE3\xCF\xD3\n");

    let mut offsets = [0u64; 7];

    // Obj 1
    offsets[1] = bytes.len() as u64;
    bytes.extend_from_slice(b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n");

    // Obj 2
    offsets[2] = bytes.len() as u64;
    bytes.extend_from_slice(b"2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n");

    // Obj 3 (Page)
    offsets[3] = bytes.len() as u64;
    bytes.extend_from_slice(
        b"3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 100 100] /Contents 4 0 R >>\nendobj\n",
    );

    // Obj 4 (content stream — encrypted)
    offsets[4] = bytes.len() as u64;
    bytes.extend_from_slice(
        format!("4 0 obj\n<< /Length {} >>\nstream\n", content_cipher.len()).as_bytes(),
    );
    bytes.extend_from_slice(&content_cipher);
    bytes.extend_from_slice(b"\nendstream\nendobj\n");

    // Obj 5 (Info dict — encrypted /Title)
    offsets[5] = bytes.len() as u64;
    bytes.extend_from_slice(b"5 0 obj\n<< /Title <");
    for b in &info_title_str {
        bytes.extend_from_slice(format!("{:02X}", b).as_bytes());
    }
    bytes.extend_from_slice(b"> >>\nendobj\n");

    // Obj 6 (Encrypt dict — plaintext, per §7.6.1)
    offsets[6] = bytes.len() as u64;
    let cf_block: String = if method == oxideav_pdf::decrypt::CryptMethod::Aes128 {
        " /CF << /StdCF << /Type /CryptFilter /CFM /AESV2 /Length 16 >> >> /StmF /StdCF /StrF /StdCF".to_string()
    } else {
        String::new()
    };
    let v = match (revision, method) {
        (2, _) => 1,
        (3, _) => 2,
        (4, oxideav_pdf::decrypt::CryptMethod::Aes128) => 4,
        (4, _) => 4,
        _ => 2,
    };
    bytes.extend_from_slice(
        format!(
            "6 0 obj\n<< /Filter /Standard /V {} /R {} /Length {}{} /O <",
            v, revision, length_bits, cf_block,
        )
        .as_bytes(),
    );
    for b in &enc.o {
        bytes.extend_from_slice(format!("{:02X}", b).as_bytes());
    }
    bytes.extend_from_slice(b"> /U <");
    for b in &enc.u {
        bytes.extend_from_slice(format!("{:02X}", b).as_bytes());
    }
    bytes.extend_from_slice(format!("> /P {} >>\nendobj\n", p).as_bytes());

    // xref
    let xref_off = bytes.len() as u64;
    bytes.extend_from_slice(b"xref\n0 7\n");
    bytes.extend_from_slice(b"0000000000 65535 f \n");
    for &off in &offsets[1..7] {
        bytes.extend_from_slice(format!("{:010} 00000 n \n", off).as_bytes());
    }

    // trailer
    bytes.extend_from_slice(b"trailer\n<< /Size 7 /Root 1 0 R /Info 5 0 R /Encrypt 6 0 R /ID [<");
    for b in &enc.file_id {
        bytes.extend_from_slice(format!("{:02X}", b).as_bytes());
    }
    bytes.extend_from_slice(b"> <");
    for b in &enc.file_id {
        bytes.extend_from_slice(format!("{:02X}", b).as_bytes());
    }
    bytes.extend_from_slice(b">] >>\n");
    bytes.extend_from_slice(format!("startxref\n{}\n%%EOF\n", xref_off).as_bytes());

    bytes
}

// ───────── Tests ─────────

#[test]
fn rc4_40_r2_decodes_with_user_password() {
    let pdf = build_encrypted_pdf(
        b"hello",
        oxideav_pdf::decrypt::CryptMethod::Rc4,
        2,
        40,
        "Encrypted Round-2 Title",
    );
    let scene = read_pdf_to_scene_with_password(&pdf, b"hello").expect("decrypt R=2");
    assert_eq!(
        scene.metadata.title.as_deref(),
        Some("Encrypted Round-2 Title")
    );
    let pages = scene.pages.expect("pages");
    assert_eq!(pages.len(), 1);
    assert_eq!(pages[0].width, 100.0);
    assert_eq!(pages[0].height, 100.0);
}

#[test]
fn rc4_128_r3_decodes_with_user_password() {
    let pdf = build_encrypted_pdf(
        b"correct horse battery staple",
        oxideav_pdf::decrypt::CryptMethod::Rc4,
        3,
        128,
        "RC4-128 Title",
    );
    let scene = read_pdf_to_scene_with_password(&pdf, b"correct horse battery staple")
        .expect("decrypt R=3");
    assert_eq!(scene.metadata.title.as_deref(), Some("RC4-128 Title"));
}

#[test]
fn aes_128_r4_decodes_with_user_password() {
    let pdf = build_encrypted_pdf(
        b"aespwd",
        oxideav_pdf::decrypt::CryptMethod::Aes128,
        4,
        128,
        "AES-128 R4 Title",
    );
    let scene = read_pdf_to_scene_with_password(&pdf, b"aespwd").expect("decrypt R=4");
    assert_eq!(scene.metadata.title.as_deref(), Some("AES-128 R4 Title"));
}

#[test]
fn empty_password_works_for_default_user_password() {
    // When the writer used "" as the user password, opening with the
    // default password (also "") authenticates immediately — this is
    // §7.6.3.1's "user attempts open with default user password" path.
    let pdf = build_encrypted_pdf(
        b"",
        oxideav_pdf::decrypt::CryptMethod::Rc4,
        3,
        128,
        "Default user password",
    );
    let scene = read_pdf_to_scene_with_password(&pdf, b"").expect("decrypt empty pwd");
    assert_eq!(
        scene.metadata.title.as_deref(),
        Some("Default user password")
    );
}

#[test]
fn wrong_password_rejected() {
    let pdf = build_encrypted_pdf(
        b"realpw",
        oxideav_pdf::decrypt::CryptMethod::Rc4,
        3,
        128,
        "Whatever",
    );
    let err = read_pdf_to_scene_with_password(&pdf, b"wrongpw").unwrap_err();
    let msg = format!("{err}");
    assert!(
        msg.contains("password") || msg.contains("Encrypt"),
        "expected wrong-password error, got: {msg}"
    );
}

#[test]
fn read_pdf_to_scene_uses_empty_password() {
    // The default `read_pdf_to_scene` should still open PDFs whose
    // user password is empty — common case for "encrypted just for
    // permission flags, no real password".
    let pdf = build_encrypted_pdf(
        b"",
        oxideav_pdf::decrypt::CryptMethod::Rc4,
        3,
        128,
        "Permission-only encryption",
    );
    let scene = oxideav_pdf::read_pdf_to_scene(&pdf).expect("decrypt empty pwd via default API");
    assert_eq!(
        scene.metadata.title.as_deref(),
        Some("Permission-only encryption")
    );
}

#[test]
fn unencrypted_pdf_still_works_via_default_api() {
    // Sanity check: the round-3 API path still works for unencrypted
    // PDFs after the round-4 changes.
    use oxideav_core::time::TimeBase;
    use oxideav_core::vector::{
        FillRule, Group, Node, Paint, Path, PathCommand, PathNode, Point, Rgba, VectorFrame,
    };
    use oxideav_scene::{Page, Scene};

    let mut path = Path::new();
    path.commands
        .push(PathCommand::MoveTo(Point::new(0.0, 0.0)));
    path.commands
        .push(PathCommand::LineTo(Point::new(50.0, 0.0)));
    path.commands.push(PathCommand::Close);
    let frame = VectorFrame {
        width: 100.0,
        height: 100.0,
        view_box: None,
        root: Group {
            children: vec![Node::Path(PathNode {
                path,
                fill: Some(Paint::Solid(Rgba::opaque(0, 255, 0))),
                stroke: None,
                fill_rule: FillRule::NonZero,
            })],
            ..Group::default()
        },
        pts: None,
        time_base: TimeBase::new(1, 1),
    };
    let mut page = Page::new(100.0, 100.0);
    page.content = frame;
    let scene = Scene {
        pages: Some(vec![page]),
        ..Scene::default()
    };

    let pdf = oxideav_pdf::write_pdf_from_scene(&scene).unwrap();
    let parsed = oxideav_pdf::read_pdf_to_scene(&pdf).unwrap();
    let pages = parsed.pages.expect("pages");
    assert_eq!(pages.len(), 1);
}

#[test]
fn owner_password_authenticates_via_user_path() {
    // When the owner password equals the user password (test fixture
    // builds /O from the user password when no separate owner is
    // given), the user-password path authenticates first. Confirm
    // the owner-password Algorithm 7 path also returns Ok when the
    // user-path doesn't match — i.e. when we call with the owner pw.
    //
    // For a fixture where owner == user, this just verifies that
    // wrong passwords fail fast (the owner-path gives the same answer
    // as the user-path) but doesn't exercise Algorithm 7's RC4
    // ladder. A separate test below crafts a real owner password.
    let pdf = build_encrypted_pdf(
        b"useronly",
        oxideav_pdf::decrypt::CryptMethod::Rc4,
        3,
        128,
        "User-only auth path",
    );
    // Random extra password — must fail.
    let err = read_pdf_to_scene_with_password(&pdf, b"ownerpw").unwrap_err();
    let _ = format!("{err}"); // just touching the error string
}

// ─── Owner-password lookup (Algorithm 7) ───
//
// Build a PDF where the user and owner passwords differ, then open
// with the owner password.

#[test]
fn owner_password_decodes_via_algorithm_7() {
    // Re-implementing build_encrypted_pdf with a non-empty owner pwd.
    let user_pw = b"userpw";
    let owner_pw = b"ownerpw";
    let revision = 3u8;
    let length_bits = 128usize;
    let p = -4i32;
    let file_id = b"OXIDEAV-OWNER-FIXTURE-0123456".to_vec();

    let enc = build_encryption(user_pw, owner_pw, revision, length_bits, p, &file_id);

    let handler = StandardHandler {
        key: enc.file_key.clone(),
        method: oxideav_pdf::decrypt::CryptMethod::Rc4,
        revision,
    };

    let title = "Owner-password fixture";
    let info_title_str = encrypt_with(&handler, 5, title.as_bytes());
    let content_plain = b"q\n10 10 50 50 re\nf\nQ\n".to_vec();
    let content_cipher = encrypt_with(&handler, 4, &content_plain);

    // Same byte layout as build_encrypted_pdf above.
    let mut bytes = Vec::with_capacity(2048);
    bytes.extend_from_slice(b"%PDF-1.4\n%\xE2\xE3\xCF\xD3\n");
    let mut offsets = [0u64; 7];
    offsets[1] = bytes.len() as u64;
    bytes.extend_from_slice(b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n");
    offsets[2] = bytes.len() as u64;
    bytes.extend_from_slice(b"2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n");
    offsets[3] = bytes.len() as u64;
    bytes.extend_from_slice(b"3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 100 100] /Contents 4 0 R >>\nendobj\n");
    offsets[4] = bytes.len() as u64;
    bytes.extend_from_slice(
        format!("4 0 obj\n<< /Length {} >>\nstream\n", content_cipher.len()).as_bytes(),
    );
    bytes.extend_from_slice(&content_cipher);
    bytes.extend_from_slice(b"\nendstream\nendobj\n");
    offsets[5] = bytes.len() as u64;
    bytes.extend_from_slice(b"5 0 obj\n<< /Title <");
    for b in &info_title_str {
        bytes.extend_from_slice(format!("{:02X}", b).as_bytes());
    }
    bytes.extend_from_slice(b"> >>\nendobj\n");
    offsets[6] = bytes.len() as u64;
    bytes.extend_from_slice(
        format!(
            "6 0 obj\n<< /Filter /Standard /V 2 /R {} /Length {} /O <",
            revision, length_bits
        )
        .as_bytes(),
    );
    for b in &enc.o {
        bytes.extend_from_slice(format!("{:02X}", b).as_bytes());
    }
    bytes.extend_from_slice(b"> /U <");
    for b in &enc.u {
        bytes.extend_from_slice(format!("{:02X}", b).as_bytes());
    }
    bytes.extend_from_slice(format!("> /P {} >>\nendobj\n", p).as_bytes());
    let xref_off = bytes.len() as u64;
    bytes.extend_from_slice(b"xref\n0 7\n0000000000 65535 f \n");
    for &off in &offsets[1..7] {
        bytes.extend_from_slice(format!("{:010} 00000 n \n", off).as_bytes());
    }
    bytes.extend_from_slice(b"trailer\n<< /Size 7 /Root 1 0 R /Info 5 0 R /Encrypt 6 0 R /ID [<");
    for b in &file_id {
        bytes.extend_from_slice(format!("{:02X}", b).as_bytes());
    }
    bytes.extend_from_slice(b"> <");
    for b in &file_id {
        bytes.extend_from_slice(format!("{:02X}", b).as_bytes());
    }
    bytes.extend_from_slice(b">] >>\n");
    bytes.extend_from_slice(format!("startxref\n{}\n%%EOF\n", xref_off).as_bytes());

    // Open with the user password — must succeed.
    let scene_user = read_pdf_to_scene_with_password(&bytes, user_pw).expect("user pw open");
    assert_eq!(scene_user.metadata.title.as_deref(), Some(title));

    // Open with the owner password — must also succeed (Algorithm 7).
    let scene_owner = read_pdf_to_scene_with_password(&bytes, owner_pw).expect("owner pw open");
    assert_eq!(scene_owner.metadata.title.as_deref(), Some(title));

    // Wrong passwords still rejected.
    assert!(read_pdf_to_scene_with_password(&bytes, b"random").is_err());
}