rgp 0.7.0

/*
Copyright (c) 2026 Ordinary Labs, LLC

SPDX-License-Identifier: AGPL-3.0-only
*/

use super::super::{
    KEY_SIZE, NONCE_SIZE, base_decrypt, base_encrypt, bytes_to_usize, usize_to_bytes,
};

use crate::ZEROED_KEY;
use anyhow::bail;
use chacha20::cipher::Key;
use chacha20::{XChaCha20, XNonce, cipher::StreamCipher};
use chacha20poly1305::{XChaCha20Poly1305, aead::Generate};
use classic_mceliece_rust::{
    CRYPTO_CIPHERTEXTBYTES as KEM_CIPHERTEXT_SIZE, CRYPTO_PUBLICKEYBYTES as KEM_PUB_KEY_SIZE,
    CRYPTO_SECRETKEYBYTES as KEM_SECRET_KEY_SIZE, Ciphertext, PublicKey as KemPublicKey,
    SecretKey as KemSecretKey, decapsulate, encapsulate, keypair as kem_keypair,
};
use rand::rngs::OsRng;
use std::io::{BufReader, Read, Seek};
#[cfg(feature = "multi-thread")]
use std::sync::mpsc::channel;
use x25519_dalek::StaticSecret;

/// #
/// **ENCRYPTED FORMAT:**
/// - nonce = 24 bytes
/// - keys count
///     - IF 0..=127
///         - is single byte = 1 bit (set)
///         - count = 7 bits
///     - ELSE
///         - is single byte = 1 bit (unset)
///         - int size = 2 bits
///         - count = 8-64 bits
/// - encrypted copies of content key + ciphertext = `pub_keys.len()` * (32 bytes + 96 bytes)
/// - encrypted content = `content.len()`
/// - signature = 64 bytes (encrypted along with the content)
/// - Poly1305 MAC = 16 bytes
/// - mode = 1 byte (set to `KEM_MODE`)
///
/// **PROCESS:**
/// 1. Generate one-time components
///     - nonce
///     - content key
/// 2. Sign plaintext to generate content signature
/// 3. Encrypt plaintext and content signature with content key
/// 4. Encrypt content key for all recipients
///     - Generate ciphertext and encapsulated key with recipient's public key
///     - Encrypt content key with encapsulated key
///     - Append ciphertext to encrypted content key
pub const KEM_MODE: u8 = 5;

/// #
/// **ENCRYPTED FORMAT:**
/// - nonce = 24 bytes
/// - keys count
///     - IF 0..=127
///         - is single byte = 1 bit (set)
///         - count = 7 bits
///     - ELSE
///         - is single byte = 1 bit (unset)
///         - int size = 2 bits
///         - count = 8-64 bits
/// - encrypted copies of content key + ciphertext = `pub_keys.len()` * (32 bytes + 96 bytes)
/// - encrypted content = `content.len()`
/// - signature = 64 bytes (encrypted along with the content)
/// - Poly1305 MAC = 16 bytes
/// - mode = 1 byte (set to `KEM_WITH_DH_HYBRID_MODE`)
///
/// **PROCESS:**
/// 1. Generate one-time components
///     - nonce
///     - content key
/// 2. Sign plaintext to generate content signature
/// 3. Encrypt plaintext and content signature with content key
/// 4. Encrypt content key for all recipients
///     - Generate ciphertext and encapsulated key with recipient's public key for `Kem`
///     - Generate shared secret with recipient's public key and sender's private key for `Dh`
///     - HMAC the `Kem` encapsulated key with the `Dh` shared secret as the key
///     - Encrypt content key with the HMAC result
///     - Append ciphertext to encrypted content key
pub const KEM_WITH_DH_HYBRID_MODE: u8 = 6;

/// generates `Kem` pub/priv key pairs.
///
///```rust
/// use rgp::generate_kem_keys;
///
/// let (secret_key, pub_key) = generate_kem_keys();
///
/// assert_eq!(secret_key.len(), 6492);
/// assert_eq!(pub_key.len(), 261120);
///```
#[allow(clippy::large_stack_arrays)]
pub fn generate_kem_keys() -> ([u8; KEM_SECRET_KEY_SIZE], [u8; KEM_PUB_KEY_SIZE]) {
    let mut public_key_buf = [0u8; KEM_PUB_KEY_SIZE];
    let mut secret_key_buf = [0u8; KEM_SECRET_KEY_SIZE];

    let (pub_key, secret_key) = kem_keypair(&mut public_key_buf, &mut secret_key_buf, &mut OsRng);

    (*secret_key.as_array(), *pub_key.as_array())
}

