crypt_guard 3.0.1

use crate::{error::CryptError, log_activity, FileMetadata, FileState, FileTypes, Key, KeyTypes};
use pqcrypto_traits::kem::{Ciphertext, PublicKey, SecretKey, SharedSecret};
use std::{
    marker::PhantomData,
    path::{Path, PathBuf},
    result::Result,
};
use zeroize::Zeroize;

/// Trait for implementing key management functions. This trait provides
/// an interface for key pair generation, encapsulation/decapsulation of secrets,
/// and key manipulation (such as setting and getting key values).
pub trait KyberKeyFunctions {
    /// Generates a new key pair.
    fn keypair() -> Result<(Vec<u8>, Vec<u8>), CryptError>;
    /// Encapsulates a secret using a public key.
    fn encap(public_key: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CryptError>;
    /// Decapsulates a secret using a secret key and a ciphertext.
    fn decap(secret_key: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>, CryptError>;
}

/// Implementation for Kyber 1024 variant.
pub struct KeyControKyber1024;
impl KyberKeyFunctions for KeyControKyber1024 {
    /// Generates a public and secret key pair using the Kyber1024 algorithm.
    fn keypair() -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        use pqcrypto_kyber::kyber1024::*;
        log_activity!(
            "Generating a new keypair.\n\tThe used KEM: ",
            format!("Kyber{}", 1024).as_str()
        );

        let (pk, sk) = keypair();
        let public_key = pk.as_bytes().to_owned();
        let secret_key = sk.as_bytes().to_owned();

        log_activity!(
            "A new keypair was created.\n\tThe used KEM: ",
            format!("Kyber{}", 1024).as_str()
        );
        Ok((public_key, secret_key))
    }

    /// Encapsulates a secret using a public key to produce a shared secret and a ciphertext.
    fn encap(public: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        use pqcrypto_kyber::kyber1024::*;
        log_activity!(
            "Generating shared_secret and ciphertext.\n\tThe used KEM: ",
            format!("Kyber{}", 1024).as_str()
        );

        let pk = PublicKey::from_bytes(public).unwrap();
        let (ss, ct) = encapsulate(&pk);

        let ciphertext = ct.as_bytes().to_vec();
        let shared_secret = ss.as_bytes().to_vec();

        log_activity!(
            "Finished generating shared_secret and ciphertext.\n\tThe used KEM: ",
            format!("Kyber{}", 1024).as_str()
        );
        Ok((shared_secret, ciphertext))
    }

    /// Decapsulates the ciphertext using a secret key to retrieve the shared secret.
    fn decap(sec: &[u8], cipher: &[u8]) -> Result<Vec<u8>, CryptError> {
        use pqcrypto_kyber::kyber1024::*;
        log_activity!("Starting decapsulation of shared_secret using secret_key and ciphertext.\n\tThe used KEM: ", format!("Kyber{}", 1024).as_str());

        let ct = Ciphertext::from_bytes(cipher).unwrap();
        let sk = SecretKey::from_bytes(sec).unwrap();
        let ss2 = decapsulate(&ct, &sk);
        let shared_secret = ss2.as_bytes().to_vec();

        log_activity!(
            "Decapsulated the shared_secret using secret_key and ciphertext.\n\tThe used KEM: ",
            format!("Kyber{}", 1024).as_str()
        );
        Ok(shared_secret)
    }
}

// Provide inherent methods so call sites can use `KeyControKyber1024::keypair()` etc.
impl KeyControKyber1024 {
    pub fn keypair() -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        <Self as KyberKeyFunctions>::keypair()
    }
    pub fn encap(public_key: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        <Self as KyberKeyFunctions>::encap(public_key)
    }
    pub fn decap(secret_key: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>, CryptError> {
        <Self as KyberKeyFunctions>::decap(secret_key, ciphertext)
    }
}

/// Implementation for Kyber 768 variant.
pub struct KeyControKyber768;
impl KyberKeyFunctions for KeyControKyber768 {
    /// Generates a public and secret key pair using the Kyber768 algorithm.
    fn keypair() -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        use pqcrypto_kyber::kyber768::*;
        log_activity!(
            "Generating a new keypair.\n\tThe used KEM: ",
            format!("Kyber{}", 768).as_str()
        );

        let (pk, sk) = keypair();
        let public_key = pk.as_bytes().to_vec();
        let secret_key = sk.as_bytes().to_vec();
        log_activity!(
            "A new keypair was created.\n\tThe used KEM: ",
            format!("Kyber{}", 768).as_str()
        );

        Ok((public_key, secret_key))
    }

