oxidize-pdf 4.1.1

//! PDF encryption detection and password handling
//!
//! This module provides functionality to detect encrypted PDFs and handle password-based
//! decryption according to ISO 32000-1 Chapter 7.6.

use super::objects::PdfDictionary;
use super::{ParseError, ParseResult};
use crate::encryption::{
    EncryptionKey, OwnerPassword, Permissions, Rc4, Rc4Key, StandardSecurityHandler, UserPassword,
};
use crate::objects::ObjectId;

/// Encryption information extracted from PDF trailer.
///
/// `#[non_exhaustive]`: this is parser output, not meant to be constructed by
/// downstream code. The attribute lets future fields (like `cfm`, added for
/// issue #364) be added without a breaking change.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct EncryptionInfo {
    /// Filter name (should be "Standard")
    pub filter: String,
    /// V entry (algorithm version)
    pub v: i32,
    /// R entry (revision)
    pub r: i32,
    /// O entry (owner password hash)
    pub o: Vec<u8>,
    /// U entry (user password hash)
    pub u: Vec<u8>,
    /// P entry (permissions)
    pub p: i32,
    /// Length entry (key length in bits)
    pub length: Option<i32>,
    /// UE entry (encrypted user key, R5/R6 only)
    pub ue: Option<Vec<u8>>,
    /// OE entry (encrypted owner key, R5/R6 only)
    pub oe: Option<Vec<u8>>,
    /// Stream crypt filter method (/CFM of the filter named by /StmF), V>=4 only.
    /// One of "V2" (RC4), "AESV2" (AES-128), "AESV3" (AES-256), "Identity".
    /// Determines the cipher for revision 4, where R alone is ambiguous.
    ///
    /// `pub(crate)`: internal cipher-selection input, not part of the public API.
    pub(crate) cfm: Option<String>,
    /// `/EncryptMetadata` flag (default true). When false and R >= 4, Algorithm 2
    /// appends 0xFFFFFFFF to the key MD5 (ISO 32000-1 §7.6.3.3, step f). Real
    /// producers often leave metadata in cleartext; honouring this is required
    /// to unlock those files, even with the empty user password (issue #379).
    ///
    /// `pub(crate)`: internal key-derivation input, not part of the public API.
    pub(crate) encrypt_metadata: bool,
}

/// PDF Encryption Handler
pub struct EncryptionHandler {
    /// Encryption information from trailer
    encryption_info: EncryptionInfo,
    /// Standard security handler
    security_handler: StandardSecurityHandler,
    /// Current encryption key (if unlocked)
    encryption_key: Option<EncryptionKey>,
    /// File ID from trailer
    file_id: Option<Vec<u8>>,
}

impl EncryptionHandler {
    /// Create encryption handler from encryption dictionary
    pub fn new(encrypt_dict: &PdfDictionary, file_id: Option<Vec<u8>>) -> ParseResult<Self> {
        let encryption_info = Self::parse_encryption_dict(encrypt_dict)?;

        // Create the security handler. For V>=4 the cipher is determined by the
        // stream crypt filter method (/CFM), not by the revision: R4 may be either
        // RC4 (/V2) or AES-128 (/AESV2). Selecting by revision alone decrypts
        // AES-128 streams with RC4 → garbage (issue #364).
        let security_handler = match encryption_info.r {
            2 => StandardSecurityHandler::rc4_40bit(),
            3 => StandardSecurityHandler::rc4_128bit(),
            4 => match encryption_info.cfm.as_deref() {
                // RC4-128 crypt filter under R4 (legacy, but valid).
                Some("V2") => StandardSecurityHandler::rc4_128bit(),
                // AES-128 is the conventional R4 cipher and what this crate writes.
                // Default to it for /AESV2, a missing /CFM, or any other value: R4
                // construction must stay infallible (it always built a handler
                // before #364), so non-RC4 cases fall back to AES rather than error.
                _ => StandardSecurityHandler::aes_128_r4(),
            },
            5 => StandardSecurityHandler::aes_256_r5(),
            6 => StandardSecurityHandler::aes_256_r6(),
            _ => {
                return Err(ParseError::SyntaxError {
                    position: 0,
                    message: format!("Encryption revision {} not supported", encryption_info.r),
                });
            }
        };

        Ok(Self {
            encryption_info,
            security_handler,
            encryption_key: None,
            file_id,
        })
    }

