saorsa-seal 0.1.3

//! Envelope encryption abstraction for recipient encryption

use crate::types::{PqcEnvelope, Recipient};
use saorsa_pqc::{MlKem768, MlKemCiphertext, MlKemOperations, MlKemPublicKey, MlKemSecretKey};

/// No-op envelope for unencrypted shares
#[derive(Clone, Debug)]
pub struct NoEnvelope;

impl NoEnvelope {
    /// Create a new no-op envelope
    pub fn new() -> Self {
        Self
    }
}

impl Default for NoEnvelope {
    fn default() -> Self {
        Self::new()
    }
}

/// Post-quantum envelope using ML-KEM-768
#[derive(Clone)]
pub struct MlKemEnvelope {
    /// Optional secret key for this instance
    secret_key: Option<MlKemSecretKey>,
    /// Optional public key for this instance
    public_key: Option<MlKemPublicKey>,
}

impl MlKemEnvelope {
    /// Create a new ML-KEM envelope
    pub fn new() -> Self {
        Self {
            secret_key: None,
            public_key: None,
        }
    }

    /// Create with a specific keypair
    pub fn with_keypair(secret_key: MlKemSecretKey, public_key: MlKemPublicKey) -> Self {
        Self {
            secret_key: Some(secret_key),
            public_key: Some(public_key),
        }
    }
}

impl PqcEnvelope for MlKemEnvelope {
    fn encap_to(&self, recipient: &Recipient, data: &[u8]) -> anyhow::Result<Vec<u8>> {
        // Get recipient's public key bytes
        let public_key_bytes = recipient
            .public_key
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("No public key for recipient"))?;

        // Deserialize public key
        let public_key = MlKemPublicKey::from_bytes(public_key_bytes)
            .map_err(|e| anyhow::anyhow!("Invalid public key: {:?}", e))?;

        // Encapsulate using ML-KEM-768
        let kem = MlKem768::new();
        let (ciphertext, shared_secret) = kem.encapsulate(&public_key)?;

        // Use shared secret to encrypt data with XChaCha20-Poly1305
        use saorsa_fec::crypto::{CryptoEngine, EncryptionKey};
        let mut key_bytes = [0u8; 32];
        key_bytes.copy_from_slice(&shared_secret.as_bytes()[..32]);
        let key = EncryptionKey::new(key_bytes);
        let mut engine = CryptoEngine::new();
        let encrypted = engine.encrypt(data, &key)?;

        // Combine ciphertext and encrypted data
        let mut result = Vec::new();
        let ct_bytes = ciphertext.as_bytes();
        result.extend_from_slice(&(ct_bytes.len() as u32).to_le_bytes());
        result.extend_from_slice(ct_bytes);
        result.extend_from_slice(&encrypted);

        Ok(result)
    }

    fn decap_from(&self, _sender: &Recipient, data: &[u8]) -> anyhow::Result<Vec<u8>> {
        // Extract ciphertext length
        if data.len() < 4 {
            return Err(anyhow::anyhow!("Invalid encapsulated data"));
        }
        let ct_len = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;

        if data.len() < 4 + ct_len {
            return Err(anyhow::anyhow!("Invalid encapsulated data length"));
        }

        // Extract ciphertext and encrypted data
        let ciphertext = MlKemCiphertext::from_bytes(&data[4..4 + ct_len])
            .map_err(|e| anyhow::anyhow!("Invalid ciphertext: {:?}", e))?;
        let encrypted = &data[4 + ct_len..];

        // Get our secret key
        let secret_key = self
            .secret_key
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("No secret key available for decapsulation"))?;

        // Decapsulate to get shared secret
        let kem = MlKem768::new();
        let shared_secret = kem.decapsulate(secret_key, &ciphertext)?;

        // Use shared secret to decrypt data
        use saorsa_fec::crypto::{CryptoEngine, EncryptionKey};
        let mut key_bytes = [0u8; 32];
        key_bytes.copy_from_slice(&shared_secret.as_bytes()[..32]);
        let key = EncryptionKey::new(key_bytes);
        let engine = CryptoEngine::new();
        let decrypted = engine.decrypt(encrypted, &key)?;

        Ok(decrypted)
    }

    fn clone_box(&self) -> Box<dyn PqcEnvelope> {
        Box::new(self.clone())
    }
}

impl Default for MlKemEnvelope {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Debug for MlKemEnvelope {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MlKemEnvelope")
            .field("has_secret_key", &self.secret_key.is_some())
            .field("has_public_key", &self.public_key.is_some())
            .finish()
    }
}

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

    #[test]
    fn test_no_envelope() {
        let _envelope = NoEnvelope::new();
    }

    #[test]
    fn test_pqc_envelope_creation() {
        let envelope = MlKemEnvelope::new();
        let boxed: Box<dyn PqcEnvelope> = Box::new(envelope);
        let _cloned = boxed.clone_box();
    }

    #[test]
    fn test_pqc_envelope_with_keypair() {
        let kem = MlKem768::new();
        let (public_key, secret_key) = kem.generate_keypair().unwrap();
        let envelope = MlKemEnvelope::with_keypair(secret_key, public_key);
        assert!(envelope.secret_key.is_some());
        assert!(envelope.public_key.is_some());
    }
}