qssl 0.1.0-alpha

//! SPHINCS+ based Key Encapsulation Mechanism
//!
//! This implements a PROPER KEM using SPHINCS+ signatures, not the broken
//! "sign random bytes" approach that QSSH uses.

use pqcrypto_sphincsplus::sphincsharaka128fsimple as sphincs;
use pqcrypto_falcon::falcon512;
use pqcrypto_traits::sign::{PublicKey as _, SecretKey as _, DetachedSignature as _};
use sha3::{Sha3_512, Digest};
use zeroize::Zeroize;
use rand::RngCore;

use crate::{QsslError, QsslResult};

/// SPHINCS+ KEM - Properly implemented
///
/// Instead of Kyber (which has vulnerabilities), we use SPHINCS+ with
/// proper ephemeral key generation and authenticated key exchange.
pub struct SphincsKem {
    /// Long-term SPHINCS+ identity key
    pub identity_sk: sphincs::SecretKey,
    pub identity_pk: sphincs::PublicKey,

    /// Ephemeral Falcon key for this session (fast)
    pub ephemeral_sk: falcon512::SecretKey,
    pub ephemeral_pk: falcon512::PublicKey,
}

impl SphincsKem {
    /// Generate new KEM keypairs
    pub fn new() -> QsslResult<Self> {
        // Long-term identity with SPHINCS+ (quantum-secure)
        let (identity_pk, identity_sk) = sphincs::keypair();

        // Ephemeral with Falcon (fast for session)
        let (ephemeral_pk, ephemeral_sk) = falcon512::keypair();

        Ok(Self {
            identity_sk,
            identity_pk,
            ephemeral_sk,
            ephemeral_pk,
        })
    }

    /// Encapsulate - Create ciphertext and shared secret
    ///
    /// This is the RIGHT way to do post-quantum key exchange:
    /// 1. Generate ephemeral keypair
    /// 2. Sign ephemeral public key with identity key
    /// 3. Create shared secret using hash of both ephemeral keys
    /// 4. Authenticate with signatures
    pub fn encapsulate(
        &self,
        peer_identity_pk: &[u8],
    ) -> QsslResult<(Vec<u8>, Vec<u8>)> {
        // Parse peer's identity key
        let peer_pk = sphincs::PublicKey::from_bytes(peer_identity_pk)
            .map_err(|e| QsslError::Crypto(format!("Invalid SPHINCS+ public key: {:?}", e)))?;

        // Generate fresh ephemeral secret
        let mut ephemeral_secret = vec![0u8; 64];
        rand::thread_rng().fill_bytes(&mut ephemeral_secret);

        // Create KEM message:
        // 1. Our ephemeral public key
        // 2. Encrypted secret
        // 3. Signature over everything

        let mut kem_message = Vec::new();

        // Add our ephemeral Falcon public key
        use pqcrypto_traits::sign::PublicKey;
        kem_message.extend_from_slice(self.ephemeral_pk.as_bytes());

        // Add the ephemeral secret (will be encrypted in transport)
        kem_message.extend_from_slice(&ephemeral_secret);

        // Sign with our SPHINCS+ identity key
        let signature = sphincs::detached_sign(&kem_message, &self.identity_sk);

        // Create ciphertext = ephemeral_pk || secret || signature
        let mut ciphertext = kem_message.clone();
        use pqcrypto_traits::sign::DetachedSignature;
        ciphertext.extend_from_slice(signature.as_bytes());

        // Derive shared secret using hash of all components
        let shared_secret = self.derive_shared_secret(
            &ephemeral_secret,
            peer_identity_pk,
            &kem_message,
        )?;

        // Zeroize ephemeral secret
        ephemeral_secret.zeroize();

        Ok((ciphertext, shared_secret))
    }