    /// Parse encryption dictionary from PDF trailer
    fn parse_encryption_dict(dict: &PdfDictionary) -> ParseResult<EncryptionInfo> {
        // Get Filter (required)
        let filter = dict
            .get("Filter")
            .and_then(|obj| obj.as_name())
            .map(|name| name.0.as_str())
            .ok_or_else(|| ParseError::MissingKey("Filter".to_string()))?;

        if filter != "Standard" {
            return Err(ParseError::SyntaxError {
                position: 0,
                message: format!("Encryption filter '{filter}' not supported"),
            });
        }

        // Get V (algorithm version)
        let v = dict
            .get("V")
            .and_then(|obj| obj.as_integer())
            .map(|i| i as i32)
            .unwrap_or(0);

        // Get R (revision)
        let r = dict
            .get("R")
            .and_then(|obj| obj.as_integer())
            .map(|i| i as i32)
            .ok_or_else(|| ParseError::MissingKey("R".to_string()))?;

        // Get O (owner password hash)
        let o = dict
            .get("O")
            .and_then(|obj| obj.as_string())
            .ok_or_else(|| ParseError::MissingKey("O".to_string()))?
            .as_bytes()
            .to_vec();

        // Get U (user password hash)
        let u = dict
            .get("U")
            .and_then(|obj| obj.as_string())
            .ok_or_else(|| ParseError::MissingKey("U".to_string()))?
            .as_bytes()
            .to_vec();

        // Get P (permissions)
        let p = dict
            .get("P")
            .and_then(|obj| obj.as_integer())
            .map(|i| i as i32)
            .ok_or_else(|| ParseError::MissingKey("P".to_string()))?;

        // Get Length (optional, defaults based on revision)
        let length = dict
            .get("Length")
            .and_then(|obj| obj.as_integer())
            .map(|i| i as i32);

        // Get UE entry (R5/R6 only — encrypted user key)
        let ue = dict
            .get("UE")
            .and_then(|obj| obj.as_string())
            .map(|s| s.as_bytes().to_vec());

        // Get OE entry (R5/R6 only — encrypted owner key)
        let oe = dict
            .get("OE")
            .and_then(|obj| obj.as_string())
            .map(|s| s.as_bytes().to_vec());

        // Get the stream crypt filter method (V>=4). For V<4 the cipher is fixed
        // by V/R (RC4), so /CF is absent and cfm stays None.
        let cfm = if v >= 4 {
            Self::parse_stream_cfm(dict)
        } else {
            None
        };

        // Get EncryptMetadata (default true). Only meaningful for R >= 4; older
        // revisions always encrypt metadata.
        let encrypt_metadata = dict
            .get("EncryptMetadata")
            .and_then(|obj| obj.as_bool())
            .unwrap_or(true);

        Ok(EncryptionInfo {
            filter: filter.to_string(),
            v,
            r,
            o,
            u,
            p,
            length,
            ue,
            oe,
            cfm,
            encrypt_metadata,
        })
    }

    /// Resolve the crypt filter method (/CFM) of the filter named by /StmF.
    ///
    /// `/Encrypt` carries `/CF << /StdCF << /CFM /AESV2 >> >>` and `/StmF /StdCF`.
    /// Returns the CFM name (e.g. "AESV2", "V2"), or "Identity" when streams are
    /// not encrypted, or None when no crypt filter is declared.
    fn parse_stream_cfm(dict: &PdfDictionary) -> Option<String> {
        let stmf = dict
            .get("StmF")
            .and_then(|o| o.as_name())
            .map(|n| n.0.as_str())
            .unwrap_or("Identity"); // ISO 32000-1 §7.6.5: default /StmF is Identity

        if stmf == "Identity" {
            return Some("Identity".to_string());
        }

        let cf = dict.get("CF").and_then(|o| o.as_dict())?;
        cf.get(stmf)
            .and_then(|o| o.as_dict())
            .and_then(|f| f.get("CFM"))
            .and_then(|o| o.as_name())
            .map(|n| n.0.clone())
    }

    /// Check if PDF is encrypted by looking for Encrypt entry in trailer
    pub fn detect_encryption(trailer: &PdfDictionary) -> bool {
        trailer.contains_key("Encrypt")
    }

    /// Try to unlock PDF with user password
    pub fn unlock_with_user_password(&mut self, password: &str) -> ParseResult<bool> {
        let user_password = UserPassword(password.to_string());

        match self.encryption_info.r {
            5 | 6 => self.unlock_user_r5_r6(&user_password),
            _ => self.unlock_user_r2_r4(&user_password),
        }
    }

    /// R2-R4 user password unlock (MD5/RC4 based)
    fn unlock_user_r2_r4(&mut self, user_password: &UserPassword) -> ParseResult<bool> {
        let permissions = Permissions::from_bits(self.encryption_info.p as u32);

        // The /EncryptMetadata 0xFFFFFFFF append (Algorithm 2, step f) applies
        // only for R >= 4; older revisions always encrypt metadata. Gate on the
        // document revision here — the security handler's revision is a cipher
        // proxy (R4-with-RC4 reports R3).
        let encrypt_metadata = self.encryption_info.encrypt_metadata || self.encryption_info.r < 4;

        let computed_u = self
            .security_handler
            .compute_user_hash_with_metadata(
                user_password,
                &self.encryption_info.o,
                permissions,
                self.file_id.as_deref(),
                encrypt_metadata,
            )
            .map_err(|e| ParseError::SyntaxError {
                position: 0,
                message: format!("Failed to compute user hash: {e}"),
            })?;

        // Compare with stored U entry (first 16 bytes for R3+)
        let comparison_length = if self.encryption_info.r >= 3 { 16 } else { 32 };

        if computed_u.len() < comparison_length || self.encryption_info.u.len() < comparison_length
        {
            return Ok(false);
        }

        let matches =
            computed_u[..comparison_length] == self.encryption_info.u[..comparison_length];

        if matches {
            let key = self
                .security_handler
                .compute_encryption_key_with_metadata(
                    user_password,
                    &self.encryption_info.o,
                    permissions,
                    self.file_id.as_deref(),
                    encrypt_metadata,
                )
                .map_err(|e| ParseError::SyntaxError {
                    position: 0,
                    message: format!("Failed to compute encryption key: {e}"),
                })?;
            self.encryption_key = Some(key);
        }

        Ok(matches)
    }

