puressh 0.2.7

//! `diffie-hellman-group{14,16,18}-sha{256,512}` (RFC 8268) and
//! `diffie-hellman-group-exchange-sha256` (RFC 4419), plus the SHA-1 forms
//! OpenSSH keeps for old peers when configured (`group1-sha1`,
//! `group14-sha1`, `group-exchange-sha1`; RFC 4253 / 4419).
//!
//! Backed by [`purecrypto::dh`]: the RFC 3526 safe-prime groups (`group14`,
//! `group16`, `group18`) and `DhPrivateKey::generate` for the ephemeral
//! exponents.

use alloc::vec::Vec;

use purecrypto::bignum::BoxedUint;
use purecrypto::dh::{DhGroup, DhPrivateKey, DhPublicKey, group14, group16, group18};
use purecrypto::hash::{Digest, Sha1, Sha256, Sha512};
use purecrypto::rng::{CryptoRng, RngCore};
use purecrypto::zeroize::Zeroizing;

use super::Kex;
use super::common::{
    KexContext, KexInitOut, KexOutput, SSH_MSG_KEX_DH_GEX_GROUP, SSH_MSG_KEX_DH_GEX_INIT,
    SSH_MSG_KEX_DH_GEX_REPLY, SSH_MSG_KEX_DH_GEX_REQUEST, SSH_MSG_KEX_ECDH_INIT,
    SSH_MSG_KEX_ECDH_REPLY,
};
use super::hash::{ExchangeHash, mpint_bytes};
use crate::error::{Error, Result};
use crate::format::{Reader, read_mpint};
use crate::hostkey::HostKeyVerify;

/// `diffie-hellman-group14-sha256` — RFC 3526 2048-bit group, SHA-256.
pub struct Group14Sha256;
impl Kex for Group14Sha256 {
    const NAME: &'static str = "diffie-hellman-group14-sha256";
    const HASH_LEN: usize = 32;
}

/// `diffie-hellman-group16-sha512` — RFC 3526 4096-bit group, SHA-512.
pub struct Group16Sha512;
impl Kex for Group16Sha512 {
    const NAME: &'static str = "diffie-hellman-group16-sha512";
    const HASH_LEN: usize = 64;
}

/// `diffie-hellman-group18-sha512` — RFC 3526 8192-bit group, SHA-512.
pub struct Group18Sha512;
impl Kex for Group18Sha512 {
    const NAME: &'static str = "diffie-hellman-group18-sha512";
    const HASH_LEN: usize = 64;
}

/// `diffie-hellman-group-exchange-sha256` — RFC 4419 GEX, SHA-256.
pub struct GexSha256;
impl Kex for GexSha256 {
    const NAME: &'static str = "diffie-hellman-group-exchange-sha256";
    const HASH_LEN: usize = 32;
}

/// `diffie-hellman-group14-sha1` — RFC 3526 2048-bit group, SHA-1 (legacy).
pub struct Group14Sha1;
impl Kex for Group14Sha1 {
    const NAME: &'static str = "diffie-hellman-group14-sha1";
    const HASH_LEN: usize = 20;
}

/// `diffie-hellman-group1-sha1` — RFC 2409 1024-bit Oakley group 2, SHA-1
/// (legacy; old devices offer nothing else).
pub struct Group1Sha1;
impl Kex for Group1Sha1 {
    const NAME: &'static str = "diffie-hellman-group1-sha1";
    const HASH_LEN: usize = 20;
}

/// `diffie-hellman-group-exchange-sha1` — RFC 4419 GEX, SHA-1 (legacy).
pub struct GexSha1;
impl Kex for GexSha1 {
    const NAME: &'static str = "diffie-hellman-group-exchange-sha1";
    const HASH_LEN: usize = 20;
}

/// RFC 2409 §6.2 Oakley group 2: the 1024-bit MODP safe prime, `g = 2`.
fn group1() -> DhGroup {
    const P: &str = "FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1\
                     29024E088A67CC74020BBEA63B139B22514A08798E3404DD\
                     EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245\
                     E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED\
                     EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE65381\
                     FFFFFFFFFFFFFFFF";
    let digits: Vec<u8> = P.bytes().filter(|b| !b.is_ascii_whitespace()).collect();
    let p: Vec<u8> = digits
        .chunks(2)
        .map(|h| {
            let n = |c: u8| (c as char).to_digit(16).unwrap_or(0) as u8;
            (n(h[0]) << 4) | n(h[1])
        })
        .collect();
    // A published safe prime (below the 2048-bit floor `from_custom`
    // enforces for peer-chosen groups); 256-bit exponents as for group14.
    DhGroup::from_custom_unchecked(BoxedUint::from_be_bytes(&p), BoxedUint::from_u64(2), 256)
        .expect("RFC 2409 group 2 is well-formed")
}