/// for reading a large volume of Classic `McEliece` public keys
pub struct KemKeyReader<R: Read + Seek> {
    pub reader: BufReader<R>,

    /// for Kem + Dh hybrid
    pub dh_priv_key: Option<[u8; KEY_SIZE]>,
}

impl<R: Read + Seek> KemKeyReader<R> {
    /// public key reader with a buffer size of 261120.
    ///
    /// for files that contain only `McEliece` public keys all in one line.
    pub fn new(source: R) -> Self {
        KemKeyReader {
            reader: BufReader::with_capacity(KEM_PUB_KEY_SIZE, source),
            dh_priv_key: None,
        }
    }

    /// public key reader with a buffer size of 261120 + 32.
    ///
    /// for files that contain `McEliece` public paired with their
    /// Diffie-Hellman counterparts (i.e Vec<...[u8; ...`mc_pub`, ...`dh_pub`]>)
    pub fn new_dh_hybrid(dh_priv_key: [u8; KEY_SIZE], source: R) -> Self {
        KemKeyReader {
            reader: BufReader::with_capacity(KEM_PUB_KEY_SIZE + KEY_SIZE, source),
            dh_priv_key: Some(dh_priv_key),
        }
    }
}

/// per-recipient content key encryption.
#[allow(clippy::type_complexity, clippy::large_stack_arrays)]
fn kem_encrypt_keys<R: Read + Seek>(
    key_reader: &mut KemKeyReader<R>,
    nonce: &[u8; 24],
    content_key: &[u8; 32],
) -> anyhow::Result<((usize, [u8; 9]), Vec<u8>)> {
    use chacha20::cipher::KeyIvInit;

    let mut key_pos = 0;
    let mut out = vec![];

    // Kem + Dh hybrid
    if let Some(dh_priv_key) = key_reader.dh_priv_key {
        let dh_priv_key = StaticSecret::from(dh_priv_key);

        let mut buf = [0u8; KEM_PUB_KEY_SIZE + KEY_SIZE];

        while key_reader.reader.read_exact(&mut buf).is_ok() {
            let mut kem_pub_key = buf[0..KEM_PUB_KEY_SIZE].try_into()?;
            let kem_pub_key = KemPublicKey::from(&mut kem_pub_key);

            let mut kem_shared_secret_buf = [0u8; KEY_SIZE];
            let (kem_ciphertext, kem_shared_secret) =
                encapsulate(&kem_pub_key, &mut kem_shared_secret_buf, &mut OsRng);

            let mut key = *kem_shared_secret.as_array();

            let dh_pub_key: [u8; KEY_SIZE] =
                buf[KEM_PUB_KEY_SIZE..KEM_PUB_KEY_SIZE + KEY_SIZE].try_into()?;

            if dh_pub_key == ZEROED_KEY {
                bail!("verifier cannot be all 0s");
            }

            let shared_secret = dh_priv_key.diffie_hellman(&dh_pub_key.into());
            if !shared_secret.was_contributory() {
                bail!("non-contributory shared secret");
            }

            // HMAC the KEM shared secret with the Diffie-Hellman shared secret
            {
                use blake2::digest::{FixedOutput, KeyInit, Mac};

                key = blake2::Blake2sMac256::new_from_slice(&shared_secret.to_bytes())?
                    .chain_update(key)
                    .finalize_fixed()
                    .into();
            }

            let mut key_cipher = XChaCha20::new(&key.into(), nonce.into());

            let mut content_key = *content_key;
            key_cipher.apply_keystream(&mut content_key);

            out.extend(content_key);
            out.extend(kem_ciphertext.as_array());

            key_pos += 1;
        }
    } else {
        // raw Kem

        let mut buf = [0u8; KEM_PUB_KEY_SIZE];

        while key_reader.reader.read_exact(&mut buf).is_ok() {
            let kem_pub_key = KemPublicKey::from(&mut buf);

            let mut kem_shared_secret_buf = [0u8; KEY_SIZE];
            let (kem_ciphertext, kem_shared_secret) =
                encapsulate(&kem_pub_key, &mut kem_shared_secret_buf, &mut OsRng);

            let key = *kem_shared_secret.as_array();

            let mut key_cipher = XChaCha20::new(&key.into(), nonce.into());

            let mut content_key = *content_key;
            key_cipher.apply_keystream(&mut content_key);

            out.extend(content_key);
            out.extend(kem_ciphertext.as_array());

            key_pos += 1;
        }
    }

    let header = usize_to_bytes(key_pos)?;

    Ok((header, out))
}

