markbase/markbase-sftp-poc/src/shell_handler.rs

use anyhow::{Error, Result};
use russh::server::{Auth, Session, Sig};
use russh::{ChannelId, SigId};
use tokio::process::Command;
use std::sync::Arc;
use log::{info, warn, error};

use crate::auth::MockAuthDb;

pub struct ShellSession {
    auth_db: Arc<MockAuthDb>,
    user: String,
}

impl ShellSession {
    pub fn new(auth_db: Arc<MockAuthDb>, user: String) -> Self {
        Self { auth_db, user }
    }
    
    // 关键方法：exec_request（rsync使用）
    async fn exec_request(&mut self, channel: ChannelId, command: &str) -> Result<()> {
        info!("Shell exec_request: user={}, command={}", self.user, command);
        
        // 1. 安全检查：只允许特定命令
        if !self.is_command_allowed(command) {
            warn!("Command not allowed: {}", command);
            return Err(Error::msg("Command not allowed"));
        }
        
        // 2. rsync命令特殊处理
        if command.starts_with("rsync --server") {
            return self.handle_rsync(channel, command).await;
        }
        
        // 3. 其他允许的shell命令执行
        let parts: Vec<&str> = command.split_whitespace().collect();
        if parts.is_empty() {
            return Err(Error::msg("Empty command"));
        }
        
        let cmd = Command::new(parts[0])
            .args(&parts[1..])
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .spawn()?;
        
        // 4. 等待命令执行完成
        let status = cmd.wait().await?;
        info!("Command exit status: {}", status);
        
        Ok(())
    }
    
    // 安全检查：命令白名单
    fn is_command_allowed(&self, command: &str) -> bool {
        // 允许的命令列表
        let allowed_commands = [
            "ls", "pwd", "cd", "echo", "cat", "rsync",
        ];
        
        let cmd_name = command.split_whitespace().next().unwrap_or("");
        allowed_commands.contains(&cmd_name)
    }
    
    // rsync命令处理（核心功能）
    async fn handle_rsync(&mut self, channel: ChannelId, command: &str) -> Result<()> {
        info!("Handling rsync command: {}", command);
        
        // 1. 解析rsync命令
        let parts: Vec<&str> = command.split_whitespace().collect();
        
        // 2. 提取路径参数（最后一个参数）
        let path = parts.last().unwrap_or(".");
        
        // 3. 用户目录限制
        let user_dir = self.auth_db.get_user_dir(&self.user);
        
        // 4. 路径安全检查
        if path.starts_with("/") && !path.starts_with(&user_dir) {
            warn!("Path access denied: user={}, path={}", self.user, path);
            return Err(Error::msg("Path access denied"));
        }
        
        // 5. 执行rsync命令
        let mut cmd = Command::new("rsync");
        cmd.args(&parts[1..parts.len()-1])  // rsync参数
           .arg(&user_dir);  // 替换为用户目录
        
        let child = cmd.spawn()?;
        let status = child.wait().await?;
        
        info!("rsync exit status: {}", status);
        
        Ok(())
    }
}

// russh server Session实现
impl Session for ShellSession {
    // shell子系统（交互式shell）
    async fn shell_request(&mut self, channel: ChannelId) -> Result<()> {
        info!("Shell request received for user: {}", self.user);
        
        // 创建交互式shell进程
        let shell = Command::new("/bin/bash")
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .spawn()?;
        
        shell.wait().await?;
        
        Ok(())
    }
    
    // exec子系统（执行命令，rsync使用）
    async fn exec_request(&mut self, channel: ChannelId, command: &str) -> Result<()> {
        self.exec_request(channel, command).await
    }
    
    // 信号处理
    async fn signal(&mut self, channel: ChannelId, signal: Sig) -> Result<()> {
        info!("Signal received: {:?}", signal);
        Ok(())
    }
}

// russh server Auth实现
impl Auth for ShellSession {
    // 密码认证
    async fn auth_password(&mut self, user: &str, password: &str) -> russh::server::AuthResult {
        info!("Auth password attempt: user={}", user);
        
        if self.auth_db.verify_password(user, password).unwrap_or(false) {
            info!("Auth success: user={}", user);
            russh::server::AuthResult::Accept
        } else {
            warn!("Auth failed: user={}", user);
            russh::server::AuthResult::Reject
        }
    }
}