use super::Encryptor;
 use rand::Rng;
 use sha2::{Digest, Sha256};
 use twofish::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
 use twofish::Twofish;
 
 type TwofishCbcEnc = cbc::Encryptor<Twofish>;
 #[allow(dead_code)]
 type TwofishCbcDec = cbc::Decryptor<Twofish>;
 
 /// Twofish cipher implementation with CBC mode
 pub struct TwofishEncryptor {
     key: [u8; 32],
 }
 
 impl TwofishEncryptor {
     /// Creates a new Twofish encryptor with the given key
     /// Uses SHA-256 to derive a proper 32-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();
 
         let mut key_bytes = [0u8; 32];
         key_bytes.copy_from_slice(&hash[..32]);
 
         println!(
             "[Twofish] Derived key from input (SHA-256): {:02x?}",
             &key_bytes[..8]
         );
 
         Self { key: key_bytes }
     }
 }
 
 impl Encryptor for TwofishEncryptor {
     fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>, String> {
         let mut rng = rand::thread_rng();
         let mut iv = [0u8; 16];
         rng.fill(&mut iv);
 
         println!(
             "[Twofish] Generated random IV for encryption: {:02x?}",
             &iv[..8]
         );
 
         let cipher = TwofishCbcEnc::new(&self.key.into(), &iv.into());
         let ciphertext = cipher.encrypt_padded_vec_mut::<Pkcs7>(data);
 
         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());
         }
 
         let iv: [u8; 16] = data[..16].try_into().unwrap();
         let ciphertext = &data[16..];
 
         println!("[Twofish] Extracted IV for decryption: {:02x?}", &iv[..8]);
 
         let cipher = TwofishCbcDec::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 {
         32
     }
 
     fn algorithm_name(&self) -> &'static str {
         "Twofish-CBC"
     }
 
     fn uses_iv(&self) -> bool {
         true
     }
 
     fn generate_decrypt_code(&self) -> &'static str {
         r#"/// Decrypts Twofish CBC encrypted data using the provided key
 ///
 /// Returns `Ok(Vec<u8>)` with decrypted data if successful, otherwise returns an error.
 pub fn decrypt_twofish(data: &[u8], key: &[u8]) -> Result<Vec<u8>, String> {
     use twofish::cipher::{block_padding::Pkcs7, BlockDecryptMut, KeyIvInit};
     use twofish::Twofish;
     
     if data.len() < 16 {
         return Err("Ciphertext too short to contain IV".to_string());
     }
     
     if key.len() != 32 {
         return Err("Key must be 32 bytes for Twofish".to_string());
     }
     
     let iv: [u8; 16] = data[..16].try_into().unwrap();
     let ciphertext = &data[16..];
     
     let key_array: [u8; 32] = key.try_into().map_err(|_| "Invalid key length".to_string())?;
     
     type TwofishCbcDec = cbc::Decryptor<Twofish>;
     let cipher = TwofishCbcDec::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![("twofish", "0.7"), ("cbc", "0.1")]
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
     fn test_twofish_encrypt_decrypt() {
         let encryptor = TwofishEncryptor::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_twofish_different_keys() {
         let encryptor1 = TwofishEncryptor::new("key1");
         let encryptor2 = TwofishEncryptor::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_twofish_empty_data() {
         let encryptor = TwofishEncryptor::new("key");
         let empty: &[u8] = &[];
 
         let encrypted = encryptor.encrypt(empty).unwrap();
         assert!(encrypted.len() >= 16);
 
         let decrypted = encryptor.decrypt(&encrypted).unwrap();
         assert_eq!(decrypted.len(), 0);
     }
 
     #[test]
     fn test_twofish_wrong_key_decryption() {
         let encryptor1 = TwofishEncryptor::new("correct-key");
         let encryptor2 = TwofishEncryptor::new("wrong-key");
         let plaintext = b"Secret message";
 
         let encrypted = encryptor1.encrypt(plaintext).unwrap();
         let result = encryptor2.decrypt(&encrypted);
 
         assert!(result.is_err() || result.unwrap() != plaintext);
     }
 
     #[test]
     fn test_twofish_long_data() {
         let encryptor = TwofishEncryptor::new("secret-key");
         let plaintext =
             b"This is a much longer message that exceeds the Twofish 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_twofish_binary_data() {
         let encryptor = TwofishEncryptor::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