//! Manage self-signed x509 certificate containing enclave remote attestation
//! endorsements.

use std::{borrow::Cow, fmt};

use anyhow::{Context, anyhow, ensure};
use asn1_rs::FromDer;
use lexe_crypto::ed25519;
use lexe_enclave::enclave;
use lexe_hex::hex;
use x509_parser::certificate::X509Certificate;
use yasna::{ASN1Error, ASN1ErrorKind, models::ObjectIdentifier};

use crate::attest_client::quote::ReportData;

/// Packages a remote attestation quote generated by
/// `lexe_tls_attest_server::quote::quote_enclave` into a x509 cert extension
/// with all the evidence a client needs to verify it.
///
/// ```asn.1
/// SgxAttestationExtension ::= SEQUENCE {
///     QUOTE      OCTET STRING
/// }
/// ```
#[derive(Eq, PartialEq)]
pub struct SgxAttestationExtension<'quote> {
    pub quote: Cow<'quote, [u8]>,
    // We don't verify the QE_REPORT because it requires calling an Intel JSON
    // API for marginal security against old SGX versions. Also, the MIN CPUSVN
    // already handles the case where the platform enclaves are too old.
    // pub qe_report: Cow<'b, [u8]>,
}

// -- impl SgxAttestationExtension -- //

impl SgxAttestationExtension<'_> {
    /// This is the Intel SGX OID prefix + 1337.7
    /// gramine uses the same but 1337.6 to embed the quote.
    pub const OID: &'static [u64] = &[1, 2, 840, 113741, 1337, 7];

    /// DER-encoded OID
    pub const OID_DER: &'static [u8] = &[
        0x06, 0x09, 0x2A, 0x86, 0x48, 0x86, 0xF8, 0x4D, 0x8A, 0x39, 0x07,
    ];

    pub fn oid_yasna() -> ObjectIdentifier {
        ObjectIdentifier::from_slice(Self::OID)
    }

    #[rustfmt::skip]
    pub const fn oid_asn1_rs() -> asn1_rs::Oid<'static> {
        // TODO(phlip9): won't parse OID_DER...
        asn1_rs::oid!(1.2.840.113741.1337.7)
    }

    /// Clients that don't understand a critical extension will immediately
    /// reject the cert. Unfortunately, setting this to true seems to break
    /// clients...
    pub const fn is_critical() -> bool {
        false
    }

    /// Serialize the attestation to DER.
    pub fn to_der_bytes(&self) -> Vec<u8> {
        yasna::construct_der(|writer| {
            writer.write_sequence(|writer| {
                writer.next().write_bytes(&self.quote);
            })
        })
    }

    pub fn to_cert_extension(&self) -> rcgen::CustomExtension {
        let mut ext = rcgen::CustomExtension::from_oid_content(
            Self::OID,
            self.to_der_bytes(),
        );
        ext.set_criticality(Self::is_critical());
        ext
    }
}

impl SgxAttestationExtension<'static> {
    /// Build a dummy attestation for testing on non-SGX platforms.
    pub fn dummy(cert_pk: &ed25519::PublicKey) -> Self {
        // Use a dummy report as the 'quote'.
        let mut report = enclave::report().clone();

        // Insert the cert pk into the first 32 bytes of the `reportdata` field,
        // See `ReportData::from_cert_pk` in `attestation::quote::quote_enclave`
        let report_data = ReportData::from_cert_pk(cert_pk);
        report.reportdata = report_data.into_inner();

        Self {
            quote: Cow::Owned(AsRef::<[u8]>::as_ref(&report).to_owned()),
        }
    }

    /// Deserialize the attestation from DER bytes.
    pub fn from_der_bytes(buf: &[u8]) -> yasna::ASN1Result<Self> {
        yasna::parse_der(buf, |reader| {
            reader.read_sequence(|reader| {
                let quote = reader.next().read_bytes()?;
                Ok(Self {
                    quote: Cow::Owned(quote),
                })
            })
        })
    }
}

impl fmt::Debug for SgxAttestationExtension<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SgxAttestationExtension")
            .field("quote", &hex::display(&self.quote))
            .finish()
    }
}

/// Additional X509 cert extensions that Intel includes in the PCK leaf cert.
//
// See: [Intel SGX PCK Certificate Specification - 1.3.5 p.11](https://download.01.org/intel-sgx/sgx-dcap/1.10/linux/docs/Intel_SGX_PCK_Certificate_CRL_Spec-1.4.pdf)
//
// ```
// X509v3 extensions:
//   X509v3 Authority Key Identifier:
//     keyid:<keyid of issuer public key>
//   X509v3 CRL Distribution Points:
//     Full Name:
//       URI: <URL>
//   X509v3 Subject Key Identifier:
//     <keyid of public key>
//   X509v3 Key Usage: critical
//     Digital Signature, Non Repudiation
//   X509v3 Basic Constraints: critical
//     CA:FALSE
//   <SGX Extensions OID>:
//     <PPID OID>: <PPID value>
//     <TCB OID>:
//       <SGX TCB Comp01 SVN OID>: <SGX TCB Comp01 SVN value>
//       <SGX TCB Comp02 SVN OID>: <SGX TCB Comp02 SVN value>
//       ...
//       <SGX TCB Comp16 SVN OID>: <SGX TCB Comp16 SVN value>
//     <PCESVN OID>: <PCESVN value>
//     <CPUSVN OID>: <CPUSVN value>
//     <PCE-ID OID>: <PCE-ID value>
//     <FMSPC OID>: <FMSPC value>
//     <SGX Type OID>: <SGX Type value>
//     <PlatformInstanceID OID>: <PlatformInstanceID value>
//     <Configuration OID>:
//       <Dynamic Platform OID>: <Dynamic Platform flag value>
//       <Cached Keys OID>: <Cached Keys flag value>
//       <SMT Enabled OID>: <SMT Enabled flag value>
// ```
//
// | name           | oid                     | type               | description                               |
// |----------------|-------------------------|--------------------|-------------------------------------------|
// | SGX Extensions | 1.2.840.113741.1.13.1   | ASN.1 Sequence     | Sequence of Intel SGX PCK cert extensions |
// | FMSPC          | 1.2.840.113741.1.13.1.4 | ASN.1 Octet String | Value of FMSPC (6 bytes)                  |
#[derive(Debug, Eq, PartialEq)]
pub(crate) struct SgxPckExtensions {
    /// This field contains the Intel CPU family of the machine that made the
    /// attestation.
    pub cpu_fmspc: CpuFmspc,
}

