enc_rust 0.2.2

use crate::{
    errors::{CrystalsError, EncryptionDecryptionError, KeyGenerationError, PackingError},
    indcpa::{
        generate_indcpa_key_pair, PrivateKey as IndcpaPrivateKey, PublicKey as IndcpaPublicKey,
    },
    params::{SecurityLevel, K, MAX_CIPHERTEXT, SHAREDSECRETBYTES, SYMBYTES},
};
use rand_chacha::ChaCha20Rng;
use rand_core::{CryptoRng, RngCore, SeedableRng};
use sha3::{
    digest::{ExtendableOutput, Update, XofReader},
    Digest, Sha3_256, Sha3_512, Shake256,
};
use subtle::{ConditionallySelectable, ConstantTimeEq};
use tinyvec::ArrayVec;

/// `PrivateKey` struct that can only be generated via the [`generate_keypair_512`], [`generate_keypair_768`], or
/// [`generate_keypair_1024`] functions and is used to [`decapsulate`](PrivateKey::decapsulate) a shared secret from a given ciphertext.
///
/// Can be accessed in byte form by packing into a `u8` array using the [`pack`](PrivateKey::pack) method,
/// and made available for use again using the [`unpack`](PrivateKey::unpack) method if `decap_key`
/// feature flag is set, or [`unpack_512`](PrivateKey::unpack_512), [`unpack_768`](PrivateKey::unpack_768), or [`unpack_1024`](PrivateKey::unpack_1024).
///
/// If using `decap_key` feature, the array used to pack must be of the correct length for the given security level, see
/// [`pack`](PrivateKey::pack) for more.
#[derive(Debug, Eq, PartialEq)]
pub struct PrivateKey {
    #[cfg(not(feature = "decap_key"))]
    key: PrivateSeed,
    #[cfg(feature = "decap_key")]
    key: PrivateKeyInner,
    sec_level: SecurityLevel,
}

#[cfg(not(feature = "decap_key"))]
#[derive(Debug, Eq, PartialEq)]
struct PrivateSeed {
    seed: [u8; 2 * SYMBYTES],
}

#[derive(Debug, Eq, PartialEq)]
struct PrivateKeyInner {
    sk: IndcpaPrivateKey,
    pk: IndcpaPublicKey,
    h_pk: [u8; SYMBYTES],
    z: [u8; SYMBYTES],
}

/// `PublicKey` struct that can only be generated via the [`generate_keypair_512`], [`generate_keypair_768`], or [`generate_keypair_1024`] functions or from the
/// corresponding [`PrivateKey`] struct using the [`get_public_key`](PrivateKey::get_public_key)
/// method and is used to [`encapsulate`](PublicKey::encapsulate) a shared secret.
///
/// Can be packed into a `u8` byte array using the [`pack`](PublicKey::pack) and
/// [`unpack`](PublicKey::unpack) methods. The array used to pack must be of the correct length
/// for the given secrity level, see [`pack`](PublicKey::pack) for more.
#[derive(Debug, Eq, PartialEq)]
pub struct PublicKey {
    pk: IndcpaPublicKey,
    h_pk: [u8; SYMBYTES],
}

/// `Ciphertext` struct that can only be generated by [`encapsulate`](PublicKey::encapsulate).
///
/// Should be converted to bytes using the [`as_bytes`](Ciphertext::as_bytes) method to be transmitted and decapsulated.
pub struct Ciphertext {
    bytes: [u8; MAX_CIPHERTEXT], // max ciphertext_bytes()
    len: usize,
}

impl Ciphertext {
    /// Returns a byte slice of the ciphertext
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// # let (pk, sk) = generate_keypair_768(None).unwrap();
    /// let (ciphertext_obj, shared_secret) = pk.encapsulate(None, None)?;
    /// let ciphertext = ciphertext_obj.as_bytes();
    ///
    /// # Ok::<(), enc_rust::errors::EncryptionDecryptionError>(())
    /// ```
    #[must_use]
    pub fn as_bytes(&self) -> &[u8] {
        &self.bytes[..self.len]
    }
}