    /// Decapsulate - Extract shared secret from ciphertext
    pub fn decapsulate(
        &self,
        ciphertext: &[u8],
        peer_identity_pk: &[u8],
    ) -> QsslResult<Vec<u8>> {
        // SPHINCS+ and Falcon sizes
        const FALCON_PK_SIZE: usize = 897;  // Falcon-512 public key
        const SPHINCS_SIG_SIZE: usize = 16976;  // SPHINCS+ haraka-128f signature (actual size)
        const EPHEMERAL_SIZE: usize = 64;

        // Parse ciphertext structure
        if ciphertext.len() < FALCON_PK_SIZE + EPHEMERAL_SIZE + SPHINCS_SIG_SIZE {
            return Err(QsslError::Crypto(
                format!("Invalid ciphertext size: got {}, expected {}",
                    ciphertext.len(),
                    FALCON_PK_SIZE + EPHEMERAL_SIZE + SPHINCS_SIG_SIZE)
            ));
        }

        let mut offset = 0;

        // Extract peer's ephemeral Falcon public key
        let peer_ephemeral_pk = &ciphertext[offset..offset + FALCON_PK_SIZE];
        offset += FALCON_PK_SIZE;

        // Extract ephemeral secret
        let ephemeral_secret = &ciphertext[offset..offset + EPHEMERAL_SIZE];
        offset += EPHEMERAL_SIZE;

        // Extract signature
        let signature = &ciphertext[offset..offset + SPHINCS_SIG_SIZE];

        // Verify signature with peer's SPHINCS+ identity
        let peer_pk = sphincs::PublicKey::from_bytes(peer_identity_pk)
            .map_err(|e| QsslError::Crypto(format!("Invalid peer public key: {:?}", e)))?;

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

        // Verify signature over ephemeral_pk || secret
        let signed_data = &ciphertext[..FALCON_PK_SIZE + EPHEMERAL_SIZE];
        if sphincs::verify_detached_signature(&sig, signed_data, &peer_pk).is_err() {
            return Err(QsslError::Crypto("Signature verification failed".to_string()));
        }

        // Derive shared secret
        let shared_secret = self.derive_shared_secret(
            ephemeral_secret,
            peer_identity_pk,
            signed_data,
        )?;

        Ok(shared_secret)
    }

    /// Derive shared secret using quantum-resistant hash
    fn derive_shared_secret(
        &self,
        ephemeral_secret: &[u8],
        peer_identity: &[u8],
        transcript: &[u8],
    ) -> QsslResult<Vec<u8>> {
        // Use SHA3-512 for quantum resistance (Grover gives 256-bit security)
        let mut hasher = Sha3_512::new();

        // Mix in all components in a deterministic order
        hasher.update(b"SPHINCS+_KEM_v1.0");
        hasher.update(ephemeral_secret);

        // Sort identities to ensure same order on both sides
        use pqcrypto_traits::sign::PublicKey;
        let our_identity = self.identity_pk.as_bytes();

        if our_identity < peer_identity {
            hasher.update(our_identity);
            hasher.update(peer_identity);
        } else {
            hasher.update(peer_identity);
            hasher.update(our_identity);
        }

        // Don't include ephemeral keys as they differ between parties
        // hasher.update(transcript);  // This also differs

        let hash = hasher.finalize();

        // Return 256-bit shared secret (quantum-secure)
        Ok(hash[..32].to_vec())
    }
}

/// Hybrid KEM combining SPHINCS+ and Falcon
///
/// This provides defense-in-depth:
/// - SPHINCS+ for long-term quantum security (hash-based)
/// - Falcon for efficiency and different mathematical assumption
pub struct HybridKem {
    sphincs_kem: SphincsKem,
    falcon_sk: falcon512::SecretKey,
    falcon_pk: falcon512::PublicKey,
}

impl HybridKem {
    /// Create new hybrid KEM
    pub fn new() -> QsslResult<Self> {
        let sphincs_kem = SphincsKem::new()?;
        let (falcon_pk, falcon_sk) = falcon512::keypair();

        Ok(Self {
            sphincs_kem,
            falcon_sk,
            falcon_pk,
        })
    }

    /// Hybrid encapsulation - combines both algorithms
    pub fn encapsulate(
        &self,
        peer_sphincs_pk: &[u8],
        peer_falcon_pk: &[u8],
    ) -> QsslResult<(Vec<u8>, Vec<u8>)> {
        // Get SPHINCS+ ciphertext and secret
        let (sphincs_ct, sphincs_ss) = self.sphincs_kem.encapsulate(peer_sphincs_pk)?;

        // Generate Falcon ephemeral and sign
        let mut falcon_ephemeral = vec![0u8; 32];
        rand::thread_rng().fill_bytes(&mut falcon_ephemeral);

        let falcon_sig = falcon512::detached_sign(&falcon_ephemeral, &self.falcon_sk);

        // Combine ciphertexts
        let mut hybrid_ct = Vec::new();
        hybrid_ct.extend_from_slice(&(sphincs_ct.len() as u32).to_be_bytes());
        hybrid_ct.extend_from_slice(&sphincs_ct);
        hybrid_ct.extend_from_slice(&falcon_ephemeral);

        use pqcrypto_traits::sign::DetachedSignature;
        let sig_bytes = falcon_sig.as_bytes();
        // Falcon-512 signature should be exactly 658 bytes (not 656)
        hybrid_ct.extend_from_slice(sig_bytes);

        // Combine shared secrets with XOR (quantum-secure combiner)
        let mut hybrid_ss = vec![0u8; 32];
        for i in 0..32 {
            hybrid_ss[i] = sphincs_ss[i] ^ falcon_ephemeral[i % falcon_ephemeral.len()];
        }

        // Additional mixing with SHA3
        let mut hasher = Sha3_512::new();
        hasher.update(b"HYBRID_KEM");
        hasher.update(&sphincs_ss);
        hasher.update(&falcon_ephemeral);
        let final_ss = hasher.finalize();

        Ok((hybrid_ct, final_ss[..32].to_vec()))
    }

