puressh 0.2.4

//! openssh-key-v1 PEM writer, key generators, and SHA-256 fingerprint.
//!
//! Mirrors the parser in `super`: encodes the binary payload, optionally
//! encrypts the inner block with `aes256-ctr` keyed by `bcrypt_pbkdf`, and
//! wraps the result in PEM at 70 columns.

use alloc::string::String;
use alloc::vec::Vec;

use purecrypto::bignum::{BoxedUint, inv_mod_boxed};
use purecrypto::cipher::{Aes256, Ctr};
use purecrypto::der::Reader as DerReader;
use purecrypto::ec::{BoxedEcdsaPrivateKey, CurveId, Ed25519PrivateKey};
use purecrypto::hash::{Digest, Sha256};
use purecrypto::kdf::bcrypt_pbkdf;
#[cfg(feature = "std")]
use purecrypto::rng::OsRng;
use purecrypto::rng::{CryptoRng, RngCore};
use purecrypto::rsa::BoxedRsaPrivateKey;
use purecrypto::zeroize::Zeroizing;

use super::{MAGIC, PrivateKey, PublicKey, base64};
use crate::error::{Error, Result};
use crate::format::{Writer, write_mpint};

/// Selectable ECDSA curve for [`PrivateKey::generate_ecdsa`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EcdsaCurve {
    /// NIST P-256.
    P256,
    /// NIST P-384.
    P384,
    /// NIST P-521.
    P521,
}

const PEM_BEGIN: &str = "-----BEGIN OPENSSH PRIVATE KEY-----";
const PEM_END: &str = "-----END OPENSSH PRIVATE KEY-----";
const PEM_WRAP: usize = 70;

/// `bcrypt_pbkdf` round count for new encrypted private keys. OpenSSH 9.x
/// defaults to 16; we bump to 24 so a single derivation costs ~50% more
/// CPU than upstream, matching `ssh-keygen -a 24`. This is a one-shot
/// hit at key-load time and well within the seconds-per-attempt range a
/// brute-forcer would face, but it noticeably slows offline cracking of
/// a leaked private key file. Existing keys keep whatever round count
/// they were originally written with — only re-encrypt operations (key
/// generation or `ssh-keygen -p`) pick up the new default.
const BCRYPT_ROUNDS: u32 = 24;
const SALT_LEN: usize = 16;
const KEY_LEN: usize = 32;
const IV_LEN: usize = 16;

impl PublicKey {
    /// SHA-256 fingerprint of the wire-format public-key blob, formatted as
    /// `"SHA256:<base64>"` with no `=` padding — matches `ssh-keygen -lf`.
    pub fn sha256_fingerprint(&self) -> String {
        let digest = Sha256::digest(&self.wire_blob());
        let mut s = base64::encode(&digest);
        while s.ends_with('=') {
            s.pop();
        }
        let mut out = String::with_capacity(7 + s.len());
        out.push_str("SHA256:");
        out.push_str(&s);
        out
    }

    /// The DNS SSHFP algorithm number (RFC 4255, 6594, 7479): RSA 1,
    /// ECDSA 3, Ed25519 4. `None` for a key type without one (the ML-DSA
    /// composite), which gets no SSHFP record.
    pub fn sshfp_algorithm(&self) -> Option<u8> {
        match self {
            PublicKey::Rsa { .. } => Some(1),
            PublicKey::EcdsaP256 { .. }
            | PublicKey::EcdsaP384 { .. }
            | PublicKey::EcdsaP521 { .. } => Some(3),
            PublicKey::Ed25519 { .. } => Some(4),
            PublicKey::MlDsa44Ed25519 { .. } => None,
        }
    }

    /// Approximate "bit length" reported by `ssh-keygen -l`: the modulus bit
    /// length for RSA, the curve bit length for ECDSA, 256 for Ed25519.
    pub fn bit_length(&self) -> u32 {
        match self {
            PublicKey::Ed25519 { .. } => 256,
            PublicKey::EcdsaP256 { .. } => 256,
            PublicKey::EcdsaP384 { .. } => 384,
            PublicKey::EcdsaP521 { .. } => 521,
            PublicKey::Rsa { n, .. } => modulus_bit_len(n),
            PublicKey::MlDsa44Ed25519 { .. } => 256,
        }
    }
}