fn sha3_256_from(input: &[u8]) -> [u8; SYMBYTES] {
    let mut hash = Sha3_256::new();
    Digest::update(&mut hash, input);

    let output: [u8; SYMBYTES] = hash.finalize().into();
    output
}

fn sha3_512_from(input: &[u8]) -> ([u8; SHAREDSECRETBYTES], [u8; SYMBYTES]) {
    let mut hash = Sha3_512::new();
    Digest::update(&mut hash, input);
    let output = hash.finalize();

    let mut o1 = [0u8; SHAREDSECRETBYTES];
    let mut o2 = [0u8; SYMBYTES];

    o1.copy_from_slice(&output[..SHAREDSECRETBYTES]);
    o2.copy_from_slice(&output[SHAREDSECRETBYTES..]);
    (o1, o2)
}

fn shake256_from(input: &[u8]) -> [u8; SHAREDSECRETBYTES] {
    let mut hash = Shake256::default();
    hash.update(input);
    let mut output = [0u8; SHAREDSECRETBYTES];
    hash.finalize_xof().read(&mut output);
    output
}

// derived new keypair deterministically from a given 64 (2 * 32) byte seed.
fn new_key_from_seed(
    seed: [u8; 2 * SYMBYTES],
    sec_level: SecurityLevel,
) -> Result<(PublicKey, PrivateKeyInner), KeyGenerationError> {
    let (sk, pk) = generate_indcpa_key_pair(&seed[..SYMBYTES], sec_level)?;

    let z: [u8; SYMBYTES] = seed[SYMBYTES..].try_into()?;

    let mut packed_pk = [0u8; MAX_CIPHERTEXT]; // max packed public key size
    pk.pack(&mut packed_pk[..sec_level.indcpa_public_key_bytes()])?;

    let h_pk: [u8; SYMBYTES] = sha3_256_from(&packed_pk[..sec_level.indcpa_public_key_bytes()]);

    Ok((PublicKey { pk, h_pk }, PrivateKeyInner { sk, pk, h_pk, z }))
}

/// Acceptable RNG to be used in encapsulation and key generation must have the
/// [`RngCore`](https://docs.rs/rand_core/latest/rand_core/trait.RngCore.html) and
/// [`CryptoRng`](https://docs.rs/rand_core/latest/rand_core/trait.CryptoRng.html) traits.
pub trait AcceptableRng: RngCore + CryptoRng {}

pub(crate) fn generate_key_pair(
    rng: Option<&mut dyn AcceptableRng>,
    k: K,
) -> Result<(PublicKey, PrivateKey), KeyGenerationError> {
    let mut seed = [0u8; 2 * SYMBYTES];

    if let Some(rng) = rng {
        rng.try_fill_bytes(&mut seed)?;
    } else {
        let mut chacha = ChaCha20Rng::from_entropy();
        chacha.try_fill_bytes(&mut seed)?;
    };

    let sec_level = SecurityLevel::new(k);

    let (pk, _sk_inner) = new_key_from_seed(seed, sec_level)?;

    Ok((
        pk,
        PrivateKey {
            #[cfg(not(feature = "decap_key"))]
            key: PrivateSeed { seed },
            #[cfg(feature = "decap_key")]
            #[allow(clippy::used_underscore_binding)]
            key: _sk_inner,
            sec_level,
        },
    ))
}

