cosmian_kms_cli 5.20.0

use std::{
    fs::File,
    io::prelude::*,
    path::{Path, PathBuf},
};

use clap::Parser;
use cosmian_kms_client::{
    ExportObjectParams, KmsClient, export_object,
    kmip_2_1::{
        kmip_attributes::Attributes,
        kmip_data_structures::KeyWrappingSpecification,
        kmip_types::{
            CryptographicAlgorithm, CryptographicParameters, EncodingOption, KeyFormatType,
        },
        requests::{create_symmetric_key_kmip_object, encrypt_request},
    },
    read_bytes_from_file,
    reexport::cosmian_kms_client_utils::symmetric_utils::DataEncryptionAlgorithm,
};
use cosmian_kms_crypto::crypto::{
    symmetric::symmetric_ciphers::{Mode, SymCipher, encrypt, random_key, random_nonce},
    wrap::wrap_object_with_key,
};
use cosmian_logger::trace;
use zeroize::Zeroizing;

use crate::{
    actions::kms::{
        console, labels::KEY_ID, shared::get_key_uid, symmetric::KeyEncryptionAlgorithm,
    },
    error::{
        KmsCliError,
        result::{KmsCliResult, KmsCliResultHelper},
    },
};

/// Encrypt a file using a symmetric cipher
///
/// Encryption can happen in two ways:
///  - server side: the data is sent to the server and encrypted server side.
///  - client side: the data is encrypted client side using a randomly generated ephemeral key
///    called the data encryption key (DEK). The DEK is then wrapped (i.e., encrypted) server side
///    using the key encryption algorithm and the key encryption key (KEK) identified by the key id.
///    The ephemeral DEK key has a size of 256 bits (512 bits for XTS).
///
/// To encrypt the data server side, do not specify the key encryption algorithm.
///
/// The bytes written to the output file are the concatenation of
///   - if client side encryption is used:
///     - the length of the encapsulated DEK as an unsigned LEB128 integer
///     - the encapsulated DEK
///   - the nonce used for data encryption (or tweak for XTS)
///   - the encrypted data (same size as the plaintext)
///   - the authentication tag generated by the data encryption algorithm (none, for XTS)
///
/// Note: server side encryption is not a streaming call:
/// the data is entirely loaded in memory before being encrypted.
#[derive(Parser, Debug, Default)]
#[clap(verbatim_doc_comment)]
pub struct EncryptAction {
    /// The file to encrypt
    #[clap(required = true, name = "FILE")]
    pub(crate) input_file: PathBuf,

    /// The symmetric key unique identifier.
    /// If not specified, tags should be specified
    #[clap(long = KEY_ID, short = 'k', group = "key-tags")]
    pub(crate) key_id: Option<String>,

    /// The data encryption algorithm.
    /// If not specified, `aes-gcm` is used.
    ///
    /// If no key encryption algorithm is specified, the data will be sent to the server
    /// and will be encrypted server side.
    #[clap(
        long = "data-encryption-algorithm",
        short = 'd',
        default_value = "aes-gcm"
    )]
    pub(crate) data_encryption_algorithm: DataEncryptionAlgorithm,

    /// The optional key encryption algorithm used to encrypt the data encryption key.
    ///
    /// If not specified:
    ///   - the encryption of the data is performed server side.
    ///   - the key id is that of the data encryption key.
    ///
    /// If specified:
    ///  - the data is encrypted client side with the data encryption algorithm, and using
    ///    a randomly generated ephemeral key called the data encryption key (DEK).
    ///  - the DEK is wrapped server side using the key encryption algorithm and the key
    ///    identified by the key id.
    #[clap(long = "key-encryption-algorithm", short = 'e', verbatim_doc_comment)]
    pub(crate) key_encryption_algorithm: Option<KeyEncryptionAlgorithm>,

    /// Tag to use to retrieve the key when no key id is specified.
    /// To specify multiple tags, use the option multiple times.
    #[clap(long = "tag", short = 't', value_name = "TAG", group = "key-tags")]
    pub(crate) tags: Option<Vec<String>>,

    /// The encrypted output file path
    #[clap(required = false, long, short = 'o')]
    pub(crate) output_file: Option<PathBuf>,

    /// Optional nonce/IV (or tweak for XTS) as a hex string.
    /// If not provided, a random value is generated.
    #[clap(required = false, long, short = 'n')]
    pub(crate) nonce: Option<String>,

    /// Optional additional authentication data as a hex string.
    /// This data needs to be provided back for decryption.
    /// This data is ignored with XTS.
    #[clap(required = false, long, short = 'a')]
    pub(crate) authentication_data: Option<String>,
}