/// kem encryption.
pub fn kem_encrypt<R: Read + Seek>(
    fingerprint: [u8; 32],
    mut content: Vec<u8>,
    key_reader: &mut KemKeyReader<R>,
) -> anyhow::Result<(Vec<u8>, [u8; KEY_SIZE])> {
    let nonce: [u8; NONCE_SIZE] = XNonce::generate().into();
    let mut out = nonce.to_vec();

    let key = Key::<XChaCha20Poly1305>::generate();

    #[cfg(feature = "multi-thread")]
    let (sender, receiver) = channel();

    #[cfg(feature = "multi-thread")]
    rayon::spawn(move || {
        let encrypted_content = base_encrypt(&fingerprint, &nonce, &key.into(), &mut content);
        sender.send(encrypted_content).expect("failed to send");
    });

    let ((size, bytes), keys) = kem_encrypt_keys(key_reader, &nonce, &key.into())?;
    out.extend_from_slice(&bytes[..size]);
    out.extend(keys);

    #[cfg(feature = "multi-thread")]
    let encrypted_content = receiver.recv()??;
    #[cfg(not(feature = "multi-thread"))]
    let encrypted_content = base_encrypt(&fingerprint, &nonce, &key.into(), &mut content)?;

    out.extend(encrypted_content);

    if key_reader.dh_priv_key.is_some() {
        out.push(KEM_WITH_DH_HYBRID_MODE);
    } else {
        out.push(KEM_MODE);
    }

    Ok((out, key.into()))
}

/// kem decryption.
pub fn kem_decrypt(
    verifier: Option<&[u8; 32]>,
    encrypted_content: &[u8],

    mut encrypted_key: [u8; KEY_SIZE],
    ciphertext: [u8; KEM_CIPHERTEXT_SIZE],
    mut secret_key: [u8; KEM_SECRET_KEY_SIZE],
    dh_components: Option<([u8; KEY_SIZE], [u8; KEY_SIZE])>,
) -> anyhow::Result<(Vec<u8>, [u8; KEY_SIZE])> {
    let nonce: [u8; NONCE_SIZE] = match encrypted_content[0..NONCE_SIZE].try_into() {
        Ok(n) => n,
        Err(_) => bail!("failed to convert to bytes"),
    };

    let encrypted_content = &encrypted_content[NONCE_SIZE..];

    let secret_key = KemSecretKey::from(&mut secret_key);
    let ciphertext = Ciphertext::from(ciphertext);

    let mut kem_shared_secret_buf = [0u8; KEY_SIZE];
    let mut key = *decapsulate(&ciphertext, &secret_key, &mut kem_shared_secret_buf).as_array();

    if let Some((pub_key, priv_key)) = dh_components {
        if pub_key == ZEROED_KEY {
            bail!("verifier cannot be all 0s");
        }

        let priv_key = StaticSecret::from(priv_key);
        let shared_secret = priv_key.diffie_hellman(&pub_key.into());

        if !shared_secret.was_contributory() {
            bail!("non-contributory shared secret");
        }

        // HMAC the KEM shared secret with the Diffie-Hellman shared secret
        {
            use blake2::digest::{FixedOutput, KeyInit, Mac};

            key = blake2::Blake2sMac256::new_from_slice(&shared_secret.to_bytes())?
                .chain_update(key)
                .finalize_fixed()
                .into();
        }
    }

    let mut key_cipher = {
        use chacha20::cipher::KeyIvInit;
        XChaCha20::new(&key.into(), &nonce.into())
    };

    key_cipher.apply_keystream(&mut encrypted_key);

    base_decrypt(verifier, &nonce, &encrypted_key, encrypted_content)
}

