markbase/markbase-core/src/ssh_server/kex.rs

// SSH密钥交换算法协商实现（Phase 2）
// 参考OpenSSH kex.c: kex_send_kexinit(), kex_choose_conf()

use crate::ssh_server::packet::{PacketType, SshPacket};
use anyhow::{anyhow, Result};
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use log::{debug, info};
use std::io::{Read, Write};

/// SSH算法类型（参考OpenSSH PROTOCOL定义）
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum AlgorithmType {
    KEX_ALGS = 0,             // 密钥交换算法
    SERVER_HOST_KEY_ALGS = 1, // 服务器主机密钥算法
    ENC_ALGS_CTOS = 2,        // 客户端到服务器加密算法
    ENC_ALGS_STOC = 3,        // 服务器到客户端加密算法
    MAC_ALGS_CTOS = 4,        // 客户端到服务器MAC算法
    MAC_ALGS_STOC = 5,        // 服务器到客户端MAC算法
    COMP_ALGS_CTOS = 6,       // 客户端到服务器压缩算法
    COMP_ALGS_STOC = 7,       // 服务器到客户端压缩算法
    LANGS_CTOS = 8,           // 客户端到服务器语言
    LANGS_STOC = 9,           // 服务器到客户端语言
}

/// SSH算法提议（参考OpenSSH kex.h: struct kex）
#[derive(Debug, Clone)]
pub struct KexProposal {
    pub kex_algorithms: String,              // 密钥交换算法列表
    pub server_host_key_algorithms: String,  // 主机密钥算法列表
    pub encryption_algorithms_ctos: String,  // 加密算法（客户端→服务器）
    pub encryption_algorithms_stoc: String,  // 加密算法（服务器→客户端）
    pub mac_algorithms_ctos: String,         // MAC算法（客户端→服务器）
    pub mac_algorithms_stoc: String,         // MAC算法（服务器→客户端）
    pub compression_algorithms_ctos: String, // 压缩算法（客户端→服务器）
    pub compression_algorithms_stoc: String, // 压缩算法（服务器→客户端）
    pub languages_ctos: String,              // 语言（客户端→服务器）
    pub languages_stoc: String,              // 语言（服务器→客户端）
    pub first_kex_packet_follows: bool,      // 是否立即发送第一个KEX packet
    pub reserved: u32,                       // 保留字段（0）
}

impl KexProposal {
    /// 创建默认算法提议（参考OpenSSH myproposal.h）
    pub fn server_default() -> Self {
        // 参考OpenSSH KEX_SERVER定义
        Self {
            // 密钥交换算法：优先Curve25519（推荐） + strict KEX extension
            kex_algorithms: "curve25519-sha256,curve25519-sha256@libssh.org,diffie-hellman-group14-sha256,ext-info-s,kex-strict-s-v00@openssh.com".to_string(),

            // 主机密钥算法：优先Ed25519
            server_host_key_algorithms: "ssh-ed25519,rsa-sha2-256,rsa-sha2-512".to_string(),

            // 加密算法：优先 AES-256-GCM（AEAD，性能优化），fallback 到 AES-CTR
            encryption_algorithms_ctos: "aes256-gcm@openssh.com,aes256-ctr,aes128-ctr".to_string(),
            encryption_algorithms_stoc: "aes256-gcm@openssh.com,aes256-ctr,aes128-ctr".to_string(),

            // MAC算法：HMAC-SHA256
            mac_algorithms_ctos: "hmac-sha2-256,hmac-sha2-512".to_string(),
            mac_algorithms_stoc: "hmac-sha2-256,hmac-sha2-512".to_string(),

            // 压缩算法：none优先
            compression_algorithms_ctos: "none,zlib".to_string(),
            compression_algorithms_stoc: "none,zlib".to_string(),

            // 语言：空
            languages_ctos: "".to_string(),
            languages_stoc: "".to_string(),

            first_kex_packet_follows: false,
            reserved: 0,
        }
    }

    /// 创建客户端默认提议（用于测试）
    pub fn client_default() -> Self {
        Self {
            kex_algorithms: "curve25519-sha256,diffie-hellman-group14-sha256".to_string(),
            server_host_key_algorithms: "ssh-ed25519,rsa-sha2-256".to_string(),
            encryption_algorithms_ctos: "aes256-gcm@openssh.com,aes256-ctr,aes128-ctr".to_string(),
            encryption_algorithms_stoc: "aes256-gcm@openssh.com,aes256-ctr,aes128-ctr".to_string(),
            mac_algorithms_ctos: "hmac-sha2-256".to_string(),
            mac_algorithms_stoc: "hmac-sha2-256".to_string(),
            compression_algorithms_ctos: "none".to_string(),
            compression_algorithms_stoc: "none".to_string(),
            languages_ctos: "".to_string(),
            languages_stoc: "".to_string(),
            first_kex_packet_follows: false,
            reserved: 0,
        }
    }