/// Client state for a fixed-group DH exchange.
pub struct DhClientState {
    group: DhGroup,
    secret: DhPrivateKey,
    e_mag: Vec<u8>,
}

/// Server reply for a fixed-group DH exchange.
pub struct DhServerReplyOut {
    /// Wire-format `SSH_MSG_KEXDH_REPLY` payload.
    pub payload: Vec<u8>,
    /// Shared secret and exchange hash.
    pub kex: KexOutput,
}

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

fn dh_client_init<R: RngCore + CryptoRng>(
    group: DhGroup,
    rng: &mut R,
) -> (DhClientState, KexInitOut) {
    let secret = DhPrivateKey::generate(group.clone(), rng);
    let pub_key = secret.public_key();
    let e_mag = pub_key.to_bytes();
    let mut payload = Vec::with_capacity(1 + 4 + e_mag.len() + 1);
    payload.push(SSH_MSG_KEX_ECDH_INIT);
    encode_mpint_into(&mut payload, &e_mag);
    (
        DhClientState {
            group,
            secret,
            e_mag,
        },
        KexInitOut { payload },
    )
}

fn encode_mpint_into(buf: &mut Vec<u8>, magnitude: &[u8]) {
    let m = strip_leading_zero(magnitude);
    if m.is_empty() {
        buf.extend_from_slice(&0u32.to_be_bytes());
        return;
    }
    if m[0] & 0x80 != 0 {
        buf.extend_from_slice(&((m.len() + 1) as u32).to_be_bytes());
        buf.push(0);
        buf.extend_from_slice(m);
    } else {
        buf.extend_from_slice(&(m.len() as u32).to_be_bytes());
        buf.extend_from_slice(m);
    }
}

fn dh_server_reply<D, R, S>(
    group: DhGroup,
    rng: &mut R,
    init_payload: &[u8],
    host_key: &S,
    ctx: &KexContext<'_>,
) -> Result<DhServerReplyOut>
where
    D: Digest,
    R: RngCore + CryptoRng,
    S: crate::hostkey::HostKey + ?Sized,
{
    let mut r = Reader::new(init_payload);
    let msg = r.read_u8()?;
    if msg != SSH_MSG_KEX_ECDH_INIT {
        return Err(Error::Protocol("expected SSH_MSG_KEXDH_INIT"));
    }
    let e_raw = read_mpint(&mut r)?;
    let peer =
        DhPublicKey::from_bytes(group.clone(), e_raw).map_err(|_| Error::Format("invalid DH e"))?;

    let secret = DhPrivateKey::generate(group.clone(), rng);
    let f_pub = secret.public_key();
    let f_mag = f_pub.to_bytes();
    let shared = secret
        .shared_secret(&peer)
        .map_err(|_| Error::Crypto("DH agreement failed"))?;
    // Locally-materialised raw DH shared secret (magnitude bytes). Wipe on
    // return; the SSH-mpint re-encoding lives on in `KexOutput.k`.
    let k_mag: Zeroizing<Vec<u8>> = Zeroizing::new(shared.into_bytes());

    let k_s = host_key.public_blob();
    let mut eh = ExchangeHash::<D>::new();
    eh.write_string(ctx.v_c);
    eh.write_string(ctx.v_s);
    eh.write_string(ctx.i_c);
    eh.write_string(ctx.i_s);
    eh.write_string(&k_s);
    eh.write_mpint(e_raw);
    eh.write_mpint(&f_mag);
    eh.write_mpint(&k_mag);
    let h = eh.finalize();

    let sig = host_key.sign(&h)?;

    let mut payload = Vec::with_capacity(1 + 4 + k_s.len() + 4 + f_mag.len() + 1 + 4 + sig.len());
    payload.push(SSH_MSG_KEX_ECDH_REPLY);
    payload.extend_from_slice(&(k_s.len() as u32).to_be_bytes());
    payload.extend_from_slice(&k_s);
    encode_mpint_into(&mut payload, &f_mag);
    payload.extend_from_slice(&(sig.len() as u32).to_be_bytes());
    payload.extend_from_slice(&sig);

    let k = mpint_bytes(&k_mag);
    Ok(DhServerReplyOut {
        payload,
        kex: KexOutput { k, h },
    })
}