    /// Encapsulates a secret using a public key to produce a shared secret and a ciphertext.
    fn encap(public: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        use pqcrypto_kyber::kyber768::*;
        log_activity!(
            "Generating shared_secret and ciphertext.\n\tThe used KEM: ",
            format!("Kyber{}", 768).as_str()
        );

        let pk = PublicKey::from_bytes(public).unwrap();
        let (ss, ct) = encapsulate(&pk);

        let ciphertext = ct.as_bytes().to_vec();
        let shared_secret = ss.as_bytes().to_vec();
        log_activity!(
            "Finished generating shared_secret and ciphertext.\n\tThe used KEM: ",
            format!("Kyber{}", 768).as_str()
        );

        Ok((shared_secret, ciphertext))
    }

    /// Decapsulates the ciphertext using a secret key to retrieve the shared secret.
    fn decap(sec: &[u8], cipher: &[u8]) -> Result<Vec<u8>, CryptError> {
        use pqcrypto_kyber::kyber768::*;
        log_activity!("Starting decapsulation of shared_secret using secret_key and ciphertext.\n\tThe used KEM: ", format!("Kyber{}", 768).as_str());

        let ct = Ciphertext::from_bytes(cipher).unwrap();
        let sk = SecretKey::from_bytes(sec).unwrap();
        let ss2 = decapsulate(&ct, &sk);
        let shared_secret = ss2.as_bytes().to_vec();
        log_activity!(
            "Decapsulated the shared_secret using secret_key and ciphertext.\n\tThe used KEM: ",
            format!("Kyber{}", 768).as_str()
        );

        Ok(shared_secret)
    }
}

// Provide inherent methods so call sites can use `KeyControKyber768::keypair()` etc.
impl KeyControKyber768 {
    pub fn keypair() -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        <Self as KyberKeyFunctions>::keypair()
    }
    pub fn encap(public_key: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        <Self as KyberKeyFunctions>::encap(public_key)
    }
    pub fn decap(secret_key: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>, CryptError> {
        <Self as KyberKeyFunctions>::decap(secret_key, ciphertext)
    }
}

/// Implementation for Kyber 512 variant.
pub struct KeyControKyber512;
impl KyberKeyFunctions for KeyControKyber512 {
    /// Generates a public and secret key pair using the Kyber512 algorithm.
    fn keypair() -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        use pqcrypto_kyber::kyber512::*;
        log_activity!(
            "Generating a new keypair.\n\tThe used KEM: ",
            format!("Kyber{}", 512).as_str()
        );

        let (pk, sk) = keypair();
        let public_key = pk.as_bytes().to_vec();
        let secret_key = sk.as_bytes().to_vec();
        log_activity!(
            "A new keypair was created.\n\tThe used KEM: ",
            format!("Kyber{}", 512).as_str()
        );

        Ok((public_key, secret_key))
    }

    /// Encapsulates a secret using a public key to produce a shared secret and a ciphertext.
    fn encap(public: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        use pqcrypto_kyber::kyber512::*;
        log_activity!(
            "Generating shared_secret and ciphertext.\n\tThe used KEM: ",
            format!("Kyber{}", 512).as_str()
        );

        let pk = PublicKey::from_bytes(public).unwrap();
        let (ss, ct) = encapsulate(&pk);

        let ciphertext = ct.as_bytes().to_vec();
        let shared_secret = ss.as_bytes().to_vec();
        log_activity!(
            "Finished generating shared_secret and ciphertext.\n\tThe used KEM: ",
            format!("Kyber{}", 512).as_str()
        );

        Ok((shared_secret, ciphertext))
    }

    /// Decapsulates the ciphertext using a secret key to retrieve the shared secret.
    fn decap(sec: &[u8], cipher: &[u8]) -> Result<Vec<u8>, CryptError> {
        use pqcrypto_kyber::kyber512::*;
        log_activity!("Starting decapsulation of shared_secret using secret_key and ciphertext.\n\tThe used KEM: ", format!("Kyber{}", 512).as_str());

        let ct = Ciphertext::from_bytes(cipher).unwrap();
        let sk = SecretKey::from_bytes(sec).unwrap();
        let ss2 = decapsulate(&ct, &sk);
        let shared_secret = ss2.as_bytes().to_vec();
        log_activity!(
            "Decapsulated the shared_secret using secret_key and ciphertext.\n\tThe used KEM: ",
            format!("Kyber{}", 512).as_str()
        );

        Ok(shared_secret)
    }
}

// Provide inherent methods so call sites can use `KeyControKyber512::keypair()` etc.
impl KeyControKyber512 {
    pub fn keypair() -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        <Self as KyberKeyFunctions>::keypair()
    }
    pub fn encap(public_key: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        <Self as KyberKeyFunctions>::encap(public_key)
    }
    pub fn decap(secret_key: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>, CryptError> {
        <Self as KyberKeyFunctions>::decap(secret_key, ciphertext)
    }
}

/// A structure to manage cryptographic keys and operations for the Kyber algorithm.
/// It encapsulates the public key, secret key, ciphertext, and shared secret.
pub struct KeyControl<T: KyberKeyFunctions> {
    pub public_key: Vec<u8>,
    pub secret_key: Vec<u8>,
    pub ciphertext: Vec<u8>,
    pub shared_secret: Vec<u8>,
    _marker: std::marker::PhantomData<T>,
}