    /// 序列化到SSH_MSG_KEXINIT packet（参考OpenSSH kex_send_kexinit()）
    pub fn to_kexinit_packet(&self) -> Result<SshPacket> {
        let mut payload = Vec::new();

        // Packet type
        payload.write_u8(PacketType::SSH_MSG_KEXINIT as u8)?;

        // Cookie（16字节随机数，OpenSSH要求）
        let mut cookie = [0u8; 16];
        use rand::Rng;
        rand::thread_rng().fill(&mut cookie);
        payload.write_all(&cookie)?;

        // 10个算法列表（SSH string格式：length + data）
        write_ssh_string(&mut payload, &self.kex_algorithms)?;
        write_ssh_string(&mut payload, &self.server_host_key_algorithms)?;
        write_ssh_string(&mut payload, &self.encryption_algorithms_ctos)?;
        write_ssh_string(&mut payload, &self.encryption_algorithms_stoc)?;
        write_ssh_string(&mut payload, &self.mac_algorithms_ctos)?;
        write_ssh_string(&mut payload, &self.mac_algorithms_stoc)?;
        write_ssh_string(&mut payload, &self.compression_algorithms_ctos)?;
        write_ssh_string(&mut payload, &self.compression_algorithms_stoc)?;
        write_ssh_string(&mut payload, &self.languages_ctos)?;
        write_ssh_string(&mut payload, &self.languages_stoc)?;

        // first_kex_packet_follows（boolean）
        payload.write_u8(if self.first_kex_packet_follows { 1 } else { 0 })?;

        // reserved（u32）
        payload.write_u32::<BigEndian>(self.reserved)?;

        Ok(SshPacket::new(payload))
    }

    /// 从SSH_MSG_KEXINIT packet解析（参考OpenSSH kex_input_kexinit()）
    pub fn from_kexinit_packet(packet: &SshPacket) -> Result<Self> {
        let mut cursor = std::io::Cursor::new(packet.payload.as_slice()); // 使用as_slice()（Rust标准）

        // Packet type
        let packet_type = cursor.read_u8()?;
        if packet_type != PacketType::SSH_MSG_KEXINIT as u8 {
            return Err(anyhow!("Invalid packet type for KEXINIT"));
        }

        // Cookie（16字节，忽略）
        cursor.read_exact(&mut [0u8; 16])?;

        // 10个算法列表
        let kex_algorithms = read_ssh_string(&mut cursor)?;
        let server_host_key_algorithms = read_ssh_string(&mut cursor)?;
        let encryption_algorithms_ctos = read_ssh_string(&mut cursor)?;
        let encryption_algorithms_stoc = read_ssh_string(&mut cursor)?;
        let mac_algorithms_ctos = read_ssh_string(&mut cursor)?;
        let mac_algorithms_stoc = read_ssh_string(&mut cursor)?;
        let compression_algorithms_ctos = read_ssh_string(&mut cursor)?;
        let compression_algorithms_stoc = read_ssh_string(&mut cursor)?;
        let languages_ctos = read_ssh_string(&mut cursor)?;
        let languages_stoc = read_ssh_string(&mut cursor)?;

        // first_kex_packet_follows
        let first_kex_packet_follows = cursor.read_u8()? != 0;

        // reserved
        let reserved = cursor.read_u32::<BigEndian>()?;

        Ok(Self {
            kex_algorithms,
            server_host_key_algorithms,
            encryption_algorithms_ctos,
            encryption_algorithms_stoc,
            mac_algorithms_ctos,
            mac_algorithms_stoc,
            compression_algorithms_ctos,
            compression_algorithms_stoc,
            languages_ctos,
            languages_stoc,
            first_kex_packet_follows,
            reserved,
        })
    }
}

/// SSH算法协商结果（参考OpenSSH struct kex）
#[derive(Debug, Clone)]
pub struct KexResult {
    pub kex_algorithm: String,      // 选定的密钥交换算法
    pub host_key_algorithm: String, // 选定的主机密钥算法
    pub encryption_ctos: String,    // 选定的加密算法（客户端→服务器）
    pub encryption_stoc: String,    // 选定的加密算法（服务器→客户端）
    pub mac_ctos: String,           // 选定的MAC算法（客户端→服务器）
    pub mac_stoc: String,           // 选定的MAC算法（服务器→客户端）
    pub compression_ctos: String,   // 选定的压缩算法（客户端→服务器）
    pub compression_stoc: String,   // 选定的压缩算法（服务器→客户端）
}

