// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! This module provide API to communicate with attestation service (AS) to get
//! attestation report endorsed by AS.

use crate::platform;
use crate::AttestationAlgorithm;
use crate::AttestationServiceConfig;
use crate::EndorsedAttestationReport;

use std::collections::HashMap;
use std::io::{ErrorKind, Read, Write};
use std::net::TcpStream;
use std::sync::Arc;

use anyhow::{anyhow, bail, Result};
use log::{debug, trace, warn};
use serde_json::json;
use sgx_crypto::ecc::EcPublicKey;

/// Root certification of the DCAP attestation service provider.
#[cfg(dcap)]
const DCAP_ROOT_CA_CERT: &str = include_str!("../../config/keys/dcap_root_ca_cert.pem");

/// URL path to get the report from the attestation service.
const AS_REPORT_URL: &str = "/sgx/dev/attestation/v4/report";

#[derive(thiserror::Error, Debug)]
pub(crate) enum AttestationServiceError {
    #[error("Invalid attestation service address.")]
    InvalidAddress,
    #[error("Attestation service responds an malformed response.")]
    InvalidResponse,
    #[error("{0} is missing in HTTP header.")]
    MissingHeader(String),
    #[error(
        "Invalid Attestation Evidence Payload. The client should not repeat the
        request without modifications."
    )]
    BadRequest,
    #[error("Failed to authenticate or authorize request.")]
    Unauthorized,
    #[error("Internal error occurred.")]
    InternalServerError,
    #[error(
        "Service is currently not able to process the request (due to a
        temporary overloading or maintenance). This is a temporary state –the
        same request can be repeated after some time."
    )]
    ServiceUnavailable,
    #[error("TLS connection error.")]
    TlsError,
    #[error("Attestation service responds an unknown error.")]
    Unknown,
}

#[cfg(all(feature = "mesalock_sgx", not(feature = "libos")))]
impl EndorsedAttestationReport {
    pub fn new(
        att_service_cfg: &AttestationServiceConfig,
        pub_k: EcPublicKey,
    ) -> anyhow::Result<Self> {
        let (mut ak_id, qe_target_info) = platform::init_sgx_quote()?;

        // For IAS-based attestation, we need to fill our SPID (obtained from Intel)
        // into the attestation key id. For DCAP-based attestation, SPID should be 0
        const SPID_OFFSET: usize = std::mem::size_of::<sgx_types::types::QlAttKeyId>();
        ak_id.att_key_id[SPID_OFFSET..(SPID_OFFSET + att_service_cfg.spid.id.len())]
            .clone_from_slice(&att_service_cfg.spid.id);

        let sgx_report = platform::create_sgx_isv_enclave_report(pub_k, qe_target_info)?;
        let quote = platform::get_sgx_quote(&ak_id, sgx_report)?;
        let as_report = get_report(
            &att_service_cfg.algo,
            &att_service_cfg.as_url,
            &att_service_cfg.api_key,
            &quote,
        )?;

        Ok(as_report)
    }
}

#[cfg(all(feature = "libos", not(feature = "mesalock_sgx")))]
impl EndorsedAttestationReport {
    pub fn new(att_service_cfg: &AttestationServiceConfig, pub_k: EcPublicKey) -> Result<Self> {
        let report_data = platform::create_sgx_report_data(pub_k);
        let quote = match &att_service_cfg.algo {
            crate::AttestationAlgorithm::SgxEpid => {
                platform::get_sgx_epid_quote(&att_service_cfg.spid, report_data)?
            }
            crate::AttestationAlgorithm::SgxEcdsa => {
                platform::get_sgx_dcap_quote(&att_service_cfg.spid, report_data)?
            }
        };
        crate::service::get_report(
            &att_service_cfg.algo,
            &att_service_cfg.as_url,
            &att_service_cfg.api_key,
            &quote,
        )
    }
}

fn new_tls_stream(
    url: &url::Url,
) -> Result<rustls::StreamOwned<rustls::client::ClientConnection, TcpStream>> {
    let host_str = url
        .host_str()
        .ok_or(AttestationServiceError::InvalidAddress)?;
    let mut root_certs = rustls::RootCertStore::empty();
    #[cfg(dcap)]
    {
        let certs = rustls_pemfile::certs(&mut DCAP_ROOT_CA_CERT.to_string().as_bytes())
            .map_err(|_| AttestationServiceError::TlsError)?;
        let (valid_count, _) = root_certs.add_parsable_certificates(&certs);
        anyhow::ensure!(valid_count >= 1, "DCAP_ROOT_CA_CERT error");
    }
    root_certs.add_server_trust_anchors(webpki_roots::TLS_SERVER_ROOTS.0.iter().map(
        |trust_anchor| {
            rustls::OwnedTrustAnchor::from_subject_spki_name_constraints(
                trust_anchor.subject,
                trust_anchor.spki,
                trust_anchor.name_constraints,
            )
        },
    ));
    let config = rustls::ClientConfig::builder()
        .with_safe_defaults()
        .with_root_certificates(root_certs)
        .with_no_client_auth();
    let client = rustls::client::ClientConnection::new(Arc::new(config), host_str.try_into()?)?;
    let addrs = url.socket_addrs(|| match url.scheme() {
        "https" => Some(443),
        _ => None,
    })?;
    let socket = TcpStream::connect(&*addrs)?;
    let stream = rustls::StreamOwned::new(client, socket);

    Ok(stream)
}