impl EncryptAction {
    pub(crate) async fn run(&self, kms_rest_client: KmsClient) -> KmsCliResult<()> {
        // Recover the unique identifier or set of tags
        let id = get_key_uid(self.key_id.as_ref(), self.tags.as_ref(), KEY_ID)?;

        let nonce = self
            .nonce
            .as_deref()
            .map(hex::decode)
            .transpose()
            .with_context(|| "failed to decode the nonce")?;

        let authentication_data = self
            .authentication_data
            .as_deref()
            .map(hex::decode)
            .transpose()
            .with_context(|| "failed to decode the authentication data")?;

        let output_file_name = self
            .output_file
            .clone()
            .unwrap_or_else(|| self.input_file.with_extension("enc"));
        let mut output_file = File::create(&output_file_name)
            .with_context(|| "failed to write the encrypted file")?;

        if let Some(key_encryption_algorithm) = self.key_encryption_algorithm {
            self.client_side_encrypt_with_file(
                kms_rest_client,
                &id,
                key_encryption_algorithm,
                self.data_encryption_algorithm,
                nonce,
                &self.input_file,
                &mut output_file,
                authentication_data,
            )
            .await?;
        } else {
            // Read the file to encrypt
            let plaintext = read_bytes_from_file(&self.input_file)
                .with_context(|| "Cannot read bytes from the file to encrypt")?;
            let (nonce, data, tag) = self
                .server_side_encrypt(
                    &kms_rest_client,
                    &id,
                    self.data_encryption_algorithm.into(),
                    nonce,
                    plaintext,
                    authentication_data,
                )
                .await?;
            if let Some(nonce) = nonce {
                output_file
                    .write_all(&nonce)
                    .with_context(|| "failed to write the nonce")?;
            }
            output_file
                .write_all(&data)
                .context("failed to write the ciphertext")?;
            if let Some(tag) = tag {
                output_file
                    .write_all(&tag)
                    .context("failed to write the authentication tag")?;
            }
        }

        let stdout = format!(
            "The encrypted file is available at {}",
            output_file_name.display()
        );
        console::Stdout::new(&stdout).write()?;

        Ok(())
    }

    /// Encrypt the data using the specified key server side
    /// Returns the nonce, the encrypted data, and the authentication tag
    ///
    /// # Errors
    /// - If the query to the KMS fails
    /// - If the nonce is empty
    /// - If the encrypted data is empty
    /// - If the authentication tag is empty
    pub async fn server_side_encrypt(
        &self,
        kms_rest_client: &KmsClient,
        data_encryption_key_id: &str,
        cryptographic_parameters: CryptographicParameters,
        nonce: Option<Vec<u8>>,
        plaintext: Vec<u8>,
        authenticated_data: Option<Vec<u8>>,
    ) -> Result<(Option<Vec<u8>>, Vec<u8>, Option<Vec<u8>>), KmsCliError> {
        // Create the kmip query
        let encrypt_request = encrypt_request(
            data_encryption_key_id,
            None,
            plaintext,
            nonce,
            authenticated_data,
            Some(cryptographic_parameters),
        )?;

        // Query the KMS with your kmip data and get the key pair ids
        let encrypt_response = kms_rest_client
            .encrypt(encrypt_request)
            .await
            .with_context(|| "Can't execute the query on the kms server")?;

        // extract the nonce and write it
        let nonce = encrypt_response.i_v_counter_nonce; // Some encryption modes (like RFC5649) don't use nonces

        // extract the ciphertext and write it
        let data = encrypt_response
            .data
            .context("The encrypted data is empty")?;

        // extract the authentication tag (only mandatory for AEAD modes: GCM, GCM-SIV, ChaCha20-Poly1305)
        // Note: For server_side_encrypt, the algorithm used is determined by the cryptographic_parameters
        // which may be different from self.data_encryption_algorithm (e.g., for key wrapping)
        let authentication_tag = encrypt_response.authenticated_encryption_tag; // Some encryption modes don't use authentication tags
        Ok((nonce, data, authentication_tag))
    }

