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

// SSH服务器完整实现（Phase 1-7集成版 + Phase 13端口转发）
// 参考OpenSSH sshd.c: complete SSH/SFTP flow + port forwarding

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 crate::ssh_server::ssh_security_config::SshSecurityConfig;  // Phase 13.1
use crate::ssh_server::port_forward::PortForwardManager;  // Phase 13
use anyhow::{Result, anyhow};
use log::{info, warn, error, debug};
use std::net::{TcpListener, TcpStream};
use std::thread;
use std::io::{Read, Write};
use std::sync::{Arc, Mutex};  // Phase 13: 端口转发线程同步

/// SSH服务器配置（Phase 13.1企业级安全配置）
pub struct SshServerConfig {
    pub port: u16,
    pub bind_address: String,
    pub security_config: SshSecurityConfig,  // Phase 13.1: 企业级安全配置
}

impl Default for SshServerConfig {
    fn default() -> Self {
        Self {
            port: 2024,
            bind_address: "127.0.0.1".to_string(),
            security_config: SshSecurityConfig::enterprise_default(),  // Phase 13.1
        }
    }
}

impl SshServerConfig {
    /// 从配置文件加载（Phase 13.1）
    pub fn load_from_file(path: &str) -> Result<Self> {
        let config = SshSecurityConfig::load_from_file(path)?;
        Ok(Self {
            port: 2024,
            bind_address: "127.0.0.1".to_string(),
            security_config: config,
        })
    }
}

/// SSH服务器主结构（Phase 1-13完整版）
pub struct SshServer {
    config: SshServerConfig,
    security_config: Arc<Mutex<SshSecurityConfig>>,  // Phase 13.1: 共享安全配置
}

impl SshServer {
    pub fn new(config: SshServerConfig) -> Self {
        let security_config = Arc::new(Mutex::new(config.security_config.clone()));  // Phase 13.1: 先clone
        Self {
            config,
            security_config,  // Phase 13.1
        }
    }
    
    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 + Port Forwarding (Phase 1-13)");
        info!("Security config: GatewayPorts={}, PermitOpen={:?}, MaxSessions={}", 
            self.config.security_config.gateway_ports,
            self.config.security_config.permit_open,
            self.config.security_config.max_sessions);
        
        let security_config = self.security_config.clone();  // Phase 13.1: 共享安全配置
        
        for stream in listener.incoming() {
            match stream {
                Ok(stream) => {
                    let client_addr = stream.peer_addr()?;
                    info!("New SSH connection from {}", client_addr);
                    
                    let security_config_clone = security_config.clone();  // Phase 13.1
                    
                    thread::spawn(move || {
                        if let Err(e) = handle_connection_complete(stream, security_config_clone) {  // Phase 13.1
                            error!("Connection error: {}", e);
                        }
                    });
                }
                Err(e) => {
                    warn!("Failed to accept connection: {}", e);
                }
            }
        }
        
        Ok(())
    }
}

