//! Tools for automatically generating certificates
//!
//! Really, the only thing you will probably ever need from this module if you aren't
//! developing the project is the [`CertGenMode`] enum, which can be passed to
//! [`Server::set_certificate_generation_mode()`].  You won't even need to call any
//! methods on it or anything.
//!
//! [`Server::set_certificate_generation_mode()`]: crate::Server::set_certificate_generation_mode()
use anyhow::{anyhow, Context, Result, ensure};
#[cfg(feature = "certgen")]
use anyhow::bail;

use rustls::ServerConfig;

use std::{path::PathBuf, sync::Arc};
#[cfg(feature = "certgen")]
use std::{
    fs,
    io::{stdin, stdout, Write},
    path::Path,
};

use rustls::internal::msgs::handshake::DigitallySignedStruct;
use tokio_rustls::rustls;
use rustls::*;

#[derive(Clone, Debug)]
#[cfg(feature = "certgen")]
/// The mode to use for determining the domains to use for a new certificate.
///
/// Used to configure a [`Server`] using [`set_certificate_generation_mode()`]
///
/// [`Server`]: crate::Server
/// [`set_certificate_generation_mode()`]: crate::Server::set_certificate_generation_mode
pub enum CertGenMode {

    /// Do not generate any certificates.  Error if not available.
    None,

    /// Use a provided set of domains
    Preset(Vec<String>),

    /// Prompt the user using stdin/stdout to enter domains.
    Interactive,
}

#[cfg(feature = "certgen")]
impl CertGenMode {
    /// Generate a new self-signed certificate
    ///
    /// Assumes that certificates do not already exist, and will overwrite anything at the
    /// provided paths.  The paths provided should be paths to non-existant files which
    /// the program has access to write to.
    ///
    /// With very rare exceptions, end users should not ever need to call this method.
    /// Instead, just pass the [`CertGenMode`] to [`set_certificate_generation_mode()`]
    ///
    /// ## Errors
    ///
    /// Returns an error if [`CertGenMode::None`], or if there is an error generating the
    /// certificate, or writing to either of the provided files.
    ///
    /// [`set_certificate_generation_mode()`]: crate::Server::set_certificate_generation_mode
    pub fn gencert(self, cert: impl AsRef<Path>, key: impl AsRef<Path>) -> Result<rcgen::Certificate> {
        let (domains, interactive) = match self {
            Self::None => bail!("Automatic certificate generation disabled"),
            Self::Preset(domains) => (domains, false),
            Self::Interactive => (prompt_domains(), true),
        };

        let certificate = rcgen::generate_simple_self_signed(domains)
            .context("Could not generate a certificate with the given domains")?;

        fs::create_dir_all(
            cert.as_ref()
                .parent()
                .expect("Received directory as certificate path, should be a file.")
        )?;

        fs::create_dir_all(
            key.as_ref()
               .parent()
               .expect("Received directory as key path, should be a file.")
        )?;

        fs::write(cert.as_ref(), &certificate.serialize_pem()?.as_bytes())
            .context("Failed to write newly generated certificate to file")?;
        fs::write(key.as_ref(), &certificate.serialize_private_key_pem().as_bytes())
            .context("Failed to write newly generated private key to file")?;

        if interactive {
            println!("Certificate generated successfully!");
        }

        Ok(certificate)
    }

    /// Attempts to load a certificate/key from a file, or generate it if not found
    ///
    /// The produced certificate & key is immediately fed to a [`rustls::ServerConfig`]
    ///
    /// See [`CertGenMode::gencert()`] for more info.
    ///
    /// ## Errors
    ///
    /// Returns an error if a certificate is not found **and** cannot be generated.
    fn load_or_generate(self, to: &mut ServerConfig, cert: impl AsRef<Path>, key: impl AsRef<Path>) -> Result<()> {
        match (load_cert_chain(&cert.as_ref().into()), load_key(&key.as_ref().into())) {
            (Ok(cert_chain), Ok(key)) => {
                to.set_single_cert(cert_chain, key)
                    .context("Failed to use loaded TLS certificate")?;
            },
            (Err(e), _) | (_, Err(e)) => {
                warn!("Failed to load certificate from file: {}, now trying automatic generation", e);
                let cert = self.gencert(cert, key).context("Could not generate certificate")?;
                to.set_single_cert(
                    vec![rustls::Certificate(cert.serialize_der()?)],
                    rustls::PrivateKey(cert.serialize_private_key_der())
                )?;
            }
        }
        Ok(())
    }
}