    /// R5/R6 user password unlock (SHA-256/AES-256 based per ISO 32000-2)
    fn unlock_user_r5_r6(&mut self, user_password: &UserPassword) -> ParseResult<bool> {
        let u_entry = &self.encryption_info.u;

        // Validate using the proper R5/R6 algorithm
        let is_valid = if self.encryption_info.r == 5 {
            self.security_handler
                .validate_r5_user_password(user_password, u_entry)
        } else {
            self.security_handler
                .validate_r6_user_password(user_password, u_entry)
        }
        .map_err(|e| ParseError::SyntaxError {
            position: 0,
            message: format!(
                "Failed to validate R{} user password: {e}",
                self.encryption_info.r
            ),
        })?;

        if is_valid {
            // Recover encryption key from UE entry
            let ue_entry =
                self.encryption_info
                    .ue
                    .as_deref()
                    .ok_or_else(|| ParseError::SyntaxError {
                        position: 0,
                        message: format!(
                            "R{} encryption requires UE entry but it is missing",
                            self.encryption_info.r
                        ),
                    })?;

            let key = if self.encryption_info.r == 5 {
                self.security_handler
                    .recover_r5_encryption_key(user_password, u_entry, ue_entry)
            } else {
                self.security_handler
                    .recover_r6_encryption_key(user_password, u_entry, ue_entry)
            }
            .map_err(|e| ParseError::SyntaxError {
                position: 0,
                message: format!(
                    "Failed to recover R{} encryption key: {e}",
                    self.encryption_info.r
                ),
            })?;

            self.encryption_key = Some(key);
        }

        Ok(is_valid)
    }

    /// R5/R6 owner password unlock (SHA-256/AES-256 based per ISO 32000-2).
    ///
    /// Mirrors [`unlock_user_r5_r6`](Self::unlock_user_r5_r6) but validates the
    /// owner password against the `/O` entry and recovers the file key from `/OE`.
    fn unlock_owner_r5_r6(&mut self, owner_password: &OwnerPassword) -> ParseResult<bool> {
        let o_entry = self.encryption_info.o.clone();
        let u_entry = self.encryption_info.u.clone();

        let is_valid = if self.encryption_info.r == 5 {
            self.security_handler
                .validate_r5_owner_password(owner_password, &o_entry, &u_entry)
        } else {
            self.security_handler
                .validate_r6_owner_password(owner_password, &o_entry, &u_entry)
        }
        .map_err(|e| ParseError::SyntaxError {
            position: 0,
            message: format!(
                "Failed to validate R{} owner password: {e}",
                self.encryption_info.r
            ),
        })?;

        if is_valid {
            let oe_entry =
                self.encryption_info
                    .oe
                    .as_deref()
                    .ok_or_else(|| ParseError::SyntaxError {
                        position: 0,
                        message: format!(
                            "R{} encryption requires OE entry but it is missing",
                            self.encryption_info.r
                        ),
                    })?;

            let key = if self.encryption_info.r == 5 {
                self.security_handler.recover_r5_owner_encryption_key(
                    owner_password,
                    &o_entry,
                    &u_entry,
                    oe_entry,
                )
            } else {
                self.security_handler.recover_r6_owner_encryption_key(
                    owner_password,
                    &o_entry,
                    &u_entry,
                    oe_entry,
                )
            }
            .map_err(|e| ParseError::SyntaxError {
                position: 0,
                message: format!(
                    "Failed to recover R{} owner encryption key: {e}",
                    self.encryption_info.r
                ),
            })?;

            self.encryption_key = Some(EncryptionKey::new(key));
        }

        Ok(is_valid)
    }

    /// Try to unlock PDF with owner password
    ///
    /// Owner password authentication works by:
    /// 1. Deriving an RC4 key from the owner password
    /// 2. Decrypting the O entry to recover the user password
    /// 3. Using the recovered user password to compute the encryption key
    pub fn unlock_with_owner_password(&mut self, password: &str) -> ParseResult<bool> {
        // R5/R6 (AES-256) use a completely different owner-password algorithm.
        // The MD5/RC4 path below assumes key_length <= 16 (MD5 output); for R5/R6
        // key_length is 32, so `hash[..key_length]` would panic. Dispatch first.
        if self.encryption_info.r >= 5 {
            return self.unlock_owner_r5_r6(&OwnerPassword(password.to_string()));
        }

        // Standard 32-byte padding from PDF spec
        const PADDING: [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,
        ];

        // Step 1: Pad owner password to 32 bytes
        let mut padded = [0u8; 32];
        let password_bytes = password.as_bytes();
        let len = password_bytes.len().min(32);
        padded[..len].copy_from_slice(&password_bytes[..len]);
        if len < 32 {
            padded[len..].copy_from_slice(&PADDING[..32 - len]);
        }

        // Step 2: MD5 hash the padded password
        let mut hash = md5::compute(&padded).to_vec();

        // Step 3: For R3+, do 50 additional MD5 iterations
        let key_length = self.security_handler.key_length;
        if self.encryption_info.r >= 3 {
            for _ in 0..50 {
                hash = md5::compute(&hash).to_vec();
            }
        }

        // Step 4: Decrypt O entry to get user password. `/O` is a 32-byte string
        // by spec; a short one from a malformed file would panic the slice below.
        if self.encryption_info.o.len() < 32 {
            return Ok(false);
        }
        let mut decrypted = self.encryption_info.o[..32].to_vec();

        if self.encryption_info.r >= 3 {
            // For R3+, decrypt with keys XOR'd with 19, 18, ..., 0
            for i in (0..20).rev() {
                let mut key_bytes = hash[..key_length].to_vec();
                for byte in &mut key_bytes {
                    *byte ^= i as u8;
                }
                let rc4_key = Rc4Key::from_slice(&key_bytes);
                let mut cipher = Rc4::new(&rc4_key);
                decrypted = cipher.process(&decrypted);
            }
        } else {
            // For R2, single RC4 decryption
            let rc4_key = Rc4Key::from_slice(&hash[..key_length]);
            let mut cipher = Rc4::new(&rc4_key);
            decrypted = cipher.process(&decrypted);
        }

        // Step 5: `decrypted` IS the padded user password (Algorithm 3, step e).
        // Authenticate it exactly as the user path does, straight from these 32
        // bytes — never reconstruct a `&str` and re-pad it. The old code searched
        // for where the standard padding began (byte 0x28, '(') and truncated
        // there; with a wrong owner password the bytes are garbage, and whenever
        // the first byte was 0x28 the "password" collapsed to "", which re-pads
        // to the exact standard padding and then authenticates ANY document with
        // an empty user password. That granted owner access on ~1/256 of wrong
        // attempts (owner-unlock fail-open). Running the raw bytes through the
        // standard verifier removes the truncation and the bypass.
        let permissions = Permissions::from_bits(self.encryption_info.p as u32);
        let encrypt_metadata = self.encryption_info.encrypt_metadata || self.encryption_info.r < 4;

        let computed_u = self
            .security_handler
            .compute_user_hash_from_padded(
                &decrypted,
                &self.encryption_info.o,
                permissions,
                self.file_id.as_deref(),
                encrypt_metadata,
            )
            .map_err(|e| ParseError::SyntaxError {
                position: 0,
                message: format!("Failed to compute user hash from /O: {e}"),
            })?;

        // R3+ compares the first 16 bytes of /U; R2 compares all 32.
        let comparison_length = if self.encryption_info.r >= 3 { 16 } else { 32 };
        if computed_u.len() < comparison_length || self.encryption_info.u.len() < comparison_length
        {
            return Ok(false);
        }
        if computed_u[..comparison_length] != self.encryption_info.u[..comparison_length] {
            return Ok(false);
        }

        // Owner authenticated: derive and retain the file key from the same bytes.
        let key = self
            .security_handler
            .compute_key_from_padded(
                &decrypted,
                &self.encryption_info.o,
                permissions,
                self.file_id.as_deref(),
                encrypt_metadata,
            )
            .map_err(|e| ParseError::SyntaxError {
                position: 0,
                message: format!("Failed to compute key from /O: {e}"),
            })?;
        self.encryption_key = Some(key);
        Ok(true)
    }

