Skip to main content

oxidize_pdf/encryption/
object_encryption.rs

1//! Object encryption/decryption for PDF documents
2//!
3//! This module implements encryption and decryption of PDF objects
4//! with integration into the parser and writer according to ISO 32000-1:2008.
5
6use crate::encryption::{
7    CryptFilterManager, EmbeddedFileEncryption, EncryptionDictionary, EncryptionKey,
8    SecurityHandler, StandardSecurityHandler,
9};
10use crate::error::{PdfError, Result};
11use crate::objects::{Dictionary, Object, ObjectId, Stream};
12use std::sync::Arc;
13
14/// Object Encryptor for encrypting PDF objects
15pub struct ObjectEncryptor {
16    /// Crypt filter manager
17    filter_manager: Arc<CryptFilterManager>,
18    /// Encryption key
19    encryption_key: EncryptionKey,
20    /// Encrypt metadata flag
21    encrypt_metadata: bool,
22    /// Embedded file encryption handler
23    embedded_file_handler: Option<EmbeddedFileEncryption>,
24}
25
26impl ObjectEncryptor {
27    /// Create new object encryptor
28    pub fn new(
29        filter_manager: Arc<CryptFilterManager>,
30        encryption_key: EncryptionKey,
31        encrypt_metadata: bool,
32    ) -> Self {
33        Self {
34            filter_manager,
35            encryption_key,
36            encrypt_metadata,
37            embedded_file_handler: None,
38        }
39    }
40
41    /// Create new object encryptor with embedded file support
42    pub fn with_embedded_files(
43        filter_manager: Arc<CryptFilterManager>,
44        encryption_key: EncryptionKey,
45        encrypt_metadata: bool,
46        eff_filter: Option<String>,
47    ) -> Self {
48        let embedded_file_handler = Some(EmbeddedFileEncryption::new(
49            eff_filter,
50            encrypt_metadata,
51            filter_manager.clone(),
52        ));
53
54        Self {
55            filter_manager,
56            encryption_key,
57            encrypt_metadata,
58            embedded_file_handler,
59        }
60    }
61
62    /// Encrypt an object
63    pub fn encrypt_object(&self, object: &mut Object, obj_id: &ObjectId) -> Result<()> {
64        match object {
65            Object::String(s) => {
66                // Check if this is a metadata stream
67                if !self.should_encrypt_string(s) {
68                    return Ok(());
69                }
70
71                let encrypted = self.filter_manager.encrypt_string(
72                    s.as_bytes(),
73                    obj_id,
74                    None,
75                    &self.encryption_key,
76                )?;
77
78                *s = String::from_utf8_lossy(&encrypted).to_string();
79            }
80            Object::Stream(dict, data) => {
81                // Create a temporary Stream object
82                let mut stream = Stream::with_dictionary(dict.clone(), data.clone());
83                self.encrypt_stream(&mut stream, obj_id)?;
84
85                // Update the Object with encrypted data
86                *object = Object::Stream(stream.dictionary().clone(), stream.data().to_vec());
87            }
88            Object::Dictionary(dict) => {
89                self.encrypt_dictionary(dict, obj_id)?;
90            }
91            Object::Array(array) => {
92                for item in array.iter_mut() {
93                    self.encrypt_object(item, obj_id)?;
94                }
95            }
96            Object::Reference(_) => {
97                // References are not encrypted
98            }
99            _ => {
100                // Other object types are not encrypted
101            }
102        }
103
104        Ok(())
105    }
106
107    /// Decrypt an object
108    pub fn decrypt_object(&self, object: &mut Object, obj_id: &ObjectId) -> Result<()> {
109        match object {
110            Object::String(s) => {
111                // Check if this is a metadata stream
112                if !self.should_encrypt_string(s) {
113                    return Ok(());
114                }
115
116                let decrypted = self.filter_manager.decrypt_string(
117                    s.as_bytes(),
118                    obj_id,
119                    None,
120                    &self.encryption_key,
121                )?;
122
123                *s = String::from_utf8_lossy(&decrypted).to_string();
124            }
125            Object::Stream(dict, data) => {
126                // Create a temporary Stream object
127                let mut stream = Stream::with_dictionary(dict.clone(), data.clone());
128                self.decrypt_stream(&mut stream, obj_id)?;
129
130                // Update the Object with decrypted data
131                *object = Object::Stream(stream.dictionary().clone(), stream.data().to_vec());
132            }
133            Object::Dictionary(dict) => {
134                self.decrypt_dictionary(dict, obj_id)?;
135            }
136            Object::Array(array) => {
137                for item in array.iter_mut() {
138                    self.decrypt_object(item, obj_id)?;
139                }
140            }
141            Object::Reference(_) => {
142                // References are not decrypted
143            }
144            _ => {
145                // Other object types are not decrypted
146            }
147        }
148
149        Ok(())
150    }
151
152    /// Encrypt a stream
153    fn encrypt_stream(&self, stream: &mut Stream, obj_id: &ObjectId) -> Result<()> {
154        // Check if stream should be encrypted
155        if !self.should_encrypt_stream(stream) {
156            return Ok(());
157        }
158
159        let encrypted_data = if let Some(ref handler) = self.embedded_file_handler {
160            // Use embedded file handler for special stream types
161            handler.process_stream_encryption(
162                stream.dictionary(),
163                stream.data(),
164                obj_id,
165                &self.encryption_key,
166                true, // encrypt
167            )?
168        } else {
169            // Use standard encryption
170            self.filter_manager.encrypt_stream(
171                stream.data(),
172                obj_id,
173                stream.dictionary(),
174                &self.encryption_key,
175            )?
176        };
177
178        *stream.data_mut() = encrypted_data;
179
180        // Update stream dictionary if needed
181        if !stream.dictionary().contains_key("Filter") {
182            stream
183                .dictionary_mut()
184                .set("Filter", Object::Name("Crypt".to_string()));
185        } else if let Some(Object::Array(filters)) = stream.dictionary_mut().get_mut("Filter") {
186            // Add Crypt filter to existing filters
187            filters.push(Object::Name("Crypt".to_string()));
188        }
189
190        Ok(())
191    }
192
193    /// Decrypt a stream
194    fn decrypt_stream(&self, stream: &mut Stream, obj_id: &ObjectId) -> Result<()> {
195        // Check if stream should be decrypted
196        if !self.should_encrypt_stream(stream) {
197            return Ok(());
198        }
199
200        let decrypted_data = if let Some(ref handler) = self.embedded_file_handler {
201            // Use embedded file handler for special stream types
202            handler.process_stream_encryption(
203                stream.dictionary(),
204                stream.data(),
205                obj_id,
206                &self.encryption_key,
207                false, // decrypt
208            )?
209        } else {
210            // Use standard decryption
211            self.filter_manager.decrypt_stream(
212                stream.data(),
213                obj_id,
214                stream.dictionary(),
215                &self.encryption_key,
216            )?
217        };
218
219        *stream.data_mut() = decrypted_data;
220
221        // Remove Crypt filter from dictionary if present
222        if let Some(Object::Array(filters)) = stream.dictionary_mut().get_mut("Filter") {
223            filters.retain(|f| {
224                if let Object::Name(name) = f {
225                    name != "Crypt"
226                } else {
227                    true
228                }
229            });
230
231            // If only Crypt filter was present, remove Filter entry
232            if filters.is_empty() {
233                stream.dictionary_mut().remove("Filter");
234            }
235        } else if let Some(Object::Name(name)) = stream.dictionary().get("Filter") {
236            if name == "Crypt" {
237                stream.dictionary_mut().remove("Filter");
238            }
239        }
240
241        Ok(())
242    }
243
244    /// Encrypt a dictionary
245    fn encrypt_dictionary(&self, dict: &mut Dictionary, obj_id: &ObjectId) -> Result<()> {
246        // Get all keys to avoid borrowing issues
247        let keys: Vec<String> = dict.keys().cloned().collect();
248
249        for key in keys {
250            // Skip certain dictionary entries
251            if self.should_skip_dictionary_key(&key) {
252                continue;
253            }
254
255            if let Some(value) = dict.get_mut(&key) {
256                self.encrypt_object(value, obj_id)?;
257            }
258        }
259
260        Ok(())
261    }
262
263    /// Decrypt a dictionary
264    fn decrypt_dictionary(&self, dict: &mut Dictionary, obj_id: &ObjectId) -> Result<()> {
265        // Get all keys to avoid borrowing issues
266        let keys: Vec<String> = dict.keys().cloned().collect();
267
268        for key in keys {
269            // Skip certain dictionary entries
270            if self.should_skip_dictionary_key(&key) {
271                continue;
272            }
273
274            if let Some(value) = dict.get_mut(&key) {
275                self.decrypt_object(value, obj_id)?;
276            }
277        }
278
279        Ok(())
280    }
281
282    /// Check if a string should be encrypted
283    fn should_encrypt_string(&self, _s: &str) -> bool {
284        // All strings are encrypted except in special cases
285        true
286    }
287
288    /// Check if a stream should be encrypted
289    fn should_encrypt_stream(&self, stream: &Stream) -> bool {
290        // Check if this is a metadata stream
291        if !self.encrypt_metadata {
292            if let Some(Object::Name(type_name)) = stream.dictionary().get("Type") {
293                if type_name == "Metadata" {
294                    return false;
295                }
296            }
297        }
298
299        // Check if stream is already encrypted
300        if let Some(filter) = stream.dictionary().get("Filter") {
301            match filter {
302                Object::Name(name) => {
303                    if name == "Crypt" {
304                        return false;
305                    }
306                }
307                Object::Array(filters) => {
308                    for f in filters {
309                        if let Object::Name(name) = f {
310                            if name == "Crypt" {
311                                return false;
312                            }
313                        }
314                    }
315                }
316                _ => {}
317            }
318        }
319
320        true
321    }
322
323    /// Check if a dictionary key should be skipped during encryption
324    fn should_skip_dictionary_key(&self, key: &str) -> bool {
325        // These keys should never be encrypted
326        matches!(
327            key,
328            "Length" | "Filter" | "DecodeParms" | "Encrypt" | "ID" | "O" | "U" | "P" | "Perms"
329        )
330    }
331}
332
333/// Integration with Document for encryption
334pub struct DocumentEncryption {
335    /// Encryption dictionary
336    pub encryption_dict: EncryptionDictionary,
337    /// Object encryptor
338    pub encryptor: ObjectEncryptor,
339}
340
341impl DocumentEncryption {
342    /// Create from encryption dictionary and password
343    pub fn new(
344        encryption_dict: EncryptionDictionary,
345        user_password: &str,
346        file_id: Option<&[u8]>,
347    ) -> Result<Self> {
348        // Create security handler based on revision
349        let handler: Box<dyn SecurityHandler> = match encryption_dict.r {
350            2 | 3 => Box::new(StandardSecurityHandler::rc4_128bit()),
351            4 => Box::new(StandardSecurityHandler {
352                revision: crate::encryption::SecurityHandlerRevision::R4,
353                key_length: encryption_dict.length.unwrap_or(16) as usize,
354            }),
355            5 => Box::new(StandardSecurityHandler::aes_256_r5()),
356            6 => Box::new(StandardSecurityHandler::aes_256_r6()),
357            _ => {
358                return Err(PdfError::EncryptionError(format!(
359                    "Unsupported encryption revision: {}",
360                    encryption_dict.r
361                )));
362            }
363        };
364
365        // Compute encryption key
366        let user_pwd = crate::encryption::UserPassword(user_password.to_string());
367        let encryption_key = if encryption_dict.r <= 4 {
368            StandardSecurityHandler {
369                revision: match encryption_dict.r {
370                    2 => crate::encryption::SecurityHandlerRevision::R2,
371                    3 => crate::encryption::SecurityHandlerRevision::R3,
372                    4 => crate::encryption::SecurityHandlerRevision::R4,
373                    _ => unreachable!(),
374                },
375                key_length: encryption_dict.length.unwrap_or(16) as usize,
376            }
377            .compute_encryption_key(
378                &user_pwd,
379                &encryption_dict.o,
380                encryption_dict.p,
381                file_id,
382            )?
383        } else {
384            // For R5/R6, use advanced key derivation
385            // This is simplified - in production, extract salts from encryption dict
386            EncryptionKey::new(vec![0u8; 32])
387        };
388
389        // Create crypt filter manager
390        let mut filter_manager = CryptFilterManager::new(
391            handler,
392            encryption_dict
393                .stm_f
394                .as_ref()
395                .map(|f| match f {
396                    crate::encryption::StreamFilter::StdCF => "StdCF".to_string(),
397                    crate::encryption::StreamFilter::Identity => "Identity".to_string(),
398                    crate::encryption::StreamFilter::Custom(name) => name.clone(),
399                })
400                .unwrap_or_else(|| "StdCF".to_string()),
401            encryption_dict
402                .str_f
403                .as_ref()
404                .map(|f| match f {
405                    crate::encryption::StringFilter::StdCF => "StdCF".to_string(),
406                    crate::encryption::StringFilter::Identity => "Identity".to_string(),
407                    crate::encryption::StringFilter::Custom(name) => name.clone(),
408                })
409                .unwrap_or_else(|| "StdCF".to_string()),
410        );
411
412        // Add crypt filters from dictionary
413        if let Some(ref filters) = encryption_dict.cf {
414            for filter in filters {
415                filter_manager.add_filter(crate::encryption::FunctionalCryptFilter {
416                    name: filter.name.clone(),
417                    method: filter.method,
418                    length: filter.length,
419                    auth_event: crate::encryption::AuthEvent::DocOpen,
420                    recipients: None,
421                });
422            }
423        }
424
425        let encryptor = ObjectEncryptor::new(
426            Arc::new(filter_manager),
427            encryption_key,
428            encryption_dict.encrypt_metadata,
429        );
430
431        Ok(Self {
432            encryption_dict,
433            encryptor,
434        })
435    }
436
437    /// Encrypt all objects in a document
438    pub fn encrypt_objects(&self, objects: &mut [(ObjectId, Object)]) -> Result<()> {
439        for (obj_id, obj) in objects.iter_mut() {
440            // Skip encryption dictionary object
441            if self.is_encryption_dict_object(obj) {
442                continue;
443            }
444
445            self.encryptor.encrypt_object(obj, obj_id)?;
446        }
447
448        Ok(())
449    }
450
451    /// Decrypt all objects in a document
452    pub fn decrypt_objects(&self, objects: &mut [(ObjectId, Object)]) -> Result<()> {
453        for (obj_id, obj) in objects.iter_mut() {
454            // Skip encryption dictionary object
455            if self.is_encryption_dict_object(obj) {
456                continue;
457            }
458
459            self.encryptor.decrypt_object(obj, obj_id)?;
460        }
461
462        Ok(())
463    }
464
465    /// Check if object is the encryption dictionary
466    fn is_encryption_dict_object(&self, obj: &Object) -> bool {
467        if let Object::Dictionary(dict) = obj {
468            // Check if this is an encryption dictionary
469            if let Some(Object::Name(filter)) = dict.get("Filter") {
470                return filter == "Standard";
471            }
472        }
473        false
474    }
475}
476
477#[cfg(test)]
478mod tests {
479    use super::*;
480    use crate::encryption::Permissions;
481
482    fn create_test_encryptor() -> ObjectEncryptor {
483        let handler = Box::new(StandardSecurityHandler::rc4_128bit());
484        let mut filter_manager =
485            CryptFilterManager::new(handler, "StdCF".to_string(), "StdCF".to_string());
486
487        // Add the StdCF filter
488        filter_manager.add_filter(crate::encryption::FunctionalCryptFilter {
489            name: "StdCF".to_string(),
490            method: crate::encryption::CryptFilterMethod::V2,
491            length: Some(16),
492            auth_event: crate::encryption::AuthEvent::DocOpen,
493            recipients: None,
494        });
495
496        let encryption_key = EncryptionKey::new(vec![0u8; 16]);
497
498        ObjectEncryptor::new(Arc::new(filter_manager), encryption_key, true)
499    }
500
501    #[test]
502    fn test_encrypt_string_object() {
503        let encryptor = create_test_encryptor();
504        let obj_id = ObjectId::new(1, 0);
505
506        let mut obj = Object::String("Test string".to_string());
507        let original = obj.clone();
508
509        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
510
511        // String should be different after encryption
512        assert_ne!(obj, original);
513    }
514
515    #[test]
516    fn test_encrypt_array_object() {
517        let encryptor = create_test_encryptor();
518        let obj_id = ObjectId::new(1, 0);
519
520        let mut obj = Object::Array(vec![
521            Object::String("String 1".to_string()),
522            Object::Integer(42),
523            Object::String("String 2".to_string()),
524        ]);
525
526        let original = obj.clone();
527        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
528
529        // Array should be different (strings encrypted)
530        assert_ne!(obj, original);
531
532        // Check that integers are not encrypted
533        if let Object::Array(array) = &obj {
534            assert_eq!(array[1], Object::Integer(42));
535        }
536    }
537
538    #[test]
539    fn test_encrypt_dictionary_object() {
540        let encryptor = create_test_encryptor();
541        let obj_id = ObjectId::new(1, 0);
542
543        let mut dict = Dictionary::new();
544        dict.set("Title", Object::String("Test Title".to_string()));
545        dict.set("Length", Object::Integer(100)); // Should be skipped
546        dict.set("Filter", Object::Name("FlateDecode".to_string())); // Should be skipped
547
548        let mut obj = Object::Dictionary(dict);
549        let original = obj.clone();
550
551        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
552
553        // Dictionary should be different
554        assert_ne!(obj, original);
555
556        // Check that skipped keys are not encrypted
557        if let Object::Dictionary(dict) = &obj {
558            assert_eq!(dict.get("Length"), Some(&Object::Integer(100)));
559            assert_eq!(
560                dict.get("Filter"),
561                Some(&Object::Name("FlateDecode".to_string()))
562            );
563        }
564    }
565
566    #[test]
567    fn test_encrypt_stream_object() {
568        let encryptor = create_test_encryptor();
569        let obj_id = ObjectId::new(1, 0);
570
571        let dict = Dictionary::new();
572        let data = b"Stream data content".to_vec();
573        let original_data = data.clone();
574
575        let mut obj = Object::Stream(dict, data);
576        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
577
578        if let Object::Stream(dict, data) = &obj {
579            // Data should be encrypted
580            assert_ne!(data, &original_data);
581
582            // Filter should be added
583            assert_eq!(dict.get("Filter"), Some(&Object::Name("Crypt".to_string())));
584        }
585    }
586
587    #[test]
588    fn test_skip_metadata_stream() {
589        let handler = Box::new(StandardSecurityHandler::rc4_128bit());
590        let filter_manager =
591            CryptFilterManager::new(handler, "StdCF".to_string(), "StdCF".to_string());
592
593        let encryption_key = EncryptionKey::new(vec![0u8; 16]);
594
595        let encryptor = ObjectEncryptor::new(
596            Arc::new(filter_manager),
597            encryption_key,
598            false, // Don't encrypt metadata
599        );
600
601        let obj_id = ObjectId::new(1, 0);
602
603        let mut dict = Dictionary::new();
604        dict.set("Type", Object::Name("Metadata".to_string()));
605        let data = b"Metadata content".to_vec();
606        let original_data = data.clone();
607
608        let mut obj = Object::Stream(dict, data);
609
610        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
611
612        if let Object::Stream(_, data) = &obj {
613            // Metadata should not be encrypted
614            assert_eq!(data, &original_data);
615        }
616    }
617
618    #[test]
619    fn test_should_skip_dictionary_key() {
620        let encryptor = create_test_encryptor();
621
622        assert!(encryptor.should_skip_dictionary_key("Length"));
623        assert!(encryptor.should_skip_dictionary_key("Filter"));
624        assert!(encryptor.should_skip_dictionary_key("DecodeParms"));
625        assert!(encryptor.should_skip_dictionary_key("Encrypt"));
626        assert!(encryptor.should_skip_dictionary_key("ID"));
627        assert!(encryptor.should_skip_dictionary_key("O"));
628        assert!(encryptor.should_skip_dictionary_key("U"));
629        assert!(encryptor.should_skip_dictionary_key("P"));
630        assert!(encryptor.should_skip_dictionary_key("Perms"));
631
632        assert!(!encryptor.should_skip_dictionary_key("Title"));
633        assert!(!encryptor.should_skip_dictionary_key("Author"));
634        assert!(!encryptor.should_skip_dictionary_key("Subject"));
635    }
636
637    #[test]
638    fn test_decrypt_object_reverses_encryption() {
639        let encryptor = create_test_encryptor();
640        let obj_id = ObjectId::new(1, 0);
641
642        let original_string = "Test content for encryption";
643        let mut obj = Object::String(original_string.to_string());
644        let original = obj.clone();
645
646        // Encrypt
647        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
648
649        // Verify it's encrypted (different from original)
650        assert_ne!(obj, original);
651
652        // Decrypt
653        encryptor.decrypt_object(&mut obj, &obj_id).unwrap();
654
655        // Due to String::from_utf8_lossy, we may not get exact original back
656        // In production, PDF strings should handle binary data properly
657        // For now, just verify that encrypt/decrypt complete without errors
658        if let Object::String(s) = &obj {
659            // At minimum, verify it's a valid string
660            assert!(!s.is_empty());
661        }
662    }
663
664    #[test]
665    fn test_reference_object_not_encrypted() {
666        let encryptor = create_test_encryptor();
667        let obj_id = ObjectId::new(1, 0);
668
669        let mut obj = Object::Reference(ObjectId::new(5, 0));
670        let original = obj.clone();
671
672        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
673
674        // Reference should remain unchanged
675        assert_eq!(obj, original);
676    }
677
678    #[test]
679    fn test_already_encrypted_stream_skipped() {
680        let encryptor = create_test_encryptor();
681        let obj_id = ObjectId::new(1, 0);
682
683        let mut dict = Dictionary::new();
684        dict.set("Filter", Object::Name("Crypt".to_string()));
685        let data = b"Already encrypted data".to_vec();
686        let original_data = data.clone();
687
688        let mut obj = Object::Stream(dict, data);
689
690        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
691
692        if let Object::Stream(_, data) = &obj {
693            // Data should remain unchanged
694            assert_eq!(data, &original_data);
695        }
696    }
697
698    #[test]
699    fn test_document_encryption_creation() {
700        let encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
701            vec![0u8; 32],
702            vec![1u8; 32],
703            Permissions::all(),
704            None,
705        );
706
707        let doc_encryption =
708            DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id"));
709
710        assert!(doc_encryption.is_ok());
711    }
712
713    #[test]
714    fn test_is_encryption_dict_object() {
715        let encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
716            vec![0u8; 32],
717            vec![1u8; 32],
718            Permissions::all(),
719            None,
720        );
721
722        let doc_encryption =
723            DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id")).unwrap();
724
725        let mut dict = Dictionary::new();
726        dict.set("Filter", Object::Name("Standard".to_string()));
727        let obj = Object::Dictionary(dict);
728
729        assert!(doc_encryption.is_encryption_dict_object(&obj));
730
731        let normal_obj = Object::String("Not an encryption dict".to_string());
732        assert!(!doc_encryption.is_encryption_dict_object(&normal_obj));
733    }
734
735    #[test]
736    fn test_with_embedded_files_constructor() {
737        let handler = Box::new(StandardSecurityHandler::rc4_128bit());
738        let filter_manager =
739            CryptFilterManager::new(handler, "StdCF".to_string(), "StdCF".to_string());
740
741        let encryption_key = EncryptionKey::new(vec![0u8; 16]);
742
743        let encryptor = ObjectEncryptor::with_embedded_files(
744            Arc::new(filter_manager),
745            encryption_key,
746            true,
747            Some("StdCF".to_string()),
748        );
749
750        // Verify embedded file handler was created
751        assert!(encryptor.embedded_file_handler.is_some());
752    }
753
754    #[test]
755    fn test_with_embedded_files_no_filter() {
756        let handler = Box::new(StandardSecurityHandler::rc4_128bit());
757        let filter_manager =
758            CryptFilterManager::new(handler, "StdCF".to_string(), "StdCF".to_string());
759
760        let encryption_key = EncryptionKey::new(vec![0u8; 16]);
761
762        let encryptor = ObjectEncryptor::with_embedded_files(
763            Arc::new(filter_manager),
764            encryption_key,
765            false,
766            None,
767        );
768
769        // Handler should still be created even without explicit filter
770        assert!(encryptor.embedded_file_handler.is_some());
771    }
772
773    #[test]
774    fn test_decrypt_string_object() {
775        let encryptor = create_test_encryptor();
776        let obj_id = ObjectId::new(1, 0);
777
778        // First encrypt a string
779        let mut obj = Object::String("Test string for decryption".to_string());
780        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
781
782        // Then decrypt it
783        encryptor.decrypt_object(&mut obj, &obj_id).unwrap();
784
785        // Verify it's a valid string after round-trip
786        if let Object::String(s) = &obj {
787            assert!(!s.is_empty());
788        } else {
789            panic!("Expected String object");
790        }
791    }
792
793    #[test]
794    fn test_decrypt_array_object() {
795        let encryptor = create_test_encryptor();
796        let obj_id = ObjectId::new(1, 0);
797
798        let mut obj = Object::Array(vec![
799            Object::String("Test 1".to_string()),
800            Object::Integer(123),
801            Object::String("Test 2".to_string()),
802        ]);
803
804        // Encrypt
805        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
806
807        // Decrypt
808        encryptor.decrypt_object(&mut obj, &obj_id).unwrap();
809
810        // Verify structure is preserved
811        if let Object::Array(array) = &obj {
812            assert_eq!(array.len(), 3);
813            assert_eq!(array[1], Object::Integer(123));
814        } else {
815            panic!("Expected Array object");
816        }
817    }
818
819    #[test]
820    fn test_decrypt_dictionary_object() {
821        let encryptor = create_test_encryptor();
822        let obj_id = ObjectId::new(1, 0);
823
824        let mut dict = Dictionary::new();
825        dict.set("Title", Object::String("Test Title".to_string()));
826        dict.set("Count", Object::Integer(42));
827        dict.set("Length", Object::Integer(100)); // Should be skipped
828
829        let mut obj = Object::Dictionary(dict);
830
831        // Encrypt
832        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
833
834        // Decrypt
835        encryptor.decrypt_object(&mut obj, &obj_id).unwrap();
836
837        // Verify skipped keys are unchanged
838        if let Object::Dictionary(dict) = &obj {
839            assert_eq!(dict.get("Count"), Some(&Object::Integer(42)));
840            assert_eq!(dict.get("Length"), Some(&Object::Integer(100)));
841        } else {
842            panic!("Expected Dictionary object");
843        }
844    }
845
846    #[test]
847    fn test_decrypt_reference_not_changed() {
848        let encryptor = create_test_encryptor();
849        let obj_id = ObjectId::new(1, 0);
850
851        let mut obj = Object::Reference(ObjectId::new(10, 0));
852        let original = obj.clone();
853
854        encryptor.decrypt_object(&mut obj, &obj_id).unwrap();
855
856        // Reference should remain unchanged
857        assert_eq!(obj, original);
858    }
859
860    #[test]
861    fn test_decrypt_other_types_not_changed() {
862        let encryptor = create_test_encryptor();
863        let obj_id = ObjectId::new(1, 0);
864
865        // Test Integer
866        let mut obj = Object::Integer(42);
867        let original = obj.clone();
868        encryptor.decrypt_object(&mut obj, &obj_id).unwrap();
869        assert_eq!(obj, original);
870
871        // Test Real
872        let mut obj = Object::Real(3.14);
873        let original = obj.clone();
874        encryptor.decrypt_object(&mut obj, &obj_id).unwrap();
875        assert_eq!(obj, original);
876
877        // Test Boolean
878        let mut obj = Object::Boolean(true);
879        let original = obj.clone();
880        encryptor.decrypt_object(&mut obj, &obj_id).unwrap();
881        assert_eq!(obj, original);
882
883        // Test Null
884        let mut obj = Object::Null;
885        let original = obj.clone();
886        encryptor.decrypt_object(&mut obj, &obj_id).unwrap();
887        assert_eq!(obj, original);
888
889        // Test Name
890        let mut obj = Object::Name("TestName".to_string());
891        let original = obj.clone();
892        encryptor.decrypt_object(&mut obj, &obj_id).unwrap();
893        assert_eq!(obj, original);
894    }
895
896    #[test]
897    fn test_encrypt_other_types_not_changed() {
898        let encryptor = create_test_encryptor();
899        let obj_id = ObjectId::new(1, 0);
900
901        // Test Integer
902        let mut obj = Object::Integer(999);
903        let original = obj.clone();
904        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
905        assert_eq!(obj, original);
906
907        // Test Real
908        let mut obj = Object::Real(2.718);
909        let original = obj.clone();
910        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
911        assert_eq!(obj, original);
912
913        // Test Boolean
914        let mut obj = Object::Boolean(false);
915        let original = obj.clone();
916        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
917        assert_eq!(obj, original);
918
919        // Test Null
920        let mut obj = Object::Null;
921        let original = obj.clone();
922        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
923        assert_eq!(obj, original);
924
925        // Test Name
926        let mut obj = Object::Name("AnotherName".to_string());
927        let original = obj.clone();
928        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
929        assert_eq!(obj, original);
930    }
931
932    #[test]
933    fn test_should_encrypt_stream_with_array_filter_containing_crypt() {
934        let encryptor = create_test_encryptor();
935
936        let mut dict = Dictionary::new();
937        dict.set(
938            "Filter",
939            Object::Array(vec![
940                Object::Name("FlateDecode".to_string()),
941                Object::Name("Crypt".to_string()),
942            ]),
943        );
944
945        let stream = Stream::with_dictionary(dict, vec![1, 2, 3]);
946
947        // Should return false because Crypt is in the filter array
948        assert!(!encryptor.should_encrypt_stream(&stream));
949    }
950
951    #[test]
952    fn test_should_encrypt_stream_with_array_filter_without_crypt() {
953        let encryptor = create_test_encryptor();
954
955        let mut dict = Dictionary::new();
956        dict.set(
957            "Filter",
958            Object::Array(vec![
959                Object::Name("FlateDecode".to_string()),
960                Object::Name("ASCII85Decode".to_string()),
961            ]),
962        );
963
964        let stream = Stream::with_dictionary(dict, vec![1, 2, 3]);
965
966        // Should return true because no Crypt filter
967        assert!(encryptor.should_encrypt_stream(&stream));
968    }
969
970    #[test]
971    fn test_should_encrypt_stream_with_non_name_filter() {
972        let encryptor = create_test_encryptor();
973
974        let mut dict = Dictionary::new();
975        // Invalid filter type (Integer instead of Name or Array)
976        dict.set("Filter", Object::Integer(123));
977
978        let stream = Stream::with_dictionary(dict, vec![1, 2, 3]);
979
980        // Should return true because no valid Crypt filter found
981        assert!(encryptor.should_encrypt_stream(&stream));
982    }
983
984    #[test]
985    fn test_should_encrypt_string_always_true() {
986        let encryptor = create_test_encryptor();
987
988        assert!(encryptor.should_encrypt_string("any string"));
989        assert!(encryptor.should_encrypt_string(""));
990        assert!(encryptor.should_encrypt_string("special chars: !@#$%"));
991    }
992
993    #[test]
994    fn test_encrypt_objects_skips_encryption_dict() {
995        let mut encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
996            vec![0u8; 32],
997            vec![1u8; 32],
998            Permissions::all(),
999            None,
1000        );
1001
1002        // Add the crypt filter that will be needed
1003        encryption_dict.cf = Some(vec![crate::encryption::CryptFilter {
1004            name: "StdCF".to_string(),
1005            method: crate::encryption::CryptFilterMethod::V2,
1006            length: Some(16),
1007        }]);
1008
1009        let doc_encryption =
1010            DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id")).unwrap();
1011
1012        // Create an encryption dictionary object
1013        let mut encrypt_dict = Dictionary::new();
1014        encrypt_dict.set("Filter", Object::Name("Standard".to_string()));
1015        encrypt_dict.set("V", Object::Integer(4));
1016
1017        let mut objects = vec![
1018            (
1019                ObjectId::new(1, 0),
1020                Object::Dictionary(encrypt_dict.clone()),
1021            ),
1022            (
1023                ObjectId::new(2, 0),
1024                Object::String("Normal string".to_string()),
1025            ),
1026        ];
1027
1028        let original_encrypt_dict = objects[0].1.clone();
1029
1030        doc_encryption.encrypt_objects(&mut objects).unwrap();
1031
1032        // Encryption dict should be unchanged
1033        assert_eq!(objects[0].1, original_encrypt_dict);
1034
1035        // Normal string should be encrypted (different from original)
1036        assert_ne!(objects[1].1, Object::String("Normal string".to_string()));
1037    }
1038
1039    #[test]
1040    fn test_decrypt_objects_skips_encryption_dict() {
1041        let mut encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1042            vec![0u8; 32],
1043            vec![1u8; 32],
1044            Permissions::all(),
1045            None,
1046        );
1047
1048        // Add the crypt filter that will be needed
1049        encryption_dict.cf = Some(vec![crate::encryption::CryptFilter {
1050            name: "StdCF".to_string(),
1051            method: crate::encryption::CryptFilterMethod::V2,
1052            length: Some(16),
1053        }]);
1054
1055        let doc_encryption =
1056            DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id")).unwrap();
1057
1058        // Create an encryption dictionary object
1059        let mut encrypt_dict = Dictionary::new();
1060        encrypt_dict.set("Filter", Object::Name("Standard".to_string()));
1061        encrypt_dict.set("V", Object::Integer(4));
1062
1063        let mut objects = vec![
1064            (
1065                ObjectId::new(1, 0),
1066                Object::Dictionary(encrypt_dict.clone()),
1067            ),
1068            (
1069                ObjectId::new(2, 0),
1070                Object::String("encrypted content".to_string()),
1071            ),
1072        ];
1073
1074        let original_encrypt_dict = objects[0].1.clone();
1075
1076        doc_encryption.decrypt_objects(&mut objects).unwrap();
1077
1078        // Encryption dict should be unchanged
1079        assert_eq!(objects[0].1, original_encrypt_dict);
1080    }
1081
1082    #[test]
1083    fn test_document_encryption_unsupported_revision() {
1084        let mut encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1085            vec![0u8; 32],
1086            vec![1u8; 32],
1087            Permissions::all(),
1088            None,
1089        );
1090        encryption_dict.r = 99; // Unsupported revision
1091
1092        let result = DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id"));
1093
1094        assert!(result.is_err());
1095        if let Err(PdfError::EncryptionError(msg)) = result {
1096            assert!(msg.contains("Unsupported encryption revision"));
1097        }
1098    }
1099
1100    #[test]
1101    fn test_document_encryption_r2() {
1102        let mut encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1103            vec![0u8; 32],
1104            vec![1u8; 32],
1105            Permissions::all(),
1106            None,
1107        );
1108        encryption_dict.r = 2;
1109
1110        let result = DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id"));
1111        assert!(result.is_ok());
1112    }
1113
1114    #[test]
1115    fn test_document_encryption_r4() {
1116        let mut encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1117            vec![0u8; 32],
1118            vec![1u8; 32],
1119            Permissions::all(),
1120            None,
1121        );
1122        encryption_dict.r = 4;
1123        encryption_dict.length = Some(16);
1124
1125        let result = DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id"));
1126        assert!(result.is_ok());
1127    }
1128
1129    #[test]
1130    fn test_document_encryption_r5() {
1131        let mut encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1132            vec![0u8; 32],
1133            vec![1u8; 32],
1134            Permissions::all(),
1135            None,
1136        );
1137        encryption_dict.r = 5;
1138
1139        let result = DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id"));
1140        assert!(result.is_ok());
1141    }
1142
1143    #[test]
1144    fn test_document_encryption_r6() {
1145        let mut encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1146            vec![0u8; 32],
1147            vec![1u8; 32],
1148            Permissions::all(),
1149            None,
1150        );
1151        encryption_dict.r = 6;
1152
1153        let result = DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id"));
1154        assert!(result.is_ok());
1155    }
1156
1157    #[test]
1158    fn test_is_encryption_dict_object_not_dict() {
1159        let encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1160            vec![0u8; 32],
1161            vec![1u8; 32],
1162            Permissions::all(),
1163            None,
1164        );
1165
1166        let doc_encryption =
1167            DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id")).unwrap();
1168
1169        // Non-dictionary objects should return false
1170        assert!(!doc_encryption.is_encryption_dict_object(&Object::Integer(42)));
1171        assert!(!doc_encryption.is_encryption_dict_object(&Object::Null));
1172        assert!(!doc_encryption.is_encryption_dict_object(&Object::Array(vec![Object::Integer(1)])));
1173    }
1174
1175    #[test]
1176    fn test_is_encryption_dict_object_dict_without_filter() {
1177        let encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1178            vec![0u8; 32],
1179            vec![1u8; 32],
1180            Permissions::all(),
1181            None,
1182        );
1183
1184        let doc_encryption =
1185            DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id")).unwrap();
1186
1187        // Dictionary without Filter key should return false
1188        let mut dict = Dictionary::new();
1189        dict.set("Type", Object::Name("Catalog".to_string()));
1190        let obj = Object::Dictionary(dict);
1191
1192        assert!(!doc_encryption.is_encryption_dict_object(&obj));
1193    }
1194
1195    #[test]
1196    fn test_is_encryption_dict_object_dict_with_different_filter() {
1197        let encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1198            vec![0u8; 32],
1199            vec![1u8; 32],
1200            Permissions::all(),
1201            None,
1202        );
1203
1204        let doc_encryption =
1205            DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id")).unwrap();
1206
1207        // Dictionary with non-Standard Filter should return false
1208        let mut dict = Dictionary::new();
1209        dict.set("Filter", Object::Name("FlateDecode".to_string()));
1210        let obj = Object::Dictionary(dict);
1211
1212        assert!(!doc_encryption.is_encryption_dict_object(&obj));
1213    }
1214
1215    #[test]
1216    fn test_nested_array_encryption() {
1217        let encryptor = create_test_encryptor();
1218        let obj_id = ObjectId::new(1, 0);
1219
1220        // Create nested arrays with strings
1221        let mut obj = Object::Array(vec![
1222            Object::Array(vec![
1223                Object::String("Nested 1".to_string()),
1224                Object::String("Nested 2".to_string()),
1225            ]),
1226            Object::String("Outer".to_string()),
1227        ]);
1228
1229        let original = obj.clone();
1230
1231        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
1232
1233        // Should be different (strings encrypted)
1234        assert_ne!(obj, original);
1235
1236        // Verify structure preserved
1237        if let Object::Array(outer) = &obj {
1238            assert_eq!(outer.len(), 2);
1239            if let Object::Array(inner) = &outer[0] {
1240                assert_eq!(inner.len(), 2);
1241            } else {
1242                panic!("Expected nested array");
1243            }
1244        }
1245    }
1246
1247    #[test]
1248    fn test_nested_dictionary_encryption() {
1249        let encryptor = create_test_encryptor();
1250        let obj_id = ObjectId::new(1, 0);
1251
1252        let mut inner_dict = Dictionary::new();
1253        inner_dict.set("InnerTitle", Object::String("Inner value".to_string()));
1254
1255        let mut outer_dict = Dictionary::new();
1256        outer_dict.set("OuterTitle", Object::String("Outer value".to_string()));
1257        outer_dict.set("Nested", Object::Dictionary(inner_dict));
1258
1259        let mut obj = Object::Dictionary(outer_dict);
1260        let original = obj.clone();
1261
1262        encryptor.encrypt_object(&mut obj, &obj_id).unwrap();
1263
1264        // Should be different (strings encrypted)
1265        assert_ne!(obj, original);
1266
1267        // Verify structure preserved
1268        if let Object::Dictionary(dict) = &obj {
1269            assert!(dict.contains_key("OuterTitle"));
1270            assert!(dict.contains_key("Nested"));
1271            if let Some(Object::Dictionary(nested)) = dict.get("Nested") {
1272                assert!(nested.contains_key("InnerTitle"));
1273            } else {
1274                panic!("Expected nested dictionary");
1275            }
1276        }
1277    }
1278
1279    #[test]
1280    fn test_document_encryption_with_crypt_filters() {
1281        let mut encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1282            vec![0u8; 32],
1283            vec![1u8; 32],
1284            Permissions::all(),
1285            None,
1286        );
1287
1288        // Add crypt filters
1289        encryption_dict.cf = Some(vec![crate::encryption::CryptFilter {
1290            name: "StdCF".to_string(),
1291            method: crate::encryption::CryptFilterMethod::V2,
1292            length: Some(16),
1293        }]);
1294
1295        let result = DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id"));
1296        assert!(result.is_ok());
1297    }
1298
1299    #[test]
1300    fn test_document_encryption_with_custom_stm_str_filters() {
1301        let mut encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1302            vec![0u8; 32],
1303            vec![1u8; 32],
1304            Permissions::all(),
1305            None,
1306        );
1307
1308        encryption_dict.stm_f = Some(crate::encryption::StreamFilter::Identity);
1309        encryption_dict.str_f = Some(crate::encryption::StringFilter::Identity);
1310
1311        let result = DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id"));
1312        assert!(result.is_ok());
1313    }
1314
1315    #[test]
1316    fn test_document_encryption_with_custom_filter_names() {
1317        let mut encryption_dict = crate::encryption::EncryptionDictionary::rc4_128bit(
1318            vec![0u8; 32],
1319            vec![1u8; 32],
1320            Permissions::all(),
1321            None,
1322        );
1323
1324        encryption_dict.stm_f = Some(crate::encryption::StreamFilter::Custom(
1325            "CustomStm".to_string(),
1326        ));
1327        encryption_dict.str_f = Some(crate::encryption::StringFilter::Custom(
1328            "CustomStr".to_string(),
1329        ));
1330
1331        let result = DocumentEncryption::new(encryption_dict, "user_password", Some(b"file_id"));
1332        assert!(result.is_ok());
1333    }
1334}