qssh 0.0.2-alpha

//! Post-quantum cryptography module using Falcon + SPHINCS+

use crate::{QsshError, Result};
use pqcrypto_sphincsplus::sphincssha256128ssimple as sphincs;
use pqcrypto_falcon::falcon512;
use pqcrypto_traits::sign::{PublicKey as SignPublicKey, SignedMessage as SignedMessageTrait};
use aes_gcm::{
    aead::{Aead, AeadCore, KeyInit, OsRng},
    Aes256Gcm, Nonce,
};
use rand::RngCore;

pub mod kdf;
pub mod quantum_cipher;
pub mod cipher_choice;
pub mod qrng_integration;
#[cfg(test)]
pub mod test_helpers;
pub use kdf::{SessionKeyDerivation, SessionKeys};
pub use quantum_cipher::QuantumCipher;
pub use cipher_choice::{CipherAlgorithm, UniversalCipher};
pub use qrng_integration::{QuantumRng, QrngConfig};

/// Post-quantum signature using SPHINCS+
pub struct PqSignature {
    pub sphincs_sk: sphincs::SecretKey,
    pub sphincs_pk: sphincs::PublicKey,
}

impl PqSignature {
    /// Generate new SPHINCS+ keypair
    pub fn new() -> Result<Self> {
        let (sphincs_pk, sphincs_sk) = sphincs::keypair();
        Ok(Self { sphincs_sk, sphincs_pk })
    }
    
    /// Get secret key bytes
    pub fn secret_bytes(&self) -> Vec<u8> {
        use pqcrypto_traits::sign::SecretKey;
        self.sphincs_sk.as_bytes().to_vec()
    }
    
    /// Get public key bytes
    pub fn public_bytes(&self) -> Vec<u8> {
        use pqcrypto_traits::sign::PublicKey;
        self.sphincs_pk.as_bytes().to_vec()
    }
}

/// Post-quantum key exchange using Falcon + SPHINCS+
/// Following QuantumHarmony's approach
pub struct PqKeyExchange {
    /// SPHINCS+ for long-term identity
    pub sphincs_sk: sphincs::SecretKey,
    pub sphincs_pk: sphincs::PublicKey,
    
    /// Falcon-512 for ephemeral operations (fast)
    pub falcon_sk: falcon512::SecretKey,
    pub falcon_pk: falcon512::PublicKey,
}

impl PqKeyExchange {
    /// Generate new post-quantum keys
    pub fn new() -> Result<Self> {
        // Check for QRNG availability
        let qrng_available = std::env::var("QSSH_QRNG_ENDPOINT").is_ok() ||
                            std::env::var("QSSH_KIRQ_ENDPOINT").is_ok();

        if qrng_available {
            log::info!("QRNG endpoint detected - keys will use quantum entropy");
        }

        // Generate keys (pqcrypto will use OS entropy, but we log QRNG availability)
        // In future, we could modify pqcrypto to accept custom RNG
        let (sphincs_pk, sphincs_sk) = sphincs::keypair();
        let (falcon_pk, falcon_sk) = falcon512::keypair();

        Ok(Self {
            sphincs_sk,
            sphincs_pk,
            falcon_sk,
            falcon_pk,
        })
    }
    
    /// Create ephemeral key share (replaces Kyber encapsulation)
    pub fn create_key_share(&self) -> Result<(Vec<u8>, Vec<u8>)> {
        // Generate ephemeral secret with best available entropy
        let mut ephemeral_secret = vec![0u8; 32];

        // Check if QRNG is configured
        if std::env::var("QSSH_QRNG_ENDPOINT").is_ok() {
            log::debug!("Using QRNG-enhanced entropy for ephemeral key share");
        }

        // Use best available RNG (OS entropy, potentially enhanced by hardware)
        rand::thread_rng().fill_bytes(&mut ephemeral_secret);
        
        // Sign with Falcon for authenticity using detached signature
        let signature = falcon512::detached_sign(&ephemeral_secret, &self.falcon_sk);

        use pqcrypto_traits::sign::DetachedSignature as DetachedSignatureTrait;
        Ok((ephemeral_secret, signature.as_bytes().to_vec()))
    }
    