    /// Try to unlock with empty password (common case)
    pub fn try_empty_password(&mut self) -> ParseResult<bool> {
        self.unlock_with_user_password("")
    }

    /// Check if the PDF is currently unlocked
    pub fn is_unlocked(&self) -> bool {
        self.encryption_key.is_some()
    }

    /// Get the current encryption key (if unlocked)
    pub fn encryption_key(&self) -> Option<&EncryptionKey> {
        self.encryption_key.as_ref()
    }

    /// Decrypt a string object.
    ///
    /// Uses the error-propagating `try_*` variant so an AES decryption failure
    /// surfaces as an error rather than silently yielding empty content (#364).
    pub fn decrypt_string(&self, data: &[u8], obj_id: &ObjectId) -> ParseResult<Vec<u8>> {
        match &self.encryption_key {
            Some(key) => self
                .security_handler
                .try_decrypt_string(data, key, obj_id)
                .map_err(|e| ParseError::SyntaxError {
                    position: 0,
                    message: format!("Failed to decrypt string for object {obj_id:?}: {e}"),
                }),
            None => Err(ParseError::EncryptionNotSupported),
        }
    }

    /// Decrypt a stream object.
    ///
    /// Uses the error-propagating `try_*` variant so an AES decryption failure
    /// surfaces as an error rather than silently yielding empty content (#364).
    pub fn decrypt_stream(&self, data: &[u8], obj_id: &ObjectId) -> ParseResult<Vec<u8>> {
        match &self.encryption_key {
            Some(key) => self
                .security_handler
                .try_decrypt_stream(data, key, obj_id)
                .map_err(|e| ParseError::SyntaxError {
                    position: 0,
                    message: format!("Failed to decrypt stream for object {obj_id:?}: {e}"),
                }),
            None => Err(ParseError::EncryptionNotSupported),
        }
    }

    /// Get encryption algorithm information
    pub fn algorithm_info(&self) -> String {
        match (
            self.encryption_info.r,
            self.encryption_info.length.unwrap_or(40),
        ) {
            (2, _) => "RC4 40-bit".to_string(),
            (3, len) => format!("RC4 {len}-bit"),
            // R4 may be RC4 or AES-128 depending on the /CFM crypt filter (#364).
            (4, len) => match self.encryption_info.cfm.as_deref() {
                Some("V2") => format!("RC4 {len}-bit with metadata control"),
                _ => "AES-128 (Revision 4)".to_string(),
            },
            (5, _) => "AES-256 (Revision 5)".to_string(),
            (6, _) => "AES-256 (Revision 6, Unicode passwords)".to_string(),
            (r, len) => format!("Unknown revision {r} with {len}-bit key"),
        }
    }

    /// Get permissions information
    pub fn permissions(&self) -> Permissions {
        Permissions::from_bits(self.encryption_info.p as u32)
    }

    /// Check if file ID is available
    pub fn has_file_id(&self) -> bool {
        self.file_id.is_some()
    }

    /// Get the encryption revision
    pub fn revision(&self) -> i32 {
        self.encryption_info.r
    }

    /// Check if strings should be encrypted
    pub fn encrypt_strings(&self) -> bool {
        // For standard security handler, strings are always encrypted
        true
    }

    /// Check if streams should be encrypted  
    pub fn encrypt_streams(&self) -> bool {
        // For standard security handler, streams are always encrypted
        true
    }

    /// Whether document metadata is encrypted (`/EncryptMetadata`, default true).
    pub fn encrypt_metadata(&self) -> bool {
        self.encryption_info.encrypt_metadata
    }
}

/// Password prompt result
#[derive(Debug)]
pub enum PasswordResult {
    /// Password was accepted
    Success,
    /// Password was rejected
    Rejected,
    /// User cancelled password entry
    Cancelled,
}

/// Trait for password prompting
pub trait PasswordProvider {
    /// Prompt for user password
    fn prompt_user_password(&self) -> ParseResult<Option<String>>;

    /// Prompt for owner password
    fn prompt_owner_password(&self) -> ParseResult<Option<String>>;
}

/// Console-based password provider
pub struct ConsolePasswordProvider;

