1use crate::{Error, Result, PublicKeys, PrivateKeys, KeyPair, HybridCiphertext, HybridSignature};
6use crate::constants::*;
7use crate::security::{SecurityManager, EntropyMonitor, TimingProtection, SecureMemory, constant_time_compare};
8use base64::{engine::general_purpose, Engine};
9use rsa::{RsaPrivateKey, RsaPublicKey, Oaep};
10use rsa::pkcs8::{EncodePrivateKey, EncodePublicKey, DecodePrivateKey, DecodePublicKey};
11use sha2::Sha256;
12use sha3::{Sha3_256, Digest};
13use aes_gcm::{
14 aead::{Aead, KeyInit},
15 Aes256Gcm, Nonce, Key
16};
17use pqcrypto_kyber::kyber1024;
18use pqcrypto_dilithium::dilithium5;
19use pqcrypto_traits::kem::{PublicKey as KemPublicKey, SecretKey as KemSecretKey, Ciphertext as _, SharedSecret as _};
20use pqcrypto_traits::sign::{PublicKey as SignPublicKey, SecretKey as SignSecretKey, DetachedSignature as _};
21use zeroize::{Zeroize, ZeroizeOnDrop};
22use std::sync::Arc;
23
24pub struct HybridCrypto {
26 keypair: KeyPair,
27 quantum_mode: bool,
28 security_manager: Arc<SecurityManager>,
29}
30
31impl HybridCrypto {
32 pub fn generate_keypair() -> Result<Self> {
42 let security_manager = Arc::new(SecurityManager::new());
43
44 if !security_manager.entropy_monitor().has_sufficient_entropy() {
46 security_manager.entropy_monitor().collect_entropy()?;
47 }
48
49 security_manager.timing_protection().add_jitter()?;
51
52 let mut rng = rand::thread_rng();
54 let rsa_private = RsaPrivateKey::new(&mut rng, 4096)
55 .map_err(|e| Error::KeyGenerationFailed(format!("RSA generation failed: {}", e)))?;
56 let rsa_public = RsaPublicKey::from(&rsa_private);
57
58 let rsa_private_pem = rsa_private.to_pkcs8_pem(rsa::pkcs8::LineEnding::LF)
59 .map_err(|e| Error::KeyGenerationFailed(format!("RSA PEM encoding failed: {}", e)))?;
60 let rsa_public_pem = rsa_public.to_public_key_pem(rsa::pkcs8::LineEnding::LF)
61 .map_err(|e| Error::KeyGenerationFailed(format!("RSA PEM encoding failed: {}", e)))?;
62
63 let (kyber_pub, kyber_priv) = kyber1024::keypair();
65
66 let (dilithium_pub, dilithium_priv) = dilithium5::keypair();
68
69 let keypair = KeyPair {
70 public: PublicKeys::new(
71 rsa_public_pem,
72 general_purpose::STANDARD.encode(kyber_pub.as_bytes()),
73 general_purpose::STANDARD.encode(dilithium_pub.as_bytes()),
74 ),
75 private: PrivateKeys::new(
76 rsa_private_pem.to_string(),
77 kyber_priv.as_bytes().to_vec(),
78 dilithium_priv.as_bytes().to_vec(),
79 ),
80 };
81
82 Ok(Self {
83 keypair,
84 quantum_mode: true,
85 security_manager,
86 })
87 }
88
89 pub fn public_keys(&self) -> &PublicKeys {
91 self.keypair.public_keys()
92 }
93
94 pub fn security_manager(&self) -> &Arc<SecurityManager> {
96 &self.security_manager
97 }
98
99 pub fn audit_security(&self) -> crate::security::SecurityAuditResult {
101 self.security_manager.audit_security()
102 }
103
104 pub fn encrypt(&self, data: &[u8], recipient_pubkeys: &PublicKeys) -> Result<HybridCiphertext> {
122 self.security_manager.timing_protection().add_jitter()?;
124
125 let symmetric_key_mem = self.security_manager.secure_key_generation()?;
127 let symmetric_key: [u8; 32] = symmetric_key_mem.as_slice()[..32].try_into()
128 .map_err(|_| Error::EncryptionFailed("Invalid symmetric key length".to_string()))?;
129
130 let key = Key::<Aes256Gcm>::from_slice(&symmetric_key);
132 let cipher = Aes256Gcm::new(key);
133 let nonce_bytes: [u8; 12] = rand::random();
134 let nonce = Nonce::from_slice(&nonce_bytes);
135
136 let encrypted = cipher.encrypt(nonce, data)
137 .map_err(|e| Error::EncryptionFailed(format!("AES encryption failed: {:?}", e)))?;
138
139 let mut ciphertext = nonce_bytes.to_vec();
141 ciphertext.extend_from_slice(&encrypted);
142
143 let rsa_pub = RsaPublicKey::from_public_key_pem(&recipient_pubkeys.rsa_pem)
145 .map_err(|e| Error::InvalidKey(format!("RSA public key parse failed: {}", e)))?;
146
147 let mut rng = rand::thread_rng();
148 let padding = Oaep::new::<Sha256>();
149 let encrypted_key_rsa = rsa_pub.encrypt(&mut rng, padding, &symmetric_key)
150 .map_err(|e| Error::EncryptionFailed(format!("RSA encryption failed: {}", e)))?;
151
152 let kyber_pub_bytes = general_purpose::STANDARD.decode(&recipient_pubkeys.kyber_base64)
154 .map_err(|e| Error::InvalidKey(format!("Kyber key decode failed: {}", e)))?;
155 let kyber_pub = kyber1024::PublicKey::from_bytes(&kyber_pub_bytes)
156 .map_err(|_| Error::InvalidKey("Invalid Kyber public key".to_string()))?;
157
158 let (kyber_ciphertext, kyber_shared_secret) = kyber1024::encapsulate(&kyber_pub);
159 let shared_secret_bytes = kyber_ciphertext.as_bytes();
160 let ciphertext_bytes = kyber_shared_secret.as_bytes();
161
162 let kyber_key = Key::<Aes256Gcm>::from_slice(&shared_secret_bytes[..32.min(shared_secret_bytes.len())]);
163 let kyber_cipher = Aes256Gcm::new(kyber_key);
164 let kyber_nonce_bytes: [u8; 12] = rand::random();
165 let kyber_nonce = Nonce::from_slice(&kyber_nonce_bytes);
166
167 let encrypted_key_kyber_inner = kyber_cipher.encrypt(kyber_nonce, &symmetric_key[..])
168 .map_err(|e| Error::EncryptionFailed(format!("Kyber key encryption failed: {:?}", e)))?;
169
170 let mut encrypted_key_kyber = ciphertext_bytes.to_vec();
171 encrypted_key_kyber.extend_from_slice(&kyber_nonce_bytes);
172 encrypted_key_kyber.extend_from_slice(&encrypted_key_kyber_inner);
173
174 Ok(HybridCiphertext::new(
175 general_purpose::STANDARD.encode(&ciphertext),
176 general_purpose::STANDARD.encode(&encrypted_key_rsa),
177 general_purpose::STANDARD.encode(&encrypted_key_kyber),
178 ))
179 }
180
181 pub fn decrypt(&self, ciphertext: &HybridCiphertext) -> Result<Vec<u8>> {
192 self.security_manager.timing_protection().add_jitter()?;
194
195 let encrypted_data = general_purpose::STANDARD.decode(&ciphertext.ciphertext)
196 .map_err(|e| Error::InvalidCiphertext(format!("Ciphertext decode failed: {}", e)))?;
197
198 let rsa_result = self.decrypt_rsa_key(&ciphertext.encrypted_key_rsa);
200 let kyber_result = self.decrypt_kyber_key(&ciphertext.encrypted_key_kyber);
201
202 let symmetric_key = match rsa_result {
203 Ok(key) => key,
204 Err(_) => kyber_result?,
205 };
206
207 if encrypted_data.len() < 12 {
209 return Err(Error::InvalidCiphertext("Ciphertext too short".to_string()));
210 }
211
212 let nonce = Nonce::from_slice(&encrypted_data[..12]);
213 let encrypted = &encrypted_data[12..];
214
215 let key = Key::<Aes256Gcm>::from_slice(&symmetric_key);
216 let cipher = Aes256Gcm::new(key);
217
218 let plaintext = cipher.decrypt(nonce, encrypted)
219 .map_err(|e| Error::DecryptionFailed(format!("AES decryption failed: {:?}", e)))?;
220
221 Ok(plaintext)
222 }
223
224 fn decrypt_rsa_key(&self, encrypted_key_b64: &str) -> Result<Vec<u8>> {
225 let encrypted_key = general_purpose::STANDARD.decode(encrypted_key_b64)?;
226
227 let rsa_priv = RsaPrivateKey::from_pkcs8_pem(&self.keypair.private.rsa_pem)
228 .map_err(|e| Error::InvalidKey(format!("RSA private key parse failed: {}", e)))?;
229
230 let padding = Oaep::new::<Sha256>();
231 let decrypted = rsa_priv.decrypt(padding, &encrypted_key)
232 .map_err(|e| Error::DecryptionFailed(format!("RSA decryption failed: {}", e)))?;
233
234 Ok(decrypted)
235 }
236
237 fn decrypt_kyber_key(&self, encrypted_key_b64: &str) -> Result<Vec<u8>> {
238 let encrypted_data = general_purpose::STANDARD.decode(encrypted_key_b64)?;
239
240 if encrypted_data.len() < KYBER1024_CIPHERTEXT_BYTES + NONCE_BYTES {
243 return Err(Error::InvalidCiphertext("Invalid Kyber encrypted key".to_string()));
244 }
245
246 let kyber_priv = kyber1024::SecretKey::from_bytes(&self.keypair.private.kyber_bytes)
247 .map_err(|_| Error::InvalidKey("Invalid Kyber private key".to_string()))?;
248
249 let kyber_ciphertext = kyber1024::Ciphertext::from_bytes(&encrypted_data[..KYBER1024_CIPHERTEXT_BYTES])
250 .map_err(|_| Error::InvalidCiphertext("Invalid Kyber ciphertext".to_string()))?;
251 let nonce_bytes = &encrypted_data[KYBER1024_CIPHERTEXT_BYTES..KYBER1024_CIPHERTEXT_BYTES + NONCE_BYTES];
252 let encrypted_key = &encrypted_data[KYBER1024_CIPHERTEXT_BYTES + NONCE_BYTES..];
253
254 let kyber_shared_secret = kyber1024::decapsulate(&kyber_ciphertext, &kyber_priv);
255 let shared_secret_bytes = kyber_shared_secret.as_bytes();
256
257 let key = Key::<Aes256Gcm>::from_slice(&shared_secret_bytes[..32.min(shared_secret_bytes.len())]);
258 let cipher = Aes256Gcm::new(key);
259 let nonce = Nonce::from_slice(nonce_bytes);
260
261 let symmetric_key = cipher.decrypt(nonce, encrypted_key)
262 .map_err(|e| Error::DecryptionFailed(format!("Kyber key decryption failed: {:?}", e)))?;
263
264 Ok(symmetric_key)
265 }
266
267 pub fn sign(&self, message: &[u8]) -> Result<HybridSignature> {
277 self.security_manager.timing_protection().add_jitter()?;
279
280 let rsa_sig = self.sign_rsa(message)?;
281 let dilithium_sig = if self.quantum_mode {
282 Some(self.sign_dilithium(message)?)
283 } else {
284 None
285 };
286
287 Ok(HybridSignature::new(rsa_sig, dilithium_sig))
288 }
289
290 fn sign_rsa(&self, message: &[u8]) -> Result<String> {
291 let mut hasher = Sha3_256::new();
292 hasher.update(message);
293 let message_hash = hasher.finalize();
294
295 let rsa_priv = RsaPrivateKey::from_pkcs8_pem(&self.keypair.private.rsa_pem)
296 .map_err(|e| Error::InvalidKey(format!("RSA private key parse failed: {}", e)))?;
297
298 let padding = rsa::Pss::new::<Sha256>();
299 let mut rng = rand::thread_rng();
300
301 let signature = rsa_priv.sign_with_rng(&mut rng, padding, &message_hash)
302 .map_err(|e| Error::SigningFailed(format!("RSA signing failed: {}", e)))?;
303
304 Ok(general_purpose::STANDARD.encode(signature))
305 }
306
307 fn sign_dilithium(&self, message: &[u8]) -> Result<String> {
308 let dilithium_priv = dilithium5::SecretKey::from_bytes(&self.keypair.private.dilithium_bytes)
309 .map_err(|_| Error::InvalidKey("Invalid Dilithium private key".to_string()))?;
310
311 let mut hasher = Sha3_256::new();
312 hasher.update(message);
313 let message_hash = hasher.finalize();
314
315 let signature = dilithium5::detached_sign(&message_hash, &dilithium_priv);
316
317 Ok(general_purpose::STANDARD.encode(signature.as_bytes()))
318 }
319
320 pub fn verify(message: &[u8], signature: &HybridSignature, pubkeys: &PublicKeys) -> Result<bool> {
333 let rsa_valid = Self::verify_rsa(message, &signature.rsa_signature, &pubkeys.rsa_pem)?;
335
336 if !rsa_valid {
338 return Ok(false);
339 }
340
341 if let Some(dilithium_sig) = &signature.dilithium_signature {
343 return Self::verify_dilithium(message, dilithium_sig, &pubkeys.dilithium_base64);
344 }
345
346 Ok(true)
347 }
348
349 fn verify_rsa(message: &[u8], signature_b64: &str, pubkey_pem: &str) -> Result<bool> {
350 let rsa_pub = RsaPublicKey::from_public_key_pem(pubkey_pem)
351 .map_err(|e| Error::InvalidKey(format!("RSA public key parse failed: {}", e)))?;
352
353 let signature_bytes = general_purpose::STANDARD.decode(signature_b64)?;
354
355 let mut hasher = Sha3_256::new();
356 hasher.update(message);
357 let message_hash = hasher.finalize();
358
359 let padding = rsa::Pss::new::<Sha256>();
360
361 let verification_result = rsa_pub.verify(padding, &message_hash, &signature_bytes);
363 Ok(verification_result.is_ok())
364 }
365
366 fn verify_dilithium(message: &[u8], signature_b64: &str, pubkey_b64: &str) -> Result<bool> {
367 let public_key_bytes = general_purpose::STANDARD.decode(pubkey_b64)?;
368 let signature_bytes = general_purpose::STANDARD.decode(signature_b64)?;
369
370 let public_key = dilithium5::PublicKey::from_bytes(&public_key_bytes)
371 .map_err(|_| Error::InvalidKey("Invalid Dilithium public key".to_string()))?;
372 let signature = dilithium5::DetachedSignature::from_bytes(&signature_bytes)
373 .map_err(|_| Error::InvalidSignature("Invalid Dilithium signature".to_string()))?;
374
375 let mut hasher = Sha3_256::new();
376 hasher.update(message);
377 let message_hash = hasher.finalize();
378
379 let verification_result = dilithium5::verify_detached_signature(&signature, &message_hash, &public_key);
381 Ok(verification_result.is_ok())
382 }
383}
384
385#[cfg(test)]
386mod tests {
387 use super::*;
388
389 #[test]
390 fn test_key_generation() {
391 let crypto = HybridCrypto::generate_keypair().unwrap();
392 assert!(!crypto.public_keys().rsa_pem.is_empty());
393 assert!(!crypto.public_keys().kyber_base64.is_empty());
394 assert!(!crypto.public_keys().dilithium_base64.is_empty());
395 }
396
397 #[test]
398 fn test_encryption_decryption() {
399 let alice = HybridCrypto::generate_keypair().unwrap();
400 let bob = HybridCrypto::generate_keypair().unwrap();
401
402 let message = b"Test message for hybrid encryption";
403 let encrypted = alice.encrypt(message, &bob.public_keys()).unwrap();
404 let decrypted = bob.decrypt(&encrypted).unwrap();
405
406 assert_eq!(message, &decrypted[..]);
407 }
408
409 #[test]
410 fn test_signing_verification() {
411 let alice = HybridCrypto::generate_keypair().unwrap();
412
413 let message = b"Message to sign";
414 let signature = alice.sign(message).unwrap();
415 let valid = HybridCrypto::verify(message, &signature, &alice.public_keys()).unwrap();
416
417 assert!(valid);
418 }
419
420 #[test]
421 fn test_tampered_message() {
422 let alice = HybridCrypto::generate_keypair().unwrap();
423
424 let message = b"Original message";
425 let signature = alice.sign(message).unwrap();
426
427 let tampered = b"Tampered message";
428 let valid = HybridCrypto::verify(tampered, &signature, &alice.public_keys()).unwrap();
429
430 assert!(!valid);
431 }
432
433 #[test]
434 fn test_failover() {
435 let alice = HybridCrypto::generate_keypair().unwrap();
436 let bob = HybridCrypto::generate_keypair().unwrap();
437
438 let message = b"Test failover";
439 let mut encrypted = alice.encrypt(message, &bob.public_keys()).unwrap();
440
441 encrypted.encrypted_key_rsa = "CORRUPTED".to_string();
443
444 let decrypted = bob.decrypt(&encrypted).unwrap();
446 assert_eq!(message, &decrypted[..]);
447 }
448}
449