//! Users' keys as SSH has them: Ed25519 public keys as an `authorized_keys` line has
//! them (`PublicKey`), and signatures as `ssh-keygen -Y sign` makes them (OpenSSH's
//! `PROTOCOL.sshsig`), checked (`verify`) and made (`sign`) here, so a user signs with
//! the keys and the tools it has, and a client of ours can sign alike without them.
//!
//! A signature is armored (`-----BEGIN SSH SIGNATURE-----`) around `SSHSIG`, its
//! version (1), the public key, a namespace (what it is for: a signature for one can't
//! be taken for another), a reserved string, the hash algorithm (`sha512` or `sha256`)
//! and the signature, over `SSHSIG`, the namespace, the reserved string, the algorithm
//! and the message's hash. Strings are SSH's: a big-endian `u32` length, then the bytes.

use base64::{Engine, engine::general_purpose::STANDARD as BASE64};
use ed25519_dalek::{Signature, Signer, SigningKey, VerifyingKey};
use sha2::{Digest, Sha256, Sha512};

const MAGIC: &[u8] = b"SSHSIG";
const KEY_TYPE: &str = "ssh-ed25519";
const BEGIN: &str = "-----BEGIN SSH SIGNATURE-----";
const END: &str = "-----END SSH SIGNATURE-----";

/// An Ed25519 public key.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct PublicKey(pub(crate) [u8; 32]);

/// A key's comment kept at most, in bytes (`PublicKey::kept_line`): what says whose and
/// where the key is, not a place for text.
pub(crate) const MAX_COMMENT: usize = 256;

impl PublicKey {
    /// From an `authorized_keys` line: `ssh-ed25519 <base64> [comment]`.
    pub(crate) fn parse(line: &str) -> Result<Self, String> {
        let mut parts = line.split_whitespace();
        let (Some(kind), Some(encoded)) = (parts.next(), parts.next()) else {
            return Err("not `ssh-ed25519 <base64> [comment]`".to_owned());
        };
        if kind != KEY_TYPE {
            return Err(format!("a {kind} key: only {KEY_TYPE} keys are taken"));
        }
        let blob = BASE64
            .decode(encoded)
            .map_err(|error| format!("its base64 doesn't decode: {error}"))?;
        Self::from_blob(&blob)
    }

    /// From SSH's encoding: the key type, then the key, each a string.
    fn from_blob(blob: &[u8]) -> Result<Self, String> {
        let mut reader = Reader(blob);
        let kind = reader.string()?;
        if kind != KEY_TYPE.as_bytes() {
            return Err(format!("not an {KEY_TYPE} key"));
        }
        let key: [u8; 32] = reader
            .string()?
            .try_into()
            .map_err(|_| "a key that isn't 32 bytes".to_owned())?;
        reader.end()?;
        Ok(PublicKey(key))
    }

    pub(crate) fn blob(&self) -> Vec<u8> {
        let mut blob = Vec::new();
        string(&mut blob, KEY_TYPE.as_bytes());
        string(&mut blob, &self.0);
        blob
    }

    /// As an `authorized_keys` line has it, without a comment.
    pub(crate) fn line(&self) -> String {
        format!("{KEY_TYPE} {}", BASE64.encode(self.blob()))
    }

    /// `line` (`ssh-ed25519 <base64> [comment]`) as a registry keeps it: the key as
    /// `line` makes it, then its comment, if any, as it was (its spaces between words
    /// made one). The comment is for whoever reads it: the key alone checks a signature.
    pub(crate) fn kept_line(line: &str) -> Result<String, String> {
        let key = Self::parse(line)?;
        let comment = line
            .split_whitespace()
            .skip(2)
            .collect::<Vec<&str>>()
            .join(" ");
        if comment.len() > MAX_COMMENT {
            return Err(format!(
                "the key's comment is longer than {MAX_COMMENT} bytes"
            ));
        }
        Ok(if comment.is_empty() {
            key.line()
        } else {
            format!("{} {comment}", key.line())
        })
    }

    /// Its fingerprint, as `ssh-keygen -l` shows it: `SHA256:` and the base64 of its
    /// encoding's SHA-256, unpadded.
    pub(crate) fn fingerprint(&self) -> String {
        let digest = Sha256::digest(self.blob());
        format!(
            "SHA256:{}",
            base64::engine::general_purpose::STANDARD_NO_PAD.encode(digest)
        )
    }
}