    /// Verify and process peer's key share (replaces Kyber decapsulation)
    pub fn process_key_share(
        &self,
        peer_falcon_pk: &[u8],
        peer_share: &[u8],
        peer_signature: &[u8],
    ) -> Result<Vec<u8>> {
        // Safety check: Validate input sizes
        use pqcrypto_traits::sign::PublicKey as SignPublicKeyTrait;
        
        // Check key share size (should be 32 bytes for our protocol)
        if peer_share.len() != 32 {
            return Err(QsshError::Crypto(format!(
                "Invalid key share size: got {}, expected 32",
                peer_share.len()
            )));
        }
        
        // Verify Falcon signature using detached signature verification
        use pqcrypto_traits::sign::{PublicKey as PubKey, DetachedSignature};

        let pk = match falcon512::PublicKey::from_bytes(peer_falcon_pk) {
            Ok(key) => key,
            Err(e) => {
                log::error!("Failed to parse Falcon public key of size {}: {:?}",
                    peer_falcon_pk.len(), e);
                return Err(QsshError::Crypto(format!("Invalid Falcon public key: {:?}", e)));
            }
        };

        // Parse as detached signature (signature only, not signed message)
        let sig = match falcon512::DetachedSignature::from_bytes(peer_signature) {
            Ok(signature) => signature,
            Err(e) => {
                log::error!("Failed to parse Falcon detached signature of size {}: {:?}",
                    peer_signature.len(), e);
                return Err(QsshError::Crypto(format!("Invalid Falcon signature: {:?}", e)));
            }
        };

        // Verify the detached signature against the share data
        match falcon512::verify_detached_signature(&sig, peer_share, &pk) {
            Ok(_) => {
                log::debug!("Falcon signature verified successfully");
            }
            Err(_) => {
                log::error!("Falcon signature verification failed");
                return Err(QsshError::Crypto("Invalid Falcon signature".into()));
            }
        }
        
        // Return the verified share
        Ok(peer_share.to_vec())
    }
    
    /// Compute shared secret from both key shares
    pub fn compute_shared_secret(
        &self,
        our_share: &[u8],
        peer_share: &[u8],
        client_random: &[u8],
        server_random: &[u8],
    ) -> Vec<u8> {
        use sha3::{Sha3_256, Digest};
        
        let mut hasher = Sha3_256::new();
        hasher.update(b"QSSH-FALCON-v2");
        
        // Sort shares lexicographically to ensure both parties get same result
        if our_share < peer_share {
            hasher.update(our_share);
            hasher.update(peer_share);
        } else {
            hasher.update(peer_share);
            hasher.update(our_share);
        }
        
        hasher.update(client_random);
        hasher.update(server_random);
        hasher.finalize().to_vec()
    }
    
    /// Sign a message with SPHINCS+ (for identity/long-term)
    pub fn sign(&self, message: &[u8]) -> Result<Vec<u8>> {
        let signature = sphincs::sign(message, &self.sphincs_sk);
        use SignedMessageTrait;
        Ok(signature.as_bytes().to_vec())
    }
    
    /// Sign with Falcon (for ephemeral/performance)
    pub fn sign_falcon(&self, message: &[u8]) -> Result<Vec<u8>> {
        let signature = falcon512::sign(message, &self.falcon_sk);
        use SignedMessageTrait;
        Ok(signature.as_bytes().to_vec())
    }
    
    /// Get Falcon secret key bytes
    pub fn secret_bytes(&self) -> Vec<u8> {
        use pqcrypto_traits::sign::SecretKey;
        self.falcon_sk.as_bytes().to_vec()
    }
    
    /// Get Falcon public key bytes
    pub fn public_bytes(&self) -> Vec<u8> {
        use pqcrypto_traits::sign::PublicKey;
        self.falcon_pk.as_bytes().to_vec()
    }
    
    /// Verify a SPHINCS+ signature
    pub fn verify(&self, message: &[u8], signature: &[u8], public_key: &[u8]) -> Result<bool> {
        let sig = sphincs::SignedMessage::from_bytes(signature)
            .map_err(|e| QsshError::Crypto(format!("Invalid signature: {:?}", e)))?;
        
        let pk = sphincs::PublicKey::from_bytes(public_key)
            .map_err(|e| QsshError::Crypto(format!("Invalid public key: {:?}", e)))?;
        
        match sphincs::open(&sig, &pk) {
            Ok(msg) => Ok(msg == message),
            Err(_) => Ok(false)
        }
    }
    