fn modulus_bit_len(n: &[u8]) -> u32 {
    let mut i = 0usize;
    while i < n.len() && n[i] == 0 {
        i += 1;
    }
    if i == n.len() {
        return 0;
    }
    let leading = n[i].leading_zeros();
    ((n.len() - i) as u32) * 8 - leading
}

impl PrivateKey {
    /// Generate a new `ssh-mldsa44-ed25519@openssh.com` keypair.
    pub fn generate_mldsa44_ed25519<R: CryptoRng + RngCore>(rng: &mut R, comment: String) -> Self {
        let mut seeds = [0u8; 64];
        rng.fill_bytes(&mut seeds);
        let public = crate::hostkey::mldsa44_ed25519_public(&seeds);
        PrivateKey::MlDsa44Ed25519 {
            seeds,
            public,
            comment,
        }
    }

    /// Generate a new ssh-ed25519 keypair.
    pub fn generate_ed25519<R: CryptoRng + RngCore>(rng: &mut R, comment: String) -> Self {
        let sk = Ed25519PrivateKey::generate(rng);
        let seed = sk.to_bytes();
        let public = sk.public_key().to_bytes();
        PrivateKey::Ed25519 {
            seed,
            public,
            comment,
        }
    }

    /// Generate a new ECDSA keypair on the chosen NIST curve.
    pub fn generate_ecdsa<R: CryptoRng + RngCore>(
        rng: &mut R,
        curve: EcdsaCurve,
        comment: String,
    ) -> Self {
        let (cid, order_len) = match curve {
            EcdsaCurve::P256 => (CurveId::P256, 32usize),
            EcdsaCurve::P384 => (CurveId::P384, 48),
            EcdsaCurve::P521 => (CurveId::P521, 66),
        };
        let sk = BoxedEcdsaPrivateKey::generate(cid, rng);
        let point = sk.public_key().to_sec1();
        let d = scalar_from_ecdsa(&sk, order_len);
        match curve {
            EcdsaCurve::P256 => PrivateKey::EcdsaP256 { d, point, comment },
            EcdsaCurve::P384 => PrivateKey::EcdsaP384 { d, point, comment },
            EcdsaCurve::P521 => PrivateKey::EcdsaP521 { d, point, comment },
        }
    }

    /// Generate a new RSA keypair (public exponent 65537, Miller-Rabin = 40).
    pub fn generate_rsa<R: CryptoRng + RngCore>(
        rng: &mut R,
        bits: usize,
        comment: String,
    ) -> Result<Self> {
        if !(2048..=16384).contains(&bits) || !bits.is_multiple_of(256) {
            return Err(Error::Crypto("rsa: invalid key size"));
        }
        let e_u = BoxedUint::from_u64(65537);
        let sk = BoxedRsaPrivateKey::generate(bits, e_u, rng, 40);
        let RsaComponents { n, e, d, p, q } = rsa_components_from_pkcs1(&sk)?;
        let p_u = BoxedUint::from_be_bytes(strip_leading_zeros(&p));
        let q_u = BoxedUint::from_be_bytes(strip_leading_zeros(&q));
        let iqmp_u = inv_mod_boxed(&q_u, &p_u).ok_or(Error::Crypto("rsa: gcd(p,q) != 1"))?;
        let iqmp_len = strip_leading_zeros(&p).len();
        let iqmp = iqmp_u.to_be_bytes(iqmp_len);
        Ok(PrivateKey::Rsa {
            n,
            e,
            d,
            p,
            q,
            iqmp,
            comment,
        })
    }