/// extract kem components.
pub fn kem_extract(
    position: usize,
    encrypted_content: &mut Vec<u8>,
) -> anyhow::Result<([u8; KEY_SIZE], [u8; KEM_CIPHERTEXT_SIZE])> {
    let (keys_count_size, keys_count) =
        bytes_to_usize(&encrypted_content[NONCE_SIZE..NONCE_SIZE + 9])?;

    let keys_start = NONCE_SIZE + keys_count_size;
    let encrypted_key_start = keys_start + (position * (KEY_SIZE + KEM_CIPHERTEXT_SIZE));

    let content_key: [u8; KEY_SIZE] =
        encrypted_content[encrypted_key_start..encrypted_key_start + KEY_SIZE].try_into()?;

    let ciphertext: [u8; KEM_CIPHERTEXT_SIZE] = encrypted_content
        [encrypted_key_start + KEY_SIZE..encrypted_key_start + (KEY_SIZE + KEM_CIPHERTEXT_SIZE)]
        .try_into()?;

    let encrypted_content_start = keys_start + (keys_count * (KEY_SIZE + KEM_CIPHERTEXT_SIZE));

    encrypted_content.copy_within(encrypted_content_start.., NONCE_SIZE);
    encrypted_content.truncate(
        encrypted_content.len()
            - keys_count_size
            - (keys_count * (KEY_SIZE + KEM_CIPHERTEXT_SIZE))
            - 1,
    );

    Ok((content_key, ciphertext))
}

#[cfg(test)]
mod tests {
    use super::{KemKeyReader, kem_decrypt, kem_encrypt, kem_extract};
    use std::fs::{File, OpenOptions, remove_file};
    use std::io::Write;
    use std::path::Path;

    #[test]
    fn test_kem() -> anyhow::Result<()> {
        let pub_key_path = Path::new("test_kem_pub_keys");
        if pub_key_path.exists() {
            remove_file(pub_key_path)?;
        }

        let (fingerprint, verifier) = crate::generate_fingerprint();

        let (secret_key, pub_key) = crate::generate_kem_keys();

        let content = vec![0u8; 1024];

        let mut pub_keys_file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(pub_key_path)?;

        pub_keys_file.write_all(&pub_key)?;
        pub_keys_file.flush()?;

        let mut key_reader = KemKeyReader::new(File::open(pub_key_path)?);

        let (mut encrypted_content, content_key) =
            kem_encrypt(fingerprint, content.clone(), &mut key_reader)?;

        let (encrypted_key, ciphertext) = kem_extract(0, &mut encrypted_content)?;

        let (decrypted_content, decrypted_content_key) = kem_decrypt(
            Some(&verifier),
            &encrypted_content,
            encrypted_key,
            ciphertext,
            secret_key,
            None,
        )?;

        assert_eq!(content, decrypted_content);
        assert_eq!(content_key, decrypted_content_key);

        remove_file(pub_key_path)?;

        Ok(())
    }

    #[test]
    fn test_kem_with_dh_hybrid() -> anyhow::Result<()> {
        let pub_key_path = Path::new("test_kem_with_dh_hybrid_pub_keys");
        if pub_key_path.exists() {
            remove_file(pub_key_path)?;
        }

        let (fingerprint, verifier) = crate::generate_fingerprint();

        let (kem_secret_key, kem_pub_key) = crate::generate_kem_keys();

        let (sender_dh_priv_key, sender_dh_pub_key) = crate::generate_dh_keys();
        let (recipient_dh_priv_key, recipient_dh_pub_key) = crate::generate_dh_keys();

        let content = vec![0u8; 1024];

        let mut pub_keys_file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(pub_key_path)?;

        pub_keys_file.write_all(&kem_pub_key)?;
        pub_keys_file.write_all(&recipient_dh_pub_key)?;
        pub_keys_file.flush()?;

        let mut key_reader =
            KemKeyReader::new_dh_hybrid(sender_dh_priv_key, File::open(pub_key_path)?);

        let (mut encrypted_content, content_key) =
            kem_encrypt(fingerprint, content.clone(), &mut key_reader)?;

        let (encrypted_key, ciphertext) = kem_extract(0, &mut encrypted_content)?;

        let (decrypted_content, decrypted_content_key) = kem_decrypt(
            Some(&verifier),
            &encrypted_content,
            encrypted_key,
            ciphertext,
            kem_secret_key,
            Some((sender_dh_pub_key, recipient_dh_priv_key)),
        )?;

        assert_eq!(content, decrypted_content);
        assert_eq!(content_key, decrypted_content_key);

        remove_file(pub_key_path)?;

        Ok(())
    }
}