/// Generates a new keypair for the 512 Security Parameters.
///
/// # Inputs
/// - `rng`: (Optional) RNG to be used when generating the keypair. Must satisfy the
///   [`RngCore`](https://docs.rs/rand_core/latest/rand_core/trait.RngCore.html) and
///   [`CryptoRng`](https://docs.rs/rand_core/latest/rand_core/trait.CryptoRng.html) traits.
///   If RNG is not present, then
///   [`ChaCha20`](https://docs.rs/rand_chacha/latest/rand_chacha/struct.ChaCha20Rng.html)
///   will be used.
///
/// # Outputs
/// - [`PublicKey`] object
/// - [`PrivateKey`] object
///
/// # Errors
/// Will return a [`KeyGenerationError`] if:
/// - Given invalid K value
/// - RNG fails
///
/// # Example
/// ```
/// # use enc_rust::kem::*;
/// let (pk, sk) = generate_keypair_512(None)?;
///
/// # Ok::<(), enc_rust::errors::KeyGenerationError>(())
/// ```
pub fn generate_keypair_512(
    rng: Option<&mut dyn AcceptableRng>,
) -> Result<(PublicKey, PrivateKey), KeyGenerationError> {
    generate_key_pair(rng, K::Two)
}

/// Generates a new keypair for the 768 Security Parameters.
///
/// # Inputs
/// - `rng`: (Optional) RNG to be used when generating the keypair. Must satisfy the
///   [`RngCore`](https://docs.rs/rand_core/latest/rand_core/trait.RngCore.html) and
///   [`CryptoRng`](https://docs.rs/rand_core/latest/rand_core/trait.CryptoRng.html) traits.
///   If RNG is not present, then
///   [`ChaCha20`](https://docs.rs/rand_chacha/latest/rand_chacha/struct.ChaCha20Rng.html)
///   will be used.
///
/// # Outputs
/// - [`PublicKey`] object
/// - [`PrivateKey`] object
///
/// # Errors
/// Will return a [`KeyGenerationError`] if:
/// - Given invalid K value
/// - RNG fails
///
/// # Example
/// ```
/// # use enc_rust::kem::*;
/// let (pk, sk) = generate_keypair_768(None)?;
///
/// # Ok::<(), enc_rust::errors::KeyGenerationError>(())
/// ```
pub fn generate_keypair_768(
    rng: Option<&mut dyn AcceptableRng>,
) -> Result<(PublicKey, PrivateKey), KeyGenerationError> {
    generate_key_pair(rng, K::Three)
}

/// Generates a new keypair for the 1024 Security Parameters.
///
/// # Inputs
/// - `rng`: (Optional) RNG to be used when generating the keypair. Must satisfy the
///   [`RngCore`](https://docs.rs/rand_core/latest/rand_core/trait.RngCore.html) and
///   [`CryptoRng`](https://docs.rs/rand_core/latest/rand_core/trait.CryptoRng.html) traits.
///   If RNG is not present, then
///   [`ChaCha20`](https://docs.rs/rand_chacha/latest/rand_chacha/struct.ChaCha20Rng.html)
///   will be used.
///
/// # Outputs
/// - [`PublicKey`] object
/// - [`PrivateKey`] object
///
/// # Errors
/// Will return a [`KeyGenerationError`] if:
/// - Given invalid K value
/// - RNG fails
///
/// # Example
/// ```
/// # use enc_rust::kem::*;
/// let (pk, sk) = generate_keypair_1024(None)?;
///
/// # Ok::<(), enc_rust::errors::KeyGenerationError>(())
/// ```
pub fn generate_keypair_1024(
    rng: Option<&mut dyn AcceptableRng>,
) -> Result<(PublicKey, PrivateKey), KeyGenerationError> {
    generate_key_pair(rng, K::Four)
}

impl PrivateKey {
    #[cfg(feature = "decap_key")]
    pub(crate) const fn sec_level(&self) -> SecurityLevel {
        self.key.sk.sec_level()
    }

    /// Returns the corresponding public key for a given private key
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// let (_, sk) = generate_keypair_768(None)?;
    /// let pk = sk.get_public_key();
    ///
    /// # Ok::<(), enc_rust::errors::KeyGenerationError>(())
    /// ```
    #[allow(clippy::missing_panics_doc, clippy::unwrap_used)]
    #[must_use]
    pub fn get_public_key(&self) -> PublicKey {
        #[cfg(not(feature = "decap_key"))]
        {
            let (pk, _) = new_key_from_seed(self.key.seed, self.sec_level).unwrap();

            pk
        }
        #[cfg(feature = "decap_key")]
        {
            PublicKey {
                pk: self.key.pk,
                h_pk: self.key.h_pk,
            }
        }
    }

