/*
 * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
 *
 * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
 */

use common::config::smtp::resolver::Tlsa;
use rustls_pki_types::CertificateDer;
use sha1::Digest;
use sha2::{Sha256, Sha512};
use trc::DaneEvent;
use x509_parser::prelude::{FromDer, X509Certificate};

use crate::queue::{Error, ErrorDetails, HostResponse, Status};

pub trait TlsaVerify {
    fn verify(
        &self,
        session_id: u64,
        hostname: &str,
        certificates: Option<&[CertificateDer<'_>]>,
    ) -> Result<(), Status<HostResponse<Box<str>>, ErrorDetails>>;
}

impl TlsaVerify for Tlsa {
    fn verify(
        &self,
        session_id: u64,
        hostname: &str,
        certificates: Option<&[CertificateDer<'_>]>,
    ) -> Result<(), Status<HostResponse<Box<str>>, ErrorDetails>> {
        let certificates = if let Some(certificates) = certificates {
            certificates
        } else {
            trc::event!(
                Dane(DaneEvent::NoCertificatesFound),
                SpanId = session_id,
                Hostname = hostname.to_string(),
            );

            return Err(Status::TemporaryFailure(ErrorDetails {
                entity: hostname.into(),
                details: Error::DaneError("No certificates were provided by host".into()),
            }));
        };

        let mut matched_end_entity = false;
        let mut matched_intermediate = false;
        'outer: for (pos, der_certificate) in certificates.iter().enumerate() {
            // Parse certificate
            let certificate = match X509Certificate::from_der(der_certificate.as_ref()) {
                Ok((_, certificate)) => certificate,
                Err(err) => {
                    trc::event!(
                        Dane(DaneEvent::CertificateParseError),
                        SpanId = session_id,
                        Hostname = hostname.to_string(),
                        Reason = err.to_string(),
                    );

                    return Err(Status::TemporaryFailure(ErrorDetails {
                        entity: hostname.into(),
                        details: Error::DaneError("Failed to parse X.509 certificate".into()),
                    }));
                }
            };

            // Match against TLSA records
            let is_end_entity = pos == 0;
            let mut sha256 = [None, None];
            let mut sha512 = [None, None];
            for record in self.entries.iter() {
                if record.is_end_entity == is_end_entity {
                    let hash: &[u8] = if record.is_sha256 {
                        &sha256[usize::from(record.is_spki)].get_or_insert_with(|| {
                            let mut hasher = Sha256::new();
                            hasher.update(if record.is_spki {
                                certificate.public_key().raw
                            } else {
                                der_certificate.as_ref()
                            });
                            hasher.finalize()
                        })[..]
                    } else {
                        &sha512[usize::from(record.is_spki)].get_or_insert_with(|| {
                            let mut hasher = Sha512::new();
                            hasher.update(if record.is_spki {
                                certificate.public_key().raw
                            } else {
                                der_certificate.as_ref()
                            });
                            hasher.finalize()
                        })[..]
                    };

                    if hash == record.data {
                        trc::event!(
                            Dane(DaneEvent::TlsaRecordMatch),
                            SpanId = session_id,
                            Hostname = hostname.to_string(),
                            Type = if is_end_entity {
                                "end-entity"
                            } else {
                                "intermediate"
                            },
                            Details = format!("{:x?}", hash),
                        );

                        if is_end_entity {
                            matched_end_entity = true;
                            if !self.has_intermediates {
                                break 'outer;
                            }
                        } else {
                            matched_intermediate = true;
                            break 'outer;
                        }
                    }
                }
            }
        }

        // DANE is valid if:
        // - EE matched even if no TA matched
        // - Both EE and TA matched
        // - EE is not present and TA matched
        if (self.has_end_entities && matched_end_entity)
            || ((self.has_end_entities == matched_end_entity)
                && (self.has_intermediates == matched_intermediate))
        {
            trc::event!(
                Dane(DaneEvent::AuthenticationSuccess),
                SpanId = session_id,
                Hostname = hostname.to_string(),
            );

            Ok(())
        } else {
            trc::event!(
                Dane(DaneEvent::AuthenticationFailure),
                SpanId = session_id,
                Hostname = hostname.to_string(),
            );

            Err(Status::PermanentFailure(ErrorDetails {
                entity: hostname.into(),
                details: Error::DaneError("No matching certificates found in TLSA records".into()),
            }))
        }
    }
}
