use super::Encryptor;
 use chacha20::cipher::{KeyIvInit, StreamCipher};
 use chacha20::ChaCha20;
 use rand::Rng;
 use sha2::{Digest, Sha256};
 
 /// ChaCha20 stream cipher implementation
 pub struct ChaCha20Encryptor {
     key: [u8; 32],
 }
 
 impl ChaCha20Encryptor {
     /// Creates a new ChaCha20 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();
 
         // Use full 32 bytes of SHA-256 hash for ChaCha20 key
         let mut key_bytes = [0u8; 32];
         key_bytes.copy_from_slice(&hash[..32]);
 
         println!(
             "[ChaCha20] Derived key from input (SHA-256): {:02x?}",
             &key_bytes[..8]
         );
 
         Self { key: key_bytes }
     }
 }
 
 impl Encryptor for ChaCha20Encryptor {
     fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>, String> {
         // Generate random nonce (12 bytes for ChaCha20)
         let mut rng = rand::thread_rng();
         let mut nonce = [0u8; 12];
         rng.fill(&mut nonce);
 
         println!(
             "[ChaCha20] Generated random nonce for encryption: {:02x?}",
             &nonce
         );
 
         let mut cipher = ChaCha20::new(&self.key.into(), &nonce.into());
         let mut buffer = data.to_vec();
         cipher.apply_keystream(&mut buffer);
 
         // Prepend nonce to ciphertext
         let mut result = nonce.to_vec();
         result.extend_from_slice(&buffer);
 
         Ok(result)
     }
 
     fn decrypt(&self, data: &[u8]) -> Result<Vec<u8>, String> {
         if data.len() < 12 {
             return Err("Ciphertext too short to contain nonce".to_string());
         }
 
         // Extract nonce from beginning of data
         let nonce: [u8; 12] = data[..12].try_into().unwrap();
         let ciphertext = &data[12..];
 
         println!("[ChaCha20] Extracted nonce for decryption: {:02x?}", &nonce);
 
         // ChaCha20 is symmetric - same operation for decrypt
         let mut cipher = ChaCha20::new(&self.key.into(), &nonce.into());
         let mut buffer = ciphertext.to_vec();
         cipher.apply_keystream(&mut buffer);
 
         Ok(buffer)
     }
 
     fn key_size(&self) -> usize {
         32 // ChaCha20 uses 32-byte keys
     }
 
     fn algorithm_name(&self) -> &'static str {
         "ChaCha20"
     }
 
     fn uses_iv(&self) -> bool {
         true // Uses nonce
     }
 
     fn generate_decrypt_code(&self) -> &'static str {
         r#"/// Decrypts ChaCha20 encrypted data using the provided key
 ///
 /// Returns `Ok(Vec<u8>)` with decrypted data if successful, otherwise returns an error.
 pub fn decrypt_chacha20(data: &[u8], key: &[u8]) -> Result<Vec<u8>, String> {
     use chacha20::cipher::{KeyIvInit, StreamCipher};
     use chacha20::ChaCha20;
     
     if data.len() < 12 {
         return Err("Ciphertext too short to contain nonce".to_string());
     }
     
     if key.len() != 32 {
         return Err("Key must be 32 bytes for ChaCha20".to_string());
     }
     
     let nonce: [u8; 12] = data[..12].try_into().unwrap();
     let ciphertext = &data[12..];
     
     let key_array: [u8; 32] = key.try_into().map_err(|_| "Invalid key length".to_string())?;
     
     let mut cipher = ChaCha20::new(&key_array.into(), &nonce.into());
     let mut buffer = ciphertext.to_vec();
     cipher.apply_keystream(&mut buffer);
     
     Ok(buffer)
 }"#
     }
 
     fn required_dependencies(&self) -> Vec<(&'static str, &'static str)> {
         vec![("chacha20", "0.9")]
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
     fn test_chacha20_encrypt_decrypt() {
         let encryptor = ChaCha20Encryptor::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_chacha20_different_keys() {
         let encryptor1 = ChaCha20Encryptor::new("key1");
         let encryptor2 = ChaCha20Encryptor::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_chacha20_empty_data() {
         let encryptor = ChaCha20Encryptor::new("key");
         let empty: &[u8] = &[];
 
         let encrypted = encryptor.encrypt(empty).unwrap();
         // Encrypted data includes 12-byte nonce
         assert_eq!(encrypted.len(), 12);
 
         let decrypted = encryptor.decrypt(&encrypted).unwrap();
         assert_eq!(decrypted.len(), 0);
     }
 
     #[test]
     fn test_chacha20_long_data() {
         let encryptor = ChaCha20Encryptor::new("secret-key");
         let plaintext = b"This is a much longer message that tests the stream cipher capabilities";
 
         let encrypted = encryptor.encrypt(plaintext).unwrap();
         assert_ne!(encrypted, plaintext);
         // Encrypted data includes 12-byte nonce + ciphertext
         assert_eq!(encrypted.len(), plaintext.len() + 12);
 
         let decrypted = encryptor.decrypt(&encrypted).unwrap();
         assert_eq!(decrypted, plaintext);
     }
 
     #[test]
     fn test_chacha20_binary_data() {
         let encryptor = ChaCha20Encryptor::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