// SPDX-License-Identifier: MIT
//
// Copyright IBM Corp. 2024

use super::AttNonce;
use crate::{
    attest::AttestationMeasAlg, brcb::BootHdrTags, crypto::calculate_hmac, request::Confidential,
    uv::ConfigUid, Result,
};
use openssl::{
    hash::MessageDigest,
    pkey::{PKeyRef, Private},
};
use std::mem::size_of;
use zerocopy::{AsBytes, BigEndian, U16, U32};

#[cfg(doc)]
use crate::attest::AttestationRequest;

/// Holds the data to be measured.
///
/// The Attestation measurement is an authentication code of the following data:
///
/// ```none
/// |-------------------------------|
/// | From SE-header:               |
/// |   Page List Digest (64)       |
/// |   Address List Digest (64)    |
/// |   Tweak List Digest (64)      |
/// |   SE Header Tag (16)          |
/// | Configuration Unique Id (16)  |
/// | user-data length (2)          |
/// | zeros (2)                     |
/// | additional data length (4)    |
/// | user-data (0-256)             |
/// | optional nonce (0 or 16)      |
/// | additional data (0+)          |
/// |-------------------------------|
/// ```
#[derive(Debug)]
pub struct AttestationItems(Confidential<Vec<u8>>);

// tags: BootHdrTags,
// cuid: ConfigUid,
// user_data_len: U16<BigEndian>,
// res: u16,
// additional_len: U32<BigEndian>,
// user_data: Vec<u8>,
// nonce: Option<[u8; 16]>,
// additional: Vec<u8>,
impl AttestationItems {
    /// Create a new attestation item struct.
    ///
    /// * `tags`: The tags from the SE header
    /// * `cuid`: The Configuration Unique Id from the SE guest for which the Measurement was
    ///   calculated
    /// * `user`: up to 256 bytes of arbitrary data generated on the SE-guest before measuring
    /// * `nonce`: technically optional nonce, but [`AttestationRequest`] enforces it
    /// * `additional`: additional data generated by the Firmware depending on the Attestation flags
    ///
    /// If size values of `user` or `additional` are longer than 16/32 bit they are silently
    /// truncated. `user-data` is limited to 256 bytes architecture wise, and additional data is
    /// limited to 8 pages by the uvdevice. Larger sizes will produce invalid measurements
    /// anyhow.
    pub fn new(
        tags: &BootHdrTags,
        cuid: &ConfigUid,
        user: Option<&[u8]>,
        nonce: Option<&AttNonce>,
        additional: Option<&[u8]>,
    ) -> Self {
        // expectations are ensured by ExchangeCtx invariants
        let user = user.unwrap_or(&[]);
        let user_data_len: U16<BigEndian> = (user.len() as u16).into();

        let additional = additional.unwrap_or(&[]);
        let additional_len: U32<BigEndian> = (additional.len() as u32).into();

        let size = size_of::<BootHdrTags>() // PLD ALD TLD TAG
            + size_of::<ConfigUid>()
            + size_of::<u16>() // user_len
            + size_of::<u16>() // reserved
            + size_of::<u32>() // additional_len
            + user.len()
            + match nonce {
                Some(_) => size_of::<AttNonce>(),
                None => 0,
            }
            + additional.len();

        let mut items = Vec::with_capacity(size);
        items.extend_from_slice(tags.as_bytes());
        items.extend_from_slice(cuid.as_bytes());
        items.extend_from_slice(user_data_len.as_bytes());
        items.extend_from_slice(&[0, 0]);
        items.extend_from_slice(additional_len.as_bytes());
        items.extend_from_slice(user);
        if let Some(nonce) = nonce {
            items.extend_from_slice(nonce);
        }
        items.extend_from_slice(additional);
        assert!(items.len() == size);
        Self(items.into())
    }
}

/// Holds an attestation measurement
#[derive(Debug)]
#[allow(clippy::len_without_is_empty)]
pub struct AttestationMeasurement(Vec<u8>);
impl AttestationMeasurement {
    /// Calculate an attestation measurement
    pub fn calculate(
        items: AttestationItems,
        mai: AttestationMeasAlg,
        meas_key: &PKeyRef<Private>,
    ) -> Result<Self> {
        match mai {
            AttestationMeasAlg::HmacSha512 => {
                calculate_hmac(meas_key, MessageDigest::sha512(), items.0.value()).map(Self)
            }
        }
    }

    /// Returns the length of the [`AttestationMeasurement`].
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Securely compares the calculated measurement with a given one
    ///
    /// Exists early when sizes do not match
    pub fn eq_secure(&self, other: &[u8]) -> bool {
        if self.len() != other.len() {
            return false;
        }
        openssl::memcmp::eq(&self.0, other)
    }
}

