Skip to main content

oxidize_pdf/encryption/
aes.rs

1//! AES encryption implementation for PDF
2//!
3//! This module provides AES-128 and AES-256 encryption support according to
4//! ISO 32000-1 Section 7.6 (PDF 1.6+ and PDF 2.0).
5//!
6//! Implementation uses production-grade RustCrypto crates (aes, cbc).
7
8use aes::{Aes128, Aes256};
9use cbc::{Decryptor, Encryptor};
10use cipher::block_padding::Pkcs7;
11use cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit};
12
13/// AES key sizes supported by PDF
14#[derive(Debug, Clone, Copy, PartialEq)]
15pub enum AesKeySize {
16    /// AES-128 (16 bytes)
17    Aes128,
18    /// AES-256 (32 bytes)
19    Aes256,
20}
21
22impl AesKeySize {
23    /// Get key size in bytes
24    pub fn key_length(&self) -> usize {
25        match self {
26            AesKeySize::Aes128 => 16,
27            AesKeySize::Aes256 => 32,
28        }
29    }
30
31    /// Get block size (always 16 bytes for AES)
32    pub fn block_size(&self) -> usize {
33        16
34    }
35}
36
37/// AES encryption key
38#[derive(Debug, Clone)]
39pub struct AesKey {
40    /// Key bytes
41    key: Vec<u8>,
42    /// Key size
43    size: AesKeySize,
44}
45
46impl AesKey {
47    /// Create new AES-128 key
48    pub fn new_128(key: Vec<u8>) -> Result<Self, AesError> {
49        if key.len() != 16 {
50            return Err(AesError::InvalidKeyLength {
51                expected: 16,
52                actual: key.len(),
53            });
54        }
55
56        Ok(Self {
57            key,
58            size: AesKeySize::Aes128,
59        })
60    }
61
62    /// Create new AES-256 key
63    pub fn new_256(key: Vec<u8>) -> Result<Self, AesError> {
64        if key.len() != 32 {
65            return Err(AesError::InvalidKeyLength {
66                expected: 32,
67                actual: key.len(),
68            });
69        }
70
71        Ok(Self {
72            key,
73            size: AesKeySize::Aes256,
74        })
75    }
76
77    /// Get key bytes
78    pub fn key(&self) -> &[u8] {
79        &self.key
80    }
81
82    /// Get key size
83    pub fn size(&self) -> AesKeySize {
84        self.size
85    }
86
87    /// Get key length in bytes
88    pub fn len(&self) -> usize {
89        self.key.len()
90    }
91
92    /// Check if key is empty (should never happen)
93    pub fn is_empty(&self) -> bool {
94        self.key.is_empty()
95    }
96}
97
98/// AES-related errors
99#[derive(Debug, Clone, PartialEq)]
100pub enum AesError {
101    /// Invalid key length
102    InvalidKeyLength { expected: usize, actual: usize },
103    /// Invalid IV length (must be 16 bytes)
104    InvalidIvLength { expected: usize, actual: usize },
105    /// Encryption failed
106    EncryptionFailed(String),
107    /// Decryption failed
108    DecryptionFailed(String),
109    /// PKCS#7 padding error
110    PaddingError(String),
111}
112
113impl std::fmt::Display for AesError {
114    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
115        match self {
116            AesError::InvalidKeyLength { expected, actual } => {
117                write!(f, "Invalid key length: expected {expected}, got {actual}")
118            }
119            AesError::InvalidIvLength { expected, actual } => {
120                write!(f, "Invalid IV length: expected {expected}, got {actual}")
121            }
122            AesError::EncryptionFailed(msg) => write!(f, "Encryption failed: {msg}"),
123            AesError::DecryptionFailed(msg) => write!(f, "Decryption failed: {msg}"),
124            AesError::PaddingError(msg) => write!(f, "Padding error: {msg}"),
125        }
126    }
127}
128
129impl std::error::Error for AesError {}
130
131/// AES cipher implementation using RustCrypto (production-grade)
132///
133/// This implementation uses the audited `aes` and `cbc` crates from the
134/// RustCrypto project for secure AES-128 and AES-256 encryption.
135pub struct Aes {
136    key: AesKey,
137}
138
139impl Aes {
140    /// Create new AES cipher
141    pub fn new(key: AesKey) -> Self {
142        Self { key }
143    }
144
145    /// Encrypt data using AES-CBC mode with PKCS#7 padding
146    ///
147    /// Uses RustCrypto's `cbc` crate for production-grade encryption.
148    pub fn encrypt_cbc(&self, data: &[u8], iv: &[u8]) -> Result<Vec<u8>, AesError> {
149        if iv.len() != 16 {
150            return Err(AesError::InvalidIvLength {
151                expected: 16,
152                actual: iv.len(),
153            });
154        }
155
156        let iv_array: [u8; 16] = iv.try_into().map_err(|_| AesError::InvalidIvLength {
157            expected: 16,
158            actual: iv.len(),
159        })?;
160
161        match self.key.size() {
162            AesKeySize::Aes128 => {
163                let key_array: [u8; 16] =
164                    self.key
165                        .key()
166                        .try_into()
167                        .map_err(|_| AesError::InvalidKeyLength {
168                            expected: 16,
169                            actual: self.key.len(),
170                        })?;
171                let encryptor = Encryptor::<Aes128>::new(&key_array.into(), &iv_array.into());
172                // Buffer size: plaintext + padding (up to 16 bytes)
173                let mut buffer = vec![0u8; data.len() + 16];
174                buffer[..data.len()].copy_from_slice(data);
175                let ciphertext = encryptor
176                    .encrypt_padded_mut::<Pkcs7>(&mut buffer, data.len())
177                    .map_err(|e| AesError::EncryptionFailed(format!("Padding error: {e}")))?;
178                Ok(ciphertext.to_vec())
179            }
180            AesKeySize::Aes256 => {
181                let key_array: [u8; 32] =
182                    self.key
183                        .key()
184                        .try_into()
185                        .map_err(|_| AesError::InvalidKeyLength {
186                            expected: 32,
187                            actual: self.key.len(),
188                        })?;
189                let encryptor = Encryptor::<Aes256>::new(&key_array.into(), &iv_array.into());
190                let mut buffer = vec![0u8; data.len() + 16];
191                buffer[..data.len()].copy_from_slice(data);
192                let ciphertext = encryptor
193                    .encrypt_padded_mut::<Pkcs7>(&mut buffer, data.len())
194                    .map_err(|e| AesError::EncryptionFailed(format!("Padding error: {e}")))?;
195                Ok(ciphertext.to_vec())
196            }
197        }
198    }
199
200    /// Decrypt data using AES-CBC mode with PKCS#7 padding removal
201    ///
202    /// Uses RustCrypto's `cbc` crate for production-grade decryption.
203    pub fn decrypt_cbc(&self, data: &[u8], iv: &[u8]) -> Result<Vec<u8>, AesError> {
204        if iv.len() != 16 {
205            return Err(AesError::InvalidIvLength {
206                expected: 16,
207                actual: iv.len(),
208            });
209        }
210
211        if data.len() % 16 != 0 {
212            return Err(AesError::DecryptionFailed(
213                "Data length must be multiple of 16 bytes".to_string(),
214            ));
215        }
216
217        let iv_array: [u8; 16] = iv.try_into().map_err(|_| AesError::InvalidIvLength {
218            expected: 16,
219            actual: iv.len(),
220        })?;
221
222        match self.key.size() {
223            AesKeySize::Aes128 => {
224                let key_array: [u8; 16] =
225                    self.key
226                        .key()
227                        .try_into()
228                        .map_err(|_| AesError::InvalidKeyLength {
229                            expected: 16,
230                            actual: self.key.len(),
231                        })?;
232                let decryptor = Decryptor::<Aes128>::new(&key_array.into(), &iv_array.into());
233                let mut buffer = data.to_vec();
234                let plaintext = decryptor
235                    .decrypt_padded_mut::<Pkcs7>(&mut buffer)
236                    .map_err(|e| AesError::PaddingError(format!("Unpadding error: {e}")))?;
237                Ok(plaintext.to_vec())
238            }
239            AesKeySize::Aes256 => {
240                let key_array: [u8; 32] =
241                    self.key
242                        .key()
243                        .try_into()
244                        .map_err(|_| AesError::InvalidKeyLength {
245                            expected: 32,
246                            actual: self.key.len(),
247                        })?;
248                let decryptor = Decryptor::<Aes256>::new(&key_array.into(), &iv_array.into());
249                let mut buffer = data.to_vec();
250                let plaintext = decryptor
251                    .decrypt_padded_mut::<Pkcs7>(&mut buffer)
252                    .map_err(|e| AesError::PaddingError(format!("Unpadding error: {e}")))?;
253                Ok(plaintext.to_vec())
254            }
255        }
256    }
257
258    /// Encrypt data using AES-ECB mode (for Perms entry in R6)
259    ///
260    /// Note: ECB mode is generally insecure but required by PDF spec for Perms entry.
261    pub fn encrypt_ecb(&self, data: &[u8]) -> Result<Vec<u8>, AesError> {
262        use aes::cipher::{BlockEncrypt, KeyInit};
263
264        if data.len() % 16 != 0 {
265            return Err(AesError::EncryptionFailed(
266                "Data length must be multiple of 16 bytes for ECB mode".to_string(),
267            ));
268        }
269
270        let mut encrypted = Vec::with_capacity(data.len());
271
272        match self.key.size() {
273            AesKeySize::Aes128 => {
274                let key_array: [u8; 16] =
275                    self.key
276                        .key()
277                        .try_into()
278                        .map_err(|_| AesError::InvalidKeyLength {
279                            expected: 16,
280                            actual: self.key.len(),
281                        })?;
282                let cipher = Aes128::new(&key_array.into());
283
284                for chunk in data.chunks(16) {
285                    let mut block =
286                        aes::cipher::generic_array::GenericArray::clone_from_slice(chunk);
287                    cipher.encrypt_block(&mut block);
288                    encrypted.extend_from_slice(&block);
289                }
290            }
291            AesKeySize::Aes256 => {
292                let key_array: [u8; 32] =
293                    self.key
294                        .key()
295                        .try_into()
296                        .map_err(|_| AesError::InvalidKeyLength {
297                            expected: 32,
298                            actual: self.key.len(),
299                        })?;
300                let cipher = Aes256::new(&key_array.into());
301
302                for chunk in data.chunks(16) {
303                    let mut block =
304                        aes::cipher::generic_array::GenericArray::clone_from_slice(chunk);
305                    cipher.encrypt_block(&mut block);
306                    encrypted.extend_from_slice(&block);
307                }
308            }
309        }
310
311        Ok(encrypted)
312    }
313
314    /// Decrypt data using AES-ECB mode (for Perms entry verification in R6)
315    #[allow(dead_code)] // Will be used in R6 implementation
316    pub fn decrypt_ecb(&self, data: &[u8]) -> Result<Vec<u8>, AesError> {
317        use aes::cipher::{BlockDecrypt, KeyInit};
318
319        if data.len() % 16 != 0 {
320            return Err(AesError::DecryptionFailed(
321                "Data length must be multiple of 16 bytes for ECB mode".to_string(),
322            ));
323        }
324
325        let mut decrypted = Vec::with_capacity(data.len());
326
327        match self.key.size() {
328            AesKeySize::Aes128 => {
329                let key_array: [u8; 16] =
330                    self.key
331                        .key()
332                        .try_into()
333                        .map_err(|_| AesError::InvalidKeyLength {
334                            expected: 16,
335                            actual: self.key.len(),
336                        })?;
337                let cipher = Aes128::new(&key_array.into());
338
339                for chunk in data.chunks(16) {
340                    let mut block =
341                        aes::cipher::generic_array::GenericArray::clone_from_slice(chunk);
342                    cipher.decrypt_block(&mut block);
343                    decrypted.extend_from_slice(&block);
344                }
345            }
346            AesKeySize::Aes256 => {
347                let key_array: [u8; 32] =
348                    self.key
349                        .key()
350                        .try_into()
351                        .map_err(|_| AesError::InvalidKeyLength {
352                            expected: 32,
353                            actual: self.key.len(),
354                        })?;
355                let cipher = Aes256::new(&key_array.into());
356
357                for chunk in data.chunks(16) {
358                    let mut block =
359                        aes::cipher::generic_array::GenericArray::clone_from_slice(chunk);
360                    cipher.decrypt_block(&mut block);
361                    decrypted.extend_from_slice(&block);
362                }
363            }
364        }
365
366        Ok(decrypted)
367    }
368
369    /// Encrypt data using AES-CBC mode WITHOUT padding
370    ///
371    /// Used for R5/R6 UE entry encryption where data is already block-aligned (32 bytes).
372    /// Unlike encrypt_cbc, this does not add PKCS#7 padding.
373    ///
374    /// # Requirements
375    /// - Data length must be a multiple of 16 bytes
376    /// - IV must be exactly 16 bytes
377    ///
378    /// # Security Note
379    /// The XOR operations in CBC chaining use simple loops. While Rust/LLVM typically
380    /// produces constant-time code for fixed-size array XOR operations, this is not
381    /// formally guaranteed. PDF encryption is designed for offline brute-force attacks
382    /// where timing side-channels are not in the threat model.
383    pub fn encrypt_cbc_raw(&self, data: &[u8], iv: &[u8]) -> Result<Vec<u8>, AesError> {
384        use aes::cipher::{BlockEncrypt, KeyInit};
385
386        if iv.len() != 16 {
387            return Err(AesError::InvalidIvLength {
388                expected: 16,
389                actual: iv.len(),
390            });
391        }
392
393        if data.len() % 16 != 0 {
394            return Err(AesError::EncryptionFailed(
395                "Data length must be multiple of 16 bytes for raw CBC mode".to_string(),
396            ));
397        }
398
399        let mut encrypted = Vec::with_capacity(data.len());
400        let mut prev_block: [u8; 16] = iv.try_into().map_err(|_| AesError::InvalidIvLength {
401            expected: 16,
402            actual: iv.len(),
403        })?;
404
405        match self.key.size() {
406            AesKeySize::Aes128 => {
407                let key_array: [u8; 16] =
408                    self.key
409                        .key()
410                        .try_into()
411                        .map_err(|_| AesError::InvalidKeyLength {
412                            expected: 16,
413                            actual: self.key.len(),
414                        })?;
415                let cipher = Aes128::new(&key_array.into());
416
417                for chunk in data.chunks(16) {
418                    // XOR with previous ciphertext (or IV for first block)
419                    let mut block: [u8; 16] = [0u8; 16];
420                    for i in 0..16 {
421                        block[i] = chunk[i] ^ prev_block[i];
422                    }
423                    let mut block_ga =
424                        aes::cipher::generic_array::GenericArray::clone_from_slice(&block);
425                    cipher.encrypt_block(&mut block_ga);
426                    prev_block.copy_from_slice(&block_ga);
427                    encrypted.extend_from_slice(&block_ga);
428                }
429            }
430            AesKeySize::Aes256 => {
431                let key_array: [u8; 32] =
432                    self.key
433                        .key()
434                        .try_into()
435                        .map_err(|_| AesError::InvalidKeyLength {
436                            expected: 32,
437                            actual: self.key.len(),
438                        })?;
439                let cipher = Aes256::new(&key_array.into());
440
441                for chunk in data.chunks(16) {
442                    // XOR with previous ciphertext (or IV for first block)
443                    let mut block: [u8; 16] = [0u8; 16];
444                    for i in 0..16 {
445                        block[i] = chunk[i] ^ prev_block[i];
446                    }
447                    let mut block_ga =
448                        aes::cipher::generic_array::GenericArray::clone_from_slice(&block);
449                    cipher.encrypt_block(&mut block_ga);
450                    prev_block.copy_from_slice(&block_ga);
451                    encrypted.extend_from_slice(&block_ga);
452                }
453            }
454        }
455
456        Ok(encrypted)
457    }
458
459    /// Decrypt data using AES-CBC mode WITHOUT padding
460    ///
461    /// Used for R5/R6 UE entry decryption where data is already block-aligned (32 bytes).
462    /// Unlike decrypt_cbc, this does not expect or remove PKCS#7 padding.
463    ///
464    /// # Requirements
465    /// - Data length must be a multiple of 16 bytes
466    /// - IV must be exactly 16 bytes
467    ///
468    /// # Security Note
469    /// The XOR operations in CBC chaining use simple loops. While Rust/LLVM typically
470    /// produces constant-time code for fixed-size array XOR operations, this is not
471    /// formally guaranteed. PDF encryption is designed for offline brute-force attacks
472    /// where timing side-channels are not in the threat model.
473    pub fn decrypt_cbc_raw(&self, data: &[u8], iv: &[u8]) -> Result<Vec<u8>, AesError> {
474        use aes::cipher::{BlockDecrypt, KeyInit};
475
476        if iv.len() != 16 {
477            return Err(AesError::InvalidIvLength {
478                expected: 16,
479                actual: iv.len(),
480            });
481        }
482
483        if data.len() % 16 != 0 {
484            return Err(AesError::DecryptionFailed(
485                "Data length must be multiple of 16 bytes for raw CBC mode".to_string(),
486            ));
487        }
488
489        let mut decrypted = Vec::with_capacity(data.len());
490        let mut prev_block: [u8; 16] = iv.try_into().map_err(|_| AesError::InvalidIvLength {
491            expected: 16,
492            actual: iv.len(),
493        })?;
494
495        match self.key.size() {
496            AesKeySize::Aes128 => {
497                let key_array: [u8; 16] =
498                    self.key
499                        .key()
500                        .try_into()
501                        .map_err(|_| AesError::InvalidKeyLength {
502                            expected: 16,
503                            actual: self.key.len(),
504                        })?;
505                let cipher = Aes128::new(&key_array.into());
506
507                for chunk in data.chunks(16) {
508                    let mut block =
509                        aes::cipher::generic_array::GenericArray::clone_from_slice(chunk);
510                    cipher.decrypt_block(&mut block);
511                    // XOR with previous ciphertext (or IV for first block)
512                    for i in 0..16 {
513                        block[i] ^= prev_block[i];
514                    }
515                    prev_block.copy_from_slice(chunk);
516                    decrypted.extend_from_slice(&block);
517                }
518            }
519            AesKeySize::Aes256 => {
520                let key_array: [u8; 32] =
521                    self.key
522                        .key()
523                        .try_into()
524                        .map_err(|_| AesError::InvalidKeyLength {
525                            expected: 32,
526                            actual: self.key.len(),
527                        })?;
528                let cipher = Aes256::new(&key_array.into());
529
530                for chunk in data.chunks(16) {
531                    let mut block =
532                        aes::cipher::generic_array::GenericArray::clone_from_slice(chunk);
533                    cipher.decrypt_block(&mut block);
534                    // XOR with previous ciphertext (or IV for first block)
535                    for i in 0..16 {
536                        block[i] ^= prev_block[i];
537                    }
538                    prev_block.copy_from_slice(chunk);
539                    decrypted.extend_from_slice(&block);
540                }
541            }
542        }
543
544        Ok(decrypted)
545    }
546}
547
548/// Generate random IV for AES encryption
549pub fn generate_iv() -> Vec<u8> {
550    // In production, use a cryptographically secure random number generator
551    // For now, use a simple approach with multiple entropy sources
552    use std::collections::hash_map::DefaultHasher;
553    use std::hash::{Hash, Hasher};
554    use std::sync::atomic::{AtomicUsize, Ordering};
555    use std::thread;
556    use std::time::SystemTime;
557
558    static COUNTER: AtomicUsize = AtomicUsize::new(0);
559
560    let mut hasher = DefaultHasher::new();
561
562    // Hash multiple entropy sources to ensure uniqueness
563    SystemTime::now().hash(&mut hasher);
564    thread::current().id().hash(&mut hasher);
565    std::process::id().hash(&mut hasher);
566    COUNTER.fetch_add(1, Ordering::SeqCst).hash(&mut hasher);
567
568    let seed = hasher.finish();
569    let mut iv = Vec::new();
570
571    for i in 0..16 {
572        iv.push(((seed >> (i * 4)) as u8) ^ (i as u8));
573    }
574
575    iv
576}
577
578#[cfg(test)]
579mod tests {
580    use super::*;
581
582    // ===== AesKey Tests =====
583
584    #[test]
585    fn test_aes_key_creation() {
586        // Test AES-128 key
587        let key_128 = vec![0u8; 16];
588        let aes_key = AesKey::new_128(key_128.clone()).unwrap();
589        assert_eq!(aes_key.key(), &key_128);
590        assert_eq!(aes_key.size(), AesKeySize::Aes128);
591        assert_eq!(aes_key.len(), 16);
592
593        // Test AES-256 key
594        let key_256 = vec![1u8; 32];
595        let aes_key = AesKey::new_256(key_256.clone()).unwrap();
596        assert_eq!(aes_key.key(), &key_256);
597        assert_eq!(aes_key.size(), AesKeySize::Aes256);
598        assert_eq!(aes_key.len(), 32);
599    }
600
601    #[test]
602    fn test_aes_key_invalid_length() {
603        // Test invalid AES-128 key length
604        let key_short = vec![0u8; 15];
605        assert!(AesKey::new_128(key_short).is_err());
606
607        let key_long = vec![0u8; 17];
608        assert!(AesKey::new_128(key_long).is_err());
609
610        // Test invalid AES-256 key length
611        let key_short = vec![0u8; 31];
612        assert!(AesKey::new_256(key_short).is_err());
613
614        let key_long = vec![0u8; 33];
615        assert!(AesKey::new_256(key_long).is_err());
616    }
617
618    #[test]
619    fn test_aes_key_size() {
620        assert_eq!(AesKeySize::Aes128.key_length(), 16);
621        assert_eq!(AesKeySize::Aes256.key_length(), 32);
622        assert_eq!(AesKeySize::Aes128.block_size(), 16);
623        assert_eq!(AesKeySize::Aes256.block_size(), 16);
624    }
625
626    #[test]
627    fn test_aes_key_size_equality() {
628        assert_eq!(AesKeySize::Aes128, AesKeySize::Aes128);
629        assert_eq!(AesKeySize::Aes256, AesKeySize::Aes256);
630        assert_ne!(AesKeySize::Aes128, AesKeySize::Aes256);
631    }
632
633    #[test]
634    fn test_aes_key_size_debug() {
635        assert_eq!(format!("{:?}", AesKeySize::Aes128), "Aes128");
636        assert_eq!(format!("{:?}", AesKeySize::Aes256), "Aes256");
637    }
638
639    #[test]
640    fn test_aes_key_is_empty() {
641        let key = AesKey::new_128(vec![0u8; 16]).unwrap();
642        assert!(!key.is_empty());
643    }
644
645    #[test]
646    fn test_aes_key_debug() {
647        let key = AesKey::new_128(vec![1u8; 16]).unwrap();
648        let debug_str = format!("{key:?}");
649        assert!(debug_str.contains("AesKey"));
650        assert!(debug_str.contains("key:"));
651        assert!(debug_str.contains("size:"));
652    }
653
654    #[test]
655    fn test_aes_key_clone() {
656        let key = AesKey::new_128(vec![1u8; 16]).unwrap();
657        let cloned = key.clone();
658        assert_eq!(key.key(), cloned.key());
659        assert_eq!(key.size(), cloned.size());
660    }
661
662    #[test]
663    fn test_aes_key_various_patterns() {
664        let patterns = vec![
665            vec![0xFF; 16],                     // All 1s
666            vec![0x00; 16],                     // All 0s
667            (0..16).map(|i| i as u8).collect(), // Sequential
668            vec![0xA5; 16],                     // Alternating bits
669        ];
670
671        for pattern in patterns {
672            let key = AesKey::new_128(pattern.clone()).unwrap();
673            assert_eq!(key.key(), &pattern);
674            assert_eq!(key.len(), 16);
675        }
676    }
677
678    #[test]
679    fn test_aes_key_256_various_patterns() {
680        let patterns = vec![
681            vec![0xFF; 32],
682            vec![0x00; 32],
683            (0..32).map(|i| i as u8).collect(),
684            vec![0x5A; 32],
685        ];
686
687        for pattern in patterns {
688            let key = AesKey::new_256(pattern.clone()).unwrap();
689            assert_eq!(key.key(), &pattern);
690            assert_eq!(key.len(), 32);
691        }
692    }
693
694    // ===== AesError Tests =====
695
696    #[test]
697    fn test_aes_error_display() {
698        let error1 = AesError::InvalidKeyLength {
699            expected: 16,
700            actual: 15,
701        };
702        assert!(error1.to_string().contains("Invalid key length"));
703
704        let error2 = AesError::EncryptionFailed("test".to_string());
705        assert!(error2.to_string().contains("Encryption failed"));
706
707        let error3 = AesError::PaddingError("bad padding".to_string());
708        assert!(error3.to_string().contains("Padding error"));
709    }
710
711    #[test]
712    fn test_aes_error_equality() {
713        let err1 = AesError::InvalidKeyLength {
714            expected: 16,
715            actual: 15,
716        };
717        let err2 = AesError::InvalidKeyLength {
718            expected: 16,
719            actual: 15,
720        };
721        let err3 = AesError::InvalidKeyLength {
722            expected: 16,
723            actual: 17,
724        };
725
726        assert_eq!(err1, err2);
727        assert_ne!(err1, err3);
728    }
729
730    #[test]
731    fn test_aes_error_clone() {
732        let errors = vec![
733            AesError::InvalidKeyLength {
734                expected: 16,
735                actual: 15,
736            },
737            AesError::InvalidIvLength {
738                expected: 16,
739                actual: 15,
740            },
741            AesError::EncryptionFailed("test".to_string()),
742            AesError::DecryptionFailed("test".to_string()),
743            AesError::PaddingError("test".to_string()),
744        ];
745
746        for error in errors {
747            let cloned = error.clone();
748            assert_eq!(error, cloned);
749        }
750    }
751
752    #[test]
753    fn test_aes_error_debug() {
754        let error = AesError::InvalidKeyLength {
755            expected: 16,
756            actual: 15,
757        };
758        let debug_str = format!("{error:?}");
759        assert!(debug_str.contains("InvalidKeyLength"));
760        assert!(debug_str.contains("expected: 16"));
761        assert!(debug_str.contains("actual: 15"));
762    }
763
764    #[test]
765    fn test_aes_error_display_all_variants() {
766        let errors = vec![
767            (
768                AesError::InvalidKeyLength {
769                    expected: 16,
770                    actual: 15,
771                },
772                "Invalid key length",
773            ),
774            (
775                AesError::InvalidIvLength {
776                    expected: 16,
777                    actual: 15,
778                },
779                "Invalid IV length",
780            ),
781            (
782                AesError::EncryptionFailed("custom error".to_string()),
783                "Encryption failed: custom error",
784            ),
785            (
786                AesError::DecryptionFailed("custom error".to_string()),
787                "Decryption failed: custom error",
788            ),
789            (
790                AesError::PaddingError("custom error".to_string()),
791                "Padding error: custom error",
792            ),
793        ];
794
795        for (error, expected_substring) in errors {
796            let display = error.to_string();
797            assert!(display.contains(expected_substring));
798        }
799    }
800
801    #[test]
802    fn test_aes_error_is_std_error() {
803        let error: Box<dyn std::error::Error> =
804            Box::new(AesError::PaddingError("test".to_string()));
805        assert_eq!(error.to_string(), "Padding error: test");
806    }
807
808    // ===== AES Encryption/Decryption Tests (Using RustCrypto) =====
809
810    #[test]
811    fn test_aes_128_encrypt_decrypt_roundtrip() {
812        // Now with RustCrypto, roundtrip should work perfectly
813        let key = AesKey::new_128(vec![
814            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
815            0x4f, 0x3c,
816        ])
817        .unwrap();
818        let aes = Aes::new(key);
819
820        let data = b"Hello, AES World!";
821        let iv = vec![
822            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
823            0x0e, 0x0f,
824        ];
825
826        let encrypted = aes.encrypt_cbc(data, &iv).unwrap();
827        assert_ne!(encrypted, data);
828        assert!(encrypted.len() >= data.len());
829        assert_eq!(encrypted.len() % 16, 0);
830
831        // With RustCrypto, decryption should produce exact original
832        let decrypted = aes.decrypt_cbc(&encrypted, &iv).unwrap();
833        assert_eq!(decrypted.as_slice(), data.as_slice());
834    }
835
836    #[test]
837    fn test_aes_256_encrypt_decrypt_roundtrip() {
838        let key = AesKey::new_256(vec![0x42; 32]).unwrap();
839        let aes = Aes::new(key);
840
841        let data = b"This is a test for AES-256 encryption!";
842        let iv = vec![0x33; 16];
843
844        let encrypted = aes.encrypt_cbc(data, &iv).unwrap();
845        assert_ne!(encrypted.as_slice(), data.as_slice());
846
847        let decrypted = aes.decrypt_cbc(&encrypted, &iv).unwrap();
848        assert_eq!(decrypted.as_slice(), data.as_slice());
849    }
850
851    #[test]
852    fn test_aes_empty_data() {
853        let key = AesKey::new_128(vec![0u8; 16]).unwrap();
854        let aes = Aes::new(key);
855        let iv = vec![0u8; 16];
856
857        let data = b"";
858        let encrypted = aes.encrypt_cbc(data, &iv).unwrap();
859        assert_eq!(encrypted.len(), 16); // One block due to PKCS#7 padding
860
861        let decrypted = aes.decrypt_cbc(&encrypted, &iv).unwrap();
862        assert_eq!(decrypted.as_slice(), data.as_slice());
863    }
864
865    #[test]
866    fn test_aes_invalid_iv() {
867        let key = AesKey::new_128(vec![0u8; 16]).unwrap();
868        let aes = Aes::new(key);
869
870        let data = b"test data";
871        let iv_short = vec![0u8; 15];
872        let iv_long = vec![0u8; 17];
873
874        assert!(aes.encrypt_cbc(data, &iv_short).is_err());
875        assert!(aes.encrypt_cbc(data, &iv_long).is_err());
876
877        let encrypted = aes.encrypt_cbc(data, &[0u8; 16]).unwrap();
878        assert!(aes.decrypt_cbc(&encrypted, &iv_short).is_err());
879        assert!(aes.decrypt_cbc(&encrypted, &iv_long).is_err());
880    }
881
882    #[test]
883    fn test_aes_multiple_blocks() {
884        let key = AesKey::new_128(vec![0x42; 16]).unwrap();
885        let aes = Aes::new(key);
886        let iv = vec![0x37; 16];
887
888        // Test data that spans multiple blocks
889        let data = vec![0x55; 48]; // 3 blocks exactly
890        let encrypted = aes.encrypt_cbc(&data, &iv).unwrap();
891        assert_eq!(encrypted.len(), 64); // PKCS#7 adds padding even for exact blocks
892
893        let decrypted = aes.decrypt_cbc(&encrypted, &iv).unwrap();
894        assert_eq!(decrypted, data);
895    }
896
897    #[test]
898    fn test_aes_large_data() {
899        let key = AesKey::new_128(vec![0x11; 16]).unwrap();
900        let aes = Aes::new(key);
901        let iv = vec![0x22; 16];
902
903        // Test with larger data (1KB)
904        let data = vec![0x33; 1024];
905        let encrypted = aes.encrypt_cbc(&data, &iv).unwrap();
906        assert!(encrypted.len() >= 1024);
907        assert_eq!(encrypted.len() % 16, 0);
908
909        let decrypted = aes.decrypt_cbc(&encrypted, &iv).unwrap();
910        assert_eq!(decrypted, data);
911    }
912
913    #[test]
914    fn test_aes_various_data_sizes() {
915        let key = AesKey::new_128(vec![0xAA; 16]).unwrap();
916        let aes = Aes::new(key);
917        let iv = vec![0xBB; 16];
918
919        for size in [1, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129] {
920            let data = vec![0xCC; size];
921            let encrypted = aes.encrypt_cbc(&data, &iv).unwrap();
922
923            // PKCS#7 always adds padding
924            let expected_size = if size % 16 == 0 {
925                size + 16
926            } else {
927                (size / 16 + 1) * 16
928            };
929            assert_eq!(encrypted.len(), expected_size, "size={size}");
930
931            let decrypted = aes.decrypt_cbc(&encrypted, &iv).unwrap();
932            assert_eq!(decrypted, data, "size={size}");
933        }
934    }
935
936    #[test]
937    fn test_decrypt_invalid_data_length() {
938        let key = AesKey::new_128(vec![0u8; 16]).unwrap();
939        let aes = Aes::new(key);
940        let iv = vec![0u8; 16];
941
942        // Data not multiple of block size
943        let invalid_data = vec![0u8; 17];
944        let result = aes.decrypt_cbc(&invalid_data, &iv);
945        assert!(result.is_err());
946        match result.unwrap_err() {
947            AesError::DecryptionFailed(msg) => {
948                assert!(msg.contains("multiple of 16"));
949            }
950            _ => panic!("Expected DecryptionFailed error"),
951        }
952    }
953
954    #[test]
955    fn test_encrypt_with_different_ivs() {
956        let key = AesKey::new_128(vec![0x42; 16]).unwrap();
957        let aes = Aes::new(key);
958
959        let data = b"Same data encrypted with different IVs";
960        let iv1 = vec![0x00; 16];
961        let iv2 = vec![0xFF; 16];
962
963        let encrypted1 = aes.encrypt_cbc(data, &iv1).unwrap();
964        let encrypted2 = aes.encrypt_cbc(data, &iv2).unwrap();
965
966        // Same data with different IVs should produce different ciphertexts
967        assert_ne!(encrypted1, encrypted2);
968        assert_eq!(encrypted1.len(), encrypted2.len());
969
970        // Both should decrypt correctly with their respective IVs
971        let decrypted1 = aes.decrypt_cbc(&encrypted1, &iv1).unwrap();
972        let decrypted2 = aes.decrypt_cbc(&encrypted2, &iv2).unwrap();
973        assert_eq!(decrypted1.as_slice(), data.as_slice());
974        assert_eq!(decrypted2.as_slice(), data.as_slice());
975    }
976
977    #[test]
978    fn test_cbc_mode_no_patterns() {
979        let key = AesKey::new_128(vec![0x11; 16]).unwrap();
980        let aes = Aes::new(key);
981
982        // Two identical plaintext blocks
983        let data = vec![0x44; 32];
984        let iv = vec![0x55; 16];
985
986        let encrypted = aes.encrypt_cbc(&data, &iv).unwrap();
987
988        // In CBC mode, the two encrypted blocks should be different
989        // even though plaintext blocks are identical
990        let block1 = &encrypted[0..16];
991        let block2 = &encrypted[16..32];
992        assert_ne!(block1, block2);
993    }
994
995    #[test]
996    fn test_error_propagation() {
997        let key = AesKey::new_128(vec![0u8; 16]).unwrap();
998        let aes = Aes::new(key);
999
1000        // Test encryption with invalid IV
1001        let result = aes.encrypt_cbc(b"test", &[0u8; 15]);
1002        assert!(matches!(result, Err(AesError::InvalidIvLength { .. })));
1003
1004        // Test decryption with invalid IV
1005        let valid_encrypted = vec![0u8; 16];
1006        let result = aes.decrypt_cbc(&valid_encrypted, &[0u8; 17]);
1007        assert!(matches!(result, Err(AesError::InvalidIvLength { .. })));
1008    }
1009
1010    // ===== ECB Mode Tests =====
1011
1012    #[test]
1013    fn test_ecb_encrypt_decrypt_roundtrip() {
1014        let key = AesKey::new_256(vec![0x55; 32]).unwrap();
1015        let aes = Aes::new(key);
1016
1017        // ECB requires data multiple of 16
1018        let data = vec![0xAB; 32]; // Two blocks
1019
1020        let encrypted = aes.encrypt_ecb(&data).unwrap();
1021        assert_eq!(encrypted.len(), 32);
1022        assert_ne!(encrypted, data);
1023
1024        let decrypted = aes.decrypt_ecb(&encrypted).unwrap();
1025        assert_eq!(decrypted, data);
1026    }
1027
1028    #[test]
1029    fn test_ecb_invalid_data_length() {
1030        let key = AesKey::new_128(vec![0u8; 16]).unwrap();
1031        let aes = Aes::new(key);
1032
1033        // ECB requires data multiple of 16
1034        let invalid_data = vec![0u8; 17];
1035        assert!(aes.encrypt_ecb(&invalid_data).is_err());
1036        assert!(aes.decrypt_ecb(&invalid_data).is_err());
1037    }
1038
1039    #[test]
1040    fn test_ecb_shows_patterns() {
1041        // ECB mode produces same ciphertext for same plaintext blocks
1042        let key = AesKey::new_128(vec![0x11; 16]).unwrap();
1043        let aes = Aes::new(key);
1044
1045        let data = vec![0x44; 32]; // Two identical blocks
1046
1047        let encrypted = aes.encrypt_ecb(&data).unwrap();
1048
1049        // In ECB mode, both blocks should be identical (unlike CBC)
1050        let block1 = &encrypted[0..16];
1051        let block2 = &encrypted[16..32];
1052        assert_eq!(
1053            block1, block2,
1054            "ECB should produce same ciphertext for identical blocks"
1055        );
1056    }
1057
1058    // ===== generate_iv Tests =====
1059
1060    #[test]
1061    fn test_generate_iv() {
1062        let iv1 = generate_iv();
1063        let iv2 = generate_iv();
1064
1065        assert_eq!(iv1.len(), 16);
1066        assert_eq!(iv2.len(), 16);
1067    }
1068
1069    #[test]
1070    fn test_generate_iv_properties() {
1071        let ivs: Vec<Vec<u8>> = (0..10).map(|_| generate_iv()).collect();
1072
1073        // All should be 16 bytes
1074        for iv in &ivs {
1075            assert_eq!(iv.len(), 16);
1076        }
1077
1078        // Check that not all IVs are identical (with counter, they should differ)
1079        let first = &ivs[0];
1080        let all_same = ivs.iter().all(|iv| iv == first);
1081        assert!(
1082            !all_same || ivs.len() == 1,
1083            "IVs should not all be identical"
1084        );
1085    }
1086
1087    // ===== NIST Test Vector =====
1088
1089    #[test]
1090    fn test_aes_256_cbc_nist_vector() {
1091        // NIST test vector for AES-256-CBC from SP 800-38A
1092        let key = AesKey::new_256(vec![
1093            0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
1094            0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
1095            0x09, 0x14, 0xdf, 0xf4,
1096        ])
1097        .unwrap();
1098        let aes = Aes::new(key);
1099
1100        let iv = vec![
1101            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
1102            0x0e, 0x0f,
1103        ];
1104
1105        let plaintext = vec![
1106            0x6b, 0xc1, 0xbe, 0xe2, 0x2e, 0x40, 0x9f, 0x96, 0xe9, 0x3d, 0x7e, 0x11, 0x73, 0x93,
1107            0x17, 0x2a,
1108        ];
1109
1110        let expected_ciphertext = vec![
1111            0xf5, 0x8c, 0x4c, 0x04, 0xd6, 0xe5, 0xf1, 0xba, 0x77, 0x9e, 0xab, 0xfb, 0x5f, 0x7b,
1112            0xfb, 0xd6,
1113        ];
1114
1115        let encrypted = aes.encrypt_cbc(&plaintext, &iv).unwrap();
1116
1117        // First 16 bytes should match NIST expected ciphertext
1118        assert_eq!(
1119            &encrypted[..16],
1120            expected_ciphertext.as_slice(),
1121            "AES-256-CBC encryption should match NIST test vector"
1122        );
1123    }
1124}