impl PasswordProvider for ConsolePasswordProvider {
    fn prompt_user_password(&self) -> ParseResult<Option<String>> {
        tracing::debug!(
            "PDF is encrypted. Enter user password (or press Enter for empty password):"
        );

        let mut input = String::new();
        std::io::stdin()
            .read_line(&mut input)
            .map_err(|e| ParseError::SyntaxError {
                position: 0,
                message: format!("Failed to read password: {e}"),
            })?;

        // Remove trailing newline
        input.truncate(input.trim_end().len());
        Ok(Some(input))
    }

    fn prompt_owner_password(&self) -> ParseResult<Option<String>> {
        tracing::debug!("User password failed. Enter owner password:");

        let mut input = String::new();
        std::io::stdin()
            .read_line(&mut input)
            .map_err(|e| ParseError::SyntaxError {
                position: 0,
                message: format!("Failed to read password: {e}"),
            })?;

        // Remove trailing newline
        input.truncate(input.trim_end().len());
        Ok(Some(input))
    }
}

/// Interactive decryption helper
pub struct InteractiveDecryption<P: PasswordProvider> {
    password_provider: P,
}

impl<P: PasswordProvider> InteractiveDecryption<P> {
    /// Create new interactive decryption helper
    pub fn new(password_provider: P) -> Self {
        Self { password_provider }
    }

