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

// SSH服务器核心实现（Phase 3完整版）
// 参考OpenSSH sshd.c: complete KEX flow

use crate::ssh_server::version::VersionExchange;
use crate::ssh_server::packet::{SshPacket, PacketType};
use crate::ssh_server::kex::{KexProposal, KexResult};
use crate::ssh_server::kex_exchange::KexExchangeHandler;
use crate::ssh_server::kex_complete::{KexState};
use crate::ssh_server::crypto::SessionKeys;
use anyhow::Result;
use log::{info, warn, error, debug};
use std::net::{TcpListener, TcpStream};
use std::thread;
use std::io::Write;  // 导入Write trait（OpenSSH标准）

/// SSH服务器配置
pub struct SshServerConfig {
    pub port: u16,
    pub bind_address: String,
}

impl Default for SshServerConfig {
    fn default() -> Self {
        Self {
            port: 2024,
            bind_address: "127.0.0.1".to_string(),
        }
    }
}

/// SSH服务器主结构（Phase 3完整版）
pub struct SshServer {
    config: SshServerConfig,
}

impl SshServer {
    pub fn new(config: SshServerConfig) -> Self {
        Self { config }
    }
    
    pub fn run(&self) -> Result<()> {
        let bind_addr = format!("{}:{}", self.config.bind_address, self.config.port);
        let listener = TcpListener::bind(&bind_addr)?;
        
        info!("MarkBaseSSH server listening on {}", bind_addr);
        info!("Implementation: Complete SSH handshake (Phase 1-3)");
        
        for stream in listener.incoming() {
            match stream {
                Ok(stream) => {
                    let client_addr = stream.peer_addr()?;
                    info!("New SSH connection from {}", client_addr);
                    
                    thread::spawn(move || {
                        if let Err(e) = handle_connection_complete(stream) {
                            error!("Connection error: {}", e);
                        }
                    });
                }
                Err(e) => {
                    warn!("Failed to accept connection: {}", e);
                }
            }
        }
        
        Ok(())
    }
}

/// 处理完整SSH连接（Phase 1-3完整流程）
fn handle_connection_complete(stream: TcpStream) -> Result<()> {
    info!("Handling client connection (Phase 1-3 complete flow)");
    
    let mut stream = stream;
    
    // Phase 1: 版本交换
    let client_version = VersionExchange::exchange(&mut stream)?;
    info!("Version exchange: client={}, server=SSH-2.0-MarkBaseSSH_1.0", client_version);
    
    // Phase 2: 算法协商
    let (kex_result, server_kexinit, client_kexinit) = perform_kex_negotiation_complete(&mut stream)?;
    info!("KEX negotiation: KEX={}, Cipher={}", kex_result.kex_algorithm, kex_result.encryption_ctos);
    
    // Phase 3: 密钥交换完整流程
    perform_complete_kex_exchange(&mut stream, client_version, kex_result, server_kexinit, client_kexinit)?;
    info!("Key exchange completed, encryption channel ready");
    
    // 测试：发送disconnect
    send_disconnect(&mut stream, "Phase 3 test complete")?;
    
    info!("Phase 3 test completed successfully");
    Ok(())
}

/// 完整算法协商（返回KEXINIT payloads）
fn perform_kex_negotiation_complete(stream: &mut TcpStream) -> Result<(KexResult, SshPacket, SshPacket)> {
    info!("Starting complete KEX negotiation");
    
    // 1. 发送服务器KEXINIT
    let server_proposal = KexProposal::server_default();
    let server_kexinit = server_proposal.to_kexinit_packet()?;
    server_kexinit.write(stream)?;
    
    info!("Sent server KEXINIT (payload size: {} bytes)", server_kexinit.payload.len());
    
    // 2. 接收客户端KEXINIT
    let client_kexinit = SshPacket::read(stream)?;
    let client_proposal = KexProposal::from_kexinit_packet(&client_kexinit)?;
    
    info!("Received client KEXINIT (payload size: {} bytes)", client_kexinit.payload.len());
    
    // 3. 算法匹配
    let kex_result = KexResult::choose_algorithms(&server_proposal, &client_proposal)?;
    
    Ok((kex_result, server_kexinit, client_kexinit))
}

/// 完整密钥交换流程（Phase 3核心）
fn perform_complete_kex_exchange(
    stream: &mut TcpStream,
    client_version: String,
    kex_result: KexResult,
    server_kexinit: SshPacket,
    client_kexinit: SshPacket,
) -> Result<()> {
    info!("Starting complete key exchange flow");
    
    // 1. 创建密钥交换状态
    let mut kex_state = KexState::new(
        client_version,
        "SSH-2.0-MarkBaseSSH_1.0".to_string(),
        kex_result,
    )?;
    
    // 2. 保存KEXINIT payloads（用于Exchange Hash）
    kex_state.save_kexinit_payloads(&client_kexinit, &server_kexinit);
    
    // 3. 接收SSH_MSG_KEX_ECDH_INIT
    let kexdh_init = SshPacket::read(stream)?;
    info!("Received SSH_MSG_KEX_ECDH_INIT");
    
    // 4. 处理KEXDH_INIT并生成KEXDH_REPLY
    let kexdh_reply = kex_state.exchange_handler.handle_kexdh_init(&kexdh_init)?;
    kexdh_reply.write(stream)?;
    info!("Sent SSH_MSG_KEX_ECDH_REPLY");
    
    // 5. 发送SSH_MSG_NEWKEYS
    let newkeys_packet = KexState::send_newkeys()?;
    newkeys_packet.write(stream)?;
    kex_state.newkeys_sent = true;
    info!("Sent SSH_MSG_NEWKEYS");
    
    // 6. 接收SSH_MSG_NEWKEYS
    let client_newkeys = SshPacket::read(stream)?;
    kex_state.handle_newkeys(&client_newkeys)?;
    info!("Received SSH_MSG_NEWKEYS");
    
    // 7. 验证加密通道建立
    if kex_state.is_encryption_ready() {
        info!("Encryption channel established successfully");
    } else {
        return Err(anyhow::anyhow!("Encryption channel not ready"));
    }
    
    Ok(())
}

/// 发送SSH_MSG_DISCONNECT
fn send_disconnect(stream: &mut TcpStream, message: &str) -> Result<()> {
    let disconnect_packet = build_disconnect_packet(2, message, "en")?;
    disconnect_packet.write(stream)?;
    Ok(())
}

/// 构建SSH_MSG_DISCONNECT packet
fn build_disconnect_packet(reason_code: u32, description: &str, language: &str) -> Result<SshPacket> {
    use byteorder::{BigEndian, WriteBytesExt};
    
    let mut payload = Vec::new();
    payload.write_u8(PacketType::SSH_MSG_DISCONNECT as u8)?;
    payload.write_u32::<BigEndian>(reason_code)?;
    payload.write_u32::<BigEndian>(description.len() as u32)?;
    payload.write_all(description.as_bytes())?;
    payload.write_u32::<BigEndian>(language.len() as u32)?;
    payload.write_all(language.as_bytes())?;
    
    Ok(SshPacket::new(payload))
}

/// SSH服务器CLI入口
pub fn run_ssh_server(port: Option<u16>) -> Result<()> {
    let config = SshServerConfig {
        port: port.unwrap_or(2024),
        bind_address: "127.0.0.1".to_string(),
    };
    
    let server = SshServer::new(config);
    server.run()
}