use super::Encryptor;
 use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
 use aes::Aes128;
 use rand::Rng;
 use sha2::{Digest, Sha256};
 
 type Aes128CbcEnc = cbc::Encryptor<Aes128>;
 #[allow(dead_code)]
 type Aes128CbcDec = cbc::Decryptor<Aes128>;
 
 /// AES-128 cipher implementation with CBC mode
 pub struct AesEncryptor {
     key: [u8; 16],
 }
 
 impl AesEncryptor {
     /// Creates a new AES encryptor with the given key
     /// Uses SHA-256 to derive a proper 16-byte key from the input
     pub fn new(key: &str) -> Self {
         let mut hasher = Sha256::new();
         hasher.update(key.as_bytes());
         let hash = hasher.finalize();
 
         // Take first 16 bytes of hash for AES-128
         let mut key_bytes = [0u8; 16];
         key_bytes.copy_from_slice(&hash[..16]);
 
         println!(
             "[AES] Derived key from input (SHA-256): {:02x?}",
             &key_bytes[..8]
         );
 
         Self { key: key_bytes }
     }
 }
 
 impl Encryptor for AesEncryptor {
     fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>, String> {
         // Generate random IV
         let mut rng = rand::thread_rng();
         let mut iv = [0u8; 16];
         rng.fill(&mut iv);
 
         println!(
             "[AES] Generated random IV for encryption: {:02x?}",
             &iv[..8]
         );
 
         let cipher = Aes128CbcEnc::new(&self.key.into(), &iv.into());
         let ciphertext = cipher.encrypt_padded_vec_mut::<Pkcs7>(data);
 
         // Prepend IV to ciphertext
         let mut result = iv.to_vec();
         result.extend_from_slice(&ciphertext);
 
         Ok(result)
     }
 
     fn decrypt(&self, data: &[u8]) -> Result<Vec<u8>, String> {
         if data.len() < 16 {
             return Err("Ciphertext too short to contain IV".to_string());
         }
 
         // Extract IV from beginning of data
         let iv: [u8; 16] = data[..16].try_into().unwrap();
         let ciphertext = &data[16..];
 
         println!("[AES] Extracted IV for decryption: {:02x?}", &iv[..8]);
 
         let cipher = Aes128CbcDec::new(&self.key.into(), &iv.into());
         cipher
             .decrypt_padded_vec_mut::<Pkcs7>(ciphertext)
             .map_err(|e| format!("Decryption error: {:?}", e))
     }
 
     fn key_size(&self) -> usize {
         16 // AES-128 uses 16-byte keys
     }
 
     fn algorithm_name(&self) -> &'static str {
         "AES-128-CBC"
     }
 
     fn uses_iv(&self) -> bool {
         true
     }
 
     fn generate_decrypt_code(&self) -> &'static str {
         r#"/// Decrypts AES-128 CBC encrypted data using the provided key
 ///
 /// Returns `Ok(Vec<u8>)` with decrypted data if successful, otherwise returns an error.
 pub fn decrypt_aes(data: &[u8], key: &[u8]) -> Result<Vec<u8>, String> {
     use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, KeyIvInit};
     use aes::Aes128;
     
     if data.len() < 16 {
         return Err("Ciphertext too short to contain IV".to_string());
     }
     
     if key.len() != 16 {
         return Err("Key must be 16 bytes for AES-128".to_string());
     }
     
     let iv: [u8; 16] = data[..16].try_into().unwrap();
     let ciphertext = &data[16..];
     
     let key_array: [u8; 16] = key.try_into().map_err(|_| "Invalid key length".to_string())?;
     
     type Aes128CbcDec = cbc::Decryptor<Aes128>;
     let cipher = Aes128CbcDec::new(&key_array.into(), &iv.into());
     cipher
         .decrypt_padded_vec_mut::<Pkcs7>(ciphertext)
         .map_err(|e| format!("Decryption error: {:?}", e))
 }"#
     }
 
     fn required_dependencies(&self) -> Vec<(&'static str, &'static str)> {
         vec![("aes", "0.8"), ("cbc", "0.1")]
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
     fn test_aes_encrypt_decrypt() {
         let encryptor = AesEncryptor::new("my-secret-key-123");
         let plaintext = b"Hello, World!";
 
         let encrypted = encryptor.encrypt(plaintext).unwrap();
         assert_ne!(encrypted, plaintext);
 
         let decrypted = encryptor.decrypt(&encrypted).unwrap();
         assert_eq!(decrypted, plaintext);
     }
 
     #[test]
     fn test_aes_different_keys() {
         let encryptor1 = AesEncryptor::new("key1");
         let encryptor2 = AesEncryptor::new("key2");
         let plaintext = b"Test data";
 
         let encrypted1 = encryptor1.encrypt(plaintext).unwrap();
         let encrypted2 = encryptor2.encrypt(plaintext).unwrap();
 
         assert_ne!(encrypted1, encrypted2);
     }
 
     #[test]
     fn test_aes_empty_data() {
         let encryptor = AesEncryptor::new("key");
         let empty: &[u8] = &[];
 
         let encrypted = encryptor.encrypt(empty).unwrap();
         // AES with padding will have at least one block
         assert!(encrypted.len() >= 16);
 
         let decrypted = encryptor.decrypt(&encrypted).unwrap();
         assert_eq!(decrypted.len(), 0);
     }
 
     #[test]
     fn test_aes_wrong_key_decryption() {
         let encryptor1 = AesEncryptor::new("correct-key");
         let encryptor2 = AesEncryptor::new("wrong-key");
         let plaintext = b"Secret message";
 
         let encrypted = encryptor1.encrypt(plaintext).unwrap();
         let result = encryptor2.decrypt(&encrypted);
 
         // Decryption with wrong key should fail or produce garbage
         assert!(result.is_err() || result.unwrap() != plaintext);
     }
 
     #[test]
     fn test_aes_long_data() {
         let encryptor = AesEncryptor::new("secret-key");
         let plaintext =
             b"This is a much longer message that exceeds the AES block size of 16 bytes";
 
         let encrypted = encryptor.encrypt(plaintext).unwrap();
         assert_ne!(encrypted, plaintext);
 
         let decrypted = encryptor.decrypt(&encrypted).unwrap();
         assert_eq!(decrypted, plaintext);
     }
 
     #[test]
     fn test_aes_binary_data() {
         let encryptor = AesEncryptor::new("key");
         let plaintext: Vec<u8> = (0..=255).collect();
 
         let encrypted = encryptor.encrypt(&plaintext).unwrap();
         assert_ne!(encrypted, plaintext);
 
         let decrypted = encryptor.decrypt(&encrypted).unwrap();
         assert_eq!(decrypted, plaintext);
     }
 }

© 2024 - 2026 djnn, inc • generated with stagit