    /// Attempt to unlock PDF interactively
    pub fn unlock_pdf(&self, handler: &mut EncryptionHandler) -> ParseResult<PasswordResult> {
        // First try empty password
        if handler.try_empty_password()? {
            return Ok(PasswordResult::Success);
        }

        // Try user password
        if let Some(password) = self.password_provider.prompt_user_password()? {
            if handler.unlock_with_user_password(&password)? {
                return Ok(PasswordResult::Success);
            }
        } else {
            return Ok(PasswordResult::Cancelled);
        }

        // Try owner password
        if let Some(password) = self.password_provider.prompt_owner_password()? {
            if handler.unlock_with_owner_password(&password)? {
                return Ok(PasswordResult::Success);
            }
        } else {
            return Ok(PasswordResult::Cancelled);
        }

        Ok(PasswordResult::Rejected)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::objects::{PdfDictionary, PdfName, PdfObject, PdfString};

    fn create_test_encryption_dict() -> PdfDictionary {
        let mut dict = PdfDictionary::new();
        dict.insert(
            "Filter".to_string(),
            PdfObject::Name(PdfName("Standard".to_string())),
        );
        dict.insert("V".to_string(), PdfObject::Integer(1));
        dict.insert("R".to_string(), PdfObject::Integer(2));
        dict.insert(
            "O".to_string(),
            PdfObject::String(PdfString::new(vec![0u8; 32])),
        );
        dict.insert(
            "U".to_string(),
            PdfObject::String(PdfString::new(vec![0u8; 32])),
        );
        dict.insert("P".to_string(), PdfObject::Integer(-4));
        dict
    }

    #[test]
    fn test_encryption_detection() {
        let mut trailer = PdfDictionary::new();
        assert!(!EncryptionHandler::detect_encryption(&trailer));

        trailer.insert("Encrypt".to_string(), PdfObject::Reference(1, 0));
        assert!(EncryptionHandler::detect_encryption(&trailer));
    }

    #[test]
    fn test_encryption_info_parsing() {
        let dict = create_test_encryption_dict();
        let info = EncryptionHandler::parse_encryption_dict(&dict).unwrap();

        assert_eq!(info.filter, "Standard");
        assert_eq!(info.v, 1);
        assert_eq!(info.r, 2);
        assert_eq!(info.o.len(), 32);
        assert_eq!(info.u.len(), 32);
        assert_eq!(info.p, -4);
        // V<4 has no crypt filters → cfm is None.
        assert_eq!(info.cfm, None);
    }

    /// Build a V=4/R=4 encryption dict whose StdCF crypt filter uses `cfm`.
    fn create_v4_encryption_dict(cfm: &str) -> PdfDictionary {
        let mut dict = PdfDictionary::new();
        dict.insert(
            "Filter".to_string(),
            PdfObject::Name(PdfName("Standard".to_string())),
        );
        dict.insert("V".to_string(), PdfObject::Integer(4));
        dict.insert("R".to_string(), PdfObject::Integer(4));
        dict.insert("Length".to_string(), PdfObject::Integer(128));
        dict.insert(
            "O".to_string(),
            PdfObject::String(PdfString::new(vec![0u8; 32])),
        );
        dict.insert(
            "U".to_string(),
            PdfObject::String(PdfString::new(vec![0u8; 32])),
        );
        dict.insert("P".to_string(), PdfObject::Integer(-4));

        let mut std_cf = PdfDictionary::new();
        std_cf.insert("CFM".to_string(), PdfObject::Name(PdfName(cfm.to_string())));
        let mut cf = PdfDictionary::new();
        cf.insert("StdCF".to_string(), PdfObject::Dictionary(std_cf));
        dict.insert("CF".to_string(), PdfObject::Dictionary(cf));
        dict.insert(
            "StmF".to_string(),
            PdfObject::Name(PdfName("StdCF".to_string())),
        );
        dict
    }

    #[test]
    fn test_cfm_parsed_from_crypt_filter() {
        // Issue #364: R4 cipher is decided by /CFM, not the revision.
        let aes =
            EncryptionHandler::parse_encryption_dict(&create_v4_encryption_dict("AESV2")).unwrap();
        assert_eq!(aes.cfm.as_deref(), Some("AESV2"));

        let rc4 =
            EncryptionHandler::parse_encryption_dict(&create_v4_encryption_dict("V2")).unwrap();
        assert_eq!(rc4.cfm.as_deref(), Some("V2"));
    }

    #[test]
    fn test_r4_algorithm_info_reflects_cipher() {
        // AESV2 under R4 must report AES-128, not RC4 (#364).
        let aes = EncryptionHandler::new(&create_v4_encryption_dict("AESV2"), None).unwrap();
        assert_eq!(aes.algorithm_info(), "AES-128 (Revision 4)");

        // A legacy R4 PDF using the V2 (RC4) crypt filter still reports RC4.
        let rc4 = EncryptionHandler::new(&create_v4_encryption_dict("V2"), None).unwrap();
        assert!(
            rc4.algorithm_info().starts_with("RC4"),
            "got: {}",
            rc4.algorithm_info()
        );
    }

    #[test]
    fn test_encryption_handler_creation() {
        let dict = create_test_encryption_dict();
        let handler = EncryptionHandler::new(&dict, None).unwrap();

        assert_eq!(handler.encryption_info.r, 2);
        assert!(!handler.is_unlocked());
        assert_eq!(handler.algorithm_info(), "RC4 40-bit");
    }

    #[test]
    fn test_empty_password_attempt() {
        let dict = create_test_encryption_dict();
        let mut handler = EncryptionHandler::new(&dict, None).unwrap();

        // Empty password should not work with test data
        let result = handler.try_empty_password().unwrap();
        assert!(!result);
        assert!(!handler.is_unlocked());
    }

    #[test]
    fn test_permissions() {
        let dict = create_test_encryption_dict();
        let handler = EncryptionHandler::new(&dict, None).unwrap();

        let permissions = handler.permissions();
        // P value of -4 should result in specific permissions
        assert!(permissions.bits() != 0);
    }

    #[test]
    fn test_encryption_flags() {
        let dict = create_test_encryption_dict();
        let handler = EncryptionHandler::new(&dict, None).unwrap();

        assert!(handler.encrypt_strings());
        assert!(handler.encrypt_streams());
        assert!(handler.encrypt_metadata());
    }

    #[test]
    fn test_decrypt_without_key() {
        let dict = create_test_encryption_dict();
        let handler = EncryptionHandler::new(&dict, None).unwrap();

        let obj_id = ObjectId::new(1, 0);
        let result = handler.decrypt_string(b"test", &obj_id);
        assert!(result.is_err());
    }

    #[test]
    fn test_unsupported_filter() {
        let mut dict = PdfDictionary::new();
        dict.insert(
            "Filter".to_string(),
            PdfObject::Name(PdfName("UnsupportedFilter".to_string())),
        );
        dict.insert("R".to_string(), PdfObject::Integer(2));
        dict.insert(
            "O".to_string(),
            PdfObject::String(PdfString::new(vec![0u8; 32])),
        );
        dict.insert(
            "U".to_string(),
            PdfObject::String(PdfString::new(vec![0u8; 32])),
        );
        dict.insert("P".to_string(), PdfObject::Integer(-4));

        let result = EncryptionHandler::new(&dict, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_unsupported_revision() {
        let mut dict = create_test_encryption_dict();
        dict.insert("R".to_string(), PdfObject::Integer(99)); // Unsupported revision

        let result = EncryptionHandler::new(&dict, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_missing_required_keys() {
        let test_cases = vec![
            ("Filter", PdfObject::Name(PdfName("Standard".to_string()))),
            ("R", PdfObject::Integer(2)),
            ("O", PdfObject::String(PdfString::new(vec![0u8; 32]))),
            ("U", PdfObject::String(PdfString::new(vec![0u8; 32]))),
            ("P", PdfObject::Integer(-4)),
        ];

        for (skip_key, _) in test_cases {
            let mut dict = create_test_encryption_dict();
            dict.0.remove(&PdfName(skip_key.to_string()));

            let result = EncryptionHandler::parse_encryption_dict(&dict);
            assert!(result.is_err(), "Should fail when {skip_key} is missing");
        }
    }

    /// Mock password provider for testing
    struct MockPasswordProvider {
        user_password: Option<String>,
        owner_password: Option<String>,
    }

    impl PasswordProvider for MockPasswordProvider {
        fn prompt_user_password(&self) -> ParseResult<Option<String>> {
            Ok(self.user_password.clone())
        }

        fn prompt_owner_password(&self) -> ParseResult<Option<String>> {
            Ok(self.owner_password.clone())
        }
    }

    #[test]
    fn test_interactive_decryption_cancelled() {
        let provider = MockPasswordProvider {
            user_password: None,
            owner_password: None,
        };

        let decryption = InteractiveDecryption::new(provider);
        let dict = create_test_encryption_dict();
        let mut handler = EncryptionHandler::new(&dict, None).unwrap();

        let result = decryption.unlock_pdf(&mut handler).unwrap();
        matches!(result, PasswordResult::Cancelled);
    }

    #[test]
    fn test_interactive_decryption_rejected() {
        let provider = MockPasswordProvider {
            user_password: Some("wrong_password".to_string()),
            owner_password: Some("also_wrong".to_string()),
        };

        let decryption = InteractiveDecryption::new(provider);
        let dict = create_test_encryption_dict();
        let mut handler = EncryptionHandler::new(&dict, None).unwrap();

        let result = decryption.unlock_pdf(&mut handler).unwrap();
        matches!(result, PasswordResult::Rejected);
    }

    // ===== ADVANCED EDGE CASE TESTS =====

    #[test]
    fn test_malformed_encryption_dictionary_invalid_types() {
        // Test with wrong type for Filter
        let mut dict = PdfDictionary::new();
        dict.insert("Filter".to_string(), PdfObject::Integer(123)); // Should be Name
        dict.insert("R".to_string(), PdfObject::Integer(2));
        dict.insert(
            "O".to_string(),
            PdfObject::String(PdfString::new(vec![0u8; 32])),
        );
        dict.insert(
            "U".to_string(),
            PdfObject::String(PdfString::new(vec![0u8; 32])),
        );
        dict.insert("P".to_string(), PdfObject::Integer(-4));

        let result = EncryptionHandler::parse_encryption_dict(&dict);
        assert!(result.is_err());

        // Test with wrong type for R
        let mut dict = create_test_encryption_dict();
        dict.insert(
            "R".to_string(),
            PdfObject::Name(PdfName("not_a_number".to_string())),
        );
        let result = EncryptionHandler::parse_encryption_dict(&dict);
        assert!(result.is_err());
    }

    #[test]
    fn test_encryption_dictionary_edge_values() {
        // Test with extreme revision values
        let mut dict = create_test_encryption_dict();
        dict.insert("R".to_string(), PdfObject::Integer(0)); // Very low revision
        let result = EncryptionHandler::new(&dict, None);
        assert!(result.is_err());

        // Test with negative revision
        let mut dict = create_test_encryption_dict();
        dict.insert("R".to_string(), PdfObject::Integer(-1));
        let result = EncryptionHandler::new(&dict, None);
        assert!(result.is_err());

        // Test with very high revision
        let mut dict = create_test_encryption_dict();
        dict.insert("R".to_string(), PdfObject::Integer(1000));
        let result = EncryptionHandler::new(&dict, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_encryption_dictionary_invalid_hash_lengths() {
        // Test with O hash too short
        let mut dict = create_test_encryption_dict();
        dict.insert(
            "O".to_string(),
            PdfObject::String(PdfString::new(vec![0u8; 16])),
        ); // Should be 32
        let result = EncryptionHandler::parse_encryption_dict(&dict);
        // Should still work but be invalid data
        assert!(result.is_ok());

        // Test with U hash too long
        let mut dict = create_test_encryption_dict();
        dict.insert(
            "U".to_string(),
            PdfObject::String(PdfString::new(vec![0u8; 64])),
        ); // Should be 32
        let result = EncryptionHandler::parse_encryption_dict(&dict);
        assert!(result.is_ok());

        // Test with empty hashes
        let mut dict = create_test_encryption_dict();
        dict.insert("O".to_string(), PdfObject::String(PdfString::new(vec![])));
        dict.insert("U".to_string(), PdfObject::String(PdfString::new(vec![])));
        let result = EncryptionHandler::parse_encryption_dict(&dict);
        assert!(result.is_ok());
    }

    #[test]
    fn test_encryption_with_different_key_lengths() {
        // Test Rev 2 (40-bit)
        let mut dict = create_test_encryption_dict();
        dict.insert("R".to_string(), PdfObject::Integer(2));
        let handler = EncryptionHandler::new(&dict, None).unwrap();
        assert_eq!(handler.algorithm_info(), "RC4 40-bit");

        // Test Rev 3 (128-bit)
        let mut dict = create_test_encryption_dict();
        dict.insert("R".to_string(), PdfObject::Integer(3));
        dict.insert("Length".to_string(), PdfObject::Integer(128));
        let handler = EncryptionHandler::new(&dict, None).unwrap();
        assert_eq!(handler.algorithm_info(), "RC4 128-bit");

        // Test Rev 4 without a V2 crypt filter → AES-128 (the conventional R4
        // cipher and this crate's default). A legacy RC4 R4 PDF would carry an
        // explicit /CFM /V2 crypt filter; see test_r4_algorithm_info_reflects_cipher.
        let mut dict = create_test_encryption_dict();
        dict.insert("R".to_string(), PdfObject::Integer(4));
        dict.insert("Length".to_string(), PdfObject::Integer(128));
        let handler = EncryptionHandler::new(&dict, None).unwrap();
        assert_eq!(handler.algorithm_info(), "AES-128 (Revision 4)");

        // Test Rev 5 (AES-256)
        let mut dict = create_test_encryption_dict();
        dict.insert("R".to_string(), PdfObject::Integer(5));
        dict.insert("V".to_string(), PdfObject::Integer(5));
        let handler = EncryptionHandler::new(&dict, None).unwrap();
        assert_eq!(handler.algorithm_info(), "AES-256 (Revision 5)");

        // Test Rev 6 (AES-256 with Unicode)
        let mut dict = create_test_encryption_dict();
        dict.insert("R".to_string(), PdfObject::Integer(6));
        dict.insert("V".to_string(), PdfObject::Integer(5));
        let handler = EncryptionHandler::new(&dict, None).unwrap();
        assert_eq!(
            handler.algorithm_info(),
            "AES-256 (Revision 6, Unicode passwords)"
        );
    }

    #[test]
    fn test_file_id_handling() {
        let dict = create_test_encryption_dict();

        // Test with file ID
        let file_id = Some(b"test_file_id_12345678".to_vec());
        let handler = EncryptionHandler::new(&dict, file_id.clone()).unwrap();
        // File ID should be stored
        assert_eq!(handler.file_id, file_id);

        // Test without file ID
        let handler = EncryptionHandler::new(&dict, None).unwrap();
        assert_eq!(handler.file_id, None);

        // Test with empty file ID
        let empty_file_id = Some(vec![]);
        let handler = EncryptionHandler::new(&dict, empty_file_id.clone()).unwrap();
        assert_eq!(handler.file_id, empty_file_id);
    }

    #[test]
    fn test_permissions_edge_cases() {
        // Test with different permission values
        let permission_values = vec![0, -1, -4, -44, -100, i32::MAX, i32::MIN];

        for p_value in permission_values {
            let mut dict = create_test_encryption_dict();
            dict.insert("P".to_string(), PdfObject::Integer(p_value as i64));
            let handler = EncryptionHandler::new(&dict, None).unwrap();

            let permissions = handler.permissions();
            assert_eq!(permissions.bits(), p_value as u32);
        }
    }

    #[test]
    fn test_decrypt_with_different_object_ids() {
        let dict = create_test_encryption_dict();
        let handler = EncryptionHandler::new(&dict, None).unwrap();
        let test_data = b"test data";

        // Test with different object IDs (should all fail since not unlocked)
        let object_ids = vec![
            ObjectId::new(1, 0),
            ObjectId::new(999, 0),
            ObjectId::new(1, 999),
            ObjectId::new(u32::MAX, u16::MAX),
        ];

        for obj_id in object_ids {
            let result = handler.decrypt_string(test_data, &obj_id);
            assert!(result.is_err());

            let result = handler.decrypt_stream(test_data, &obj_id);
            assert!(result.is_err());
        }
    }

    #[test]
    fn test_password_scenarios_comprehensive() {
        let dict = create_test_encryption_dict();
        let mut handler = EncryptionHandler::new(&dict, None).unwrap();

        // Test various password scenarios (using String for uniformity)
        let test_passwords = vec![
            "".to_string(),                     // Empty
            " ".to_string(),                    // Single space
            "   ".to_string(),                  // Multiple spaces
            "password".to_string(),             // Simple
            "Password123!@#".to_string(),       // Complex
            "a".repeat(32),                     // Exactly 32 chars
            "a".repeat(50),                     // Over 32 chars
            "unicode_ñáéíóú".to_string(),       // Unicode
            "pass\nwith\nnewlines".to_string(), // Newlines
            "pass\twith\ttabs".to_string(),     // Tabs
            "pass with spaces".to_string(),     // Spaces
            "🔐🗝️📄".to_string(),               // Emojis
        ];

        for password in test_passwords {
            // All should fail with test data but not crash
            let result = handler.unlock_with_user_password(&password);
            assert!(result.is_ok());
            assert!(!result.unwrap());

            let result = handler.unlock_with_owner_password(&password);
            assert!(result.is_ok());
            assert!(!result.unwrap());
        }
    }

    #[test]
    fn test_encryption_handler_thread_safety_simulation() {
        // Simulate what would happen in multi-threaded access
        let dict = create_test_encryption_dict();
        let handler = EncryptionHandler::new(&dict, None).unwrap();

        // Test multiple read operations (safe)
        for _ in 0..100 {
            assert!(!handler.is_unlocked());
            assert_eq!(handler.algorithm_info(), "RC4 40-bit");
            assert!(handler.encrypt_strings());
            assert!(handler.encrypt_streams());
            assert!(handler.encrypt_metadata());
        }
    }

    #[test]
    fn test_encryption_state_transitions() {
        let dict = create_test_encryption_dict();
        let mut handler = EncryptionHandler::new(&dict, None).unwrap();

        // Initial state
        assert!(!handler.is_unlocked());

        // Try unlock (should fail with test data)
        let result = handler.try_empty_password().unwrap();
        assert!(!result);
        assert!(!handler.is_unlocked());

        // Try user password (should fail with test data)
        let result = handler.unlock_with_user_password("test").unwrap();
        assert!(!result);
        assert!(!handler.is_unlocked());

        // Try owner password (should fail with test data)
        let result = handler.unlock_with_owner_password("test").unwrap();
        assert!(!result);
        assert!(!handler.is_unlocked());

        // State should remain consistent
        assert!(!handler.is_unlocked());
    }

    #[test]
    fn test_interactive_decryption_edge_cases() {
        // Test provider that returns None for both passwords
        let provider = MockPasswordProvider {
            user_password: None,
            owner_password: None,
        };

        let decryption = InteractiveDecryption::new(provider);
        let dict = create_test_encryption_dict();
        let mut handler = EncryptionHandler::new(&dict, None).unwrap();

        let result = decryption.unlock_pdf(&mut handler).unwrap();
        matches!(result, PasswordResult::Cancelled);

        // Test provider that returns empty strings
        let provider = MockPasswordProvider {
            user_password: Some("".to_string()),
            owner_password: Some("".to_string()),
        };

        let decryption = InteractiveDecryption::new(provider);
        let mut handler = EncryptionHandler::new(&dict, None).unwrap();

        let result = decryption.unlock_pdf(&mut handler).unwrap();
        matches!(result, PasswordResult::Rejected);
    }

    /// Test custom MockPasswordProvider for edge cases
    struct EdgeCasePasswordProvider {
        call_count: std::cell::RefCell<usize>,
        passwords: Vec<Option<String>>,
    }

    impl EdgeCasePasswordProvider {
        fn new(passwords: Vec<Option<String>>) -> Self {
            Self {
                call_count: std::cell::RefCell::new(0),
                passwords,
            }
        }
    }

    impl PasswordProvider for EdgeCasePasswordProvider {
        fn prompt_user_password(&self) -> ParseResult<Option<String>> {
            let mut count = self.call_count.borrow_mut();
            if *count < self.passwords.len() {
                let result = self.passwords[*count].clone();
                *count += 1;
                Ok(result)
            } else {
                Ok(None)
            }
        }

        fn prompt_owner_password(&self) -> ParseResult<Option<String>> {
            self.prompt_user_password()
        }
    }

    #[test]
    fn test_interactive_decryption_with_sequence() {
        let passwords = vec![
            Some("first_attempt".to_string()),
            Some("second_attempt".to_string()),
            None, // Cancelled
        ];

        let provider = EdgeCasePasswordProvider::new(passwords);
        let decryption = InteractiveDecryption::new(provider);
        let dict = create_test_encryption_dict();
        let mut handler = EncryptionHandler::new(&dict, None).unwrap();

        let result = decryption.unlock_pdf(&mut handler).unwrap();
        matches!(result, PasswordResult::Cancelled);
    }
}