    /// Packs the private key as bytes and returns it as a 64 byte array
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// let (_, sk) = generate_keypair_768(None).unwrap();
    /// #[cfg(feature = "decap_key")]
    /// {
    ///     let mut sk_bytes = [0u8; 2400];
    ///     sk.pack(&mut sk_bytes)?;
    /// }
    /// #[cfg(not(feature = "decap_key"))]
    /// let sk_bytes = sk.pack();
    ///
    /// # Ok::<(), enc_rust::errors::PackingError>(())
    /// ```
    #[must_use]
    #[cfg(not(feature = "decap_key"))]
    pub const fn pack(&self) -> [u8; 2 * SYMBYTES] {
        self.key.seed
    }

    /// Packs private key into a given buffer
    ///
    /// # Inputs
    /// - `bytes`: Buffer for the private key to be packed into. For corresponding
    ///   security levels, `bytes` should be of length:
    ///
    /// | Security Level | Length |
    /// |----------------|--------|
    /// | 512            | 1632   |
    /// | 768            | 2400   |
    /// | 1024           | 3168   |
    ///
    /// # Errors
    /// Will return a [`PackingError`] if the buffer is of the wrong length
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// let (_, sk) = generate_keypair_768(None).unwrap();
    /// #[cfg(feature = "decap_key")]
    /// {
    ///     let mut sk_bytes = [0u8; 2400];
    ///     sk.pack(&mut sk_bytes)?;
    /// }
    /// #[cfg(not(feature = "decap_key"))]
    /// let sk_bytes = sk.pack();
    ///
    /// # Ok::<(), enc_rust::errors::PackingError>(())
    /// ```
    #[cfg(feature = "decap_key")]
    pub fn pack(&self, bytes: &mut [u8]) -> Result<(), PackingError> {
        let sec_level = self.sec_level();

        if bytes.len() != sec_level.private_key_bytes() {
            return Err(CrystalsError::IncorrectBufferLength(
                bytes.len(),
                sec_level.private_key_bytes(),
            )
            .into());
        }

        let (sk_bytes, rest) = bytes.split_at_mut(sec_level.indcpa_private_key_bytes());
        let (pk_bytes, rest) = rest.split_at_mut(sec_level.indcpa_public_key_bytes());
        let (h_pk_bytes, z_bytes) = rest.split_at_mut(SYMBYTES);
        self.key.sk.pack(sk_bytes)?;
        self.key.pk.pack(pk_bytes)?;
        h_pk_bytes.copy_from_slice(&self.key.h_pk);
        z_bytes.copy_from_slice(&self.key.z);

        Ok(())
    }

    /// Unpacks a buffer of bytes into a [`PrivateKey`]
    ///
    /// # Inputs
    /// - `bytes`: Buffer for the private key to be extracted from
    ///
    /// # Outputs
    /// - [`PrivateKey`] object
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// # let (pk, new_sk) = generate_keypair_768(None).unwrap();
    /// # #[cfg(feature = "decap_key")]
    /// # {
    /// #     let mut sk_bytes = [0u8; 2400];
    /// #     new_sk.pack(&mut sk_bytes)?;
    /// # }
    /// # #[cfg(not(feature = "decap_key"))]
    /// # let sk_bytes = new_sk.pack();
    ///
    /// # #[cfg(not(feature = "decap_key"))]
    /// let sk = PrivateKey::unpack_512(sk_bytes);
    ///
    /// # Ok::<(), enc_rust::errors::PackingError>(())
    /// ```
    #[must_use]
    #[cfg(not(feature = "decap_key"))]
    pub const fn unpack_512(bytes: [u8; 2 * SYMBYTES]) -> Self {
        Self {
            key: PrivateSeed { seed: bytes },
            sec_level: SecurityLevel::new(K::Two),
        }
    }

