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

// SSH Packet基础结构定义
// 参考OpenSSH packet.c: ssh_packet_read(), ssh_packet_write()

use anyhow::{Result, anyhow};
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use std::io::{Read, Write};

/// SSH Packet类型（参考OpenSSH SSH_MSG_*定义）
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum PacketType {
    // SSH握手相关
    SSH_MSG_DISCONNECT = 1,
    SSH_MSG_IGNORE = 2,
    SSH_MSG_UNIMPLEMENTED = 3,
    SSH_MSG_DEBUG = 4,
    SSH_MSG_SERVICE_REQUEST = 5,
    SSH_MSG_SERVICE_ACCEPT = 6,
    SSH_MSG_KEXINIT = 20,
    SSH_MSG_NEWKEYS = 21,
    
    // 密钥交换相关
    SSH_MSG_KEXDH_INIT = 30,
    SSH_MSG_KEXDH_REPLY = 31,
    // 注意：Curve25519和DH使用相同的消息类型（30/31）
    // SSH_MSG_KEX_ECDH_INIT和SSH_MSG_KEX_ECDH_REPLY已在代码中注释
    // 使用SSH_MSG_KEXDH_INIT和SSH_MSG_KEXDH_REPLY代替
    
    // 认证相关
    SSH_MSG_USERAUTH_REQUEST = 50,
    SSH_MSG_USERAUTH_FAILURE = 51,
    SSH_MSG_USERAUTH_SUCCESS = 52,
    SSH_MSG_USERAUTH_BANNER = 53,
    SSH_MSG_USERAUTH_PK_OK = 60,
    
    // Channel相关
    SSH_MSG_GLOBAL_REQUEST = 80,
    SSH_MSG_REQUEST_SUCCESS = 81,
    SSH_MSG_REQUEST_FAILURE = 82,
    SSH_MSG_CHANNEL_OPEN = 90,
    SSH_MSG_CHANNEL_OPEN_CONFIRMATION = 91,
    SSH_MSG_CHANNEL_OPEN_FAILURE = 92,
    SSH_MSG_CHANNEL_WINDOW_ADJUST = 93,
    SSH_MSG_CHANNEL_DATA = 94,
    SSH_MSG_CHANNEL_EXTENDED_DATA = 95,
    SSH_MSG_CHANNEL_EOF = 96,
    SSH_MSG_CHANNEL_CLOSE = 97,
    SSH_MSG_CHANNEL_REQUEST = 98,
    SSH_MSG_CHANNEL_SUCCESS = 99,
    SSH_MSG_CHANNEL_FAILURE = 100,
}

/// SSH Packet结构（未加密状态）
/// 参考OpenSSH packet结构：
/// - packet_length (4 bytes)
/// - padding_length (1 byte)
/// - payload (variable)
/// - padding (variable)
/// - MAC (optional, encrypted阶段)
#[derive(Debug, Clone)]
pub struct SshPacket {
    pub packet_length: u32,
    pub padding_length: u8,
    pub payload: Vec<u8>,
    pub padding: Vec<u8>,
}

impl SshPacket {
    /// 创建新的SSH packet
    pub fn new(payload: Vec<u8>) -> Self {
        // 计算padding（SSH协议要求：packet总长度必须是block_size的倍数）
        // 参考OpenSSH：block_size = 8（未加密阶段）
        let block_size = 8;
        
        // packet_length = padding_length + payload_length + 1 (type byte)
        let payload_length = payload.len();
        let min_padding = 4;  // OpenSSH要求最少4字节padding
        
        // 计算padding长度
        let total_without_mac = 1 + payload_length;  // padding_length byte + payload
        let padding_needed = (block_size - (total_without_mac % block_size)) % block_size;
        let padding_length = std::cmp::max(min_padding as u32, padding_needed as u32) as u8;
        
        // 计算packet总长度
        let packet_length = 1 + payload_length + padding_length as usize;
        
        // 生成随机padding（简化：使用0，实际应随机）
        let padding = vec![0u8; padding_length as usize];
        
        Self {
            packet_length: packet_length as u32,
            padding_length,
            payload,
            padding,
        }
    }
    
    /// 写入packet到stream（未加密阶段）
    /// 参考OpenSSH packet_write_poll()
    pub fn write<T: Write>(&self, stream: &mut T) -> Result<()> {
        // 写入packet_length（BigEndian）
        stream.write_u32::<BigEndian>(self.packet_length)?;
        
        // 写入padding_length
        stream.write_u8(self.padding_length)?;
        
        // 写入payload
        stream.write_all(&self.payload)?;
        
        // 写入padding
        stream.write_all(&self.padding)?;
        
        stream.flush()?;
        Ok(())
    }
    