fn dh_client_finish<D: Digest>(
    state: DhClientState,
    reply_payload: &[u8],
    verifier: &dyn HostKeyVerify,
    ctx: &KexContext<'_>,
) -> Result<KexOutput> {
    let mut r = Reader::new(reply_payload);
    let msg = r.read_u8()?;
    if msg != SSH_MSG_KEX_ECDH_REPLY {
        return Err(Error::Protocol("expected SSH_MSG_KEXDH_REPLY"));
    }
    let k_s = r.read_string()?;
    let f_raw = read_mpint(&mut r)?;
    let sig = r.read_string()?;

    let peer = DhPublicKey::from_bytes(state.group.clone(), f_raw)
        .map_err(|_| Error::Format("invalid DH f"))?;
    let shared = state
        .secret
        .shared_secret(&peer)
        .map_err(|_| Error::Crypto("DH agreement failed"))?;
    // Wipe the locally-materialised raw DH shared secret on return.
    let k_mag: Zeroizing<Vec<u8>> = Zeroizing::new(shared.into_bytes());

    let mut eh = ExchangeHash::<D>::new();
    eh.write_string(ctx.v_c);
    eh.write_string(ctx.v_s);
    eh.write_string(ctx.i_c);
    eh.write_string(ctx.i_s);
    eh.write_string(k_s);
    eh.write_mpint(&state.e_mag);
    eh.write_mpint(f_raw);
    eh.write_mpint(&k_mag);
    let h = eh.finalize();

    verifier.verify(&h, sig)?;

    let k = mpint_bytes(&k_mag);
    Ok(KexOutput { k, h })
}

macro_rules! dh_fixed_group_impl {
    ($ty:ident, $group:expr, $hash:ty) => {
        impl $ty {
            /// Algorithm name.
            pub const NAME: &'static str = <Self as Kex>::NAME;
            /// Exchange-hash output length in bytes.
            pub const HASH_LEN: usize = <Self as Kex>::HASH_LEN;

            /// Client side: generate `x`, compute `e = g^x mod p`, build the
            /// `SSH_MSG_KEXDH_INIT` payload.
            pub fn client_init<R: RngCore + CryptoRng>(rng: &mut R) -> (DhClientState, KexInitOut) {
                dh_client_init($group(), rng)
            }

            /// Server side.
            pub fn server_reply<R, S>(
                rng: &mut R,
                init_payload: &[u8],
                host_key: &S,
                ctx: &KexContext<'_>,
            ) -> Result<DhServerReplyOut>
            where
                R: RngCore + CryptoRng,
                S: crate::hostkey::HostKey + ?Sized,
            {
                dh_server_reply::<$hash, _, _>($group(), rng, init_payload, host_key, ctx)
            }

            /// Client side.
            pub fn client_finish(
                state: DhClientState,
                reply_payload: &[u8],
                verifier: &dyn HostKeyVerify,
                ctx: &KexContext<'_>,
            ) -> Result<KexOutput> {
                dh_client_finish::<$hash>(state, reply_payload, verifier, ctx)
            }
        }
    };
}

dh_fixed_group_impl!(Group14Sha256, group14, Sha256);
dh_fixed_group_impl!(Group16Sha512, group16, Sha512);
dh_fixed_group_impl!(Group18Sha512, group18, Sha512);
dh_fixed_group_impl!(Group14Sha1, group14, Sha1);
dh_fixed_group_impl!(Group1Sha1, group1, Sha1);

// -----------------------------------------------------------------------------
// RFC 4419 — diffie-hellman-group-exchange-sha256.
// -----------------------------------------------------------------------------

/// Client wishes for `(min, n, max)` from RFC 4419.
#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub struct GexRequest {
    /// Minimum acceptable prime bit size.
    pub min: u32,
    /// Preferred prime bit size.
    pub n: u32,
    /// Maximum acceptable prime bit size.
    pub max: u32,
}

impl GexRequest {
    /// OpenSSH's request for a connection whose symmetric algorithms need
    /// `need_bytes` of security (the largest cipher security length, block,
    /// IV or MAC key): 2048..=8192, preferring `dh_estimate(need * 8)`.
    pub fn for_strength(need_bytes: usize) -> Self {
        let bits = need_bytes * 8;
        let n = if bits <= 112 {
            2048
        } else if bits <= 128 {
            3072
        } else if bits <= 192 {
            7680
        } else {
            8192
        };
        GexRequest {
            min: 2048,
            n,
            max: 8192,
        }
    }
}

impl Default for GexRequest {
    fn default() -> Self {
        // OpenSSH defaults: 2048..=8192, preferring 8192.
        GexRequest {
            min: 2048,
            n: 8192,
            max: 8192,
        }
    }
}

/// Client state for a GEX exchange. Retains the request, the server-supplied
/// group, and the local exponent's `e` value across the three round-trips.
pub struct GexClientState {
    request: GexRequest,
    group: Option<DhGroup>,
    p_bytes: Vec<u8>,
    g_bytes: Vec<u8>,
    secret: Option<DhPrivateKey>,
    e_mag: Vec<u8>,
}

impl GexClientState {
    /// The `(min, n, max)` sent in `SSH_MSG_KEX_DH_GEX_REQUEST`.
    pub fn request(&self) -> GexRequest {
        self.request
    }