    /// Hybrid decapsulation
    pub fn decapsulate(
        &self,
        ciphertext: &[u8],
        peer_sphincs_pk: &[u8],
        peer_falcon_pk: &[u8],
    ) -> QsslResult<Vec<u8>> {
        if ciphertext.len() < 4 {
            return Err(QsslError::Crypto("Invalid hybrid ciphertext".to_string()));
        }

        // Parse hybrid ciphertext
        let sphincs_len = u32::from_be_bytes([
            ciphertext[0], ciphertext[1], ciphertext[2], ciphertext[3]
        ]) as usize;

        const FALCON_SIG_SIZE: usize = 658;  // Falcon-512 detached signature actual size

        if ciphertext.len() < 4 + sphincs_len + 32 + FALCON_SIG_SIZE {
            return Err(QsslError::Crypto(format!(
                "Invalid hybrid ciphertext size: got {}, expected at least {}",
                ciphertext.len(),
                4 + sphincs_len + 32 + FALCON_SIG_SIZE
            )));
        }

        let sphincs_ct = &ciphertext[4..4 + sphincs_len];
        let falcon_ephemeral = &ciphertext[4 + sphincs_len..4 + sphincs_len + 32];
        let falcon_sig = &ciphertext[4 + sphincs_len + 32..];

        // Verify Falcon signature
        let peer_falcon = falcon512::PublicKey::from_bytes(peer_falcon_pk)
            .map_err(|e| QsslError::Crypto(format!("Invalid Falcon key: {:?}", e)))?;

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

        if falcon512::verify_detached_signature(&sig, falcon_ephemeral, &peer_falcon).is_err() {
            return Err(QsslError::Crypto("Falcon signature verification failed".to_string()));
        }

        // Decapsulate SPHINCS+
        let sphincs_ss = self.sphincs_kem.decapsulate(sphincs_ct, peer_sphincs_pk)?;

        // Combine secrets
        let mut hasher = Sha3_512::new();
        hasher.update(b"HYBRID_KEM");
        hasher.update(&sphincs_ss);
        hasher.update(falcon_ephemeral);
        let final_ss = hasher.finalize();

        Ok(final_ss[..32].to_vec())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sphincs_kem() {
        // Alice and Bob generate their keys
        let alice = SphincsKem::new().unwrap();
        let bob = SphincsKem::new().unwrap();

        use pqcrypto_traits::sign::PublicKey;
        let alice_pk = alice.identity_pk.as_bytes();
        let bob_pk = bob.identity_pk.as_bytes();

        // Alice encapsulates for Bob
        let (ciphertext, alice_secret) = alice.encapsulate(bob_pk).unwrap();

        // Bob decapsulates
        let bob_secret = bob.decapsulate(&ciphertext, alice_pk).unwrap();

        // Shared secrets should match
        assert_eq!(alice_secret, bob_secret);
    }

    #[test]
    fn test_hybrid_kem() {
        let alice = HybridKem::new().unwrap();
        let bob = HybridKem::new().unwrap();

        use pqcrypto_traits::sign::PublicKey;
        let alice_sphincs = alice.sphincs_kem.identity_pk.as_bytes();
        let alice_falcon = alice.falcon_pk.as_bytes();
        let bob_sphincs = bob.sphincs_kem.identity_pk.as_bytes();
        let bob_falcon = bob.falcon_pk.as_bytes();

        // Exchange
        let (ct, alice_ss) = alice.encapsulate(bob_sphincs, bob_falcon).unwrap();
        let bob_ss = bob.decapsulate(&ct, alice_sphincs, alice_falcon).unwrap();

        assert_eq!(alice_ss, bob_ss);
    }
}