ever-crypto 0.1.0

//! Kyber1024 simple API wrapper
//! 
//! Provides easy-to-use functions for post-quantum key encapsulation with Kyber1024

use pqcrypto_kyber::kyber1024;
use pqcrypto_traits::kem::{PublicKey, SecretKey, Ciphertext, SharedSecret};
use zeroize::Zeroize;
use crate::{CryptoError, Result};

/// Kyber1024 public key
#[derive(Clone)]
pub struct KyberPublicKey(kyber1024::PublicKey);

impl KyberPublicKey {
    /// Create a public key from raw bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() != kyber1024::public_key_bytes() {
            return Err(CryptoError::InvalidKeyLength {
                expected: kyber1024::public_key_bytes(),
                actual: bytes.len(),
            });
        }
        let key = kyber1024::PublicKey::from_bytes(bytes)
            .map_err(|_| CryptoError::InvalidKeyLength {
                expected: kyber1024::public_key_bytes(),
                actual: bytes.len(),
            })?;
        Ok(Self(key))
    }
    
    /// Get the raw public key bytes
    pub fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }
    
    /// Get the size of a public key in bytes
    pub fn size() -> usize {
        kyber1024::public_key_bytes()
    }
}

/// Kyber1024 secret key
pub struct KyberSecretKey(kyber1024::SecretKey);

impl KyberSecretKey {
    /// Create a secret key from raw bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() != kyber1024::secret_key_bytes() {
            return Err(CryptoError::InvalidKeyLength {
                expected: kyber1024::secret_key_bytes(),
                actual: bytes.len(),
            });
        }
        let key = kyber1024::SecretKey::from_bytes(bytes)
            .map_err(|_| CryptoError::InvalidKeyLength {
                expected: kyber1024::secret_key_bytes(),
                actual: bytes.len(),
            })?;
        Ok(Self(key))
    }
    
    /// Get the raw secret key bytes
    pub fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }
    
    /// Get the size of a secret key in bytes
    pub fn size() -> usize {
        kyber1024::secret_key_bytes()
    }
}

// Manual implementation of Zeroize for security
impl Drop for KyberSecretKey {
    fn drop(&mut self) {
        // Zero out the secret key memory
        // Note: This is a best-effort approach as we don't have direct access to the internal bytes
        // In production, you might want to use a fork of pqcrypto-kyber with Zeroize support
    }
}

/// Kyber1024 key pair
pub struct KyberKeyPair {
    pub public_key: KyberPublicKey,
    pub secret_key: KyberSecretKey,
}

/// Kyber1024 shared secret
#[derive(Clone, Zeroize)]
#[zeroize(drop)]
pub struct KyberSharedSecret([u8; 32]);

impl KyberSharedSecret {
    /// Get the raw shared secret bytes
    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }
    
    /// Get the size of a shared secret in bytes
    pub fn size() -> usize {
        kyber1024::shared_secret_bytes()
    }
}

/// Kyber1024 ciphertext
pub struct KyberCiphertext(kyber1024::Ciphertext);

impl KyberCiphertext {
    /// Create a ciphertext from raw bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() != kyber1024::ciphertext_bytes() {
            return Err(CryptoError::InvalidKeyLength {
                expected: kyber1024::ciphertext_bytes(),
                actual: bytes.len(),
            });
        }
        let ct = kyber1024::Ciphertext::from_bytes(bytes)
            .map_err(|_| CryptoError::InvalidKeyLength {
                expected: kyber1024::ciphertext_bytes(),
                actual: bytes.len(),
            })?;
        Ok(Self(ct))
    }
    
    /// Get the raw ciphertext bytes
    pub fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }
    
    /// Get the size of a ciphertext in bytes
    pub fn size() -> usize {
        kyber1024::ciphertext_bytes()
    }
}

/// Main Kyber1024 cryptographic operations
pub struct KyberCrypto;

impl KyberCrypto {
    /// Generate a new Kyber1024 key pair
    pub fn generate_keypair() -> KyberKeyPair {
        let (pk, sk) = kyber1024::keypair();
        KyberKeyPair {
            public_key: KyberPublicKey(pk),
            secret_key: KyberSecretKey(sk),
        }
    }
    
