//! Generate SSH certificates.

use crate::error::CAError;
use crate::key::{KeyError, KeyPair, PublicKey, SshKey};
use crate::util::prompt_token;
use serde::{Deserialize, Serialize};
use std::process::Command;
use tempfile::tempdir;

/// An SSH host certificate.
#[derive(Debug)]
pub struct HostCertificate {
    cert: PublicKey,
}

impl HostCertificate {
    fn new(cert: PublicKey) -> Self {
        Self { cert }
    }

    // This is only used by tests in this module, so it doesn't need
    // to be public.
    #[allow(dead_code)]
    fn as_str(&self) -> &str {
        self.cert.as_str()
    }
}

impl std::fmt::Display for HostCertificate {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.cert.as_str())
    }
}

/// An SSH user certificate.
#[derive(Debug)]
pub struct UserCertificate {
    cert: PublicKey,
}

impl UserCertificate {
    fn new(cert: PublicKey) -> Self {
        Self { cert }
    }

    // This is only used by tests in this module, so it doesn't need
    // to be public.
    #[allow(dead_code)]
    fn as_str(&self) -> &str {
        self.cert.as_str()
    }
}

impl std::fmt::Display for UserCertificate {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.cert.as_str())
    }
}

/// A trait for an SSH certificate authority.
///
/// This trait provides the common interface, and the shared
/// functionality, for host and user CA. There are distinct types for
/// those, to let the compiler enforce the right kind of CA being
/// used.
pub trait CertificateAuthority {
    /// The actual type of certificate created by this CA.
    type Certificate;

    /// Additional options for ssh-keygen.
    const OPTS: &'static [&'static str];

    /// String describing the kind of certificate this CA creates.
    const KIND: &'static str;

    /// Create a new CA from its key pair.
    fn new(name: String, keypair: KeyPair) -> Self;

    /// Create a certificate from a public key.
    fn certificate(key: PublicKey) -> Self::Certificate;

    /// Return reference to the key pair for this CA.
    fn keypair(&self) -> &KeyPair;

    /// Certify a host or user. Return a certificate of the correct
    /// type, if successful.
    fn certify(
        &self,
        subject_key: &PublicKey,
        valid_for: &str,
        principals: &[String],
    ) -> Result<Self::Certificate, CAError> {
        assert!(!principals.is_empty());

        let dirname = tempdir().map_err(|e| CAError::KeyError(KeyError::TempDir(e)))?;
        let ca_key = dirname.path().join("ca");
        let sub_key_pub = dirname.path().join("sub.pub");
        let cert = dirname.path().join("sub-cert.pub");

        let pair = self.keypair();
        pair.private().write(&ca_key).map_err(CAError::KeyError)?;

        subject_key.write(&sub_key_pub).map_err(CAError::KeyError)?;

        let mut principals_list = String::new();
        for (i, p) in principals.iter().enumerate() {
            if i > 0 {
                principals_list.push(',');
            }
            principals_list.push_str(p);
        }

        if pair.public().requires_token() {
            prompt_token();
        }
        let output = Command::new("ssh-keygen")
            .arg("-s")
            .arg(&ca_key)
            .args(Self::OPTS)
            .args(["-n", &principals_list])
            .args([
                "-I",
                &format!("certificate for {} {}", Self::KIND, principals_list),
            ])
            .args(["-V", valid_for])
            .arg(&sub_key_pub)
            .output()
            .map_err(|e| CAError::KeyError(KeyError::Run(e)))?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr).into_owned();
            return Err(CAError::KeyError(KeyError::CertError(stderr)));
        }

        let cert = PublicKey::read(&cert).map_err(CAError::KeyError)?;
        Ok(Self::certificate(cert))
    }
}

/// An SSH host CA.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HostCa {
    name: String,
    keypair: KeyPair,
}

impl CertificateAuthority for HostCa {
    type Certificate = HostCertificate;
    const OPTS: &'static [&'static str] = &["-h"];
    const KIND: &'static str = "host";

    fn new(name: String, keypair: KeyPair) -> Self {
        Self { name, keypair }
    }

    fn certificate(key: PublicKey) -> HostCertificate {
        HostCertificate::new(key)
    }

    fn keypair(&self) -> &KeyPair {
        &self.keypair
    }
}

/// An SSH user CA.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserCa {
    name: String,
    keypair: KeyPair,
}

impl CertificateAuthority for UserCa {
    type Certificate = UserCertificate;
    const OPTS: &'static [&'static str] = &[];
    const KIND: &'static str = "user";

    fn new(name: String, keypair: KeyPair) -> Self {
        Self { name, keypair }
    }

    fn certificate(key: PublicKey) -> UserCertificate {
        UserCertificate::new(key)
    }

    fn keypair(&self) -> &KeyPair {
        &self.keypair
    }
}

/// Parse an option validity period given by the user.
pub fn parse_validity(v: &str) -> Result<String, CAError> {
    let (v, unit) = if let Some(v) = v.strip_suffix('d') {
        (v, 'd')
    } else if let Some(v) = v.strip_suffix('h') {
        (v, 'h')
    } else if let Some(v) = v.strip_suffix('m') {
        (v, 'm')
    } else if let Some(v) = v.strip_suffix('w') {
        (v, 'w')
    } else {
        (v, 'd')
    };
    if let Ok(n) = v.parse::<usize>() {
        Ok(format!("+{n}{unit}"))
    } else {
        Err(CAError::Validity(v.into()))
    }
}

#[cfg(test)]
mod test {
    use super::{parse_validity, CertificateAuthority, HostCa, KeyPair, UserCa};
    use crate::key::KeyKind;

    fn valdity_period_is(input: &str, expected: &str) {
        let actual = parse_validity(input);
        assert!(actual.is_ok());
        let actual = actual.unwrap();
        assert_eq!(actual, expected);
    }

    #[test]
    fn unitless_validity_is_days() {
        valdity_period_is("123", "+123d");
    }

    #[test]
    fn unit_m_is_minutes() {
        valdity_period_is("123m", "+123m");
    }

    #[test]
    fn unit_w_is_weeks() {
        valdity_period_is("123w", "+123w");
    }

    #[test]
    fn unit_d_is_days() {
        valdity_period_is("123d", "+123d");
    }

    #[test]
    fn unit_h_is_hours() {
        valdity_period_is("123h", "+123h");
    }

    #[test]
    fn unparseable_validity_returns_error() {
        assert!(parse_validity("xyzzy").is_err());
    }

    #[test]
    fn certify_host_key() {
        let ca = KeyPair::generate(KeyKind::Ed25519).unwrap();
        let ca = HostCa::new("hostca".to_string(), ca);
        let host = KeyPair::generate(KeyKind::Ed25519).unwrap();
        let names = vec!["dummy".to_string()];
        let valid_for = parse_validity("1").unwrap();
        let cert = ca.certify(host.public(), &valid_for, &names).unwrap();
        assert_ne!(cert.as_str(), "");
    }

    #[test]
    fn certify_user_key() {
        let ca = KeyPair::generate(KeyKind::Ed25519).unwrap();
        let ca = UserCa::new("userca".to_string(), ca);
        let user = KeyPair::generate(KeyKind::Ed25519).unwrap();
        let valid_for = parse_validity("1").unwrap();
        let cert = ca
            .certify(user.public(), &valid_for, &["tomjon".to_string()])
            .unwrap();
        assert_ne!(cert.as_str(), "");
    }
}
