//! Encrypt
//! ----------------
//! KMS uses this API to encrypt data using a key in an external key manager.
use std::sync::Arc;

use actix_web::{
    HttpRequest, HttpResponse, post,
    web::{Data, Json, Path},
};
use base64::{Engine, engine::general_purpose::STANDARD};
use cosmian_kms_server_database::reexport::{
    cosmian_kmip::kmip_2_1::{
        kmip_operations::Encrypt,
        kmip_types::{CryptographicAlgorithm, CryptographicParameters, UniqueIdentifier},
    },
    cosmian_kms_crypto::crypto::symmetric::symmetric_ciphers::AES_256_GCM_IV_LENGTH,
};
use openssl::{rand::rand_bytes, sha::Sha256};
use serde::{Deserialize, Serialize};
use tracing::{debug, info};
use zeroize::Zeroizing;

use crate::{
    core::KMS,
    error::KmsError,
    result::KResult,
    routes::aws_xks::{
        encrypt_decrypt::{CdivAlgorithm, EncryptionAlgorithm, RequestMetadata},
        error::{XksErrorName, XksErrorReply},
    },
};

/// KMS uses this API to encrypt data using a key in an external key manager.
///
/// Example:
/// ```json
/// {
///    "requestMetadata": {
///    "awsPrincipalArn": "arn:aws:iam::123456789012:user/Alice",
///    "kmsKeyArn": "arn:aws:kms:us-east-2:123456789012:/key/1234abcd-12ab-34cd-56ef-1234567890ab",
///    "kmsOperation": "Encrypt",
///    "kmsRequestId": "4112f4d6-db54-4af4-ae30-c55a22a8dfae",
///    "kmsViaService": "ebs"
///    },
///    "additionalAuthenticatedData": "cHJvamVjdD1uaWxlLGRlcGFydG1lbnQ9bWFya2V0aW5n",
///    "plaintext": "SGVsbG8gV29ybGQh",
///    "encryptionAlgorithm": "AES_GCM",
///    "ciphertextDataIntegrityValueAlgorithm": "SHA_256"
/// }
/// ```
#[derive(Deserialize, Debug, Serialize)]
#[allow(non_snake_case)]
pub(crate) struct EncryptRequest {
    /// The HTTP body of the request contains requestMetadata fields
    /// that provide additional context on the request being made.
    /// This information is helpful for auditing and for implementing
    /// an optional secondary layer of authorization at the XKS Proxy
    /// (see a later section on Authorization).
    /// There is no expectation for the XKS Proxy to validate any information
    /// included in the requestMetadata beyond validating the signature
    /// that covers the entire request payload.
    pub requestMetadata: RequestMetadata,
    ///  Base64-encoded plaintext provided to external key manager for encryption.
    /// The proxy MUST support the ability to process up to 4300 bytes of plaintext data.
    /// Note that Base64 encoding of 4300 bytes of binary data will result in a string
    /// that is 5736 bytes.
    /// Plaintext passed to the encrypt API MUST NOT be logged at XKS Proxy
    /// or the external key manager.
    /// This field is REQUIRED.
    pub plaintext: String,
    /// Specifies the algorithm that will be used for encryption.
    /// For the v1 specification, this MUST be `AES_GCM`.
    /// This field is REQUIRED.
    pub encryptionAlgorithm: EncryptionAlgorithm,
    /// AES-GCM is an example of an AEAD (Authenticated Encryption with Additional Data) cipher
    /// for which the encrypt operation produces an authenticationTag in addition to the ciphertext.
    /// The authenticationTag can be used to ensure the integrity of the ciphertext
    /// and additional data passed as AAD.
    /// For a decrypt call to succeed, the same AAD that was used to create the ciphertext
    /// must be supplied to the decrypt operation.
    /// This field is OPTIONAL.
    /// When present, this field MUST be specified as a Base64 encoded string
    /// and used as the Additional Authenticated Data (AAD) input to the AES-GCM operation
    /// inside the external key manager.
    /// The XKS Proxy MUST be able to handle AAD values up to 8192 bytes in length
    /// (the Base64 encoding of 8192 bytes will be 10924 bytes).
    pub additionalAuthenticatedData: Option<String>,
    /// Indicates the hashing algorithm to be used in the computation
    /// of the Ciphertext Data Integrity Value (CDIV).
    /// For the first version (v1) of this specification, this MUST be "`SHA_256`".
    /// This field is OPTIONAL.
    /// When present, the XKS Proxy MUST return a ciphertextDataIntegrityValue field
    /// in its response as described below.
    pub ciphertextDataIntegrityValueAlgorithm: Option<CdivAlgorithm>,
}