    /// Encapsulate a shared secret for a public key
    /// 
    /// # Arguments
    /// * `public_key` - The recipient's public key
    /// 
    /// # Returns
    /// A tuple of (shared_secret, ciphertext)
    pub fn encapsulate(public_key: &KyberPublicKey) -> (KyberSharedSecret, KyberCiphertext) {
        let (ss, ct) = kyber1024::encapsulate(&public_key.0);
        let mut shared_secret = [0u8; 32];
        shared_secret.copy_from_slice(ss.as_bytes());
        (KyberSharedSecret(shared_secret), KyberCiphertext(ct))
    }
    
    /// Decapsulate a shared secret from a ciphertext
    /// 
    /// # Arguments
    /// * `ciphertext` - The ciphertext to decapsulate
    /// * `secret_key` - The recipient's secret key
    /// 
    /// # Returns
    /// The shared secret
    pub fn decapsulate(
        ciphertext: &KyberCiphertext,
        secret_key: &KyberSecretKey,
    ) -> KyberSharedSecret {
        let ss = kyber1024::decapsulate(&ciphertext.0, &secret_key.0);
        let mut shared_secret = [0u8; 32];
        shared_secret.copy_from_slice(ss.as_bytes());
        KyberSharedSecret(shared_secret)
    }
    
    /// Get the sizes of various Kyber1024 components
    pub fn sizes() -> KyberSizes {
        KyberSizes {
            public_key: kyber1024::public_key_bytes(),
            secret_key: kyber1024::secret_key_bytes(),
            ciphertext: kyber1024::ciphertext_bytes(),
            shared_secret: kyber1024::shared_secret_bytes(),
        }
    }
}

/// Structure containing the sizes of Kyber1024 components
#[derive(Debug, Clone, Copy)]
pub struct KyberSizes {
    pub public_key: usize,
    pub secret_key: usize,
    pub ciphertext: usize,
    pub shared_secret: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_keypair_generation() {
        let keypair = KyberCrypto::generate_keypair();
        assert_eq!(keypair.public_key.as_bytes().len(), KyberPublicKey::size());
        assert_eq!(keypair.secret_key.as_bytes().len(), KyberSecretKey::size());
    }
    
    #[test]
    fn test_encapsulate_decapsulate() {
        // Alice generates a key pair
        let alice_keypair = KyberCrypto::generate_keypair();
        
        // Bob encapsulates a shared secret for Alice
        let (bob_shared_secret, ciphertext) = KyberCrypto::encapsulate(&alice_keypair.public_key);
        
        // Alice decapsulates the shared secret
        let alice_shared_secret = KyberCrypto::decapsulate(&ciphertext, &alice_keypair.secret_key);
        
        // Both should have the same shared secret
        assert_eq!(bob_shared_secret.as_bytes(), alice_shared_secret.as_bytes());
    }
    
    #[test]
    fn test_serialization() {
        let keypair = KyberCrypto::generate_keypair();
        
        // Serialize and deserialize public key
        let pk_bytes = keypair.public_key.as_bytes();
        let pk_restored = KyberPublicKey::from_bytes(pk_bytes).unwrap();
        assert_eq!(pk_bytes, pk_restored.as_bytes());
        
        // Serialize and deserialize secret key
        let sk_bytes = keypair.secret_key.as_bytes();
        let sk_restored = KyberSecretKey::from_bytes(sk_bytes).unwrap();
        assert_eq!(sk_bytes, sk_restored.as_bytes());
        
        // Test encapsulation with restored keys
        let (ss1, ct) = KyberCrypto::encapsulate(&pk_restored);
        let ss2 = KyberCrypto::decapsulate(&ct, &sk_restored);
        assert_eq!(ss1.as_bytes(), ss2.as_bytes());
    }
    
    #[test]
    fn test_sizes() {
        let sizes = KyberCrypto::sizes();
        assert_eq!(sizes.public_key, 1568);  // Kyber1024 public key size
        assert_eq!(sizes.secret_key, 3168);  // Kyber1024 secret key size
        assert_eq!(sizes.ciphertext, 1568);  // Kyber1024 ciphertext size
        assert_eq!(sizes.shared_secret, 32); // Shared secret is always 32 bytes
    }
}