// SPDX-License-Identifier: AGPL-3.0-or-later
//! This module contains types which contribute to the configuration of TLS.

use std::sync::Arc;

use btlib::{
    crypto::{Creds, CredsPriv, HashKind, Scheme, Sign, Verifier},
    BlockPath, Principal, Writecap,
};
use core::ops::Deref;
use log::error;
use quinn::{ClientConfig, ServerConfig};
use rustls::{
    client::{HandshakeSignatureValid, ResolvesClientCert},
    internal::msgs::base::PayloadU16,
    server::{ClientCertVerified, ResolvesServerCert},
    sign::{CertifiedKey, SigningKey},
    Certificate, ConfigBuilder, ConfigSide, SignatureAlgorithm, SignatureScheme, WantsCipherSuites,
    WantsVerifier,
};

use crate::Result;

pub(crate) fn server_config(resolver: Arc<CertResolver>) -> Result<ServerConfig> {
    let server_config = common_config(rustls::ServerConfig::builder())?
        .with_client_cert_verifier(Arc::new(ClientCertVerifier))
        .with_cert_resolver(resolver);
    Ok(ServerConfig::with_crypto(Arc::new(server_config)))
}

pub(crate) fn client_config(
    server_path: Arc<BlockPath>,
    resolver: Arc<CertResolver>,
) -> Result<ClientConfig> {
    let client_config = common_config(rustls::ClientConfig::builder())?
        .with_custom_certificate_verifier(Arc::new(ServerCertVerifier::new(server_path)))
        .with_client_cert_resolver(resolver);
    Ok(ClientConfig::new(Arc::new(client_config)))
}

fn common_config<Side: ConfigSide>(
    builder: ConfigBuilder<Side, WantsCipherSuites>,
) -> Result<ConfigBuilder<Side, WantsVerifier>> {
    builder
        .with_cipher_suites(&[rustls::cipher_suite::TLS13_AES_128_GCM_SHA256])
        .with_kx_groups(&[&rustls::kx_group::SECP256R1])
        .with_protocol_versions(&[&rustls::version::TLS13])
        .map_err(|err| err.into())
}

fn to_cert_err(err: btlib::Error) -> rustls::Error {
    rustls::Error::InvalidCertificateData(err.to_string())
}

fn verify_tls13_signature(
    message: &[u8],
    cert: &Certificate,
    dss: &rustls::DigitallySignedStruct,
) -> std::result::Result<(), rustls::Error> {
    let (_, subject_key) = Writecap::from_cert_chain(cert, &[]).map_err(to_cert_err)?;
    subject_key
        .verify(&mut std::iter::once(message), dss.signature())
        .map_err(|_| rustls::Error::InvalidCertificateSignature)?;
    Ok(())
}

/// Verifier for the certificate chain presented by the server.
struct ServerCertVerifier {
    server_path: Arc<BlockPath>,
}

impl ServerCertVerifier {
    fn new(server_path: Arc<BlockPath>) -> Self {
        Self { server_path }
    }
}

impl rustls::client::ServerCertVerifier for ServerCertVerifier {
    fn verify_server_cert(
        &self,
        end_entity: &Certificate,
        intermediates: &[Certificate],
        _server_name: &rustls::ServerName,
        _scts: &mut dyn Iterator<Item = &[u8]>,
        _ocsp_response: &[u8],
        _now: std::time::SystemTime,
    ) -> std::result::Result<rustls::client::ServerCertVerified, rustls::Error> {
        let (writecap, ..) =
            Writecap::from_cert_chain(end_entity, intermediates).map_err(to_cert_err)?;
        let path = writecap.bind_path();
        if &path != self.server_path.as_ref() {
            return Err(rustls::Error::InvalidCertificateData(format!(
                "expected writecap with path '{}' got writecap with path '{path}'",
                self.server_path
            )));
        }
        writecap.assert_valid_for(&path).map_err(to_cert_err)?;
        Ok(rustls::client::ServerCertVerified::assertion())
    }

    fn verify_tls13_signature(
        &self,
        message: &[u8],
        cert: &Certificate,
        dss: &rustls::DigitallySignedStruct,
    ) -> std::result::Result<rustls::client::HandshakeSignatureValid, rustls::Error> {
        verify_tls13_signature(message, cert, dss)?;
        Ok(HandshakeSignatureValid::assertion())
    }
}

/// Verifier for the certificate chain presented by the client.
struct ClientCertVerifier;