#[cfg(feature = "certgen")]
/// Attempt to get domains by prompting the user
///
/// Guaranteed to return at least one domain.  The user is provided `localhost` as a
/// default.
///
/// ## Panics
/// Panics if reading from stdin or writing to stdout returns an error.
fn prompt_domains() -> Vec<String> {
    let mut domains = Vec::with_capacity(1);
    let mut input = String::with_capacity(8);
    println!("Now generating self-signed certificate...");
    print!("Please enter a domain (CN) for your certificate [localhost]: ");
    stdout().flush().unwrap();
    loop {
        stdin().read_line(&mut input).unwrap();
        let domain = input.trim();
        if domain.is_empty() {
            if domains.is_empty() {
                println!("Using `localhost` as domain.");
                domains.push("localhost".to_string());
            }
            return domains;
        } else {
            domains.push(domain.to_owned());
            input.clear();
            print!("Add another domain, or finish with a blank input: ");
            stdout().flush().unwrap();
        }
    }
}

pub fn tls_config(
    cert_path: &PathBuf,
    key_path: &PathBuf,
    #[cfg(feature = "certgen")]
    mode: CertGenMode,
) -> Result<Arc<ServerConfig>> {
    let mut config = ServerConfig::new(AllowAnonOrSelfsignedClient::new());

    #[cfg(feature = "certgen")]
    mode.load_or_generate(&mut config, cert_path, key_path)?;

    #[cfg(not(feature = "certgen"))] {
        let cert_chain = load_cert_chain(cert_path)
            .context("Failed to load TLS certificate")?;
        let key = load_key(key_path)
            .context("Failed to load TLS key")?;
        config.set_single_cert(cert_chain, key)
            .context("Failed to use loaded TLS certificate")?;
    }

    Ok(config.into())
}

fn load_cert_chain(cert_path: &PathBuf) -> Result<Vec<Certificate>> {
    let certs = std::fs::File::open(cert_path)
        .with_context(|| format!("Failed to open `{:?}`", cert_path))?;
    let mut certs = std::io::BufReader::new(certs);
    let certs = rustls::internal::pemfile::certs(&mut certs)
        .map_err(|_| anyhow!("failed to load certs `{:?}`", cert_path))?;

    Ok(certs)
}

fn load_key(key_path: &PathBuf) -> Result<PrivateKey> {
    let keys = std::fs::File::open(key_path)
        .with_context(|| format!("Failed to open `{:?}`", key_path))?;
    let mut keys = std::io::BufReader::new(keys);
    let mut keys = rustls::internal::pemfile::pkcs8_private_keys(&mut keys)
        .map_err(|_| anyhow!("failed to load key `{:?}`", key_path))?;

    ensure!(!keys.is_empty(), "no key found");

    let key = keys.swap_remove(0);

    Ok(key)
}

/// A client cert verifier that accepts all connections
///
/// Unfortunately, rustls doesn't provide a ClientCertVerifier that accepts self-signed
/// certificates, so we need to implement this ourselves.
struct AllowAnonOrSelfsignedClient { }

impl AllowAnonOrSelfsignedClient {

    /// Create a new verifier
    fn new() -> Arc<Self> {
        Arc::new(Self {})
    }

}

impl ClientCertVerifier for AllowAnonOrSelfsignedClient {

    fn client_auth_root_subjects(
        &self,
        _: Option<&webpki::DNSName>
    ) -> Option<DistinguishedNames> {
        Some(Vec::new())
    }

    fn client_auth_mandatory(&self, _sni: Option<&webpki::DNSName>) -> Option<bool> {
        Some(false)
    }

    // the below methods are a hack until webpki doesn't break with certain certs

    fn verify_client_cert(
        &self,
        _: &[Certificate],
        _: Option<&webpki::DNSName>
    ) -> Result<ClientCertVerified, TLSError> {
        Ok(ClientCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &Certificate,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, TLSError> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &Certificate,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, TLSError> {
        Ok(HandshakeSignatureValid::assertion())
    }
}