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

//! Helpers for unit tests.

#![cfg_attr(not(test), expect(dead_code))]

use crate::jwt::JwtAlgorithm;
use crate::jwt::JwtHeader;
use base64::Engine;
use openhcl_attestation_protocol::igvm_attest::akv;
use openssl::hash::MessageDigest;
use openssl::pkey::PKey;
use openssl::pkey::Private;
use openssl::rsa::Padding;
use openssl::x509::X509;
use openssl::x509::X509Name;

pub const CIPHERTEXT: &str = "test";

/// Generate a self-signed X.509 certificate for testing.
pub fn generate_x509(private: &PKey<Private>) -> X509 {
    let mut x509 = X509::builder().unwrap();

    // Generate a public key from the private key and set it as the public key of the certificate
    let public = private.public_key_to_pem().unwrap();
    let public = PKey::public_key_from_pem(&public).unwrap();
    x509.set_pubkey(&public).unwrap();

    x509.set_version(2).unwrap();
    x509.set_serial_number(
        &openssl::bn::BigNum::from_u32(1)
            .unwrap()
            .to_asn1_integer()
            .unwrap(),
    )
    .unwrap();
    x509.set_not_before(&openssl::asn1::Asn1Time::days_from_now(0).unwrap())
        .unwrap();
    x509.set_not_after(&openssl::asn1::Asn1Time::days_from_now(365).unwrap())
        .unwrap();

    let mut name = X509Name::builder().unwrap();
    name.append_entry_by_text("C", "US").unwrap();
    name.append_entry_by_text("ST", "Washington").unwrap();
    name.append_entry_by_text("L", "Redmond").unwrap();
    name.append_entry_by_text("O", "Example INC").unwrap();
    name.append_entry_by_text("CN", "example.com").unwrap();
    let name = name.build();
    x509.set_subject_name(&name).unwrap();
    x509.set_issuer_name(&name).unwrap();

    x509.sign(private, MessageDigest::sha256()).unwrap();

    x509.build()
}

/// 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(private: &PKey<Private>) -> Vec<String> {
    let cert = generate_x509(private);
    let intermediate = generate_x509(private);
    let root = generate_x509(private);

    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(private: &PKey<Private>) -> (String, String, String) {
    let header = JwtHeader {
        alg: JwtAlgorithm::RS256,
        x5c: generate_x5c(private),
    };
    // 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 mut signer = openssl::sign::Signer::new(MessageDigest::sha256(), private).unwrap();
    signer.set_rsa_padding(Padding::PKCS1).unwrap();
    signer.update(message.as_bytes()).unwrap();
    let signature = signer.sign_to_vec().unwrap();
    let base64_signature = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(&signature);

    (base64_header, base64_body, base64_signature)
}