impl rustls::server::ClientCertVerifier for ClientCertVerifier {
    fn verify_client_cert(
        &self,
        end_entity: &Certificate,
        intermediates: &[Certificate],
        _now: std::time::SystemTime,
    ) -> std::result::Result<rustls::server::ClientCertVerified, rustls::Error> {
        let (writecap, ..) =
            Writecap::from_cert_chain(end_entity, intermediates).map_err(to_cert_err)?;
        writecap
            .assert_valid_for(writecap.path())
            .map_err(to_cert_err)?;
        Ok(ClientCertVerified::assertion())
    }

    fn client_auth_root_subjects(&self) -> Option<rustls::DistinguishedNames> {
        let der = match Principal::default().to_name_der() {
            Ok(der) => der,
            Err(err) => {
                error!("failed to create distinguished name from root principal: {err}");
                return None;
            }
        };
        Some(vec![PayloadU16(der)])
    }

    fn verify_tls13_signature(
        &self,
        message: &[u8],
        cert: &Certificate,
        dss: &rustls::DigitallySignedStruct,
    ) -> std::result::Result<rustls::client::HandshakeSignatureValid, rustls::Error> {
        verify_tls13_signature(message, cert, dss)?;
        Ok(rustls::client::HandshakeSignatureValid::assertion())
    }
}

pub(crate) struct CertResolver {
    cert_key: Arc<CertifiedKey>,
}

impl CertResolver {
    pub(crate) fn new<C: Creds + Send + Sync + 'static>(creds: Arc<C>) -> Result<Self> {
        let writecap = creds.writecap().ok_or(btlib::BlockError::MissingWritecap)?;
        let chain = writecap.to_cert_chain(creds.public_sign())?;
        let mut certs = Vec::with_capacity(chain.len());
        for cert in chain {
            certs.push(Certificate(cert))
        }
        let key = Arc::new(CredRef::new(creds));
        let cert_key = Arc::new(CertifiedKey {
            cert: certs,
            key,
            ocsp: None,
            sct_list: None,
        });
        Ok(Self { cert_key })
    }
}

impl ResolvesClientCert for CertResolver {
    fn resolve(
        &self,
        _acceptable_issuers: &[&[u8]],
        _sigschemes: &[rustls::SignatureScheme],
    ) -> Option<Arc<rustls::sign::CertifiedKey>> {
        Some(self.cert_key.clone())
    }

    fn has_certs(&self) -> bool {
        true
    }
}

impl ResolvesServerCert for CertResolver {
    fn resolve(
        &self,
        _client_hello: rustls::server::ClientHello,
    ) -> Option<Arc<rustls::sign::CertifiedKey>> {
        Some(self.cert_key.clone())
    }
}

trait SignExt {
    fn as_signature_scheme(&self) -> rustls::SignatureScheme;
    fn as_signature_algorithm(&self) -> rustls::SignatureAlgorithm;
}

impl SignExt for Sign {
    fn as_signature_scheme(&self) -> SignatureScheme {
        match self {
            Self::RsaSsaPss(scheme) => match scheme.hash_kind() {
                HashKind::Sha2_256 => SignatureScheme::RSA_PSS_SHA256,
                HashKind::Sha2_512 => SignatureScheme::RSA_PSS_SHA512,
            },
        }
    }

    fn as_signature_algorithm(&self) -> SignatureAlgorithm {
        match self {
            Self::RsaSsaPss(..) => SignatureAlgorithm::RSA,
        }
    }
}

/// A new type around `Arc<C>` which allows rustls' traits to be implemented.
struct CredRef<C> {
    creds: Arc<C>,
}

impl<C> CredRef<C> {
    fn new(creds: Arc<C>) -> Self {
        Self { creds }
    }
}

impl<C> Deref for CredRef<C> {
    type Target = C;
    fn deref(&self) -> &Self::Target {
        &self.creds
    }
}

impl<C: CredsPriv + Send + Sync + 'static> SigningKey for CredRef<C> {
    fn choose_scheme(
        &self,
        offered: &[rustls::SignatureScheme],
    ) -> Option<Box<dyn rustls::sign::Signer>> {
        if offered.contains(&self.sign_kind().as_signature_scheme()) {
            Some(Box::new(Self::new(self.creds.clone())))
        } else {
            None
        }
    }

    fn algorithm(&self) -> rustls::SignatureAlgorithm {
        self.sign_kind().as_signature_algorithm()
    }
}

impl<C: CredsPriv + Send + Sync> rustls::sign::Signer for CredRef<C> {
    fn sign(&self, message: &[u8]) -> std::result::Result<Vec<u8>, rustls::Error> {
        self.creds
            .sign(&mut std::iter::once(message))
            .map(|sig| sig.take_data())
            .map_err(|err| rustls::Error::General(err.to_string()))
    }

    fn scheme(&self) -> rustls::SignatureScheme {
        self.sign_kind().as_signature_scheme()
    }
}