impl<T: KyberKeyFunctions> Drop for KeyControl<T> {
    fn drop(&mut self) {
        self.secret_key.zeroize();
        self.shared_secret.zeroize();
    }
}

impl<T: KyberKeyFunctions> KeyControl<T> {
    /// Constructs a new instance of `KeyControl`.
    pub fn new() -> Self {
        KeyControl {
            public_key: Vec::new(),
            secret_key: Vec::new(),
            ciphertext: Vec::new(),
            shared_secret: Vec::new(),
            _marker: PhantomData,
        }
    }
    /// Sets the ciphertext for the `KeyControl` instance.
    pub fn set_ciphertext(&mut self, cipher: Vec<u8>) -> Result<(), CryptError> {
        self.ciphertext = cipher;
        Ok(())
    }

    /// Sets the public key for the `KeyControl` instance.
    pub fn set_public_key(&mut self, public: Vec<u8>) -> Result<(), CryptError> {
        self.public_key = public;
        Ok(())
    }

    /// Sets the secret key for the `KeyControl` instance.
    pub fn set_secret_key(&mut self, sec: Vec<u8>) -> Result<(), CryptError> {
        self.secret_key = sec;
        Ok(())
    }

    /// Retrieves a specified key based on `KeyTypes`.
    pub fn get_key(&self, key: KeyTypes) -> Result<Key, CryptError> {
        let key = match key {
            KeyTypes::None => unimplemented!(),
            KeyTypes::PublicKey => Key::new(KeyTypes::PublicKey, self.public_key.to_vec()),
            KeyTypes::SecretKey => Key::new(KeyTypes::SecretKey, self.secret_key.to_vec()),
            KeyTypes::Ciphertext => Key::new(KeyTypes::Ciphertext, self.ciphertext.to_vec()),
            KeyTypes::SharedSecret => Key::new(KeyTypes::SharedSecret, self.shared_secret.to_vec()),
        };
        Ok(key)
    }

    /// Saves a specified key to a file at the given base path.
    pub fn save(&self, key: KeyTypes, base_path: PathBuf) -> Result<(), CryptError> {
        let key = match key {
            KeyTypes::None => unimplemented!(),
            KeyTypes::PublicKey => Key::new(KeyTypes::PublicKey, self.public_key.to_vec()),
            KeyTypes::SecretKey => Key::new(KeyTypes::SecretKey, self.secret_key.to_vec()),
            KeyTypes::Ciphertext => Key::new(KeyTypes::Ciphertext, self.ciphertext.to_vec()),
            KeyTypes::SharedSecret => unimplemented!(),
        };
        key.save(base_path)
    }
    /// Loads a specified key from a file.
    pub fn load(&self, key: KeyTypes, path: &Path) -> Result<Vec<u8>, CryptError> {
        let key = match key {
            KeyTypes::None => unimplemented!(),
            KeyTypes::PublicKey => {
                FileMetadata::from(PathBuf::from(path), FileTypes::PublicKey, FileState::Other)
            }
            KeyTypes::SecretKey => {
                FileMetadata::from(PathBuf::from(path), FileTypes::SecretKey, FileState::Other)
            }
            KeyTypes::Ciphertext => {
                FileMetadata::from(PathBuf::from(path), FileTypes::Ciphertext, FileState::Other)
            }
            KeyTypes::SharedSecret => unimplemented!(),
        };
        Ok(key.load().unwrap())
    }

    /// Getter methods for public_key, secret_key, ciphertext, and shared_secret.
    pub fn public_key(&self) -> Result<Vec<u8>, CryptError> {
        let key = &self.public_key;
        Ok(key.to_vec())
    }
    pub fn secret_key(&self) -> Result<Vec<u8>, CryptError> {
        let key = &self.secret_key;
        Ok(key.to_vec())
    }
    pub fn ciphertext(&self) -> Result<Vec<u8>, CryptError> {
        let key = &self.ciphertext;
        Ok(key.to_vec())
    }
    pub fn shared_secret(&self) -> Result<Vec<u8>, CryptError> {
        let key = &self.shared_secret;
        Ok(key.to_vec())
    }

    /// Encapsulates a secret using a public key.
    pub fn encap(&self, public: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CryptError> {
        T::encap(public)
    }
    /// Decapsulates the ciphertext using a secret key to retrieve the shared secret.
    pub fn decap(&self, secret_key: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>, CryptError> {
        // Call the `decap` method on the type `T` that `KeyControl` wraps around
        T::decap(secret_key, ciphertext)
    }
}

impl<T: KyberKeyFunctions> Default for KeyControl<T> {
    fn default() -> Self {
        Self::new()
    }
}