//! SSH keys and key pairs.

use crate::util::prompt_token;

use serde::{Deserialize, Serialize};
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;

/// An SSH public key.
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct PublicKey {
    key: String,
}

impl PublicKey {
    pub fn requires_token(&self) -> bool {
        self.key.as_str().starts_with("sk-ssh-")
    }
}

impl SshKey for PublicKey {
    fn new(key: String) -> Self {
        Self { key }
    }

    fn as_str(&self) -> &str {
        &self.key
    }
}

/// An SSH private key.
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct PrivateKey {
    key: String,
}

impl SshKey for PrivateKey {
    fn new(key: String) -> Self {
        Self { key }
    }

    fn as_str(&self) -> &str {
        &self.key
    }
}

/// A trait for SSH keys.
///
/// This provides the interface and some common functionality for
/// public and private SSH keys.
pub trait SshKey {
    fn new(key: String) -> Self
    where
        Self: Sized;
    fn as_str(&self) -> &str;

    fn read(filename: &Path) -> Result<Self, KeyError>
    where
        Self: Sized,
    {
        Ok(Self::new(read_string(filename)?))
    }

    fn write(&self, filename: &Path) -> Result<(), KeyError> {
        write_string(filename, self.as_str())?;
        Ok(())
    }
}

/// Errors from this SSH keys.
#[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.
///
/// These are the types of keys OpenSSH supports, some with
/// pre-defined lengths in bits. We list them for completeness, though
/// we actually only use Ed25519 keys, which are nice and short.
#[derive(Eq, PartialEq)]
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,

    /// Ed25519 key of fixed length, tied to a security token
    Ed25519Sk,
}

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",
            Self::Ed25519Sk => "ed25519-sk",
        }
    }

    /// 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,
            Self::Ed25519Sk => 1024,
        }
    }
}

/// A public/private key pair.
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct KeyPair {
    public: PublicKey,
    private: PrivateKey,
}

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

    /// Generate a new key pair of the desired kind.
    pub fn generate(kind: KeyKind) -> Result<Self, KeyError> {
        if kind == KeyKind::Ed25519Sk {
            prompt_token();
        }
        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("")
            .arg("-C")
            .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");

        let public_key = PublicKey::read(&public_key)?;
        let private_key = PrivateKey::read(&private_key)?;

        Ok(Self::new(public_key, private_key))
    }

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

    /// Private key of the pair, as a string.
    pub fn private(&self) -> &PrivateKey {
        &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)?)
}

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(())
}