/// The key that made `armored`, a signature of `message` for `namespace`, if it is
/// one: valid, of an Ed25519 key, for that namespace.
pub(crate) fn verify(armored: &str, namespace: &str, message: &[u8]) -> Result<PublicKey, String> {
    verify_blob(&dearmor(armored)?, namespace, message)
}

/// The SSHSIG blob an armored signature holds.
pub(crate) fn dearmor(armored: &str) -> Result<Vec<u8>, String> {
    let body: String = armored
        .trim()
        .strip_prefix(BEGIN)
        .and_then(|rest| rest.trim_end().strip_suffix(END))
        .ok_or("not an armored SSH signature")?
        .split_whitespace()
        .collect();
    BASE64
        .decode(body)
        .map_err(|error| format!("its base64 doesn't decode: {error}"))
}

/// [`verify`] of a signature's SSHSIG blob (`dearmor`).
pub(crate) fn verify_blob(
    blob: &[u8],
    namespace: &str,
    message: &[u8],
) -> Result<PublicKey, String> {
    let mut reader = Reader(
        blob.strip_prefix(MAGIC)
            .ok_or("not an SSH signature (no SSHSIG)")?,
    );
    if reader.u32()? != 1 {
        return Err("an SSH signature of another version".to_owned());
    }
    let key = PublicKey::from_blob(reader.string()?)?;
    let signed_for = reader.string()?;
    if signed_for != namespace.as_bytes() {
        return Err(format!(
            "signed for {:?}, not {namespace:?}",
            String::from_utf8_lossy(signed_for)
        ));
    }
    let reserved = reader.string()?;
    let algorithm = reader.string()?;
    let hash = match algorithm {
        b"sha512" => Sha512::digest(message).to_vec(),
        b"sha256" => Sha256::digest(message).to_vec(),
        other => {
            return Err(format!("hashed with {}", String::from_utf8_lossy(other)));
        }
    };
    let mut signature = Reader(reader.string()?);
    reader.end()?;
    if signature.string()? != KEY_TYPE.as_bytes() {
        return Err(format!("not an {KEY_TYPE} signature"));
    }
    let bytes: [u8; 64] = signature
        .string()?
        .try_into()
        .map_err(|_| "a signature that isn't 64 bytes".to_owned())?;
    signature.end()?;
    let signed = signed_data(signed_for, reserved, algorithm, &hash);
    VerifyingKey::from_bytes(&key.0)
        .map_err(|_| "not a valid Ed25519 key".to_owned())?
        .verify_strict(&signed, &Signature::from_bytes(&bytes))
        .map_err(|_| "the signature doesn't match".to_owned())?;
    Ok(key)
}

/// `message` signed with `key` for `namespace`, armored, as `ssh-keygen -Y sign` makes
/// it (SHA-512).
pub(crate) fn sign(key: &SigningKey, namespace: &str, message: &[u8]) -> String {
    let public = PublicKey(key.verifying_key().to_bytes());
    let hash = Sha512::digest(message);
    let signed = signed_data(namespace.as_bytes(), b"", b"sha512", &hash);
    let mut signature = Vec::new();
    string(&mut signature, KEY_TYPE.as_bytes());
    string(&mut signature, &key.sign(&signed).to_bytes());
    let mut blob = MAGIC.to_vec();
    blob.extend_from_slice(&1u32.to_be_bytes());
    string(&mut blob, &public.blob());
    string(&mut blob, namespace.as_bytes());
    string(&mut blob, b"");
    string(&mut blob, b"sha512");
    string(&mut blob, &signature);
    let encoded = BASE64.encode(blob);
    let mut armored = format!("{BEGIN}\n");
    for line in encoded.as_bytes().chunks(70) {
        armored.push_str(std::str::from_utf8(line).expect("base64 is ASCII"));
        armored.push('\n');
    }
    armored.push_str(END);
    armored.push('\n');
    armored
}

/// What a signature is made over.
fn signed_data(namespace: &[u8], reserved: &[u8], algorithm: &[u8], hash: &[u8]) -> Vec<u8> {
    let mut signed = MAGIC.to_vec();
    string(&mut signed, namespace);
    string(&mut signed, reserved);
    string(&mut signed, algorithm);
    string(&mut signed, hash);
    signed
}