    /// Encode this key as an openssh-key-v1 PEM document, using the host
    /// `OsRng` for the checkint and (if encrypting) the bcrypt salt. Pass
    /// `None` for an unencrypted key; pass a non-empty passphrase to encrypt
    /// with `aes256-ctr` keyed by `bcrypt_pbkdf`.
    ///
    /// Available only with the `std` feature; in `no_std` builds call
    /// [`PrivateKey::to_openssh_pem_with_rng`] and supply a `CryptoRng` of
    /// your choice.
    #[cfg(feature = "std")]
    pub fn to_openssh_pem(&self, passphrase: Option<&[u8]>) -> Result<String> {
        let mut rng = OsRng;
        self.to_openssh_pem_with_rng(&mut rng, passphrase)
    }

    /// `to_openssh_pem` variant that takes the entropy source explicitly. The
    /// RNG drives the openssh-key-v1 checkint and, if `passphrase` is set, the
    /// 16-byte bcrypt salt — both quantities must be unpredictable, so pass a
    /// real `CryptoRng`.
    pub fn to_openssh_pem_with_rng<R: CryptoRng + RngCore>(
        &self,
        rng: &mut R,
        passphrase: Option<&[u8]>,
    ) -> Result<String> {
        self.to_openssh_pem_with_rng_rounds(rng, passphrase, BCRYPT_ROUNDS)
    }

    /// [`Self::to_openssh_pem`] with `rounds` bcrypt KDF rounds instead of
    /// the default 16 (`ssh-keygen -a`); only used when encrypting.
    #[cfg(feature = "std")]
    pub fn to_openssh_pem_with_rounds(
        &self,
        passphrase: Option<&[u8]>,
        rounds: u32,
    ) -> Result<String> {
        let mut rng = OsRng;
        self.to_openssh_pem_with_rng_rounds(&mut rng, passphrase, rounds)
    }

    /// [`Self::to_openssh_pem_with_rng`] with an explicit bcrypt round count.
    pub fn to_openssh_pem_with_rng_rounds<R: CryptoRng + RngCore>(
        &self,
        rng: &mut R,
        passphrase: Option<&[u8]>,
        rounds: u32,
    ) -> Result<String> {
        if rounds == 0 {
            return Err(Error::Config("bcrypt rounds must be > 0"));
        }
        // `non_empty_pass` is the single source of truth for "encrypt?": both
        // the `block` choice and the encryption branch read from it, so the
        // branches can't disagree (which was the latent risk behind the prior
        // `passphrase.expect("checked above")`).
        let non_empty_pass = passphrase.filter(|p| !p.is_empty());
        let block = if non_empty_pass.is_some() { 16 } else { 8 };

        let inner = encode_inner_block(rng, self, block);
        let pub_blob = self.public_key().wire_blob();

        let (ciphername, kdfname, kdfoptions, payload): (_, _, _, Zeroizing<Vec<u8>>) =
            if let Some(pass) = non_empty_pass {
                let mut salt = [0u8; SALT_LEN];
                rng.fill_bytes(&mut salt);
                // `derived` is a wrapping passphrase-equivalent secret: from
                // it you can reconstruct the AES key and IV that protect the
                // private key blob. Wrap in `Zeroizing` so it is wiped from
                // memory at scope exit — without this it would linger in the
                // heap until the allocator reused the slot.
                let derived: Zeroizing<Vec<u8>> = Zeroizing::new(
                    bcrypt_pbkdf(pass, &salt, rounds, KEY_LEN + IV_LEN)
                        .map_err(|_| Error::Crypto("bcrypt_pbkdf: invalid parameters"))?,
                );
                // `key` and `iv` are derived material too; wrap so the stack
                // copies are wiped when the block ends. The `Aes256` instance
                // takes a copy internally but that is owned by the Ctr cipher
                // and dropped at end of scope; the copies we hold are wiped.
                let mut key: Zeroizing<[u8; KEY_LEN]> = Zeroizing::new([0u8; KEY_LEN]);
                key.copy_from_slice(&derived[..KEY_LEN]);
                let mut iv: Zeroizing<[u8; IV_LEN]> = Zeroizing::new([0u8; IV_LEN]);
                iv.copy_from_slice(&derived[KEY_LEN..KEY_LEN + IV_LEN]);
                let mut buf = inner;
                let mut ctr = Ctr::new(Aes256::new(&key), &iv);
                ctr.apply_keystream(&mut buf);
                let mut opts = Writer::new();
                opts.write_string(&salt);
                opts.write_u32(rounds);
                ("aes256-ctr", "bcrypt", opts.into_vec(), buf)
            } else {
                ("none", "none", Vec::new(), inner)
            };

        let mut w = Writer::new();
        w.write_raw(MAGIC);
        w.write_string(ciphername.as_bytes());
        w.write_string(kdfname.as_bytes());
        w.write_string(&kdfoptions);
        w.write_u32(1);
        w.write_string(&pub_blob);
        w.write_string(&payload);
        // For an unencrypted key `bin` (and its base64 form) carry the
        // cleartext secret; wipe the intermediates, leaving only the
        // returned PEM string as the caller's responsibility.
        let bin: Zeroizing<Vec<u8>> = Zeroizing::new(w.into_vec());

        let b64: Zeroizing<String> = Zeroizing::new(base64::encode(&bin));
        Ok(wrap_pem(&b64))
    }
}