    /// The size in bits of the group the server chose (0 before it did).
    pub fn group_bits(&self) -> u32 {
        let p = self.p_bytes.iter().skip_while(|b| **b == 0);
        let lead = p.clone().next().copied().unwrap_or(0);
        let n = p.count() as u32;
        if n == 0 {
            0
        } else {
            (n - 1) * 8 + (8 - lead.leading_zeros())
        }
    }

    /// How many bits are set in our public value `e` (OpenSSH logs it).
    pub fn e_bits_set(&self) -> u32 {
        self.e_mag.iter().map(|b| b.count_ones()).sum()
    }
}

impl GexSha256 {
    /// Algorithm name.
    pub const NAME: &'static str = <Self as Kex>::NAME;
    /// Exchange-hash output length in bytes.
    pub const HASH_LEN: usize = <Self as Kex>::HASH_LEN;

    /// Step 1 (client → server): build `SSH_MSG_KEX_DH_GEX_REQUEST` carrying
    /// `(min, n, max)`.
    pub fn client_request(request: GexRequest) -> (GexClientState, KexInitOut) {
        let mut payload = Vec::with_capacity(1 + 12);
        payload.push(SSH_MSG_KEX_DH_GEX_REQUEST);
        payload.extend_from_slice(&request.min.to_be_bytes());
        payload.extend_from_slice(&request.n.to_be_bytes());
        payload.extend_from_slice(&request.max.to_be_bytes());
        (
            GexClientState {
                request,
                group: None,
                p_bytes: Vec::new(),
                g_bytes: Vec::new(),
                secret: None,
                e_mag: Vec::new(),
            },
            KexInitOut { payload },
        )
    }

    /// Step 2 (server → client): parse the request and emit
    /// `SSH_MSG_KEX_DH_GEX_GROUP` carrying `(p, g)`.
    ///
    /// `select` chooses the group to offer; the default is `group14` for
    /// `n <= 2048` and `group16` for larger requests, but callers can wrap
    /// this to return a custom safe prime per RFC 4419 §3.
    pub fn server_group(
        request_payload: &[u8],
        select: impl FnOnce(GexRequest) -> DhGroup,
    ) -> Result<(GexRequest, DhGroup, KexInitOut)> {
        let mut r = Reader::new(request_payload);
        let msg = r.read_u8()?;
        let request = match msg {
            SSH_MSG_KEX_DH_GEX_REQUEST => {
                let min = r.read_u32()?;
                let n = r.read_u32()?;
                let max = r.read_u32()?;
                GexRequest { min, n, max }
            }
            // SSH_MSG_KEX_DH_GEX_REQUEST_OLD only carries `n`; reuse it as
            // both bounds.
            30 => {
                let n = r.read_u32()?;
                GexRequest { min: n, n, max: n }
            }
            _ => return Err(Error::Protocol("expected SSH_MSG_KEX_DH_GEX_REQUEST")),
        };
        let group = select(request);
        let mut payload = Vec::new();
        payload.push(SSH_MSG_KEX_DH_GEX_GROUP);
        let p_bytes = group.p().to_be_bytes(group.byte_size());
        let g_bytes = group.g().to_be_bytes(group.byte_size());
        encode_mpint_into(&mut payload, &p_bytes);
        encode_mpint_into(&mut payload, &g_bytes);
        Ok((request, group, KexInitOut { payload }))
    }

    /// Step 3 (client → server): parse `SSH_MSG_KEX_DH_GEX_GROUP`, pick a
    /// fresh `x`, and emit `SSH_MSG_KEX_DH_GEX_INIT` carrying `e = g^x mod p`.
    pub fn client_init<R: RngCore + CryptoRng>(
        mut state: GexClientState,
        group_payload: &[u8],
        rng: &mut R,
    ) -> Result<(GexClientState, KexInitOut)> {
        let mut r = Reader::new(group_payload);
        let msg = r.read_u8()?;
        if msg != SSH_MSG_KEX_DH_GEX_GROUP {
            return Err(Error::Protocol("expected SSH_MSG_KEX_DH_GEX_GROUP"));
        }
        let p_raw = read_mpint(&mut r)?;
        let g_raw = read_mpint(&mut r)?;
        let p = BoxedUint::from_be_bytes(strip_leading_zero(p_raw));
        let g = BoxedUint::from_be_bytes(strip_leading_zero(g_raw));

        let bits = p.bit_len();
        if bits < state.request.min as usize || bits > state.request.max as usize {
            return Err(Error::Crypto("GEX group out of requested range"));
        }
        let priv_bits = bits.clamp(160, 512);
        // Size-checked above; no primality proof, as with OpenSSH's client
        // (dh_new_group + dh_pub_is_valid): the group comes from the server
        // the exchange is authenticated against, and proving an 8192-bit
        // prime safe took seconds per login.
        let group = DhGroup::from_custom_unchecked(p, g, priv_bits)
            .map_err(|_| Error::Format("invalid GEX group"))?;
        let secret = DhPrivateKey::generate(group.clone(), rng);
        let pub_key = secret.public_key();
        let e_mag = pub_key.to_bytes();

        state.p_bytes = strip_leading_zero(p_raw).to_vec();
        state.g_bytes = strip_leading_zero(g_raw).to_vec();
        state.group = Some(group);
        state.secret = Some(secret);
        state.e_mag = e_mag.clone();

        let mut payload = Vec::new();
        payload.push(SSH_MSG_KEX_DH_GEX_INIT);
        encode_mpint_into(&mut payload, &e_mag);
        Ok((state, KexInitOut { payload }))
    }

