1use aes::{Aes128, Aes256};
9use cbc::{Decryptor, Encryptor};
10use cipher::block_padding::Pkcs7;
11use cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit};
12
13#[derive(Debug, Clone, Copy, PartialEq)]
15pub enum AesKeySize {
16 Aes128,
18 Aes256,
20}
21
22impl AesKeySize {
23 pub fn key_length(&self) -> usize {
25 match self {
26 AesKeySize::Aes128 => 16,
27 AesKeySize::Aes256 => 32,
28 }
29 }
30
31 pub fn block_size(&self) -> usize {
33 16
34 }
35}
36
37#[derive(Debug, Clone)]
39pub struct AesKey {
40 key: Vec<u8>,
42 size: AesKeySize,
44}
45
46impl AesKey {
47 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 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 pub fn key(&self) -> &[u8] {
79 &self.key
80 }
81
82 pub fn size(&self) -> AesKeySize {
84 self.size
85 }
86
87 pub fn len(&self) -> usize {
89 self.key.len()
90 }
91
92 pub fn is_empty(&self) -> bool {
94 self.key.is_empty()
95 }
96}
97
98#[derive(Debug, Clone, PartialEq)]
100pub enum AesError {
101 InvalidKeyLength { expected: usize, actual: usize },
103 InvalidIvLength { expected: usize, actual: usize },
105 EncryptionFailed(String),
107 DecryptionFailed(String),
109 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
131pub struct Aes {
136 key: AesKey,
137}
138
139impl Aes {
140 pub fn new(key: AesKey) -> Self {
142 Self { key }
143 }
144
145 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 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 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 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 #[allow(dead_code)] 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 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 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 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 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 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 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
548pub fn generate_iv() -> Vec<u8> {
550 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 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 #[test]
585 fn test_aes_key_creation() {
586 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 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 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 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], vec![0x00; 16], (0..16).map(|i| i as u8).collect(), vec![0xA5; 16], ];
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 #[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 #[test]
811 fn test_aes_128_encrypt_decrypt_roundtrip() {
812 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 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); 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 let data = vec![0x55; 48]; let encrypted = aes.encrypt_cbc(&data, &iv).unwrap();
891 assert_eq!(encrypted.len(), 64); 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 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 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 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 assert_ne!(encrypted1, encrypted2);
968 assert_eq!(encrypted1.len(), encrypted2.len());
969
970 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 let data = vec![0x44; 32];
984 let iv = vec![0x55; 16];
985
986 let encrypted = aes.encrypt_cbc(&data, &iv).unwrap();
987
988 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 let result = aes.encrypt_cbc(b"test", &[0u8; 15]);
1002 assert!(matches!(result, Err(AesError::InvalidIvLength { .. })));
1003
1004 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 #[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 let data = vec![0xAB; 32]; 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 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 let key = AesKey::new_128(vec![0x11; 16]).unwrap();
1043 let aes = Aes::new(key);
1044
1045 let data = vec![0x44; 32]; let encrypted = aes.encrypt_ecb(&data).unwrap();
1048
1049 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 #[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 for iv in &ivs {
1075 assert_eq!(iv.len(), 16);
1076 }
1077
1078 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 #[test]
1090 fn test_aes_256_cbc_nist_vector() {
1091 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 assert_eq!(
1119 &encrypted[..16],
1120 expected_ciphertext.as_slice(),
1121 "AES-256-CBC encryption should match NIST test vector"
1122 );
1123 }
1124}