impl PrivateKey {
    /// Encode the body of an `SSH_AGENTC_ADD_IDENTITY` request
    /// (`PROTOCOL.agent` §4.2): `string key-type` followed by the
    /// type-specific private-key fields and the trailing `string comment`.
    ///
    /// This is the same per-type field layout as the OpenSSH private-key
    /// inner block (see `encode_inner_block`) but **without** the leading
    /// check-ints or the trailing block padding — the agent frame is not an
    /// encrypted blob, so neither applies. The caller wraps the returned
    /// bytes in an agent message frame (`encode_message`).
    ///
    /// The returned `Vec` contains the cleartext private key; callers
    /// should wrap it in [`purecrypto::zeroize::Zeroizing`] (as the agent client does)
    /// so it is wiped once the request has been written.
    pub fn to_agent_add_body(&self) -> Vec<u8> {
        let mut w = Writer::new();
        match self {
            PrivateKey::Ed25519 {
                seed,
                public,
                comment,
            } => {
                w.write_string(b"ssh-ed25519");
                w.write_string(public);
                let mut sk: Zeroizing<[u8; 64]> = Zeroizing::new([0u8; 64]);
                sk[..32].copy_from_slice(seed);
                sk[32..].copy_from_slice(public);
                w.write_string(&sk[..]);
                w.write_string(comment.as_bytes());
            }
            PrivateKey::MlDsa44Ed25519 {
                seeds,
                public,
                comment,
            } => {
                w.write_string(b"ssh-mldsa44-ed25519@openssh.com");
                w.write_string(public);
                w.write_string(seeds);
                w.write_string(comment.as_bytes());
            }
            PrivateKey::EcdsaP256 { d, point, comment } => {
                w.write_string(b"ecdsa-sha2-nistp256");
                w.write_string(b"nistp256");
                w.write_string(point);
                write_mpint(&mut w, d);
                w.write_string(comment.as_bytes());
            }
            PrivateKey::EcdsaP384 { d, point, comment } => {
                w.write_string(b"ecdsa-sha2-nistp384");
                w.write_string(b"nistp384");
                w.write_string(point);
                write_mpint(&mut w, d);
                w.write_string(comment.as_bytes());
            }
            PrivateKey::EcdsaP521 { d, point, comment } => {
                w.write_string(b"ecdsa-sha2-nistp521");
                w.write_string(b"nistp521");
                w.write_string(point);
                write_mpint(&mut w, d);
                w.write_string(comment.as_bytes());
            }
            PrivateKey::Rsa {
                n,
                e,
                d,
                p,
                q,
                iqmp,
                comment,
            } => {
                w.write_string(b"ssh-rsa");
                write_mpint(&mut w, n);
                write_mpint(&mut w, e);
                write_mpint(&mut w, d);
                write_mpint(&mut w, iqmp);
                write_mpint(&mut w, p);
                write_mpint(&mut w, q);
                w.write_string(comment.as_bytes());
            }
        }
        w.into_vec()
    }

