//! Generate SSH host keys and certificates.

use std::fs::{read, File, Permissions};
use std::io::Write;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::process::Command;
use tempfile::tempdir;

/// Errors from this module.
#[derive(Debug, thiserror::Error)]
pub enum KeyError {
    /// Could not generate a new key pair.
    #[error("ssh-keygen failed to generate a key: {0}")]
    KeyGen(String),

    /// Error creating a certificate.
    #[error("ssh-keygen failed to certify a key: {0}")]
    CertError(String),

    /// Error creating a temporary directory.
    #[error("Couldn't create temporary directory")]
    TempDir(#[source] std::io::Error),

    /// Error running ssh-keygen.
    #[error("Couldn't run ssh-keygen")]
    Run(#[source] std::io::Error),

    /// Error reading a file.
    #[error("Couldn't read file {0}")]
    Read(PathBuf, #[source] std::io::Error),

    /// Error writing a file.
    #[error("Couldn't write file {0}")]
    Write(PathBuf, #[source] std::io::Error),

    /// Error creating a file.
    #[error("Couldn't create file {0}")]
    Create(PathBuf, #[source] std::io::Error),

    /// Error setting file permissions.
    #[error("Couldn't set permissions for file {0}")]
    SetPerm(PathBuf, #[source] std::io::Error),

    /// Error parsing a string as UTF8.
    #[error(transparent)]
    Utf8Error(#[from] std::string::FromUtf8Error),
}

/// Type of SSH key.
pub enum KeyKind {
    /// RSA key of desired length in bits.
    RSA(u32),

    /// DSA of fixed length.
    DSA,

    /// ECDSA key of 256 bits.
    ECDSA256,

    /// ECDSA key of 384 bits.
    ECDSA384,

    /// ECDSA key of 521 bits.
    ECDSA521,

    /// Ed25519 key of fixed length.
    Ed25519,
}

impl KeyKind {
    /// Type of key as string for ssh-keygen -t option.
    pub fn as_str(&self) -> &str {
        match self {
            Self::RSA(_) => "rsa",
            Self::DSA => "dsa",
            Self::ECDSA256 => "ecdsa",
            Self::ECDSA384 => "ecdsa",
            Self::ECDSA521 => "ecdsa",
            Self::Ed25519 => "ed25519",
        }
    }

    /// Number of bits needed for the key.
    ///
    /// This is only really meaningful for RSA keys.
    pub fn bits(&self) -> u32 {
        match self {
            Self::RSA(bits) => *bits,
            Self::DSA => 1024,
            Self::ECDSA256 => 256,
            Self::ECDSA384 => 384,
            Self::ECDSA521 => 521,
            Self::Ed25519 => 1024,
        }
    }
}

/// A public/private key pair.
pub struct KeyPair {
    public: String,
    private: String,
}

impl KeyPair {
    /// Create pair from string representation.
    pub fn from_str(public: String, private: String) -> Self {
        Self { public, private }
    }

    /// Generate a new key pair of the desired kind.
    pub fn generate(kind: KeyKind) -> Result<Self, KeyError> {
        let dirname = tempdir().map_err(KeyError::TempDir)?;
        let private_key = dirname.path().join("key");
        let output = Command::new("ssh-keygen")
            .arg("-f")
            .arg(&private_key)
            .arg("-t")
            .arg(kind.as_str())
            .arg("-b")
            .arg(format!("{}", kind.bits()))
            .arg("-N")
            .arg("")
            .output()
            .map_err(KeyError::Run)?;

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

        let public_key = private_key.with_extension("pub");

        Ok(Self::from_str(
            read_string(&public_key)?,
            read_string(&private_key)?,
        ))
    }

    /// Public key of the pair, as a string.
    pub fn public(&self) -> &str {
        &self.public
    }

    /// Private key of the pair, as a string.
    pub fn private(&self) -> &str {
        &self.private
    }
}

fn read_string(filename: &Path) -> Result<String, KeyError> {
    let bytes = read(filename).map_err(|err| KeyError::Read(filename.to_path_buf(), err))?;
    Ok(String::from_utf8(bytes)?)
}

/// A key for SSH certificate authority.
///
/// This is used for creating host certificates.
pub struct CaKey {
    private: String,
}

impl CaKey {
    /// Create new CA key from a key pair.
    pub fn from(pair: KeyPair) -> Self {
        Self {
            private: pair.private().to_string(),
        }
    }

    /// Read CA key from a file.
    pub fn from_file(filename: &Path) -> Result<Self, KeyError> {
        let private = read_string(filename)?;
        Ok(Self { private })
    }

    /// Create a host certificate.
    ///
    /// Return as a string.
    pub fn certify_host(&self, host_key: &KeyPair, hostname: &str) -> Result<String, KeyError> {
        let dirname = tempdir().map_err(KeyError::TempDir)?;
        let ca_key = dirname.path().join("ca");
        let host_key_pub = dirname.path().join("host.pub");
        let cert = dirname.path().join("host-cert.pub");

        write_string(&ca_key, &self.private)?;
        write_string(&host_key_pub, host_key.public())?;

        let output = Command::new("ssh-keygen")
            .arg("-s")
            .arg(&ca_key)
            .arg("-h")
            .arg("-n")
            .arg(hostname)
            .arg("-I")
            .arg(format!("host key for {}", hostname))
            .arg(&host_key_pub)
            .output()
            .map_err(KeyError::Run)?;

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

        read_string(&cert)
    }
}

fn write_string(filename: &Path, s: &str) -> Result<(), KeyError> {
    let mut file =
        File::create(filename).map_err(|err| KeyError::Create(filename.to_path_buf(), err))?;
    let ro_user = Permissions::from_mode(0o600);
    file.set_permissions(ro_user)
        .map_err(|err| KeyError::SetPerm(filename.to_path_buf(), err))?;
    file.write_all(s.as_bytes())
        .map_err(|err| KeyError::Write(filename.to_path_buf(), err))?;
    Ok(())
}

#[cfg(test)]
mod keypair_test {
    use super::{CaKey, KeyKind, KeyPair};

    #[test]
    fn generate_key() {
        let pair = KeyPair::generate(KeyKind::Ed25519).unwrap();
        assert_ne!(pair.public(), "");
        assert_ne!(pair.private(), "");
        assert_ne!(pair.private(), pair.public());
    }

    #[test]
    fn certify_host_key() {
        let ca = KeyPair::generate(KeyKind::Ed25519).unwrap();
        let ca = CaKey::from(ca);
        let host = KeyPair::generate(KeyKind::Ed25519).unwrap();
        let cert = ca.certify_host(&host, "dummy").unwrap();
        assert_ne!(cert, "");
    }
}