qssl 0.2.0

//! SPHINCS+ based Key Encapsulation Mechanism
//!
//! Pure Rust via paraxiom-pqc. Zero C.

use paraxiom_pqc::sign::{self as pqc_sign, SignAlgorithm, SigningKey, VerificationKey, Signature};
use sha3::{Sha3_512, Digest};
use zeroize::Zeroize;
use rand::RngCore;

use crate::{QsslError, QsslResult};

/// SPHINCS+ KEM - uses SLH-DSA (FIPS 205) + Falcon for ephemeral keys
pub struct SphincsKem {
    /// Long-term SLH-DSA identity key
    pub identity_sk: SigningKey,
    pub identity_pk: VerificationKey,

    /// Ephemeral Falcon key for this session (fast)
    pub ephemeral_sk: SigningKey,
    pub ephemeral_pk: VerificationKey,
}

impl SphincsKem {
    /// Generate new KEM keypairs
    pub fn new() -> QsslResult<Self> {
        // Long-term identity with SLH-DSA (quantum-secure, hash-based)
        let identity = pqc_sign::keypair(SignAlgorithm::SlhDsaShake128f)
            .map_err(|e| QsslError::Crypto(format!("SLH-DSA keygen failed: {}", e)))?;

        // Ephemeral with Falcon (fast for session)
        let ephemeral = pqc_sign::keypair(SignAlgorithm::Falcon512)
            .map_err(|e| QsslError::Crypto(format!("Falcon keygen failed: {}", e)))?;

        Ok(Self {
            identity_sk: identity.sk,
            identity_pk: identity.vk,
            ephemeral_sk: ephemeral.sk,
            ephemeral_pk: ephemeral.vk,
        })
    }

    /// Encapsulate - Create ciphertext and shared secret
    pub fn encapsulate(
        &self,
        peer_identity_pk: &[u8],
    ) -> QsslResult<(Vec<u8>, Vec<u8>)> {
        // Generate fresh ephemeral secret
        let mut ephemeral_secret = vec![0u8; 64];
        rand::thread_rng().fill_bytes(&mut ephemeral_secret);

        // Create KEM message: ephemeral_pk || secret
        let mut kem_message = Vec::new();
        kem_message.extend_from_slice(&self.ephemeral_pk.bytes);
        kem_message.extend_from_slice(&ephemeral_secret);

        // Sign with our SLH-DSA identity key
        let signature = pqc_sign::sign(&self.identity_sk, &kem_message)
            .map_err(|e| QsslError::Crypto(format!("Sign failed: {}", e)))?;

        // Create ciphertext = ephemeral_pk || secret || signature
        let mut ciphertext = kem_message.clone();
        ciphertext.extend_from_slice(&signature.bytes);

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

        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>> {
        // Falcon-512 public key size
        let falcon_pk_size = self.ephemeral_pk.bytes.len();
        const EPHEMERAL_SIZE: usize = 64;

        let msg_size = falcon_pk_size + EPHEMERAL_SIZE;
        if ciphertext.len() <= msg_size {
            return Err(QsslError::Crypto(
                format!("Invalid ciphertext size: got {}", ciphertext.len())
            ));
        }

        // Extract components
        let signed_data = &ciphertext[..msg_size];
        let sig_bytes = &ciphertext[msg_size..];
        let ephemeral_secret = &ciphertext[falcon_pk_size..msg_size];

        // Verify signature with peer's SLH-DSA identity
        let peer_vk = VerificationKey {
            algorithm: SignAlgorithm::SlhDsaShake128f,
            bytes: peer_identity_pk.to_vec(),
        };
        let sig = Signature {
            algorithm: SignAlgorithm::SlhDsaShake128f,
            bytes: sig_bytes.to_vec(),
        };

        let valid = pqc_sign::verify(&peer_vk, signed_data, &sig)
            .map_err(|e| QsslError::Crypto(format!("Verify failed: {}", e)))?;
        if !valid {
            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>> {
        let mut hasher = Sha3_512::new();
        hasher.update(b"SPHINCS+_KEM_v1.0");
        hasher.update(ephemeral_secret);

        let our_identity = &self.identity_pk.bytes;
        if our_identity.as_slice() < peer_identity {
            hasher.update(our_identity);
            hasher.update(peer_identity);
        } else {
            hasher.update(peer_identity);
            hasher.update(our_identity);
        }

        let hash = hasher.finalize();
        Ok(hash[..32].to_vec())
    }
}

/// Hybrid KEM combining SLH-DSA and Falcon
pub struct HybridKem {
    sphincs_kem: SphincsKem,
    falcon_sk: SigningKey,
    falcon_pk: VerificationKey,
}

impl HybridKem {
    pub fn new() -> QsslResult<Self> {
        let sphincs_kem = SphincsKem::new()?;
        let falcon_kp = pqc_sign::keypair(SignAlgorithm::Falcon512)
            .map_err(|e| QsslError::Crypto(format!("Falcon keygen failed: {}", e)))?;

        Ok(Self {
            sphincs_kem,
            falcon_sk: falcon_kp.sk,
            falcon_pk: falcon_kp.vk,
        })
    }

    pub fn encapsulate(
        &self,
        peer_sphincs_pk: &[u8],
        _peer_falcon_pk: &[u8],
    ) -> QsslResult<(Vec<u8>, Vec<u8>)> {
        let (sphincs_ct, sphincs_ss) = self.sphincs_kem.encapsulate(peer_sphincs_pk)?;

        let mut falcon_ephemeral = vec![0u8; 32];
        rand::thread_rng().fill_bytes(&mut falcon_ephemeral);

        let falcon_sig = pqc_sign::sign(&self.falcon_sk, &falcon_ephemeral)
            .map_err(|e| QsslError::Crypto(format!("Falcon sign failed: {}", e)))?;

        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);
        hybrid_ct.extend_from_slice(&falcon_sig.bytes);

        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()))
    }

    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()));
        }

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

        if ciphertext.len() < 4 + sphincs_len + 32 {
            return Err(QsslError::Crypto("Invalid hybrid ciphertext size".into()));
        }

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

        // Verify Falcon signature
        let peer_vk = VerificationKey {
            algorithm: SignAlgorithm::Falcon512,
            bytes: peer_falcon_pk.to_vec(),
        };
        let sig = Signature {
            algorithm: SignAlgorithm::Falcon512,
            bytes: falcon_sig_bytes.to_vec(),
        };
        let valid = pqc_sign::verify(&peer_vk, falcon_ephemeral, &sig)
            .map_err(|e| QsslError::Crypto(format!("Falcon verify failed: {}", e)))?;
        if !valid {
            return Err(QsslError::Crypto("Falcon signature verification failed".into()));
        }

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

        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() {
        let alice = SphincsKem::new().unwrap();
        let bob = SphincsKem::new().unwrap();

        let alice_pk = &alice.identity_pk.bytes;
        let bob_pk = &bob.identity_pk.bytes;

        let (ciphertext, alice_secret) = alice.encapsulate(bob_pk).unwrap();
        let bob_secret = bob.decapsulate(&ciphertext, alice_pk).unwrap();

        assert_eq!(alice_secret, bob_secret);
    }

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

        let alice_sphincs = &alice.sphincs_kem.identity_pk.bytes;
        let alice_falcon = &alice.falcon_pk.bytes;
        let bob_sphincs = &bob.sphincs_kem.identity_pk.bytes;
        let bob_falcon = &bob.falcon_pk.bytes;

        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);
    }
}