/// 处理完整SSH连接（Phase 1-13完整流程）
fn handle_connection_complete(stream: TcpStream, security_config: Arc<Mutex<SshSecurityConfig>>) -> Result<()> {
    info!("Handling client connection (Phase 1-13 complete flow with port forwarding)");
    
    // Phase 13.1: 增加活动会话数
    {
        let mut security = security_config.lock().unwrap();
        security.increment_sessions()?;
    }
    
    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 13: PortForwardManager初始化
    let mut port_forward_manager = PortForwardManager::new();
    
    // Phase 6-13: SSH服务循环（处理channel请求 + 端口转发）
    let security_config_clone = security_config.clone();  // Phase 13.1: clone for service loop
    handle_ssh_service_loop(&mut stream, &mut channel_manager, &mut encryption_ctx, &mut port_forward_manager, security_config_clone)?;
    
    info!("SSH session completed successfully");
    
    // Phase 13.1: 减少活动会话数
    {
        let mut security = security_config.lock().unwrap();
        security.decrement_sessions();
    }
    
    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) => {
                // message包含可用的认证方法列表（如"password,publickey"）
                let mut failure_payload = Vec::new();
                failure_payload.write_u8(PacketType::SSH_MSG_USERAUTH_FAILURE as u8)?;
                failure_payload.write_u32::<BigEndian>(message.len() as u32)?;
                failure_payload.write_all(message.as_bytes())?;
                failure_payload.write_u8(0)?;  // partial_success = false
                
                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;
            }
            AuthResult::PublicKeyOk(algorithm, public_key_blob) => {
                // SSH_MSG_USERAUTH_PK_OK：public key acceptable
                info!("Public key acceptable, sending USERAUTH_PK_OK");
                
                let mut pk_ok_payload = Vec::new();
                pk_ok_payload.write_u8(PacketType::SSH_MSG_USERAUTH_PK_OK as u8)?;
                
                // algorithm (SSH string)
                pk_ok_payload.write_u32::<BigEndian>(algorithm.len() as u32)?;
                pk_ok_payload.write_all(algorithm.as_bytes())?;
                
                // public key blob (SSH string)
                pk_ok_payload.write_u32::<BigEndian>(public_key_blob.len() as u32)?;
                pk_ok_payload.write_all(&public_key_blob)?;
                
                let encrypted_pk_ok = EncryptedPacket::new(
                    &pk_ok_payload,
                    encryption_ctx,
                    true,
                )?;
                encrypted_pk_ok.write(stream)?;
                info!("Sent SSH_MSG_USERAUTH_PK_OK");
                
                continue;  // Wait for signed request
            }
        }
    }
}