/// Get attestation report form the attestation service (e.g., Intel Attestation
/// Service and customized DCAP attestation service).
pub(crate) fn get_report(
    algo: &AttestationAlgorithm,
    url: &url::Url,
    api_key: &str,
    quote: &[u8],
) -> Result<EndorsedAttestationReport> {
    debug!("get_report");
    let encoded_quote = base64::encode(quote);
    let encoded_json = json!({ "isvEnclaveQuote": encoded_quote }).to_string();
    let host_str = url
        .host_str()
        .ok_or(AttestationServiceError::InvalidAddress)?;

    let request = format!(
        "POST {} HTTP/1.1\r\n\
         HOST: {}\r\n\
         Ocp-Apim-Subscription-Key: {}\r\n\
         Connection: Close\r\n\
         Content-Length: {}\r\n\
         Content-Type: application/json\r\n\r\n\
         {}",
        AS_REPORT_URL,
        host_str,
        api_key,
        encoded_json.len(),
        encoded_json
    );
    trace!("{}", request);

    let mut stream = new_tls_stream(url).map_err(|_| AttestationServiceError::TlsError)?;
    stream.write_all(request.as_bytes())?;
    let mut response = Vec::new();
    if let Err(e) = stream.read_to_end(&mut response) {
        match e.kind() {
            // Server may send CloseNotify ConnectionAborted for Connection:Close request
            ErrorKind::ConnectionAborted => warn!("connection aborted: {:?}", e),
            _ => bail!("{:?} is not allowed", e),
        }
    };

    trace!("{}", String::from_utf8_lossy(&response));

    let mut headers = [httparse::EMPTY_HEADER; 16];
    let mut http_response = httparse::Response::new(&mut headers);

    debug!("http_response.parse");
    let header_len = match http_response
        .parse(&response)
        .map_err(|_| AttestationServiceError::InvalidResponse)?
    {
        httparse::Status::Complete(s) => s,
        _ => bail!(AttestationServiceError::InvalidResponse),
    };

    match http_response.code {
        Some(200) => {
            debug!("Operation successful.");
        }
        Some(400) => {
            debug!(
                "Invalid Attestation Evidence Payload. The client should not
                 repeat the request without modifications."
            );
            bail!(AttestationServiceError::BadRequest);
        }
        Some(401) => {
            debug!("Failed to authenticate or authorize request.");
            bail!(AttestationServiceError::Unauthorized);
        }
        Some(500) => {
            debug!("Internal error occurred.");
            bail!(AttestationServiceError::InternalServerError);
        }
        Some(503) => {
            debug!(
                "Service is currently not able to process the request (due to a
                 temporary overloading or maintenance). This is a temporary
                 state, the same request can be repeated after some time."
            );
            bail!(AttestationServiceError::ServiceUnavailable);
        }
        _ => {
            debug!("Attestation service responds an unknown error");
            bail!(AttestationServiceError::Unknown);
        }
    }

    let header_map = parse_headers(&http_response);
    debug!("ias header_map: {:?}", header_map);

    debug!("get_content_length");
    if !header_map.contains_key("content-length")
        || header_map
            .get("content-length")
            .ok_or_else(|| AttestationServiceError::MissingHeader("content-length".to_string()))?
            .parse::<u32>()
            .unwrap_or(0)
            == 0
    {
        bail!(AttestationServiceError::MissingHeader(
            "content-length".to_string()
        ));
    }

    debug!("get_signature");
    let signature_header = match algo {
        AttestationAlgorithm::SgxEpid => "x-iasreport-signature",
        AttestationAlgorithm::SgxEcdsa => "x-dcapreport-signature",
    };
    let signature = header_map
        .get(signature_header)
        .ok_or_else(|| AttestationServiceError::MissingHeader(signature_header.to_string()))?;
    let signature = base64::decode(signature)?;

    debug!("get_signing_cert");
    let signing_cert_header = match algo {
        AttestationAlgorithm::SgxEpid => "x-iasreport-signing-certificate",
        AttestationAlgorithm::SgxEcdsa => "x-dcapreport-signing-certificate",
    };
    let certs: Vec<Vec<u8>> = {
        let cert_str = header_map.get(signing_cert_header).ok_or_else(|| {
            AttestationServiceError::MissingHeader(signing_cert_header.to_string())
        })?;
        let decoded_cert = percent_encoding::percent_decode_str(cert_str).decode_utf8()?;
        let certs = rustls_pemfile::certs(&mut decoded_cert.as_bytes())
            .map_err(|_| anyhow!("pemfile error"))?;
        certs
    };

    debug!("return_report");
    let report = response[header_len..].to_vec();
    Ok(EndorsedAttestationReport {
        report,
        signature,
        certs,
    })
}

fn parse_headers(resp: &httparse::Response) -> HashMap<String, String> {
    debug!("parse_headers");
    let mut header_map = HashMap::new();
    for h in resp.headers.iter() {
        header_map.insert(
            // HTTP header name is case insensitive
            h.name.to_lowercase(),
            String::from_utf8_lossy(h.value).into_owned(),
        );
    }

    header_map
}