//! Small SSH binary codec and RFC4253 / OpenSSH EtM packet protection.
include!("ssh_protection_tests.inc");
use aes::cipher::{KeyIvInit, StreamCipher};
use hmac::{Hmac, Mac};
use rand_core::{OsRng, RngCore};
use sha2::{Sha256, Sha512};
use std::io::{self, Read, Write};
pub type Aes = ctr::Ctr128BE<aes::Aes128>;
pub fn fail(message: &str) -> io::Error {
    io::Error::other(format!("ssh: {message}"))
}
pub fn random(bytes: &mut [u8]) -> io::Result<()> {
    OsRng
        .try_fill_bytes(bytes)
        .map_err(|_| fail("random source unavailable"))
}
pub fn uint(out: &mut Vec<u8>, n: u32) {
    out.extend_from_slice(&n.to_be_bytes());
}
pub fn string(out: &mut Vec<u8>, bytes: &[u8]) {
    uint(out, bytes.len() as u32);
    out.extend_from_slice(bytes);
}
pub fn message(kind: u8, strings: &[&[u8]]) -> Vec<u8> {
    let mut out = vec![kind];
    for s in strings {
        string(&mut out, s);
    }
    out
}
pub struct Decoder<'a>(pub &'a [u8]);
impl<'a> Decoder<'a> {
    pub fn take(&mut self, len: usize) -> io::Result<&'a [u8]> {
        if len > self.0.len() {
            return Err(fail("truncated message"));
        }
        let (a, b) = self.0.split_at(len);
        self.0 = b;
        Ok(a)
    }
    pub fn byte(&mut self) -> io::Result<u8> {
        Ok(self.take(1)?[0])
    }
    pub fn uint(&mut self) -> io::Result<u32> {
        let b = self.take(4)?;
        Ok(u32::from_be_bytes([b[0], b[1], b[2], b[3]]))
    }
    pub fn string(&mut self) -> io::Result<&'a [u8]> {
        let len = self.uint()? as usize;
        self.take(len)
    }
    pub fn end(&self) -> io::Result<()> {
        if self.0.is_empty() {
            Ok(())
        } else {
            Err(fail("trailing message bytes"))
        }
    }
}
/// MAC key/tag sizes match OpenSSH mac.c; EtM framing follows PROTOCOL §1.5.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MacAlgorithm {
    Sha256,
    Sha256Etm,
    Sha512Etm,
}
impl MacAlgorithm {
    pub fn from_name(name: &[u8]) -> io::Result<Self> {
        match name {
            b"hmac-sha2-256" => Ok(Self::Sha256),
            b"hmac-sha2-256-etm@openssh.com" => Ok(Self::Sha256Etm),
            b"hmac-sha2-512-etm@openssh.com" => Ok(Self::Sha512Etm),
            _ => Err(fail("unsupported MAC")),
        }
    }
    pub fn key_len(self) -> usize {
        if self == Self::Sha512Etm { 64 } else { 32 }
    }
    fn etm(self) -> bool {
        self != Self::Sha256
    }
}
enum Tag {
    Sha256(Hmac<Sha256>),
    Sha512(Hmac<Sha512>),
}
impl Tag {
    fn bytes(self) -> Vec<u8> {
        match self {
            Self::Sha256(h) => h.finalize().into_bytes().to_vec(),
            Self::Sha512(h) => h.finalize().into_bytes().to_vec(),
        }
    }
    fn verify(self, tag: &[u8]) -> io::Result<()> {
        match self {
            Self::Sha256(h) => h.verify_slice(tag),
            Self::Sha512(h) => h.verify_slice(tag),
        }
        .map_err(|_| fail("packet MAC mismatch"))
    }
}
pub struct Protection {
    cipher: Aes,
    mac: Vec<u8>,
    algorithm: MacAlgorithm,
}
impl Protection {
    /// Historical non-EtM constructor, retained for existing callers.
    pub fn new(key: &[u8], iv: &[u8], mac: Vec<u8>) -> io::Result<Self> {
        Self::with_mac(key, iv, mac, MacAlgorithm::Sha256)
    }
    pub fn with_mac(
        key: &[u8],
        iv: &[u8],
        mac: Vec<u8>,
        algorithm: MacAlgorithm,
    ) -> io::Result<Self> {
        if mac.len() != algorithm.key_len() {
            return Err(fail("invalid MAC key size"));
        }
        Ok(Self {
            cipher: Aes::new_from_slices(key, iv).map_err(|_| fail("invalid key size"))?,
            mac,
            algorithm,
        })
    }
    fn tag(&self, seq: u32, bytes: &[u8]) -> io::Result<Tag> {
        if self.algorithm == MacAlgorithm::Sha512Etm {
            let mut h = <Hmac<Sha512> as Mac>::new_from_slice(&self.mac)
                .map_err(|_| fail("invalid MAC key"))?;
            h.update(&seq.to_be_bytes());
            h.update(bytes);
            Ok(Tag::Sha512(h))
        } else {
            let mut h = <Hmac<Sha256> as Mac>::new_from_slice(&self.mac)
                .map_err(|_| fail("invalid MAC key"))?;
            h.update(&seq.to_be_bytes());
            h.update(bytes);
            Ok(Tag::Sha256(h))
        }
    }
}
pub struct Packets<S> {
    pub stream: S,
    pub send: Option<Protection>,
    pub recv: Option<Protection>,
    tx: u32,
    rx: u32,
}
impl<S: Read + Write> Packets<S> {
    pub fn new(stream: S) -> Self {
        Self {
            stream,
            send: None,
            recv: None,
            tx: 0,
            rx: 0,
        }
    }
    pub fn write(&mut self, payload: &[u8]) -> io::Result<()> {
        if payload.is_empty() || payload.len() > 65536 {
            return Err(fail("outgoing packet size"));
        }
        let block = if self.send.is_some() { 16 } else { 8 };
        let etm = self.send.as_ref().is_some_and(|p| p.algorithm.etm());
        let mut padding = block - (payload.len() + if etm { 1 } else { 5 }) % block;
        if padding < 4 {
            padding += block;
        }
        let mut packet = Vec::with_capacity(payload.len() + padding + 5);
        uint(&mut packet, (payload.len() + padding + 1) as u32);
        packet.push(padding as u8);
        packet.extend_from_slice(payload);
        let start = packet.len();
        packet.resize(start + padding, 0);
        random(&mut packet[start..])?;
        // Legacy MAC authenticates plaintext; EtM authenticates length + ciphertext.
        let mut tag = None;
        if let Some(p) = &mut self.send {
            if !etm {
                tag = Some(p.tag(self.tx, &packet)?.bytes());
            }
            p.cipher
                .apply_keystream(&mut packet[if etm { 4 } else { 0 }..]);
            if etm {
                tag = Some(p.tag(self.tx, &packet)?.bytes());
            }
        }
        self.stream.write_all(&packet)?;
        if let Some(tag) = tag {
            self.stream.write_all(&tag)?;
        }
        self.stream.flush()?;
        self.tx = self.tx.wrapping_add(1);
        Ok(())
    }
    pub fn read(&mut self) -> io::Result<Vec<u8>> {
        let block = if self.recv.is_some() { 16 } else { 8 };
        let etm = self.recv.as_ref().is_some_and(|p| p.algorithm.etm());
        let prefix = if etm { 4 } else { block };
        let mut packet = vec![0; prefix];
        self.stream.read_exact(&mut packet)?;
        if !etm {
            if let Some(p) = &mut self.recv {
                p.cipher.apply_keystream(&mut packet);
            }
        }
        let len = Decoder(&packet).uint()? as usize;
        if !(5..=262144).contains(&len) {
            return Err(fail("invalid packet length"));
        }
        let aligned = len + if etm { 0 } else { 4 };
        if aligned % block != 0 || len + 4 < prefix {
            return Err(fail("invalid packet length"));
        }
        packet.resize(len + 4, 0);
        self.stream.read_exact(&mut packet[prefix..])?;
        if let Some(p) = &mut self.recv {
            let mut tag = vec![0; p.algorithm.key_len()];
            if etm {
                self.stream.read_exact(&mut tag)?;
                p.tag(self.rx, &packet)?.verify(&tag)?;
                // Never advance the cipher or inspect padding before authenticating EtM.
                p.cipher.apply_keystream(&mut packet[4..]);
            } else {
                p.cipher.apply_keystream(&mut packet[prefix..]);
                self.stream.read_exact(&mut tag)?;
                p.tag(self.rx, &packet)?.verify(&tag)?;
            }
        }
        self.rx = self.rx.wrapping_add(1);
        let padding = packet[4] as usize;
        if padding < 4 || padding + 1 >= len {
            return Err(fail("invalid packet padding"));
        }
        Ok(packet[5..len + 4 - padding].to_vec())
    }
    pub fn meaningful(&mut self) -> io::Result<Vec<u8>> {
        for _ in 0..256 {
            let p = self.read()?;
            match p.first() {
                Some(1) => return Err(fail("server disconnected")),
                Some(2 | 4) => continue, // IGNORE, DEBUG (never log remote text)
                Some(_) => return Ok(p),
                None => return Err(fail("empty packet")),
            }
        }
        Err(fail("too many non-progress messages"))
    }
}