/// The HTTP response body contains the keySpec, keyUsage, and keyStatus fields.
/// ```json
/// {
///     "authenticationTag": "vBxN2ncH1oEkR8WVXpmyYQ==",
///     "ciphertext": "ghxkK1txeDNn3q8Y",
///     "ciphertextDataIntegrityValue": "qHA/ImC9h5HsLRXqCyPmWgYx7tzyoTplzILbP0fPXsc=",
///     "ciphertextMetadata": "a2V5X3ZlcnNpb249MQ==",
///     "initializationVector": "HMrlRw85cAJUd5Ax"
/// }
/// ```
#[derive(Serialize, Default, Deserialize)]
#[allow(non_snake_case)]
pub(crate) struct EncryptResponse {
    /// Base64 encoded ciphertext generated by the external key manager from provided plaintext.
    /// Since `AES_GCM` is a stream cipher, the length of the ciphertext
    /// MUST be the same as the length of the plaintext
    pub ciphertext: String,
    /// The XKS Proxy MAY return up to 20 bytes of ciphertext metadata for internal housekeeping,
    /// e.g. an external key manager may implement automatic key rotation
    /// and use the extra bytes to encode versioning of the key material.
    /// This is an OPTIONAL, vendor-specific field.
    /// When present, the size of the field MUST NOT exceed 20 bytes
    /// and the value MUST be Base64-encoded (the encoded string will be more than 20 bytes).
    /// The XKS Proxy MUST append the ciphertextMetadata to the additionalAuthenticatedData
    /// before normal AES GCM processing to ensure that integrity protection
    /// offered by the authenticationTag extends to the ciphertextMetadata.
    ///
    /// NOTE: It is important to explicitly include the length of additionalAuthenticatedData
    /// and the length of the ciphertextMetadata to avoid unintended successful decrypts,
    /// e.g. when a caller calls encrypt with no additionalAuthenticatedData,
    /// receives a ciphertextMetadata in the response and then calls decrypt
    /// passing the ciphertextMetadata as additionalAuthenticatedData and no ciphertextMetadata.
    /// The AAD input for the external key manager should be computed as (2-byte length,
    /// before Base64 encoding, of additionalAuthenticatedData in big-endian format
    /// || additionalAuthenticatedData || 1-byte length, before Base64 encoding,
    /// of ciphertextMetadata || ciphertextMetadata) where || represents concatenation of
    /// the binary values before Base64 encoding. If the additionalAuthenticatedData
    /// or ciphertextMetadata is not present, the corresponding length MUST be set to zero.
    /// If the inclusion of the lengths represents a departure from previously implemented behavior,
    /// the XKS proxy SHOULD encode the new behavior in the ciphertextMetadata
    /// and use the encoding to follow the same behavior during decrypt as was used
    /// for the corresponding encrypt.
    /// Otherwise, previously generated ciphertext will no longer be decryptable.
    ///
    /// For example, let's say version A of an XKS proxy concatenated the ciphertextMetadata
    /// directly to additionalAuthenticatedData (without including the lengths)
    /// but Version B implements new guidance then there needs to be a mechanism to distinguish
    /// whether a decrypt call should use the old way or the new way to create the AAD
    /// for the external key manager.
    /// If Version B always implements the new behavior then ciphertext created by Version A
    /// will no longer be decryptable.
    /// The ciphertextMetadata is the natural place to encode this difference in
    /// how the authenticationTag was created.
    pub ciphertextMetadata: Option<String>,
    /// Base64 encoded initialization vector generated by the external key manager
    /// that was used during encrypt operation.
    /// The initialization vector MUST be either 12 bytes (96 bits) or 16 bytes (128 bits).
    /// The Base64 encoding will have 16 bytes or 24 bytes.
    pub initializationVector: String,
    /// Base64 encoded message authentication code generated by external key manager
    /// performing AES-GCM encryption.
    /// Authentication tag size MUST be 16 bytes (128 bits).
    /// Some key managers append the authentication tag to the ciphertext.
    /// In such cases, the XKS proxy MUST separate the two before composing the response.
    pub authenticationTag: String,
    /// This field is a Base64 encoded hash computed over the
    /// additionalAuthenticatedData (if present in the request), ciphertextMetadata (if present),
    /// initializationVector, ciphertext and authenticationTag fields in the response.
    /// It MUST be included whenever the request includes the ciphertextDataIntegrityValueAlgorithm
    /// field.
    /// The hashing algorithm used to compute this value MUST be the one specified
    /// as the ciphertextDataIntegrityValueAlgorithm in the request.
    /// KMS will independently calculate the ciphertextDataIntegrityValue (CDIV)
    /// and return an error to the caller if the computed value does not match the value
    /// in the response.
    /// KMS interprets a match as assurance from the XKS Proxy that a subsequent decrypt call
    /// where the caller passes in the same additionalAuthenticatedData
    /// (if present, in the encrypt request), initializationVector, ciphertext
    /// and authenticationTag values will succeed and return the plaintext
    /// that was passed as input to this encrypt API.
    /// See Appendix C for a complete example and specific CDIV implementation guidelines.
    pub ciphertextDataIntegrityValue: Option<String>,
}