fn string(out: &mut Vec<u8>, bytes: &[u8]) {
    out.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
    out.extend_from_slice(bytes);
}

/// Reads SSH's encoding.
struct Reader<'a>(&'a [u8]);

impl<'a> Reader<'a> {
    fn u32(&mut self) -> Result<u32, String> {
        let (bytes, rest) = self.0.split_first_chunk::<4>().ok_or("cut short")?;
        self.0 = rest;
        Ok(u32::from_be_bytes(*bytes))
    }

    fn string(&mut self) -> Result<&'a [u8], String> {
        let length = self.u32()? as usize;
        if self.0.len() < length {
            return Err("cut short".to_owned());
        }
        let (bytes, rest) = self.0.split_at(length);
        self.0 = rest;
        Ok(bytes)
    }

    fn end(&self) -> Result<(), String> {
        if self.0.is_empty() {
            Ok(())
        } else {
            Err("more than it should hold".to_owned())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // OpenSSH's test data (regress/unittests/sshsig/testdata: ed25519.pub, ed25519.sig,
    // signed-data, namespace), made by `ssh-keygen -Y sign`.
    const OPENSSH_KEY: &str = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIIlizSIN3DQV78VPNjVvvwjgf+P5HqbPYCYu3BOMdjAE ED25519 test key";
    const OPENSSH_SIGNATURE: &str = "-----BEGIN SSH SIGNATURE-----
U1NIU0lHAAAAAQAAADMAAAALc3NoLWVkMjU1MTkAAAAgiWLNIg3cNBXvxU82NW+/COB/4/
keps9gJi7cE4x2MAQAAAAIdW5pdHRlc3QAAAAAAAAABnNoYTUxMgAAAFMAAAALc3NoLWVk
MjU1MTkAAABAihQsbUzuNEFflk5Tw1+H9aLS7tZQk0RG8KW1DtOmDYYnWe3D3UKiG3fcJa
DNg4vBWp1j1gLRiBMOF+gwYNegDg==
-----END SSH SIGNATURE-----
";
    const OPENSSH_MESSAGE: &[u8] = b"This is a test, this is only a test";

    /// A signature `ssh-keygen -Y sign` made checks, by its key, for its namespace and
    /// message only.
    #[test]
    fn what_ssh_keygen_signs_checks() {
        let key = PublicKey::parse(OPENSSH_KEY).unwrap();
        assert_eq!(
            verify(OPENSSH_SIGNATURE, "unittest", OPENSSH_MESSAGE),
            Ok(key)
        );
        assert!(verify(OPENSSH_SIGNATURE, "continuum-value", OPENSSH_MESSAGE).is_err());
        assert!(verify(OPENSSH_SIGNATURE, "unittest", b"This is a test").is_err());
        assert_eq!(PublicKey::parse(&key.line()), Ok(key));
        assert!(key.fingerprint().starts_with("SHA256:"));
        // Kept with its comment, which doesn't change the key.
        let commented =
            PublicKey::kept_line(&format!("{}  alice@laptop  2026", key.line())).unwrap();
        assert_eq!(commented, format!("{} alice@laptop 2026", key.line()));
        assert_eq!(PublicKey::parse(&commented), Ok(key));
        assert_eq!(PublicKey::kept_line(&key.line()), Ok(key.line()));
        let long = format!("{} {}", key.line(), "x".repeat(MAX_COMMENT + 1));
        assert!(PublicKey::kept_line(&long).is_err());
    }

    /// What `sign` makes checks, and reads as `ssh-keygen` writes it.
    #[test]
    fn signed_here_checks() {
        let key = SigningKey::from_bytes(&[7; 32]);
        let armored = sign(&key, "continuum-value", b"a value");
        assert!(armored.starts_with(BEGIN));
        assert_eq!(
            verify(&armored, "continuum-value", b"a value"),
            Ok(PublicKey(key.verifying_key().to_bytes()))
        );
        assert!(verify(&armored, "continuum-value", b"another").is_err());
    }

    #[test]
    fn other_keys_are_refused() {
        assert!(PublicKey::parse("ssh-rsa AAAAB3NzaC1yc2E= someone").is_err());
        assert!(PublicKey::parse("ssh-ed25519 !!!").is_err());
        assert!(PublicKey::parse("").is_err());
        assert!(verify("not a signature", "unittest", b"").is_err());
    }
}