impl AsRef<[u8]> for AttestationMeasurement {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}
impl From<Vec<u8>> for AttestationMeasurement {
    fn from(value: Vec<u8>) -> Self {
        Self(value)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use openssl::pkey::PKey;

    const M_KEY: [u8; 64] = [0x41; 64];
    const BOOT_HDR_TAGS: BootHdrTags = BootHdrTags::new([1; 64], [2; 64], [3; 64], [4; 16]);
    const CUID: [u8; 16] = [5; 16];
    const USER: [u8; 256] = [7; 256];
    const NONCE: [u8; 16] = [8; 16];
    const ADDITIONAL: [u8; 128] = [9; 128];

    // just for better output in case of a test failure
    impl PartialEq<[u8]> for AttestationMeasurement {
        fn eq(&self, other: &[u8]) -> bool {
            self.eq_secure(other)
        }
    }

    #[test]
    fn measurement_all() {
        const EXP_HMAC: [u8; 64] = [
            0x88, 0x79, 0x4c, 0x62, 0xcc, 0xe7, 0xbc, 0xf2, 0x62, 0x16, 0xde, 0xb3, 0xf4, 0x8f,
            0x13, 0xfe, 0xa6, 0x37, 0x4b, 0x6d, 0x7e, 0x35, 0xbc, 0xc5, 0xc2, 0xce, 0x68, 0x12,
            0x1d, 0xb6, 0xf4, 0x5d, 0xfc, 0x8c, 0x17, 0x18, 0x56, 0x46, 0x35, 0x49, 0x40, 0x8b,
            0xf8, 0xe7, 0xd1, 0xac, 0xa1, 0x1e, 0xfa, 0xd0, 0xa8, 0x78, 0xaf, 0x97, 0xdc, 0x9e,
            0x21, 0xa1, 0xfc, 0x2a, 0x32, 0xf3, 0xa6, 0x75,
        ];
        let items = AttestationItems::new(
            &BOOT_HDR_TAGS,
            &CUID,
            Some(&USER),
            Some(&NONCE),
            Some(&ADDITIONAL),
        );
        let key = PKey::hmac(&M_KEY).unwrap();
        let meas =
            AttestationMeasurement::calculate(items, AttestationMeasAlg::HmacSha512, &key).unwrap();
        assert_eq!(meas, EXP_HMAC[..]);
        assert!(meas.eq_secure(&EXP_HMAC[..]));
    }

    #[test]
    fn measurement_user_add() {
        const EXP_HMAC: [u8; 64] = [
            0xfb, 0xd4, 0xf7, 0x38, 0xa3, 0x90, 0xed, 0xd9, 0x47, 0xcd, 0x4f, 0x11, 0xaf, 0x3a,
            0x2f, 0x3b, 0xab, 0x2f, 0xdf, 0x8b, 0xf8, 0x9b, 0xf8, 0x1b, 0xeb, 0x49, 0x51, 0x17,
            0xf4, 0x38, 0x2c, 0xf4, 0x2f, 0x07, 0x30, 0xc8, 0xc7, 0xd9, 0xe3, 0xca, 0x27, 0xfb,
            0x25, 0xad, 0xfc, 0xeb, 0x21, 0x22, 0x4f, 0x57, 0xfd, 0xb3, 0x98, 0xdc, 0xf4, 0x1a,
            0x83, 0xc1, 0x46, 0xe6, 0xa2, 0x3d, 0xb7, 0x60,
        ];
        let items =
            AttestationItems::new(&BOOT_HDR_TAGS, &CUID, Some(&USER), None, Some(&ADDITIONAL));
        let key = PKey::hmac(&M_KEY).unwrap();
        let meas =
            AttestationMeasurement::calculate(items, AttestationMeasAlg::HmacSha512, &key).unwrap();
        assert_eq!(meas, EXP_HMAC[..]);
        assert!(meas.eq_secure(&EXP_HMAC[..]));
    }

    #[test]
    fn measurement_add() {
        const EXP_HMAC: [u8; 64] = [
            0x63, 0x67, 0x1f, 0xbf, 0x29, 0x50, 0x36, 0xeb, 0x10, 0x23, 0xea, 0x71, 0xf7, 0x18,
            0x2e, 0x7d, 0x63, 0x43, 0xdc, 0x7b, 0x2d, 0xa5, 0x84, 0xe8, 0x24, 0xd0, 0xa7, 0xd1,
            0x98, 0xab, 0x9c, 0xde, 0xd7, 0x56, 0xc9, 0x3b, 0x39, 0x05, 0x0f, 0xfb, 0x76, 0x45,
            0x55, 0xb0, 0x1f, 0x88, 0xcb, 0x82, 0x01, 0x7a, 0x6a, 0x15, 0xc7, 0xe0, 0xba, 0xfc,
            0x60, 0x05, 0xf1, 0xe4, 0xf7, 0x8a, 0xa1, 0x24,
        ];
        let items = AttestationItems::new(&BOOT_HDR_TAGS, &CUID, None, None, Some(&ADDITIONAL));
        let key = PKey::hmac(&M_KEY).unwrap();
        let meas =
            AttestationMeasurement::calculate(items, AttestationMeasAlg::HmacSha512, &key).unwrap();
        assert_eq!(meas, EXP_HMAC[..]);
        assert!(meas.eq_secure(&EXP_HMAC[..]));
    }

    #[test]
    fn measurement_minimal() {
        const EXP_HMAC: [u8; 64] = [
            0xc5, 0xc3, 0x4c, 0x93, 0x83, 0x5d, 0x1e, 0xc2, 0x3f, 0x5c, 0x2d, 0x77, 0x8d, 0xfa,
            0x20, 0x12, 0x9b, 0x11, 0xb3, 0x05, 0x60, 0x17, 0x42, 0xcb, 0x2f, 0x38, 0xe0, 0xed,
            0x98, 0x94, 0xdc, 0xdb, 0x73, 0xfc, 0x86, 0x95, 0xab, 0x6a, 0x8d, 0xba, 0xd0, 0x74,
            0x40, 0x73, 0xdd, 0xc8, 0x1a, 0x5e, 0xaa, 0xfa, 0x52, 0xe4, 0xa1, 0x5a, 0xf8, 0xde,
            0xb8, 0xd7, 0x61, 0x09, 0x19, 0x22, 0x84, 0x7f,
        ];
        let items = AttestationItems::new(&BOOT_HDR_TAGS, &CUID, None, None, None);
        let key = PKey::hmac(&M_KEY).unwrap();
        let meas =
            AttestationMeasurement::calculate(items, AttestationMeasAlg::HmacSha512, &key).unwrap();
        assert_eq!(meas, EXP_HMAC[..]);
        assert!(meas.eq_secure(&EXP_HMAC[..]));
    }
}