/// SSH服务循环（Phase 6-13完整版）
fn handle_ssh_service_loop(
    stream: &mut TcpStream,
    channel_manager: &mut ChannelManager,
    encryption_ctx: &mut EncryptionContext,
    port_forward_manager: &mut PortForwardManager,  // Phase 13
    security_config: Arc<Mutex<SshSecurityConfig>>,  // Phase 13.1
) -> Result<()> {
    info!("Starting SSH service loop (channel management + port forwarding)");
    
    loop {
        // 使用EncryptedPacket读取加密packet（Phase 6）
        let encrypted_packet = EncryptedPacket::read(stream, encryption_ctx, true)?;
        let packet = SshPacket::new(encrypted_packet.payload().to_vec());
        
        match packet.payload.first() {
            // Phase 13: SSH_MSG_GLOBAL_REQUEST处理（端口转发）
            Some(&pt) if pt == PacketType::SSH_MSG_GLOBAL_REQUEST as u8 => {
                info!("Received SSH_MSG_GLOBAL_REQUEST (port forwarding)");
                
                // Phase 13.1: 安全配置验证
                let security = security_config.lock().unwrap();
                if !security.allow_tcp_forwarding {
                    warn!("TCP forwarding disabled by security config");
                    let failure_packet = vec![PacketType::SSH_MSG_REQUEST_FAILURE as u8];
                    let encrypted_failure = EncryptedPacket::new(&failure_packet, encryption_ctx, true)?;
                    encrypted_failure.write(stream)?;
                    info!("Sent SSH_MSG_REQUEST_FAILURE (TCP forwarding disabled)");
                    continue;
                }
                
                // Phase 13.2: 调用PortForwardManager处理（传递security_config）
                let (success, response) = port_forward_manager.handle_global_request(&packet.payload, &security)?;
                drop(security);  // 释放锁
                
                if success {
                    if let Some(response_data) = response {
                        let encrypted_response = EncryptedPacket::new(&response_data, encryption_ctx, true)?;
                        encrypted_response.write(stream)?;
                        info!("Sent SSH_MSG_REQUEST_SUCCESS (tcpip-forward accepted)");
                    } else {
                        // 无响应数据时，发送简单的SUCCESS
                        let success_packet = vec![PacketType::SSH_MSG_REQUEST_SUCCESS as u8];
                        let encrypted_success = EncryptedPacket::new(&success_packet, encryption_ctx, true)?;
                        encrypted_success.write(stream)?;
                        info!("Sent SSH_MSG_REQUEST_SUCCESS");
                    }
                } else {
                    let failure_packet = vec![PacketType::SSH_MSG_REQUEST_FAILURE as u8];
                    let encrypted_failure = EncryptedPacket::new(&failure_packet, encryption_ctx, true)?;
                    encrypted_failure.write(stream)?;
                    info!("Sent SSH_MSG_REQUEST_FAILURE (tcpip-forward rejected)");
                }
            }
            
            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)?;
                let encrypted_response = EncryptedPacket::new(&response.payload, encryption_ctx, true)?;
                encrypted_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)? {
                    let encrypted_response = EncryptedPacket::new(&response.payload, encryption_ctx, true)?;
                    encrypted_response.write(stream)?;
                    
                    // Phase 6: 检查是否有命令输出需要发送
                    if let Some(channel_id) = channel_manager.get_channel_with_output() {
                        if let Some(output) = channel_manager.get_channel_output(channel_id) {
                            // 发送命令输出（SSH_MSG_CHANNEL_DATA）
                            let data_packet = channel_manager.build_channel_data(channel_id, &output)?;
                            let encrypted_data = EncryptedPacket::new(&data_packet.payload, encryption_ctx, true)?;
                            encrypted_data.write(stream)?;
                            info!("Sent command output ({} bytes)", output.len());
                            
                            // 发送SSH_MSG_CHANNEL_EOF
                            let eof_packet = channel_manager.build_channel_eof(channel_id)?;
                            let encrypted_eof = EncryptedPacket::new(&eof_packet.payload, encryption_ctx, true)?;
                            encrypted_eof.write(stream)?;
                            info!("Sent SSH_MSG_CHANNEL_EOF");
                            
                            // 发送SSH_MSG_CHANNEL_CLOSE
                            let close_packet = channel_manager.build_channel_close(channel_id)?;
                            let encrypted_close = EncryptedPacket::new(&close_packet.payload, encryption_ctx, true)?;
                            encrypted_close.write(stream)?;
                            info!("Sent SSH_MSG_CHANNEL_CLOSE");
                            
                            // 移除channel
                            channel_manager.remove_channel(channel_id);
                        }
                    }
                }
            }
            Some(&pt) if pt == PacketType::SSH_MSG_CHANNEL_DATA as u8 => {
                info!("Received SSH_MSG_CHANNEL_DATA");
                if let Some(response) = channel_manager.handle_channel_data(&packet)? {
                    // Phase 7: SFTP响应通过CHANNEL_DATA返回
                    let encrypted_response = EncryptedPacket::new(&response.payload, encryption_ctx, true)?;
                    encrypted_response.write(stream)?;
                    info!("Sent SSH_MSG_CHANNEL_DATA (SFTP response)");
                }
            }
            Some(&pt) if pt == PacketType::SSH_MSG_CHANNEL_CLOSE as u8 => {
                info!("Received SSH_MSG_CHANNEL_CLOSE");
                if let Some(response) = channel_manager.handle_channel_close(&packet)? {
                    let encrypted_response = EncryptedPacket::new(&response.payload, encryption_ctx, true)?;
                    encrypted_response.write(stream)?;
                }
                break;
            }
            Some(&pt) if pt == PacketType::SSH_MSG_CHANNEL_EOF as u8 => {
                info!("Received SSH_MSG_CHANNEL_EOF");
                // EOF means client won't send more data, just acknowledge and continue
            }
            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(),
        security_config: SshSecurityConfig::enterprise_default(),  // Phase 13.1: 添加安全配置
    };
    
    let server = SshServer::new(config);
    server.run()
}