impl SgxPckExtensions {
    /// The asn1 OID for the outer "SGX Extensions" sequence.
    #[rustfmt::skip]
    const OID: asn1_rs::Oid<'static> = asn1_rs::oid!(1.2.840.113741.1.13.1);

    /// The asn1 OID for the inner FMSPC entry in the outer sequence.
    const OID_FMSPC: &[u64] = &[1, 2, 840, 113741, 1, 13, 1, 4];

    /// Try to parse from a whole DER-encoded x509 cert.
    pub(crate) fn from_cert_der(cert_der: &[u8]) -> anyhow::Result<Self> {
        let (unparsed_data, cert) =
            X509Certificate::from_der(cert_der).context("Invalid PCK cert")?;
        ensure!(unparsed_data.is_empty(), "leftover unparsed PCK cert data");

        let cert_ext = cert
            .get_extension_unique(&Self::OID)
            .context("Invalid SGX PCK cert extension")?
            .context("Missing SGX PCK cert extension")?;

        Self::from_der_bytes(cert_ext.value)
            .map_err(|err| anyhow!("Invalid SGX PCK cert extension: {err:#}"))
    }

    /// Try to parse from the x509 cert extension contents.
    ///
    /// These are the bytes contained in the `"SGX Extensions"
    /// (oid=1.2.840.113741.1.13.1)` x509 cert extension.
    fn from_der_bytes(buf: &[u8]) -> yasna::ASN1Result<Self> {
        // Inside `buf` is an ASN1 sequence. We just want the entry for FMSPC.

        let mut fmspc = None;

        yasna::parse_der(buf, |reader| {
            reader.read_sequence_of(|reader| {
                reader.read_sequence(|reader| {
                    let oid = reader.next().read_oid()?;

                    if oid.as_ref() == Self::OID_FMSPC {
                        let bytes = reader.next().read_bytes()?;
                        let bytes = <[u8; 6]>::try_from(bytes.as_slice())
                            .map_err(|_| Self::invalid_asn1())?;
                        if fmspc.replace(bytes).is_some() {
                            return Err(Self::invalid_asn1());
                        }
                    } else {
                        // Skip other extensions we don't care about
                        let _ = reader.next().read_der()?;
                    }

                    Ok(())
                })
            })?;

            let cpu_fmspc = CpuFmspc(fmspc.ok_or_else(Self::invalid_asn1)?);
            Ok(Self { cpu_fmspc })
        })
    }

    fn invalid_asn1() -> ASN1Error {
        ASN1Error::new(ASN1ErrorKind::Invalid)
    }
}

/// An Intel CPU FMSPC (Family-Model-Stepping-Platform-Config) code that
/// identifies a family of Intel CPUs
#[derive(Copy, Clone, Eq, PartialEq, Hash)]
pub(crate) struct CpuFmspc(pub [u8; 6]);

lexe_byte_array::impl_byte_array!(CpuFmspc, 6);
lexe_byte_array::impl_debug_display_as_hex!(CpuFmspc);

#[cfg(test)]
mod test {
    use std::fs;

    use lexe_common::ByteArray;
    use rustls::pki_types::pem::PemObject;

    use super::*;

    #[test]
    fn test_sgx_attestation_ext_oid_der_bytes() {
        let oid = SgxAttestationExtension::oid_yasna();
        let oid_der = yasna::encode_der(&oid);
        assert_eq!(SgxAttestationExtension::OID_DER, &oid_der);
    }

    #[test]
    fn test_sgx_attestation_ext_serde() {
        let ext = SgxAttestationExtension {
            quote: b"test".as_slice().into(),
        };

        let der_bytes = ext.to_der_bytes();
        let ext2 = SgxAttestationExtension::from_der_bytes(&der_bytes).unwrap();
        assert_eq!(ext, ext2);

        let der_bytes2 = ext2.to_der_bytes();
        assert_eq!(der_bytes, der_bytes2);

        assert_eq!(b"test".as_slice(), ext2.quote.as_ref());
    }

    #[test]
    fn test_sgx_pck_extensions_parse() {
        let cert_pem = fs::read_to_string("test_data/pck_cert.pem").unwrap();
        let cert_der = rustls::pki_types::CertificateDer::from_pem_slice(
            cert_pem.as_bytes(),
        )
        .unwrap();

        let extensions = SgxPckExtensions::from_cert_der(&cert_der).unwrap();
        let fmspc = CpuFmspc::from_hex("00606a000000").unwrap();
        assert_eq!(extensions.cpu_fmspc, fmspc);
    }
}
