// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

//! Helpers for unit tests.

#![cfg(test)]

use crate::jwt::JwtAlgorithm;
use crate::jwt::JwtHeader;
use base64::Engine;
use crypto::rsa::RsaKeyPair;
use crypto::x509::X509Certificate;
use openhcl_attestation_protocol::igvm_attest::akv;

pub const CIPHERTEXT: &str = "test";

/// Generate a self-signed X.509 certificate for testing.
pub fn generate_x509(key_pair: &RsaKeyPair) -> X509Certificate {
    let mut builder = crypto::x509::X509Builder::new().unwrap();
    builder.set_pubkey_from_rsa_key_pair(key_pair).unwrap();
    builder
        .set_subject_and_issuer_name("US", "Washington", "Redmond", "Example INC", "example.com")
        .unwrap();
    builder.set_validity_days(365).unwrap();
    builder.sign_and_build(key_pair).unwrap()
}

/// Generate an X.509 certificate chain for testing.
/// The chain consists of three certificates: cert, intermediate, and root.
/// All certs are signed by the same private key and have the same subject and issuer.
fn generate_x5c(key_pair: &RsaKeyPair) -> Vec<String> {
    let cert = generate_x509(key_pair);
    let intermediate = generate_x509(key_pair);
    let root = generate_x509(key_pair);

    let base64_cert = base64::engine::general_purpose::STANDARD.encode(cert.to_der().unwrap());
    let base64_intermediate =
        base64::engine::general_purpose::STANDARD.encode(intermediate.to_der().unwrap());
    let base64_root = base64::engine::general_purpose::STANDARD.encode(root.to_der().unwrap());

    vec![base64_cert, base64_intermediate, base64_root]
}

/// Generate the base64 encoded components of a JWT.
pub fn generate_base64_encoded_jwt_components(key_pair: &RsaKeyPair) -> (String, String, String) {
    let header = JwtHeader {
        alg: JwtAlgorithm::RS256,
        x5c: generate_x5c(key_pair),
    };
    // Header is a base64-url encoded JSON object
    let base64_header = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .encode(serde_json::to_string(&header).unwrap());

    let key_hsm = akv::AkvKeyReleaseKeyBlob {
        ciphertext: CIPHERTEXT.as_bytes().to_vec(),
    };

    let body = akv::AkvKeyReleaseJwtBody {
        response: akv::AkvKeyReleaseResponse {
            key: akv::AkvKeyReleaseKeyObject {
                key: akv::AkvJwk {
                    key_hsm: serde_json::to_string(&key_hsm).unwrap().as_bytes().to_vec(),
                },
            },
        },
    };
    // Body is a base64-url encoded JSON object
    let base64_body = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .encode(serde_json::to_string(&body).unwrap().as_bytes());

    // The signature is generated by signing the concatenation of base64_header and base64_body
    let message = format!("{}.{}", base64_header, base64_body);
    let signature = key_pair.sign_pkcs1_sha256(message.as_bytes()).unwrap();
    let base64_signature = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(&signature);

    (base64_header, base64_body, base64_signature)
}
