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

// SSH服务器完整实现（Phase 1-7集成版）
// 参考OpenSSH sshd.c: complete SSH/SFTP flow

use crate::ssh_server::version::VersionExchange;
use crate::ssh_server::packet::{SshPacket, PacketType};
use crate::ssh_server::kex::{KexResult, KexProposal};
use crate::ssh_server::kex_complete::{KexState};
use crate::ssh_server::auth::{AuthHandler, AuthResult};
use crate::ssh_server::channel::{ChannelManager};
use crate::ssh_server::cipher::{EncryptionContext, EncryptedPacket};
use anyhow::{Result, anyhow};
use log::{info, warn, error, debug};
use std::net::{TcpListener, TcpStream};
use std::thread;
use std::io::{Read, Write};

/// 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 1-7完整版）
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/SFTP (Phase 1-7)");
        
        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-7完整流程）
fn handle_connection_complete(stream: TcpStream) -> Result<()> {
    info!("Handling client connection (Phase 1-7 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: 密钥交换完整流程
    let mut encryption_ctx = perform_complete_kex_exchange(&mut stream, client_version.clone(), kex_result, server_kexinit, client_kexinit)?;
    info!("Key exchange completed, encryption channel ready");
    
    // Phase 5: SSH认证（参考OpenSSH auth2.c）
    let mut auth_handler = AuthHandler::new()?;
    let auth_user = perform_ssh_auth(&mut stream, &mut auth_handler, &mut encryption_ctx)?;
    info!("SSH authentication succeeded: user={}", auth_user);
    
    // Phase 6: SSH Channel管理（参考OpenSSH channel.c）
    let mut channel_manager = ChannelManager::new();
    
    // Phase 6-7: SSH服务循环（处理channel请求）
    handle_ssh_service_loop(&mut stream, &mut channel_manager)?;
    
    info!("SSH session 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<EncryptionContext> {
    info!("Starting complete key exchange flow");
    
    let mut kex_state = KexState::new(
        client_version,
        "SSH-2.0-MarkBaseSSH_1.0".to_string(),
        kex_result,
    )?;
    
    kex_state.save_kexinit_payloads(&client_kexinit, &server_kexinit);
    
    let kexdh_init = SshPacket::read(stream)?;
    info!("Received SSH_MSG_KEX_ECDH_INIT");
    
    let kexdh_reply = kex_state.exchange_handler.handle_kexdh_init(
        &kexdh_init,
        &kex_state.client_version,
        &kex_state.server_version,
        &kex_state.client_kexinit_payload,
        &kex_state.server_kexinit_payload,
    )?;
    kexdh_reply.write(stream)?;
    info!("Sent SSH_MSG_KEX_ECDH_REPLY");
    
    let newkeys_packet = KexState::send_newkeys()?;
    newkeys_packet.write(stream)?;
    kex_state.newkeys_sent = true;
    info!("Sent SSH_MSG_NEWKEYS");
    
    let client_newkeys = SshPacket::read(stream)?;
    kex_state.handle_newkeys(&client_newkeys)?;
    info!("Received SSH_MSG_NEWKEYS");
    
    if kex_state.is_encryption_ready() {
        info!("Encryption channel established successfully");
    } else {
        return Err(anyhow::anyhow!("Encryption channel not ready"));
    }
    
    let session_keys = kex_state.exchange_handler.compute_session_keys()?;
    let encryption_ctx = EncryptionContext::from_session_keys(&session_keys);
    
    Ok(encryption_ctx)
}

/// SSH认证流程（Phase 5）
fn perform_ssh_auth(
    stream: &mut TcpStream, 
    auth_handler: &mut AuthHandler,
    encryption_ctx: &mut EncryptionContext,
) -> Result<String> {
    info!("Starting SSH authentication");
    info!("Encryption context: key_ctos_len={}, key_stoc_len={}, iv_ctos_len={}, iv_stoc_len={}",
        encryption_ctx.encryption_key_ctos.len(),
        encryption_ctx.encryption_key_stoc.len(),
        encryption_ctx.iv_ctos.len(),
        encryption_ctx.iv_stoc.len()
    );
    
    // OpenSSH strict KEX: SSH_MSG_EXT_INFO may be sent before SSH_MSG_SERVICE_REQUEST
    let mut encrypted_request = EncryptedPacket::read(stream, encryption_ctx, true)?;
    let payload = encrypted_request.payload();
    
    if payload[0] == PacketType::SSH_MSG_EXT_INFO as u8 {
        info!("Received SSH_MSG_EXT_INFO, reading next packet");
        encrypted_request = EncryptedPacket::read(stream, encryption_ctx, true)?;
    }
    
    let payload = encrypted_request.payload();
    info!("Received packet type: {}", payload[0]);
    
    if payload[0] != PacketType::SSH_MSG_SERVICE_REQUEST as u8 {
        return Err(anyhow!("Expected SSH_MSG_SERVICE_REQUEST, got type {}", payload[0]));
    }
    
    use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
    let mut cursor = std::io::Cursor::new(&payload[1..]);
    let service_name_len = cursor.read_u32::<BigEndian>()?;
    let mut service_name = vec![0u8; service_name_len as usize];
    cursor.read_exact(&mut service_name)?;
    let service_name_str = String::from_utf8_lossy(&service_name);
    
    if service_name_str != "ssh-userauth" {
        return Err(anyhow!("Unsupported service: {}", service_name_str));
    }
    
    let mut service_accept_payload = Vec::new();
    service_accept_payload.write_u8(PacketType::SSH_MSG_SERVICE_ACCEPT as u8)?;
    service_accept_payload.write_u32::<BigEndian>(12)?;  // "ssh-userauth" length is 12, not 14!
    service_accept_payload.write_all("ssh-userauth".as_bytes())?;
    
    let encrypted_accept = EncryptedPacket::new(
        &service_accept_payload,
        encryption_ctx,
        true,
    )?;
    encrypted_accept.write(stream)?;
    info!("Sent encrypted SSH_MSG_SERVICE_ACCEPT");
    
    loop {
        let auth_packet = EncryptedPacket::read(stream, encryption_ctx, true)?;  // Reading from client, use cipher_ctos
        let auth_payload = auth_packet.payload();
        info!("Received encrypted SSH_MSG_USERAUTH_REQUEST");
        
        let auth_request = SshPacket::new(auth_payload.to_vec());
        
        match auth_handler.handle_userauth_request(&auth_request)? {
            AuthResult::Success => {
                let success_payload = vec![PacketType::SSH_MSG_USERAUTH_SUCCESS as u8];
                let encrypted_success = EncryptedPacket::new(
                    &success_payload,
                    encryption_ctx,
                    true,
                )?;
                encrypted_success.write(stream)?;
                info!("Sent encrypted SSH_MSG_USERAUTH_SUCCESS");
                
                return Ok("demo".to_string());
            }
            AuthResult::Failure(message) => {
                let mut failure_payload = Vec::new();
                failure_payload.write_u8(PacketType::SSH_MSG_USERAUTH_FAILURE as u8)?;
                failure_payload.write_u32::<BigEndian>(9)?;
                failure_payload.write_all("password".as_bytes())?;
                failure_payload.write_u8(0)?;
                
                let encrypted_failure = EncryptedPacket::new(
                    &failure_payload,
                    encryption_ctx,
                    true,
                )?;
                encrypted_failure.write(stream)?;
                warn!("Sent encrypted SSH_MSG_USERAUTH_FAILURE: {}", message);
            }
            AuthResult::PartialSuccess => {
                warn!("Partial success auth not implemented");
                continue;
            }
        }
    }
}

/// SSH服务循环（Phase 6）
fn handle_ssh_service_loop(
    stream: &mut TcpStream,
    channel_manager: &mut ChannelManager,
) -> Result<()> {
    info!("Starting SSH service loop (channel management)");
    
    loop {
        let packet = SshPacket::read(stream)?;
        
        match packet.payload.first() {
            Some(&pt) if pt == PacketType::SSH_MSG_CHANNEL_OPEN as u8 => {
                info!("Received SSH_MSG_CHANNEL_OPEN");
                let response = channel_manager.handle_channel_open(&packet)?;
                response.write(stream)?;
                info!("Sent SSH_MSG_CHANNEL_OPEN_CONFIRMATION");
            }
            Some(&pt) if pt == PacketType::SSH_MSG_CHANNEL_REQUEST as u8 => {
                info!("Received SSH_MSG_CHANNEL_REQUEST");
                if let Some(response) = channel_manager.handle_channel_request(&packet)? {
                    response.write(stream)?;
                }
            }
            Some(&pt) if pt == PacketType::SSH_MSG_CHANNEL_DATA as u8 => {
                info!("Received SSH_MSG_CHANNEL_DATA");
                channel_manager.handle_channel_data(&packet)?;
            }
            Some(&pt) if pt == PacketType::SSH_MSG_CHANNEL_CLOSE as u8 => {
                info!("Received SSH_MSG_CHANNEL_CLOSE");
                channel_manager.handle_channel_close(&packet)?;
                break;
            }
            Some(&pt) if pt == PacketType::SSH_MSG_DISCONNECT as u8 => {
                info!("Received SSH_MSG_DISCONNECT");
                break;
            }
            _ => {
                warn!("Unknown packet type: {:?}", packet.payload.first());
            }
        }
    }
    
    Ok(())
}

/// 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()
}