    /// 从stream读取packet（未加密阶段）
    /// 参考OpenSSH packet_read_poll()
    pub fn read<T: Read>(stream: &mut T) -> Result<Self> {
        // 读取packet_length（BigEndian）
        let packet_length = stream.read_u32::<BigEndian>()?;
        
        // 检查packet长度限制（OpenSSH限制：256KB）
        if packet_length > 256 * 1024 {
            return Err(anyhow!("Packet too large: {}", packet_length));
        }
        
        // 读取padding_length
        let padding_length = stream.read_u8()?;
        
        // 读取payload（packet_length - padding_length - 1）
        let payload_length = packet_length - padding_length as u32 - 1;
        let mut payload = vec![0u8; payload_length as usize];
        stream.read_exact(&mut payload)?;
        
        // 读取padding
        let mut padding = vec![0u8; padding_length as usize];
        stream.read_exact(&mut padding)?;
        
        Ok(Self {
            packet_length,
            padding_length,
            payload,
            padding,
        })
    }
    
    /// 获取payload中的packet type
    pub fn get_type(&self) -> Result<PacketType> {
        if self.payload.is_empty() {
            return Err(anyhow!("Empty payload"));
        }
        
        let type_byte = self.payload[0];
        
        // 转换为PacketType enum
        match type_byte {
            1 => Ok(PacketType::SSH_MSG_DISCONNECT),
            2 => Ok(PacketType::SSH_MSG_IGNORE),
            3 => Ok(PacketType::SSH_MSG_UNIMPLEMENTED),
            4 => Ok(PacketType::SSH_MSG_DEBUG),
            5 => Ok(PacketType::SSH_MSG_SERVICE_REQUEST),
            6 => Ok(PacketType::SSH_MSG_SERVICE_ACCEPT),
            20 => Ok(PacketType::SSH_MSG_KEXINIT),
            21 => Ok(PacketType::SSH_MSG_NEWKEYS),
            30 => Ok(PacketType::SSH_MSG_KEXDH_INIT),
            31 => Ok(PacketType::SSH_MSG_KEXDH_REPLY),
            50 => Ok(PacketType::SSH_MSG_USERAUTH_REQUEST),
            51 => Ok(PacketType::SSH_MSG_USERAUTH_FAILURE),
            52 => Ok(PacketType::SSH_MSG_USERAUTH_SUCCESS),
            53 => Ok(PacketType::SSH_MSG_USERAUTH_BANNER),
            80 => Ok(PacketType::SSH_MSG_GLOBAL_REQUEST),
            81 => Ok(PacketType::SSH_MSG_REQUEST_SUCCESS),
            82 => Ok(PacketType::SSH_MSG_REQUEST_FAILURE),
            90 => Ok(PacketType::SSH_MSG_CHANNEL_OPEN),
            91 => Ok(PacketType::SSH_MSG_CHANNEL_OPEN_CONFIRMATION),
            92 => Ok(PacketType::SSH_MSG_CHANNEL_OPEN_FAILURE),
            93 => Ok(PacketType::SSH_MSG_CHANNEL_WINDOW_ADJUST),
            94 => Ok(PacketType::SSH_MSG_CHANNEL_DATA),
            95 => Ok(PacketType::SSH_MSG_CHANNEL_EXTENDED_DATA),
            96 => Ok(PacketType::SSH_MSG_CHANNEL_EOF),
            97 => Ok(PacketType::SSH_MSG_CHANNEL_CLOSE),
            98 => Ok(PacketType::SSH_MSG_CHANNEL_REQUEST),
            99 => Ok(PacketType::SSH_MSG_CHANNEL_SUCCESS),
            100 => Ok(PacketType::SSH_MSG_CHANNEL_FAILURE),
            _ => Err(anyhow!("Unknown packet type: {}", type_byte)),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    
    #[test]
    fn test_packet_creation() {
        let payload = vec![PacketType::SSH_MSG_KEXINIT as u8];
        let packet = SshPacket::new(payload);
        
        assert!(packet.packet_length > 0);
        assert!(packet.padding_length >= 4);
    }
    
    #[test]
    fn test_packet_write_read() {
        let payload = vec![PacketType::SSH_MSG_KEXINIT as u8];
        let packet = SshPacket::new(payload);
        
        let mut buffer = Vec::new();
        packet.write(&mut buffer).unwrap();
        
        let mut cursor = Cursor::new(buffer);
        let read_packet = SshPacket::read(&mut cursor).unwrap();
        
        assert_eq!(packet.packet_length, read_packet.packet_length);
        assert_eq!(packet.payload, read_packet.payload);
    }
}