    /// Step 4 (server → client): parse `SSH_MSG_KEX_DH_GEX_INIT`, agree, sign,
    /// emit `SSH_MSG_KEX_DH_GEX_REPLY`. Returns the `(K, H)` and the wire
    /// payload to send.
    pub fn server_reply<R, S>(
        rng: &mut R,
        request: GexRequest,
        group: &DhGroup,
        init_payload: &[u8],
        host_key: &S,
        ctx: &KexContext<'_>,
    ) -> Result<DhServerReplyOut>
    where
        R: RngCore + CryptoRng,
        S: crate::hostkey::HostKey + ?Sized,
    {
        Self::server_reply_with::<Sha256, _, _>(rng, request, group, init_payload, host_key, ctx)
    }

    /// Step 5 (client side): parse `SSH_MSG_KEX_DH_GEX_REPLY`, verify the
    /// host-key signature on `H`, and return `(K, H)`.
    pub fn client_finish(
        state: GexClientState,
        reply_payload: &[u8],
        verifier: &dyn HostKeyVerify,
        ctx: &KexContext<'_>,
    ) -> Result<KexOutput> {
        Self::client_finish_with::<Sha256>(state, reply_payload, verifier, ctx)
    }

    /// Step 4 (server → client): parse `SSH_MSG_KEX_DH_GEX_INIT`, agree, sign,
    /// emit `SSH_MSG_KEX_DH_GEX_REPLY`. Returns the `(K, H)` and the wire
    /// payload to send.
    fn server_reply_with<D: Digest, R, S>(
        rng: &mut R,
        request: GexRequest,
        group: &DhGroup,
        init_payload: &[u8],
        host_key: &S,
        ctx: &KexContext<'_>,
    ) -> Result<DhServerReplyOut>
    where
        R: RngCore + CryptoRng,
        S: crate::hostkey::HostKey + ?Sized,
    {
        let mut r = Reader::new(init_payload);
        let msg = r.read_u8()?;
        if msg != SSH_MSG_KEX_DH_GEX_INIT {
            return Err(Error::Protocol("expected SSH_MSG_KEX_DH_GEX_INIT"));
        }
        let e_raw = read_mpint(&mut r)?;
        let peer = DhPublicKey::from_bytes(group.clone(), e_raw)
            .map_err(|_| Error::Format("invalid GEX e"))?;

        let secret = DhPrivateKey::generate(group.clone(), rng);
        let f_pub = secret.public_key();
        let f_mag = f_pub.to_bytes();
        let shared = secret
            .shared_secret(&peer)
            .map_err(|_| Error::Crypto("DH agreement failed"))?;
        // Wipe the locally-materialised raw GEX shared secret on return; the
        // SSH-mpint re-encoding lives on in `KexOutput.k`.
        let k_mag: Zeroizing<Vec<u8>> = Zeroizing::new(shared.into_bytes());
        let p_mag = group.p().to_be_bytes(group.byte_size());
        let g_mag = group.g().to_be_bytes(group.byte_size());

        let k_s = host_key.public_blob();
        let mut eh = ExchangeHash::<D>::new();
        eh.write_string(ctx.v_c);
        eh.write_string(ctx.v_s);
        eh.write_string(ctx.i_c);
        eh.write_string(ctx.i_s);
        eh.write_string(&k_s);
        eh.write_u32(request.min);
        eh.write_u32(request.n);
        eh.write_u32(request.max);
        eh.write_mpint(&p_mag);
        eh.write_mpint(&g_mag);
        eh.write_mpint(e_raw);
        eh.write_mpint(&f_mag);
        eh.write_mpint(&k_mag);
        let h = eh.finalize();

        let sig = host_key.sign(&h)?;

        let mut payload = Vec::new();
        payload.push(SSH_MSG_KEX_DH_GEX_REPLY);
        payload.extend_from_slice(&(k_s.len() as u32).to_be_bytes());
        payload.extend_from_slice(&k_s);
        encode_mpint_into(&mut payload, &f_mag);
        payload.extend_from_slice(&(sig.len() as u32).to_be_bytes());
        payload.extend_from_slice(&sig);

        let k = mpint_bytes(&k_mag);
        Ok(DhServerReplyOut {
            payload,
            kex: KexOutput { k, h },
        })
    }