    /// Verify a Falcon signature
    pub fn verify_falcon(&self, message: &[u8], signature: &[u8], public_key: &[u8]) -> Result<bool> {
        use pqcrypto_traits::sign::{PublicKey as PubKey, DetachedSignature};

        let sig = falcon512::DetachedSignature::from_bytes(signature)
            .map_err(|e| QsshError::Crypto(format!("Invalid Falcon signature: {:?}", e)))?;

        let pk = falcon512::PublicKey::from_bytes(public_key)
            .map_err(|e| QsshError::Crypto(format!("Invalid Falcon public key: {:?}", e)))?;

        match falcon512::verify_detached_signature(&sig, message, &pk) {
            Ok(_) => Ok(true),
            Err(_) => Ok(false)
        }
    }
}

/// Symmetric encryption using AES-256-GCM (NIST approved)
pub struct SymmetricCrypto {
    cipher: Aes256Gcm,
}

impl SymmetricCrypto {
    /// Create from shared secret
    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 = aes_gcm::Key::<Aes256Gcm>::from_slice(&shared_secret[..32]);
        let cipher = Aes256Gcm::new(key);
        
        Ok(Self { cipher })
    }
    
    /// Encrypt data
    pub fn encrypt(&self, plaintext: &[u8]) -> Result<(Vec<u8>, Vec<u8>)> {
        let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
        let ciphertext = self.cipher
            .encrypt(&nonce, plaintext)
            .map_err(|e| QsshError::Crypto(format!("Encryption failed: {}", e)))?;
        
        Ok((ciphertext, nonce.to_vec()))
    }
    
    /// Decrypt data
    pub fn decrypt(&self, ciphertext: &[u8], nonce: &[u8]) -> Result<Vec<u8>> {
        let nonce = Nonce::from_slice(nonce);
        let plaintext = self.cipher
            .decrypt(nonce, ciphertext)
            .map_err(|e| QsshError::Crypto(format!("Decryption failed: {}", e)))?;
        
        Ok(plaintext)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    // TODO: Fix segfault in Falcon tests on macOS
    // #[test]
    #[allow(dead_code)]
    fn test_falcon_key_agreement() {
        let alice = PqKeyExchange::new().unwrap();
        let bob = PqKeyExchange::new().unwrap();
        
        // Alice creates key share
        let (alice_share, alice_sig) = alice.create_key_share().unwrap();
        
        // Bob verifies and creates his share
        let verified_alice_share = bob.process_key_share(
            alice.falcon_pk.as_bytes(),
            &alice_share,
            &alice_sig
        ).unwrap();
        assert_eq!(alice_share, verified_alice_share);
        
        let (bob_share, bob_sig) = bob.create_key_share().unwrap();
        
        // Alice verifies Bob's share
        let verified_bob_share = alice.process_key_share(
            bob.falcon_pk.as_bytes(),
            &bob_share,
            &bob_sig
        ).unwrap();
        assert_eq!(bob_share, verified_bob_share);
        
        // Both compute shared secret
        let client_random = vec![0x11; 32];
        let server_random = vec![0x22; 32];
        
        let alice_secret = alice.compute_shared_secret(
            &alice_share,
            &bob_share,
            &client_random,
            &server_random
        );
        
        let bob_secret = bob.compute_shared_secret(
            &bob_share,
            &alice_share,
            &client_random,
            &server_random
        );
        
        // Secrets should match
        assert_eq!(alice_secret, bob_secret);
    }
    
    // #[test]
    #[allow(dead_code)]
    fn test_falcon_signatures() {
        let signer = PqKeyExchange::new().unwrap();
        let message = b"Test message for Falcon";
        
        let signature = signer.sign_falcon(message).unwrap();
        let valid = signer.verify_falcon(message, &signature, signer.falcon_pk.as_bytes()).unwrap();
        
        assert!(valid);
        
        // Wrong message should fail
        let wrong_message = b"Wrong message";
        let invalid = signer.verify_falcon(wrong_message, &signature, signer.falcon_pk.as_bytes()).unwrap();
        
        assert!(!invalid);
    }
    
    #[test]
    fn test_symmetric_crypto() {
        let secret = vec![0x42; 32];
        let crypto = SymmetricCrypto::from_shared_secret(&secret).expect("Failed to create crypto");

        let plaintext = b"Hello, quantum world!";
        let (ciphertext, nonce) = crypto.encrypt(plaintext).expect("Failed to encrypt");

        let decrypted = crypto.decrypt(&ciphertext, &nonce).expect("Failed to decrypt");

        assert_eq!(plaintext, &decrypted[..]);
    }
}