    /// Unpacks a buffer of bytes into a [`PrivateKey`]
    ///
    /// # Inputs
    /// - `bytes`: Buffer for the private key to be extracted from
    ///
    /// # Outputs
    /// - [`PrivateKey`] object
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// # let (pk, new_sk) = generate_keypair_768(None).unwrap();
    /// # #[cfg(feature = "decap_key")]
    /// # {
    /// #     let mut sk_bytes = [0u8; 2400];
    /// #     new_sk.pack(&mut sk_bytes)?;
    /// # }
    /// # #[cfg(not(feature = "decap_key"))]
    /// # let sk_bytes = new_sk.pack();
    ///
    /// # #[cfg(not(feature = "decap_key"))]
    /// let sk = PrivateKey::unpack_768(sk_bytes);
    ///
    /// # Ok::<(), enc_rust::errors::PackingError>(())
    /// ```
    #[must_use]
    #[cfg(not(feature = "decap_key"))]
    pub const fn unpack_768(bytes: [u8; 2 * SYMBYTES]) -> Self {
        Self {
            key: PrivateSeed { seed: bytes },
            sec_level: SecurityLevel::new(K::Three),
        }
    }

    /// Unpacks a buffer of bytes into a [`PrivateKey`]
    ///
    /// # Inputs
    /// - `bytes`: Buffer for the private key to be extracted from
    ///
    /// # Outputs
    /// - [`PrivateKey`] object
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// # let (pk, new_sk) = generate_keypair_768(None).unwrap();
    /// # #[cfg(feature = "decap_key")]
    /// # {
    /// #     let mut sk_bytes = [0u8; 2400];
    /// #     new_sk.pack(&mut sk_bytes)?;
    /// # }
    /// # #[cfg(not(feature = "decap_key"))]
    /// # let sk_bytes = new_sk.pack();
    ///
    /// # #[cfg(not(feature = "decap_key"))]
    /// let sk = PrivateKey::unpack_1024(sk_bytes);
    ///
    /// # Ok::<(), enc_rust::errors::PackingError>(())
    /// ```
    #[must_use]
    #[cfg(not(feature = "decap_key"))]
    pub const fn unpack_1024(bytes: [u8; 2 * SYMBYTES]) -> Self {
        Self {
            key: PrivateSeed { seed: bytes },
            sec_level: SecurityLevel::new(K::Four),
        }
    }

    /// Unpacks a buffer of bytes into a [`PrivateKey`]
    ///
    /// # Inputs
    /// - `bytes`: Buffer for the private key to be extracted from
    ///
    /// # Outputs
    /// - [`PrivateKey`] object
    ///
    /// # Errors
    /// Will return a [`PackingError`] if the buffer is of the wrong length
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// # let (pk, new_sk) = generate_keypair_768(None).unwrap();
    /// # #[cfg(feature = "decap_key")]
    /// # {
    /// #     let mut sk_bytes = [0u8; 2400];
    /// #     new_sk.pack(&mut sk_bytes)?;
    ///
    ///       let sk = PrivateKey::unpack(&sk_bytes)?;
    /// # }
    ///
    /// # Ok::<(), enc_rust::errors::PackingError>(())
    /// ```
    #[cfg(feature = "decap_key")]
    pub fn unpack(bytes: &[u8]) -> Result<Self, PackingError> {
        let sec_level = match bytes.len() {
            1632 => SecurityLevel::new(K::Two),
            2400 => SecurityLevel::new(K::Three),
            3168 => SecurityLevel::new(K::Four),
            _ => return Err(CrystalsError::IncorrectBufferLength(bytes.len(), 3168).into()),
        };
        let (sk_bytes, rest) = bytes.split_at(sec_level.indcpa_private_key_bytes());
        let (pk_bytes, rest) = rest.split_at(sec_level.indcpa_public_key_bytes());
        let (h_pk_bytes, z_bytes) = rest.split_at(SYMBYTES);

        let sk = IndcpaPrivateKey::unpack(sk_bytes)?;
        let pk = IndcpaPublicKey::unpack(pk_bytes)?;
        let mut h_pk = [0u8; SYMBYTES];
        h_pk.copy_from_slice(h_pk_bytes);
        let mut z = [0u8; SYMBYTES];
        z.copy_from_slice(z_bytes);

        Ok(Self {
            key: PrivateKeyInner { sk, pk, h_pk, z },
            sec_level,
        })
    }