/// 算法匹配逻辑（参考OpenSSH kex_choose_conf()）
impl KexResult {
    /// 从服务器和客户端提议中选择算法（参考OpenSSH kex_choose_conf()）
    pub fn choose_algorithms(server: &KexProposal, client: &KexProposal) -> Result<Self> {
        info!("Starting algorithm negotiation");

        // 算法匹配：优先客户端偏好（OpenSSH逻辑）
        // 参考OpenSSH：客户端列出的算法顺序为偏好顺序

        // 密钥交换算法匹配
        let kex_algorithm = match_algorithm(&client.kex_algorithms, &server.kex_algorithms)?;

        // 主机密钥算法匹配
        let host_key_algorithm = match_algorithm(
            &client.server_host_key_algorithms,
            &server.server_host_key_algorithms,
        )?;

        // 加密算法匹配
        let encryption_ctos = match_algorithm(
            &client.encryption_algorithms_ctos,
            &server.encryption_algorithms_ctos,
        )?;
        let encryption_stoc = match_algorithm(
            &client.encryption_algorithms_stoc,
            &server.encryption_algorithms_stoc,
        )?;

        // MAC算法匹配
        let mac_ctos = match_algorithm(&client.mac_algorithms_ctos, &server.mac_algorithms_ctos)?;
        let mac_stoc = match_algorithm(&client.mac_algorithms_stoc, &server.mac_algorithms_stoc)?;

        // 压缩算法匹配
        let compression_ctos = match_algorithm(
            &client.compression_algorithms_ctos,
            &server.compression_algorithms_ctos,
        )?;
        let compression_stoc = match_algorithm(
            &client.compression_algorithms_stoc,
            &server.compression_algorithms_stoc,
        )?;

        info!("Algorithm negotiation completed:");
        debug!("  KEX: {}", kex_algorithm);
        debug!("  Host key: {}", host_key_algorithm);
        debug!("  Encryption (C->S): {}", encryption_ctos);
        debug!("  Encryption (S->C): {}", encryption_stoc);
        debug!("  MAC (C->S): {}", mac_ctos);
        debug!("  MAC (S->C): {}", mac_stoc);

        Ok(Self {
            kex_algorithm,
            host_key_algorithm,
            encryption_ctos,
            encryption_stoc,
            mac_ctos,
            mac_stoc,
            compression_ctos,
            compression_stoc,
        })
    }
}

/// 算法匹配函数（参考OpenSSH match.c: match_list()）
fn match_algorithm(client_algs: &str, server_algs: &str) -> Result<String> {
    // 算法列表格式：name1,name2,name3,...
    let client_list: Vec<&str> = client_algs.split(',').collect();
    let server_list: Vec<&str> = server_algs.split(',').collect();

    // OpenSSH逻辑：按客户端偏好顺序匹配
    for client_alg in &client_list {
        if server_list.contains(client_alg) {
            return Ok(client_alg.to_string());
        }
    }

    Err(anyhow!(
        "No matching algorithm found: client={}, server={}",
        client_algs,
        server_algs
    ))
}

/// SSH string写入辅助函数（length + data）
fn write_ssh_string<W: Write>(writer: &mut W, s: &str) -> Result<()> {
    writer.write_u32::<BigEndian>(s.len() as u32)?;
    writer.write_all(s.as_bytes())?;
    Ok(())
}

/// SSH string读取辅助函数（length + data）
fn read_ssh_string<R: Read>(reader: &mut R) -> Result<String> {
    let length = reader.read_u32::<BigEndian>()?;
    let mut buffer = vec![0u8; length as usize];
    reader.read_exact(&mut buffer)?;
    Ok(String::from_utf8(buffer)?)
}

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

    #[test]
    fn test_kex_proposal_creation() {
        let proposal = KexProposal::server_default();
        assert!(proposal.kex_algorithms.contains("curve25519-sha256"));
    }

    #[test]
    fn test_kex_proposal_serialization() {
        let proposal = KexProposal::server_default();
        let packet = proposal.to_kexinit_packet().unwrap();
        assert!(packet.payload.len() > 0);
    }

    #[test]
    fn test_algorithm_matching() {
        let client = "curve25519-sha256,aes256-ctr";
        let server = "aes256-ctr,diffie-hellman-group14-sha256";

        let matched = match_algorithm(client, server).unwrap();
        assert_eq!(matched, "aes256-ctr"); // 按客户端顺序匹配
    }

    #[test]
    fn test_kex_negotiation() {
        let server = KexProposal::server_default();
        let client = KexProposal::client_default();

        let result = KexResult::choose_algorithms(&server, &client).unwrap();
        assert_eq!(result.kex_algorithm, "curve25519-sha256"); // 优先Curve25519
        assert_eq!(result.encryption_ctos, "aes256-gcm@openssh.com"); // AES-256-GCM (higher priority)
    }
}