    /// Step 5 (client side): parse `SSH_MSG_KEX_DH_GEX_REPLY`, verify the
    /// host-key signature on `H`, and return `(K, H)`.
    fn client_finish_with<D: Digest>(
        state: GexClientState,
        reply_payload: &[u8],
        verifier: &dyn HostKeyVerify,
        ctx: &KexContext<'_>,
    ) -> Result<KexOutput> {
        let group = state
            .group
            .as_ref()
            .ok_or(Error::Protocol("GEX state missing group"))?;
        let secret = state
            .secret
            .as_ref()
            .ok_or(Error::Protocol("GEX state missing secret"))?;

        let mut r = Reader::new(reply_payload);
        let msg = r.read_u8()?;
        if msg != SSH_MSG_KEX_DH_GEX_REPLY {
            return Err(Error::Protocol("expected SSH_MSG_KEX_DH_GEX_REPLY"));
        }
        let k_s = r.read_string()?;
        let f_raw = read_mpint(&mut r)?;
        let sig = r.read_string()?;

        let peer = DhPublicKey::from_bytes(group.clone(), f_raw)
            .map_err(|_| Error::Format("invalid GEX f"))?;
        let shared = secret
            .shared_secret(&peer)
            .map_err(|_| Error::Crypto("DH agreement failed"))?;
        // Wipe the locally-materialised raw GEX shared secret on return.
        let k_mag: Zeroizing<Vec<u8>> = Zeroizing::new(shared.into_bytes());

        let mut eh = ExchangeHash::<D>::new();
        eh.write_string(ctx.v_c);
        eh.write_string(ctx.v_s);
        eh.write_string(ctx.i_c);
        eh.write_string(ctx.i_s);
        eh.write_string(k_s);
        eh.write_u32(state.request.min);
        eh.write_u32(state.request.n);
        eh.write_u32(state.request.max);
        eh.write_mpint(&state.p_bytes);
        eh.write_mpint(&state.g_bytes);
        eh.write_mpint(&state.e_mag);
        eh.write_mpint(f_raw);
        eh.write_mpint(&k_mag);
        let h = eh.finalize();

        verifier.verify(&h, sig)?;

        let k = mpint_bytes(&k_mag);
        Ok(KexOutput { k, h })
    }
}

impl GexSha1 {
    /// Algorithm name.
    pub const NAME: &'static str = <Self as Kex>::NAME;
    /// Exchange-hash output length in bytes.
    pub const HASH_LEN: usize = <Self as Kex>::HASH_LEN;

    /// As [`GexSha256::server_reply`], hashing with SHA-1. (The request,
    /// group and init steps do not depend on the hash: use [`GexSha256`]'s.)
    pub fn server_reply<R, S>(
        rng: &mut R,
        request: GexRequest,
        group: &DhGroup,
        init_payload: &[u8],
        host_key: &S,
        ctx: &KexContext<'_>,
    ) -> Result<DhServerReplyOut>
    where
        R: RngCore + CryptoRng,
        S: crate::hostkey::HostKey + ?Sized,
    {
        GexSha256::server_reply_with::<Sha1, _, _>(rng, request, group, init_payload, host_key, ctx)
    }

    /// As [`GexSha256::client_finish`], hashing with SHA-1.
    pub fn client_finish(
        state: GexClientState,
        reply_payload: &[u8],
        verifier: &dyn HostKeyVerify,
        ctx: &KexContext<'_>,
    ) -> Result<KexOutput> {
        GexSha256::client_finish_with::<Sha1>(state, reply_payload, verifier, ctx)
    }
}

// -----------------------------------------------------------------------------
// moduli(5): the server's group-exchange groups.
// -----------------------------------------------------------------------------

/// One safe prime from a `moduli(5)` file (OpenSSH's `struct dhgroup`). The
/// group itself is built only for the entry a key exchange picks.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModuliEntry {
    /// The prime's size in bits (the file lists one less).
    pub bits: u32,
    /// The generator, big-endian.
    pub g: Vec<u8>,
    /// The prime, big-endian.
    pub p: Vec<u8>,
}

impl ModuliEntry {
    /// The Diffie-Hellman group, with a private exponent sized as for the
    /// RFC 3526 groups.
    pub fn group(&self) -> Option<DhGroup> {
        let p = BoxedUint::from_be_bytes(&self.p);
        let g = BoxedUint::from_be_bytes(&self.g);
        DhGroup::from_custom_unchecked(p, g, 512).ok()
    }
}