    /// Encrypt a file using a symmetric stream cipher
    /// and return the ephemeral key
    #[expect(clippy::too_many_arguments)]
    async fn client_side_encrypt_with_file(
        &self,
        kms_rest_client: KmsClient,
        kek_id: &str,
        key_encryption_algorithm: KeyEncryptionAlgorithm,
        data_encryption_algorithm: DataEncryptionAlgorithm,
        nonce: Option<Vec<u8>>,
        input_file_name: &Path,
        output_file: &mut File,
        aad: Option<Vec<u8>>,
    ) -> KmsCliResult<Zeroizing<Vec<u8>>> {
        // Additional authenticated data (AAD) for AEAD ciphers
        // (empty for XTS)
        let aad = match data_encryption_algorithm {
            DataEncryptionAlgorithm::AesXts | DataEncryptionAlgorithm::AesCbc => vec![],
            DataEncryptionAlgorithm::AesGcm => aad.unwrap_or_default(),
            #[cfg(feature = "non-fips")]
            DataEncryptionAlgorithm::Chacha20Poly1305 | DataEncryptionAlgorithm::AesGcmSiv => {
                aad.unwrap_or_default()
            }
        };

        // Generate an ephemeral key (DEK) and wrap it with the KEK.
        let (dek, encapsulation) = self
            .server_side_kem_encapsulation(
                kms_rest_client,
                kek_id,
                key_encryption_algorithm,
                data_encryption_algorithm,
            )
            .await?;

        // write the encapsulation to the output file, starting with the length of the encapsulation
        // as an unsigned LEB128 integer
        leb128::write::unsigned(output_file, u64::try_from(encapsulation.len())?)?;
        output_file.write_all(&encapsulation)?;

        // Determine the DEM parameters
        let cryptographic_parameters: CryptographicParameters = data_encryption_algorithm.into();
        let cipher = SymCipher::from_algorithm_and_key_size(
            cryptographic_parameters
                .cryptographic_algorithm
                .ok_or_else(|| {
                    KmsCliError::Default(
                        "No data encryption cryptographic algorithm specified".to_owned(),
                    )
                })?,
            cryptographic_parameters.block_cipher_mode,
            dek.len(),
        )?;

        // we need a nonce (or tweak)
        let nonce = match nonce {
            Some(n) => n,
            None => random_nonce(cipher)?,
        };
        output_file.write_all(&nonce)?;

        // instantiate the stream cipher
        let mut stream_cipher = cipher.stream_cipher(Mode::Encrypt, &dek, &nonce, &aad)?;
        // process the data read from the file by 4096 chunks and write the encrypted data
        let mut file = File::open(input_file_name)?;
        let mut chunk = vec![0; 2 ^ 16]; //64K
        loop {
            let bytes_read = file.read(&mut chunk)?;
            if bytes_read == 0 {
                break;
            }
            chunk.truncate(bytes_read);
            let ciphertext = stream_cipher.update(&chunk)?;
            output_file.write_all(&ciphertext)?;
        }
        // finalize the encryption and write the remaining bytes
        let (remaining, tag) = stream_cipher.finalize_encryption()?;
        output_file.write_all(&remaining)?;
        // write the tag
        output_file.write_all(&tag)?;
        output_file.flush()?;
        Ok(dek)
    }