    /// Encode the body of an `SSH_AGENTC_ADD_IDENTITY` request adding this
    /// key under the certificate `cert_blob` (whose embedded key it must
    /// be): `string cert-type`, `string certificate`, then the key's
    /// private fields without those the certificate already carries
    /// (OpenSSH's `sshkey_private_serialize` for a certificate key), and
    /// the `string comment`.
    pub fn to_agent_add_cert_body(&self, cert_blob: &[u8]) -> Vec<u8> {
        let cert_type = crate::format::Reader::new(cert_blob)
            .read_string()
            .unwrap_or_default();
        let mut w = Writer::new();
        w.write_string(cert_type);
        w.write_string(cert_blob);
        let comment = match self {
            PrivateKey::Ed25519 {
                seed,
                public,
                comment,
            } => {
                w.write_string(public);
                let mut sk: Zeroizing<[u8; 64]> = Zeroizing::new([0u8; 64]);
                sk[..32].copy_from_slice(seed);
                sk[32..].copy_from_slice(public);
                w.write_string(&sk[..]);
                comment
            }
            PrivateKey::MlDsa44Ed25519 { seeds, comment, .. } => {
                w.write_string(seeds);
                comment
            }
            PrivateKey::EcdsaP256 { d, comment, .. }
            | PrivateKey::EcdsaP384 { d, comment, .. }
            | PrivateKey::EcdsaP521 { d, comment, .. } => {
                write_mpint(&mut w, d);
                comment
            }
            PrivateKey::Rsa {
                d,
                p,
                q,
                iqmp,
                comment,
                ..
            } => {
                write_mpint(&mut w, d);
                write_mpint(&mut w, iqmp);
                write_mpint(&mut w, p);
                write_mpint(&mut w, q);
                comment
            }
        };
        w.write_string(comment.as_bytes());
        w.into_vec()
    }
}

/// Encode the cleartext inner block (check-ints, private fields, comment,
/// padding). The result is the secret itself, so it is returned wrapped in
/// [`Zeroizing`] and wiped once the caller has encrypted / serialised it.
fn encode_inner_block<R: CryptoRng + RngCore>(
    rng: &mut R,
    pk: &PrivateKey,
    block: usize,
) -> Zeroizing<Vec<u8>> {
    let mut check = [0u8; 4];
    rng.fill_bytes(&mut check);
    let checkint = u32::from_be_bytes(check);

    let mut w = Writer::new();
    w.write_u32(checkint);
    w.write_u32(checkint);
    match pk {
        PrivateKey::Ed25519 {
            seed,
            public,
            comment,
        } => {
            w.write_string(b"ssh-ed25519");
            w.write_string(public);
            let mut sk: Zeroizing<[u8; 64]> = Zeroizing::new([0u8; 64]);
            sk[..32].copy_from_slice(seed);
            sk[32..].copy_from_slice(public);
            w.write_string(&sk[..]);
            w.write_string(comment.as_bytes());
        }
        PrivateKey::MlDsa44Ed25519 {
            seeds,
            public,
            comment,
        } => {
            w.write_string(b"ssh-mldsa44-ed25519@openssh.com");
            w.write_string(public);
            w.write_string(seeds);
            w.write_string(comment.as_bytes());
        }
        PrivateKey::EcdsaP256 { d, point, comment } => {
            w.write_string(b"ecdsa-sha2-nistp256");
            w.write_string(b"nistp256");
            w.write_string(point);
            write_mpint(&mut w, d);
            w.write_string(comment.as_bytes());
        }
        PrivateKey::EcdsaP384 { d, point, comment } => {
            w.write_string(b"ecdsa-sha2-nistp384");
            w.write_string(b"nistp384");
            w.write_string(point);
            write_mpint(&mut w, d);
            w.write_string(comment.as_bytes());
        }
        PrivateKey::EcdsaP521 { d, point, comment } => {
            w.write_string(b"ecdsa-sha2-nistp521");
            w.write_string(b"nistp521");
            w.write_string(point);
            write_mpint(&mut w, d);
            w.write_string(comment.as_bytes());
        }
        PrivateKey::Rsa {
            n,
            e,
            d,
            p,
            q,
            iqmp,
            comment,
        } => {
            w.write_string(b"ssh-rsa");
            write_mpint(&mut w, n);
            write_mpint(&mut w, e);
            write_mpint(&mut w, d);
            write_mpint(&mut w, iqmp);
            write_mpint(&mut w, p);
            write_mpint(&mut w, q);
            w.write_string(comment.as_bytes());
        }
    }
    let mut buf: Zeroizing<Vec<u8>> = Zeroizing::new(w.into_vec());
    let pad_remainder = buf.len() % block;
    if pad_remainder != 0 {
        let pad_n = block - pad_remainder;
        let mut next = 1u8;
        let mut count = 0usize;
        while count < pad_n {
            buf.push(next);
            next = next.wrapping_add(1);
            count += 1;
        }
    }
    buf
}