#[post("/kms/xks/v1/keys/{key_id}/encrypt")]
pub(crate) async fn encrypt(
    req_http: HttpRequest,
    key_id: Path<String>,
    request: Json<EncryptRequest>,
    kms: Data<Arc<KMS>>,
) -> HttpResponse {
    let request = request.into_inner();
    let key_id = key_id.into_inner();
    info!(
        "POST /kms/xks/v1/keys/{key_id}/encrypt - operation: {} - id: {} - user: {}",
        request.requestMetadata.kmsOperation,
        request.requestMetadata.kmsRequestId,
        request.requestMetadata.awsPrincipalArn
    );
    debug!("encrypt request: {:?}", request.requestMetadata);
    let kms = kms.into_inner();
    match encrypt_inner(req_http, request, key_id, &kms)
        .await
        .map(Json)
    {
        Ok(wrap_response) => HttpResponse::Ok().json(wrap_response),
        Err(e) => match e {
            KmsError::Unauthorized(msg) => XksErrorReply {
                errorName: XksErrorName::InvalidKeyUsageException,
                errorMessage: Some(msg),
            }
            .into(),
            KmsError::ItemNotFound(msg) => XksErrorReply {
                errorName: XksErrorName::KeyNotFoundException,
                errorMessage: Some(msg),
            }
            .into(),
            KmsError::CryptographicError(msg) => XksErrorReply {
                errorName: XksErrorName::ValidationException,
                errorMessage: Some(msg),
            }
            .into(),
            _ => HttpResponse::from_error(e),
        },
    }
}

async fn encrypt_inner(
    _req_http: HttpRequest,
    request: EncryptRequest,
    key_id_or_tags: String,
    kms: &Arc<KMS>,
) -> KResult<EncryptResponse> {
    let user = request.requestMetadata.awsPrincipalArn;

    let cryptographic_parameters = match request.encryptionAlgorithm {
        EncryptionAlgorithm::AES_GCM => CryptographicParameters {
            cryptographic_algorithm: Some(CryptographicAlgorithm::AES),
            ..Default::default()
        },
    };
    let data = Zeroizing::new(STANDARD.decode(&request.plaintext)?);
    // Supplied Nonce or new one.
    let nonce: [u8; AES_256_GCM_IV_LENGTH] = {
        let mut iv = [0; AES_256_GCM_IV_LENGTH];
        rand_bytes(&mut iv)?;
        iv
    };
    let aead = match request.additionalAuthenticatedData {
        None => None,
        Some(b64) => Some(STANDARD.decode(&b64)?),
    };
    let response = kms
        .encrypt(
            Encrypt {
                unique_identifier: Some(UniqueIdentifier::TextString(key_id_or_tags.clone())),
                cryptographic_parameters: Some(cryptographic_parameters),
                data: Some(data),
                i_v_counter_nonce: Some(nonce.to_vec()),
                correlation_value: None,
                init_indicator: None,
                final_indicator: None,
                authenticated_encryption_additional_data: aead.clone(),
            },
            &user,
        )
        .await?;
    let ciphertext = response
        .data
        .ok_or_else(|| KmsError::ServerError("Missing AES GCM ciphertext".to_owned()))?;
    let returned_key_id = response.unique_identifier.to_string();
    let ciphertext_metadata = if returned_key_id == key_id_or_tags {
        None
    } else {
        // The encryption is likely performed using tags; keep track of the actual key used
        Some(returned_key_id.as_bytes().to_vec())
    };
    let tag = response.authenticated_encryption_tag.ok_or_else(|| {
        KmsError::ServerError("Missing AES GCM authenticated encryption tag".to_owned())
    })?;
    let integrity_hash = match request.ciphertextDataIntegrityValueAlgorithm {
        None => None,
        Some(CdivAlgorithm::SHA_256) => {
            let mut hasher = Sha256::new();
            if let Some(aead) = &aead {
                hasher.update(aead);
            }
            if let Some(ciphertext_metadata) = &ciphertext_metadata {
                hasher.update(ciphertext_metadata);
            }
            hasher.update(&nonce);
            hasher.update(&ciphertext);
            hasher.update(&tag);
            Some(hasher.finish())
        }
    };
    Ok(EncryptResponse {
        ciphertext: STANDARD.encode(ciphertext),
        ciphertextMetadata: ciphertext_metadata.map(|b| STANDARD.encode(&b)),
        initializationVector: STANDARD.encode(nonce),
        authenticationTag: STANDARD.encode(tag),
        ciphertextDataIntegrityValue: integrity_hash.map(|b| STANDARD.encode(b)),
    })
}