    /// Generate an ephemeral key (DEK) and wrap it with the KEK.
    /// This encapsulation has the following format:
    /// | `kem_nonce` | `kem_ciphertext` | `kem_tag` |
    ///
    /// # Errors
    /// - If the cryptographic algorithm is not specified
    pub async fn server_side_kem_encapsulation(
        &self,
        kms_rest_client: KmsClient,
        kek_id: &str,
        key_encryption_algorithm: KeyEncryptionAlgorithm,
        data_encryption_algorithm: DataEncryptionAlgorithm,
    ) -> KmsCliResult<(Zeroizing<Vec<u8>>, Vec<u8>)> {
        // Generate the ephemeral key (DEK)
        let dek = match data_encryption_algorithm {
            DataEncryptionAlgorithm::AesGcm => random_key(SymCipher::Aes256Gcm)?,
            DataEncryptionAlgorithm::AesCbc => random_key(SymCipher::Aes256Cbc)?,
            #[cfg(feature = "non-fips")]
            DataEncryptionAlgorithm::Chacha20Poly1305 => random_key(SymCipher::Chacha20Poly1305)?,
            #[cfg(feature = "non-fips")]
            DataEncryptionAlgorithm::AesGcmSiv => random_key(SymCipher::Aes256Gcm)?,
            DataEncryptionAlgorithm::AesXts => random_key(SymCipher::Aes256Xts)?,
        };

        // Wrap the DEK with the KEK
        let (kem_nonce, kem_ciphertext, kem_tag) = self
            .server_side_encrypt(
                &kms_rest_client,
                kek_id,
                key_encryption_algorithm.into(),
                None,
                dek.to_vec(),
                None,
            )
            .await?;

        let encapsulation: Vec<u8> = [
            kem_nonce.unwrap_or_default(),
            kem_ciphertext,
            kem_tag.unwrap_or_default(),
        ]
        .concat();
        Ok((dek, encapsulation))
    }

    /// Generate an ephemeral key (DEK) and wrap it with the KEK.
    /// The encapsulation depends on the Key Encryption Algorithm which is provided by the KMIP symmetric key.
    /// Either AES-GCM or RFC5649 will be used.
    /// This encapsulation has the following format:
    /// | `kem_nonce` | `kem_ciphertext` | `kem_tag` |
    ///
    /// # Errors
    /// - If the cryptographic algorithm is not specified
    pub async fn client_side_kem_encapsulation(
        &self,
        kms_rest_client: &KmsClient,
        kek_id: &str,
        data_encryption_algorithm: DataEncryptionAlgorithm,
    ) -> KmsCliResult<(Zeroizing<Vec<u8>>, Vec<u8>)> {
        trace!("data_encryption_algorithm: {data_encryption_algorithm}");
        // Generate the ephemeral key (DEK)
        let dek: Zeroizing<Vec<u8>> = match data_encryption_algorithm {
            DataEncryptionAlgorithm::AesCbc => random_key(SymCipher::Aes256Cbc)?,
            DataEncryptionAlgorithm::AesGcm => random_key(SymCipher::Aes256Gcm)?,
            #[cfg(feature = "non-fips")]
            DataEncryptionAlgorithm::Chacha20Poly1305 => random_key(SymCipher::Chacha20Poly1305)?,
            #[cfg(feature = "non-fips")]
            DataEncryptionAlgorithm::AesGcmSiv => random_key(SymCipher::Aes256Gcm)?,
            DataEncryptionAlgorithm::AesXts => random_key(SymCipher::Aes256Xts)?,
        };
        trace!("dek (len={}): {dek:?}", dek.len());

        // First export the KEK locally
        let wrapping_key = export_object(
            kms_rest_client,
            kek_id,
            ExportObjectParams {
                key_format_type: Some(KeyFormatType::TransparentSymmetricKey),
                ..ExportObjectParams::default()
            },
        )
        .await?
        .1;
        // Create the KMIP object corresponding to the DEK
        let mut dek_object = create_symmetric_key_kmip_object(
            kms_rest_client.config.vendor_id.as_str(),
            &dek,
            &Attributes {
                cryptographic_algorithm: Some(CryptographicAlgorithm::AES),
                ..Default::default()
            },
        )?;

        // Wrap the DEK with the KEK
        wrap_object_with_key(
            &mut dek_object,
            &wrapping_key,
            &KeyWrappingSpecification {
                encoding_option: Some(EncodingOption::NoEncoding),
                ..Default::default()
            },
        )?;

        let encapsulation = dek_object.key_block()?.wrapped_key_bytes()?;

        Ok((dek, encapsulation.to_vec()))
    }