    /// Decapsulates a ciphertext (given as a byte slice) into the shared secret
    ///
    /// # Inputs
    /// - `ciphertext`: Byte slice containing the ciphertext to be decapsulated
    ///
    /// # Outputs
    /// - `[u8; 32]`: The shared secret, a 32 byte array
    ///
    /// # Errors
    /// Will return an [`EncryptionDecryptionError`] if:
    /// - Given invalid ciphertext length
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// # let (pk, sk) = generate_keypair_768(None).unwrap();
    /// # let (ciphertext_obj, secret) = pk.encapsulate(None, None).unwrap();
    /// # let ciphertext = ciphertext_obj.as_bytes();
    /// let shared_secret = sk.decapsulate(ciphertext)?;
    ///
    /// # Ok::<(), enc_rust::errors::EncryptionDecryptionError>(())
    /// ```
    pub fn decapsulate(
        &self,
        ciphertext: &[u8],
    ) -> Result<[u8; SHAREDSECRETBYTES], EncryptionDecryptionError> {
        let valid_bytes = [
            SecurityLevel::new(K::Two).ciphertext_bytes(),
            SecurityLevel::new(K::Three).ciphertext_bytes(),
            SecurityLevel::new(K::Four).ciphertext_bytes(),
        ];

        let sec_level = match ciphertext.len() {
            len if len == valid_bytes[0] => {
                Ok::<SecurityLevel, CrystalsError>(SecurityLevel::new(K::Two))
            }
            len if len == valid_bytes[1] => {
                Ok::<SecurityLevel, CrystalsError>(SecurityLevel::new(K::Three))
            }
            len if len == valid_bytes[2] => {
                Ok::<SecurityLevel, CrystalsError>(SecurityLevel::new(K::Four))
            }
            _ => Err(CrystalsError::InvalidCiphertextLength(ciphertext.len())),
        }?;

        #[cfg(not(feature = "decap_key"))]
        let (_, inner) = new_key_from_seed(self.key.seed, sec_level)?;
        #[cfg(feature = "decap_key")]
        let inner = &self.key;

        let m = inner.sk.decrypt(ciphertext)?;

        let (k, r) = sha3_512_from(&[m, inner.h_pk].concat());

        let k_bar = shake256_from(&[&inner.z, ciphertext].concat());

        let mut ct = [0u8; MAX_CIPHERTEXT]; // max indcpa_bytes()
        inner
            .pk
            .encrypt(&m, &r, &mut ct[..sec_level.indcpa_bytes()])?;

        let equal = ct.ct_eq(ciphertext);

        Ok(k.iter()
            .zip(k_bar.iter())
            .map(|(x, y)| u8::conditional_select(x, y, equal))
            .collect::<ArrayVec<[u8; SHAREDSECRETBYTES]>>()
            .into_inner())
    }
}

impl PublicKey {
    pub(crate) const fn sec_level(&self) -> SecurityLevel {
        self.pk.sec_level()
    }

    /// Packs [`PublicKey`] into a given buffer
    ///
    /// # Inputs
    /// - `bytes`: Buffer for the public key to be packed into. For corresponding
    ///   security levels, `bytes` should be of length:
    ///
    /// | Security Level | Length |
    /// |----------------|--------|
    /// | 512            | 800    |
    /// | 768            | 1184   |
    /// | 1024           | 1568   |
    ///
    /// # Errors
    /// Will return a [`PackingError`] if the buffer is of the wrong length
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// # let (pk, sk) = generate_keypair_768(None).unwrap();
    /// let mut pk_bytes = [0u8; 1184];
    /// pk.pack(&mut pk_bytes)?;
    ///
    /// # Ok::<(), enc_rust::errors::PackingError>(())
    /// ```
    pub fn pack(&self, bytes: &mut [u8]) -> Result<(), PackingError> {
        if bytes.len() != self.sec_level().public_key_bytes() {
            return Err(CrystalsError::IncorrectBufferLength(
                bytes.len(),
                self.sec_level().public_key_bytes(),
            )
            .into());
        }

        self.pk.pack(bytes)?;

        Ok(())
    }