fn hex_bytes(s: &str) -> Option<Vec<u8>> {
    let digit = |c: u8| (c as char).to_digit(16).map(|d| d as u8);
    let b = s.as_bytes();
    let mut out = Vec::with_capacity(b.len().div_ceil(2));
    // An odd count of digits: the first stands alone.
    let (head, rest) = b.split_at(b.len() % 2);
    for &c in head {
        out.push(digit(c)?);
    }
    for pair in rest.chunks(2) {
        out.push((digit(pair[0])? << 4) | digit(pair[1])?);
    }
    Some(out)
}

/// Bits in a big-endian magnitude.
fn bit_len(b: &[u8]) -> u32 {
    let b = strip_leading_zero(b);
    match b.first() {
        None => 0,
        Some(lead) => (b.len() as u32 - 1) * 8 + (8 - lead.leading_zeros()),
    }
}

/// Parse a `moduli(5)` file as OpenSSH's `parse_prime` does: fields `time
/// type tests tries size generator modulus`; only tested (not composite)
/// safe primes whose stated size matches the modulus count. Comments, blank
/// and malformed lines are skipped.
pub fn parse_moduli(text: &str) -> Vec<ModuliEntry> {
    let mut out = Vec::new();
    for line in text.lines() {
        let line = line.trim_start();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }
        let f: Vec<&str> = line.split_whitespace().collect();
        let [_time, kind, tests, tries, size, g, p] = f[..] else {
            continue;
        };
        const MODULI_TYPE_SAFE: u32 = 2;
        const MODULI_TESTS_COMPOSITE: u32 = 0x01;
        let (Ok(kind), Ok(tests), Ok(tries), Ok(size)) = (
            kind.parse::<u32>(),
            tests.parse::<u32>(),
            tries.parse::<u32>(),
            size.parse::<u32>(),
        ) else {
            continue;
        };
        if kind != MODULI_TYPE_SAFE
            || tests > 0x1f
            || tests & MODULI_TESTS_COMPOSITE != 0
            || tests & !MODULI_TESTS_COMPOSITE == 0
            || tries == 0
            || tries > 1 << 30
            || size == 0
            || size > 64 * 1024
        {
            continue;
        }
        let (Some(g), Some(p)) = (hex_bytes(g), hex_bytes(p)) else {
            continue;
        };
        // The whole group is one bit larger than the listed size.
        let bits = size + 1;
        if bit_len(&p) != bits || bit_len(&g) < 2 {
            continue;
        }
        out.push(ModuliEntry {
            bits,
            g: strip_leading_zero(&g).to_vec(),
            p: strip_leading_zero(&p).to_vec(),
        });
    }
    out
}

/// Smallest and largest group sizes OpenSSH's server agrees to
/// (`DH_GRP_MIN`, `DH_GRP_MAX`).
pub const GEX_GROUP_MIN: u32 = 2048;
/// See [`GEX_GROUP_MIN`].
pub const GEX_GROUP_MAX: u32 = 8192;

/// OpenSSH's `choose_dh`: of the entries within the client's `[min, max]`
/// (clamped to [`GEX_GROUP_MIN`]..=[`GEX_GROUP_MAX`]), the size closest to
/// `n` from above (else the largest below it), and one of that size chosen
/// by `pick(count)` (a uniform index below `count`). `None` when no entry
/// fits, or the request is inconsistent (OpenSSH then refuses the exchange).
pub fn choose_moduli(
    entries: &[ModuliEntry],
    req: GexRequest,
    pick: impl FnOnce(usize) -> usize,
) -> Option<&ModuliEntry> {
    if req.max < req.min || req.n < req.min || req.max < req.n || req.max < GEX_GROUP_MIN {
        return None;
    }
    let min = req.min.max(GEX_GROUP_MIN);
    let max = req.max.min(GEX_GROUP_MAX);
    let want = req.n.clamp(GEX_GROUP_MIN, GEX_GROUP_MAX);
    let mut best = 0u32;
    let mut count = 0usize;
    for e in entries {
        if e.bits > max || e.bits < min {
            continue;
        }
        if (e.bits > want && e.bits < best) || (e.bits > best && best < want) {
            best = e.bits;
            count = 0;
        }
        if e.bits == best {
            count += 1;
        }
    }
    if count == 0 {
        return None;
    }
    let which = pick(count).min(count - 1);
    entries
        .iter()
        .filter(|e| e.bits >= min && e.bits <= max && e.bits == best)
        .nth(which)
}

