1  use rand::RngCore;
  2  use sha3::{Sha3_512, Digest};
  3  use zeroize::Zeroize;
  4  use chacha20poly1305::{
  5      aead::{Aead, KeyInit, OsRng},
  6      ChaCha20Poly1305, Key, Nonce,
  7  };
  8  use argon2::{Argon2, password_hash::SaltString, PasswordHasher};
  9  
 10  // Derive a salt from the password itself
 11  pub fn derive_salt_from_password(password: &str) -> [u8; 16] {
 12      let mut hasher = Sha3_512::new();
 13      hasher.update(password.as_bytes());
 14      let hash_result = hasher.finalize();
 15  
 16      let mut salt = [0u8; 16];
 17      salt.copy_from_slice(&hash_result[..16]); // Use the first 16 bytes of the hash as the salt
 18      salt
 19  }
 20  
 21  // Derive encryption key using Argon2
 22  pub fn derive_key(password: &str, salt: &[u8]) -> [u8; 32] {
 23      let argon2 = Argon2::default();
 24      let salt = SaltString::encode_b64(salt).expect("Failed to generate salt string");
 25      let hash = argon2
 26          .hash_password(password.as_bytes(), &salt)
 27          .expect("Failed to hash password");
 28      let hash_bytes = hash.hash.expect("Hash missing in PasswordHash structure");
 29  
 30      let mut key = [0u8; 32];
 31      key.copy_from_slice(hash_bytes.as_bytes());
 32      key
 33  }
 34  
 35  pub fn combine_shared_secrets(
 36      kyber_secret: &str,
 37      ecdh_secret: &str,
 38  ) -> Result<String, Box<dyn std::error::Error>> {
 39      use sha3::{Digest, Sha3_512};
 40      use hex; // For hexadecimal encoding
 41  
 42      // Concatenate the secrets
 43      let combined = [kyber_secret.as_bytes(), ecdh_secret.as_bytes()].concat();
 44  
 45      // Hash the combined secrets to produce a fixed-length shared secret
 46      let mut hasher = Sha3_512::new();
 47      hasher.update(combined);
 48  
 49      // Convert the hash result to a hexadecimal string
 50      Ok(hex::encode(hasher.finalize()))
 51  }
 52  
 53  // Encrypt the data using ChaCha20Poly1305
 54  pub fn encrypt_data(plain_text: &str, password: &str) -> Result<String, Box<dyn std::error::Error>> {
 55      // Generate random salt for key derivation
 56      let mut salt = [0u8; 16];
 57      OsRng.fill_bytes(&mut salt);
 58  
 59      // Derive encryption key using Argon2
 60      let mut key = derive_key(password, &salt);
 61      let cipher = ChaCha20Poly1305::new(&Key::from_slice(&key));
 62  
 63      // Generate random nonce (12 bytes)
 64      let mut nonce_bytes = [0u8; 12];
 65      OsRng.fill_bytes(&mut nonce_bytes);
 66      let nonce = Nonce::from_slice(&nonce_bytes);
 67  
 68      // Encrypt the data
 69      let encrypted_data = cipher
 70          .encrypt(nonce, plain_text.as_bytes())
 71          .map_err(|_| "Encryption error")?;
 72  
 73      // Clear the key from memory after usage
 74      key.zeroize();
 75  
 76      // Return the formatted encrypted message with salt, nonce, and encrypted data
 77      Ok(format!(
 78          "{}:{}:{}",
 79          hex::encode(salt),
 80          hex::encode(nonce_bytes),
 81          hex::encode(encrypted_data)
 82      ))
 83  }
 84  
 85  // Decrypt the data using ChaCha20Poly1305
 86  pub fn decrypt_data(encrypted_text: &str, password: &str) -> Result<String, Box<dyn std::error::Error>> {
 87      // Split the encrypted data into salt, nonce, and encrypted part
 88      let parts: Vec<&str> = encrypted_text.split(':').collect();
 89      if parts.len() != 3 {
 90          return Err("Invalid encrypted data format".into());
 91      }
 92  
 93      // Decode hex-encoded salt, nonce, and encrypted data
 94      let salt = hex::decode(parts[0]).map_err(|_| "Decryption error: Invalid salt format")?;
 95      let nonce_bytes = hex::decode(parts[1]).map_err(|_| "Decryption error: Invalid nonce format")?;
 96      let encrypted_data = hex::decode(parts[2]).map_err(|_| "Decryption error: Invalid encrypted data format")?;
 97  
 98      // Derive the decryption key using the password and salt
 99      let mut key = derive_key(password, &salt);
100      let cipher = ChaCha20Poly1305::new(&Key::from_slice(&key));
101  
102      // Ensure nonce is of the correct length (12 bytes for ChaCha20Poly1305)
103      let nonce = Nonce::from_slice(&nonce_bytes);
104  
105      // Decrypt the data
106      let decrypted_data = cipher
107          .decrypt(nonce, encrypted_data.as_ref())
108          .map_err(|_| "Decryption error: Failed to decrypt")?;
109  
110      // Clear the key from memory after usage
111      key.zeroize();
112  
113      // Convert decrypted bytes into a string
114      Ok(String::from_utf8(decrypted_data).map_err(|_| "Decryption error: Invalid UTF-8 data")?)
115  }