    /// Unpacks a buffer of bytes into a [`PublicKey`]
    ///
    /// # Inputs
    /// - `bytes`: Buffer for the public key to be extracted from
    ///
    /// # Outputs
    /// - [`PublicKey`] object
    ///
    /// # Errors
    /// Will return a [`PackingError`] if the buffer is of the wrong length
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// # let (new_pk, sk) = generate_keypair_768(None).unwrap();
    /// # let mut pk_bytes = [0u8; 1184];
    /// # new_pk.pack(&mut pk_bytes)?;
    /// let pk = PublicKey::unpack(&pk_bytes)?;
    ///
    /// # Ok::<(), enc_rust::errors::PackingError>(())
    /// ```
    pub fn unpack(bytes: &[u8]) -> Result<Self, PackingError> {
        let pk = IndcpaPublicKey::unpack(bytes)?;
        let h_pk = sha3_256_from(bytes);

        Ok(Self { pk, h_pk })
    }

    /// Encapsulates a generated shared secret into a ciphertext to be shared
    ///
    /// # Inputs
    /// - `seed`: (Optional) a 64 byte slice used as a seed for randomness
    /// - `rng`: (Optional) RNG to be used during encapsulation. Must satisfy the
    ///   [`RngCore`](https://docs.rs/rand_core/latest/rand_core/trait.RngCore.html) and
    ///   [`CryptoRng`](https://docs.rs/rand_core/latest/rand_core/trait.CryptoRng.html) traits.
    ///   If RNG is not present, then
    ///   [`ChaCha20`](https://docs.rs/rand_chacha/latest/rand_chacha/struct.ChaCha20Rng.html)
    ///   will be used.
    ///
    /// # Outputs
    /// - [`Ciphertext`] object
    /// - `[u8; 32]`: The shared secret, a 32 byte array
    ///
    /// # Errors
    /// Will return an [`EncryptionDecryptionError`] if:
    /// - Given invalid seed length
    /// - RNG fails
    ///
    /// # Example
    /// ```
    /// # use enc_rust::kem::*;
    /// # let (pk, sk) = generate_keypair_768(None).unwrap();
    /// let (ciphertext_obj, shared_secret) = pk.encapsulate(None, None)?;
    ///
    /// # Ok::<(), enc_rust::errors::EncryptionDecryptionError>(())
    /// ```
    pub fn encapsulate(
        &self,
        seed: Option<&[u8]>,
        rng: Option<&mut dyn AcceptableRng>,
    ) -> Result<(Ciphertext, [u8; SHAREDSECRETBYTES]), EncryptionDecryptionError> {
        let sec_level = self.pk.sec_level();

        let mut m = [0u8; SYMBYTES];
        if let Some(seed) = seed {
            if seed.len() != SYMBYTES {
                return Err(CrystalsError::InvalidSeedLength(seed.len(), SYMBYTES).into());
            }
            m.copy_from_slice(seed);
        } else if let Some(rng) = rng {
            rng.try_fill_bytes(&mut m)?;
        } else {
            let mut chacha = ChaCha20Rng::from_entropy();
            chacha.try_fill_bytes(&mut m)?;
        }

        let (k, r) = sha3_512_from(&[m, self.h_pk].concat());
        let mut bytes = [0u8; MAX_CIPHERTEXT]; // max ciphertext_bytes
        self.pk
            .encrypt(&m, &r, &mut bytes[..sec_level.ciphertext_bytes()])?;

        Ok((
            Ciphertext {
                bytes,
                len: sec_level.ciphertext_bytes(),
            },
            k,
        ))
    }
}