use sha3::{Sha3_256, Digest};
use crate::{Result, QsshError};
pub struct QuantumCipher {
key: Vec<u8>,
nonce_counter: u64,
}
impl QuantumCipher {
pub fn from_shared_secret(shared_secret: &[u8]) -> Result<Self> {
if shared_secret.len() < 32 {
return Err(QsshError::Crypto("Shared secret too short".into()));
}
let key = shared_secret[..32].to_vec();
Ok(Self {
key,
nonce_counter: 0,
})
}
fn generate_keystream(&self, nonce: &[u8], counter: u64, length: usize) -> Vec<u8> {
let mut keystream = Vec::with_capacity(length);
let mut block_counter = counter;
while keystream.len() < length {
let mut hasher = Sha3_256::new();
hasher.update(&self.key);
hasher.update(nonce);
hasher.update(block_counter.to_le_bytes());
let block = hasher.finalize();
keystream.extend_from_slice(&block);
block_counter += 1;
}
keystream.truncate(length);
keystream
}
pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<(Vec<u8>, Vec<u8>)> {
let mut hasher = Sha3_256::new();
hasher.update(&self.key);
hasher.update(b"nonce");
hasher.update(self.nonce_counter.to_le_bytes());
let nonce_hash = hasher.finalize();
let nonce = nonce_hash[..16].to_vec();
self.nonce_counter += 1;
let keystream = self.generate_keystream(&nonce, 0, plaintext.len());
let ciphertext: Vec<u8> = plaintext.iter()
.zip(keystream.iter())
.map(|(p, k)| p ^ k)
.collect();
Ok((ciphertext, nonce))
}
pub fn decrypt(&self, ciphertext: &[u8], nonce: &[u8]) -> Result<Vec<u8>> {
let keystream = self.generate_keystream(nonce, 0, ciphertext.len());
let plaintext: Vec<u8> = ciphertext.iter()
.zip(keystream.iter())
.map(|(c, k)| c ^ k)
.collect();
Ok(plaintext)
}
pub fn encrypt_authenticated(&mut self, plaintext: &[u8]) -> Result<(Vec<u8>, Vec<u8>, Vec<u8>)> {
let (ciphertext, nonce) = self.encrypt(plaintext)?;
let mut hasher = Sha3_256::new();
hasher.update(&self.key);
hasher.update(&nonce);
hasher.update(&ciphertext);
let tag = hasher.finalize().to_vec();
Ok((ciphertext, nonce, tag))
}
pub fn decrypt_authenticated(&self, ciphertext: &[u8], nonce: &[u8], tag: &[u8]) -> Result<Vec<u8>> {
let mut hasher = Sha3_256::new();
hasher.update(&self.key);
hasher.update(nonce);
hasher.update(ciphertext);
let computed_tag = hasher.finalize();
if tag.len() != computed_tag.len() {
return Err(QsshError::Crypto("Authentication failed".into()));
}
let mut diff = 0u8;
for (a, b) in tag.iter().zip(computed_tag.iter()) {
diff |= a ^ b;
}
if diff != 0 {
return Err(QsshError::Crypto("Authentication failed".into()));
}
self.decrypt(ciphertext, nonce)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_quantum_cipher() {
let shared_secret = vec![0x42u8; 32];
let mut cipher = QuantumCipher::from_shared_secret(&shared_secret).expect("Failed to create cipher");
let plaintext = b"Hello, quantum world!";
let (ciphertext, nonce) = cipher.encrypt(plaintext).expect("Failed to encrypt");
assert_ne!(ciphertext, plaintext);
let decrypted = cipher.decrypt(&ciphertext, &nonce).expect("Failed to decrypt");
assert_eq!(decrypted, plaintext);
}
#[test]
fn test_authenticated_encryption() {
let shared_secret = vec![0x42u8; 32];
let mut cipher = QuantumCipher::from_shared_secret(&shared_secret).expect("Failed to create cipher");
let plaintext = b"Secret quantum message";
let (ciphertext, nonce, tag) = cipher.encrypt_authenticated(plaintext).expect("Failed to encrypt with auth");
let decrypted = cipher.decrypt_authenticated(&ciphertext, &nonce, &tag).expect("Failed to decrypt with auth");
assert_eq!(decrypted, plaintext);
let mut bad_tag = tag.clone();
bad_tag[0] ^= 1;
assert!(cipher.decrypt_authenticated(&ciphertext, &nonce, &bad_tag).is_err());
}
}