    /// Encrypt a buffer using a symmetric stream cipher
    /// and return the ciphertext. Ciphertext format is:
    /// | `kem_nonce` | `kem_ciphertext` | `kem_tag` | ciphertext |
    /// # Errors
    /// - If the cryptographic algorithm is not specified
    /// - If the nonce cannot be generated
    /// - If the ciphertext cannot be generated
    /// - If the tag cannot be generated
    pub fn client_side_encrypt_with_buffer(
        &self,
        dek: &Zeroizing<Vec<u8>>,
        encapsulation: &[u8],
        data_encryption_algorithm: DataEncryptionAlgorithm,
        nonce: Option<Vec<u8>>,
        plaintext: &[u8],
        aad: Option<Vec<u8>>,
    ) -> KmsCliResult<Vec<u8>> {
        // Additional authenticated data (AAD) for AEAD ciphers
        // (empty for XTS or CBC modes)
        let aad = match data_encryption_algorithm {
            DataEncryptionAlgorithm::AesXts | DataEncryptionAlgorithm::AesCbc => vec![],
            DataEncryptionAlgorithm::AesGcm => aad.unwrap_or_default(),
            #[cfg(feature = "non-fips")]
            DataEncryptionAlgorithm::Chacha20Poly1305 | DataEncryptionAlgorithm::AesGcmSiv => {
                aad.unwrap_or_default()
            }
        };

        // write the encapsulation to the output file, starting with the length of the encapsulation
        // as an unsigned LEB128 integer
        let mut output_buffer = Vec::with_capacity(encapsulation.len() + 2 * plaintext.len());
        let encapsulation_len = u64::try_from(encapsulation.len())?;
        leb128::write::unsigned(&mut output_buffer, encapsulation_len)?;
        output_buffer.write_all(encapsulation)?;

        // Determine the DEM parameters
        let cryptographic_parameters: CryptographicParameters = data_encryption_algorithm.into();
        let sym_cipher = SymCipher::from_algorithm_and_key_size(
            cryptographic_parameters
                .cryptographic_algorithm
                .ok_or_else(|| {
                    KmsCliError::Default(
                        "No data encryption cryptographic algorithm specified".to_owned(),
                    )
                })?,
            cryptographic_parameters.block_cipher_mode,
            dek.len(),
        )?;

        // we need a nonce (or tweak)
        let nonce = match nonce {
            Some(n) => n,
            None => random_nonce(sym_cipher)?,
        };
        output_buffer.write_all(&nonce)?;

        let (ciphertext, tag) = encrypt(sym_cipher, dek, &nonce, &aad, plaintext, None)?;
        // write the ciphertext and tag
        output_buffer.write_all(&ciphertext)?;
        output_buffer.write_all(&tag)?;
        output_buffer.flush()?;
        Ok(output_buffer)
    }
}