//! RFC8731 curve25519-sha256, RFC4253 key derivation, ed25519 host signatures.
include!("ssh_negotiation_tests.inc");
use super::ssh_wire::*;
use sha2::{Digest, Sha256};
use signature::Verifier;
use std::io::{self, Read, Write};
use x25519_dalek::{PublicKey, StaticSecret};

pub const VERSION: &[u8] = b"SSH-2.0-SausageGit";
const KEX: &[u8] = b"curve25519-sha256";
const HOST: &[u8] = b"ssh-ed25519";
const CIPHER: &[u8] = b"aes128-ctr";
const MAC: &[u8] = b"hmac-sha2-256-etm@openssh.com,hmac-sha2-512-etm@openssh.com,hmac-sha2-256";

pub fn banner<S: Read + Write>(stream: &mut S) -> io::Result<Vec<u8>> {
    stream.write_all(VERSION)?;
    stream.write_all(b"\r\n")?;
    stream.flush()?;
    // RFC4253 permits pre-banner lines, bounded here to 8KiB.
    let mut line = Vec::new();
    for _ in 0..8192 {
        let mut b = [0];
        stream.read_exact(&mut b)?;
        if b[0] == b'\n' {
            if line.last() == Some(&b'\r') {
                line.pop();
            }
            if line.starts_with(b"SSH-") {
                if !line.starts_with(b"SSH-2.0-") || line.len() > 253 {
                    return Err(fail("unsupported server version"));
                }
                return Ok(line);
            }
            line.clear();
        } else {
            line.push(b[0]);
            if line.len() > 1024 {
                return Err(fail("server banner too long"));
            }
        }
    }
    Err(fail("server banner missing"))
}
fn init() -> io::Result<Vec<u8>> {
    let mut p = vec![20];
    let mut cookie = [0; 16];
    random(&mut cookie)?;
    p.extend(cookie);
    for list in [
        KEX, HOST, CIPHER, CIPHER, MAC, MAC, b"none", b"none", b"", b"",
    ] {
        string(&mut p, list);
    }
    p.push(0);
    uint(&mut p, 0);
    Ok(p)
}
fn contains(list: &[u8], name: &[u8]) -> bool {
    list.split(|b| *b == b',').any(|n| n == name)
}
fn select<'a>(kind: &str, ours: &'a [u8], theirs: &[u8]) -> io::Result<&'a [u8]> {
    ours.split(|b| *b == b',')
        .find(|name| contains(theirs, name))
        .ok_or_else(|| {
            fail(&format!(
                "no matching {kind}; client offered {:?}; server offered {:?}",
                String::from_utf8_lossy(ours),
                String::from_utf8_lossy(theirs)
            ))
        })
}
#[derive(Debug)]
struct Negotiated {
    send: MacAlgorithm,
    recv: MacAlgorithm,
    wrong_guess: bool,
}
fn negotiate(server: &[u8]) -> io::Result<Negotiated> {
    let mut d = Decoder(server);
    if d.byte()? != 20 {
        return Err(fail("expected KEXINIT"));
    }
    d.take(16)?;
    let kex = d.string()?;
    let host = d.string()?;
    select("kex", KEX, kex)?;
    select("host-key", HOST, host)?;
    select("cipher client-to-server", CIPHER, d.string()?)?;
    select("cipher server-to-client", CIPHER, d.string()?)?;
    let send = MacAlgorithm::from_name(select("MAC client-to-server", MAC, d.string()?)?)?;
    let recv = MacAlgorithm::from_name(select("MAC server-to-client", MAC, d.string()?)?)?;
    select("compression client-to-server", b"none", d.string()?)?;
    select("compression server-to-client", b"none", d.string()?)?;
    d.string()?;
    d.string()?;
    let follows = d.byte()? != 0;
    d.uint()?;
    d.end()?;
    Ok(Negotiated {
        send,
        recv,
        wrong_guess: follows
            && (kex.split(|b| *b == b',').next() != Some(KEX)
                || host.split(|b| *b == b',').next() != Some(HOST)),
    })
}
/// RFC8731 §3.1 reinterprets the raw X25519 output as network-order unsigned K.
pub fn mpint(shared: &[u8; 32]) -> Vec<u8> {
    let mut value: Vec<_> = shared.iter().copied().skip_while(|b| *b == 0).collect();
    if value.first().is_some_and(|b| b & 128 != 0) {
        value.insert(0, 0);
    }
    let mut out = Vec::new();
    string(&mut out, &value);
    out
}
pub fn derive(k: &[u8], hash: &[u8], letter: u8, session: &[u8]) -> Vec<u8> {
    let mut h = Sha256::new();
    h.update(k);
    h.update(hash);
    h.update([letter]);
    h.update(session);
    h.finalize().to_vec()
}
/// RFC4253 §7.2 expansion: SHA-512 MAC needs 64 bytes even with SHA-256 KEX.
pub fn derive_len(k: &[u8], hash: &[u8], letter: u8, session: &[u8], len: usize) -> Vec<u8> {
    let mut key = derive(k, hash, letter, session);
    while key.len() < len {
        let mut h = Sha256::new();
        h.update(k);
        h.update(hash);
        h.update(&key);
        key.extend_from_slice(&h.finalize());
    }
    key.truncate(len);
    key
}
/// Authenticate a key exchange and switch protection in each direction at NEWKEYS.
pub fn exchange<S: Read + Write>(
    packets: &mut Packets<S>,
    server_version: &[u8],
    session: &mut Vec<u8>,
    server_init: Option<Vec<u8>>,
    mut verify: impl FnMut(&[u8]) -> io::Result<()>,
) -> io::Result<()> {
    let ours = init()?;
    packets.write(&ours)?;
    let theirs = match server_init {
        Some(p) => p,
        None => packets.meaningful()?,
    };
    let negotiated = negotiate(&theirs)?;
    let mut secret = [0; 32];
    random(&mut secret)?;
    let secret = StaticSecret::from(secret);
    let public = PublicKey::from(&secret);
    packets.write(&message(30, &[public.as_bytes()]))?;
    if negotiated.wrong_guess {
        packets.meaningful()?;
    }
    let reply = packets.meaningful()?;
    let mut d = Decoder(&reply);
    if d.byte()? != 31 {
        return Err(fail("expected curve25519 reply"));
    }
    let host = d.string()?;
    let remote = d.string()?;
    let signature = d.string()?;
    d.end()?;
    let remote: [u8; 32] = remote
        .try_into()
        .map_err(|_| fail("invalid curve25519 public key"))?;
    let shared = secret.diffie_hellman(&PublicKey::from(remote));
    if !shared.was_contributory() {
        return Err(fail("invalid all-zero shared secret"));
    }
    let k = mpint(shared.as_bytes());
    let mut transcript = Vec::new();
    for item in [
        VERSION,
        server_version,
        &ours,
        &theirs,
        host,
        public.as_bytes(),
        &remote,
    ] {
        string(&mut transcript, item);
    }
    transcript.extend_from_slice(&k);
    let hash = Sha256::digest(&transcript);
    let key = ssh_key::PublicKey::from_bytes(host).map_err(|_| fail("invalid host public key"))?;
    if key.algorithm() != ssh_key::Algorithm::Ed25519 {
        return Err(fail("unsupported host key"));
    }
    let signature =
        ssh_key::Signature::try_from(signature).map_err(|_| fail("invalid host signature"))?;
    Verifier::verify(&key, &hash, &signature)
        .map_err(|_| fail("host key exchange signature mismatch"))?;
    verify(host)?;
    if session.is_empty() {
        session.extend_from_slice(&hash);
    }
    let send = Protection::with_mac(
        &derive(&k, &hash, b'C', session)[..16],
        &derive(&k, &hash, b'A', session)[..16],
        derive_len(&k, &hash, b'E', session, negotiated.send.key_len()),
        negotiated.send,
    )?;
    let recv = Protection::with_mac(
        &derive(&k, &hash, b'D', session)[..16],
        &derive(&k, &hash, b'B', session)[..16],
        derive_len(&k, &hash, b'F', session, negotiated.recv.key_len()),
        negotiated.recv,
    )?;
    packets.write(&[21])?;
    packets.send = Some(send);
    if packets.meaningful()? != [21] {
        return Err(fail("expected NEWKEYS"));
    }
    packets.recv = Some(recv);
    Ok(())
}