fn wrap_pem(b64: &str) -> String {
    let mut out = String::with_capacity(b64.len() + b64.len() / PEM_WRAP + 80);
    out.push_str(PEM_BEGIN);
    out.push('\n');
    let bytes = b64.as_bytes();
    let mut i = 0usize;
    let max_iter = bytes.len() / PEM_WRAP + 2;
    let mut guard = 0usize;
    while i < bytes.len() {
        let end = (i + PEM_WRAP).min(bytes.len());
        out.push_str(core::str::from_utf8(&bytes[i..end]).expect("base64 is ASCII"));
        out.push('\n');
        i = end;
        guard += 1;
        if guard > max_iter {
            break;
        }
    }
    out.push_str(PEM_END);
    out.push('\n');
    out
}

fn scalar_from_ecdsa(sk: &BoxedEcdsaPrivateKey, order_len: usize) -> Vec<u8> {
    // BoxedEcdsaPrivateKey exposes its scalar only through to_sec1_der, which
    // wraps the scalar as an OCTET STRING inside an ECPrivateKey SEQUENCE.
    let der = sk.to_sec1_der();
    if let Some(d) = parse_sec1_priv_octet(&der, order_len) {
        return d;
    }
    alloc::vec![0]
}

fn parse_sec1_priv_octet(der: &[u8], order_len: usize) -> Option<Vec<u8>> {
    let mut outer = DerReader::new(der);
    let mut seq = outer.read_sequence().ok()?;
    let _ver = seq.read_integer_bytes().ok()?;
    let priv_bytes = seq.read_octet_string().ok()?;
    if priv_bytes.len() < order_len {
        let mut v = alloc::vec![0u8; order_len];
        v[order_len - priv_bytes.len()..].copy_from_slice(priv_bytes);
        return Some(v);
    }
    Some(priv_bytes.to_vec())
}

fn strip_leading_zeros(b: &[u8]) -> &[u8] {
    let mut i = 0;
    while i < b.len() && b[i] == 0 {
        i += 1;
    }
    &b[i..]
}

struct RsaComponents {
    n: Vec<u8>,
    e: Vec<u8>,
    d: Vec<u8>,
    p: Vec<u8>,
    q: Vec<u8>,
}

fn rsa_components_from_pkcs1(sk: &BoxedRsaPrivateKey) -> Result<RsaComponents> {
    let der = sk.to_pkcs1_der();
    let mut outer = DerReader::new(&der);
    let mut seq = outer
        .read_sequence()
        .map_err(|_| Error::Format("rsa: bad pkcs1 sequence"))?;
    let _ver = seq
        .read_integer_bytes()
        .map_err(|_| Error::Format("rsa: bad version"))?;
    let n = seq
        .read_integer_bytes()
        .map_err(|_| Error::Format("rsa: bad n"))?
        .to_vec();
    let e = seq
        .read_integer_bytes()
        .map_err(|_| Error::Format("rsa: bad e"))?
        .to_vec();
    let d = seq
        .read_integer_bytes()
        .map_err(|_| Error::Format("rsa: bad d"))?
        .to_vec();
    let p = seq
        .read_integer_bytes()
        .map_err(|_| Error::Format("rsa: bad p"))?
        .to_vec();
    let q = seq
        .read_integer_bytes()
        .map_err(|_| Error::Format("rsa: bad q"))?
        .to_vec();
    Ok(RsaComponents { n, e, d, p, q })
}