/// OpenSSH's `dh_new_group_fallback`: the fixed group a server without a
/// usable `moduli(5)` file offers for a request up to `max` bits.
pub fn fallback_gex_group(max: u32) -> DhGroup {
    let max = max.min(GEX_GROUP_MAX);
    if max < 3072 {
        group14()
    } else if max < 6144 {
        group16()
    } else {
        group18()
    }
}

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

    #[test]
    fn moduli_parse_and_choose_as_openssh() {
        // group14's prime, listed as OpenSSH's moduli lists it.
        let p14 = group14().p().to_be_bytes(group14().byte_size());
        let hex: alloc::string::String = p14.iter().map(|b| alloc::format!("{b:02X}")).collect();
        let text = alloc::format!(
            "# Time Type Tests Tries Size Generator Modulus\n\
             20260511084302 2 6 100 2047 2 {hex}\n\
             20260511084303 2 6 100 2047 5 {hex}\n\
             20260511084304 1 6 100 2047 2 {hex}\n\
             20260511084305 2 7 100 2047 2 {hex}\n\
             20260511084306 2 6 100 2048 2 {hex}\n\
             20260511084307 2 6 100 2047 1 {hex}\n\
             garbage\n"
        );
        let entries = parse_moduli(&text);
        // Not safe, composite, wrong size and g = 1 are all dropped.
        assert_eq!(entries.len(), 2);
        assert_eq!(entries[0].bits, 2048);
        assert_eq!(entries[1].g, [5]);
        let req = |min, n, max| GexRequest { min, n, max };
        let chosen = choose_moduli(&entries, req(2048, 3072, 8192), |n| n - 1).unwrap();
        assert_eq!(chosen.g, [5]);
        let g = chosen.group().unwrap();
        assert_eq!(g.bit_size(), 2048);
        assert!(choose_moduli(&entries, req(3072, 4096, 8192), |_| 0).is_none());
        assert!(choose_moduli(&entries, req(4096, 2048, 8192), |_| 0).is_none());
        // Below the floor, but the floor itself fits.
        assert!(choose_moduli(&entries, req(1024, 1024, 2048), |_| 0).is_some());

        // Closest above `n` wins over closer below it.
        let mk = |bits| ModuliEntry {
            bits,
            g: alloc::vec![2],
            p: alloc::vec![],
        };
        let set = [mk(2048), mk(3072), mk(4096), mk(8192)];
        assert_eq!(
            choose_moduli(&set, req(2048, 3500, 8192), |_| 0).map(|e| e.bits),
            Some(4096)
        );
        assert_eq!(
            choose_moduli(&set, req(2048, 8192, 8192), |_| 0).map(|e| e.bits),
            Some(8192)
        );
        assert_eq!(
            choose_moduli(&set, req(2048, 7000, 8000), |_| 0).map(|e| e.bits),
            Some(4096)
        );
        assert_eq!(
            choose_moduli(&set, req(2048, 3072, 3500), |_| 0).map(|e| e.bits),
            Some(3072)
        );
        assert_eq!(fallback_gex_group(2048).bit_size(), 2048);
        assert_eq!(fallback_gex_group(4096).bit_size(), 4096);
        assert_eq!(fallback_gex_group(8192).bit_size(), 8192);
    }

    #[test]
    fn algorithm_names() {
        assert_eq!(Group14Sha256::NAME, "diffie-hellman-group14-sha256");
        assert_eq!(Group14Sha256::HASH_LEN, 32);
        assert_eq!(Group16Sha512::NAME, "diffie-hellman-group16-sha512");
        assert_eq!(Group16Sha512::HASH_LEN, 64);
        assert_eq!(Group18Sha512::NAME, "diffie-hellman-group18-sha512");
        assert_eq!(Group18Sha512::HASH_LEN, 64);
        assert_eq!(GexSha256::NAME, "diffie-hellman-group-exchange-sha256");
        assert_eq!(GexSha256::HASH_LEN, 32);
    }

    #[test]
    fn gex_request_payload_layout() {
        let req = GexRequest {
            min: 1024,
            n: 2048,
            max: 8192,
        };
        let (_state, out) = GexSha256::client_request(req);
        assert_eq!(out.payload[0], SSH_MSG_KEX_DH_GEX_REQUEST);
        assert_eq!(&out.payload[1..5], &1024u32.to_be_bytes());
        assert_eq!(&out.payload[5..9], &2048u32.to_be_bytes());
        assert_eq!(&out.payload[9..13], &8192u32.to_be_bytes());
    }

    #[test]
    fn mpint_helpers() {
        let mut v = Vec::new();
        encode_mpint_into(&mut v, &[0x00, 0x80, 0x01]);
        assert_eq!(v, &[0, 0, 0, 3, 0x00, 0x80, 0x01]);
        let mut v = Vec::new();
        encode_mpint_into(&mut v, &[0x00, 0x00, 0x01]);
        assert_eq!(v, &[0, 0, 0, 1, 0x01]);
        let mut v = Vec::new();
        encode_mpint_into(&mut v, &[]);
        assert_eq!(v, &[0, 0, 0, 0]);
    }
}