markbase/markbase-core/src/ssh2_mod/scp_handler.rs

// SCP Handler实现（ssh2辅助模块）
// 支持 scp -f（从服务器下载）和 scp -t（上传到服务器）

use anyhow::{Result, anyhow};
use ssh2::Channel;
use std::path::{Path, PathBuf};
use std::fs::{File, OpenOptions};
use std::io::{Read, Write, BufReader, BufWriter};
use log::{info, warn, error, debug};

/// SCP Handler
pub struct ScpHandler {
    base_path: PathBuf,
    user_id: String,
}

impl ScpHandler {
    pub fn new(base_path: PathBuf, user_id: String) -> Self {
        Self { base_path, user_id }
    }
    
    /// 处理SCP命令
    pub fn handle_scp_command(&self, channel: &mut Channel, command: &str) -> Result<()> {
        info!("SCP command: {}", command);
        
        // 解析SCP命令
        if command.contains("-t") {
            // scp -t：接收文件（上传）
            self.handle_scp_receive(channel, command)?;
        } else if command.contains("-f") {
            // scp -f：发送文件（下载）
            self.handle_scp_send(channel, command)?;
        } else if command.contains("-r") {
            // scp -r：递归目录
            if command.contains("-t") {
                self.handle_scp_receive_dir(channel, command)?;
            } else if command.contains("-f") {
                self.handle_scp_send_dir(channel, command)?;
            }
        } else {
            warn!("Unsupported SCP command: {}", command);
            return Err(anyhow!("Unsupported SCP command"));
        }
        
        Ok(())
    }
    
    /// SCP接收文件（scp -t，客户端上传）
    fn handle_scp_receive(&self, channel: &mut Channel, command: &str) -> Result<()> {
        info!("SCP receive mode: {}", command);
        
        // 发送确认（0x00）
        channel.write_all(&[0x00])?;
        
        // 读取文件信息（C0644 <size> <filename>)
        let mut buf = vec![0u8; 8192];
        let len = channel.read(&mut buf)?;
        let header = String::from_utf8_lossy(&buf[..len]);
        debug!("SCP header: {}", header);
        
        // 解析header：C0644 <size> <filename>
        let parts: Vec<&str> = header.trim().split_whitespace().collect();
        if parts.len() < 3 || !parts[0].starts_with("C") {
            return Err(anyhow!("Invalid SCP header: {}", header));
        }
        
        let mode = parts[0];  // C0644
        let size: u64 = parts[1].parse()?;
        let filename = parts[2].trim_matches('\n');
        
        // 发送确认
        channel.write_all(&[0x00])?;
        
        // 构建文件路径
        let file_path = self.base_path.join(&self.user_id).join(filename);
        info!("SCP receive file: {} ({})", file_path.display(), size);
        
        // 创建文件
        let mut file = BufWriter::new(File::create(&file_path)?);
        
        // 读取文件内容
        let mut received = 0;
        while received < size {
            let to_read = std::cmp::min(8192, (size - received) as usize);
            let len = channel.read(&mut buf[..to_read])?;
            if len == 0 {
                break;
            }
            file.write_all(&buf[..len])?;
            received += len as u64;
        }
        
        file.flush()?;
        
        // 发送确认
        channel.write_all(&[0x00])?;
        
        // 读取结束标志（E）
        let len = channel.read(&mut buf)?;
        if len > 0 && buf[0] == 'E' as u8 {
            channel.write_all(&[0x00])?;
        }
        
        info!("SCP receive completed: {} bytes", received);
        Ok(())
    }
    
    /// SCP发送文件（scp -f，客户端下载）
    fn handle_scp_send(&self, channel: &mut Channel, command: &str) -> Result<()> {
        info!("SCP send mode: {}", command);
        
        // 解析路径
        let path_str = command.split_whitespace().last().unwrap_or("");
        let file_path = self.base_path.join(&self.user_id).join(path_str);
        
        if !file_path.exists() {
            warn!("SCP file not found: {}", file_path.display());
            channel.write_all(b"SCP error: file not found\n")?;
            return Err(anyhow!("File not found: {}", file_path.display()));
        }
        
        // 获取文件信息
        let metadata = std::fs::metadata(&file_path)?;
        let size = metadata.len();
        let filename = file_path.file_name().unwrap().to_str().unwrap();
        
        info!("SCP send file: {} ({})", file_path.display(), size);
        
        // 等待客户端确认
        let mut buf = vec![0u8; 1];
        channel.read(&mut buf)?;
        
        // 发送文件信息
        let header = format!("C0644 {} {}\n", size, filename);
        channel.write_all(header.as_bytes())?;
        
        // 等待客户端确认
        channel.read(&mut buf)?;
        
        // 发送文件内容
        let mut file = BufReader::new(File::open(&file_path)?);
        let mut chunk = vec![0u8; 8192];
        while let Ok(len) = file.read(&mut chunk) {
            if len == 0 {
                break;
            }
            channel.write_all(&chunk[..len])?;
        }
        
        // 发送结束确认
        channel.write_all(&[0x00])?;
        
        // 等待客户端确认
        channel.read(&mut buf)?;
        
        // 发送结束标志
        channel.write_all("E\n".as_bytes())?;
        
        info!("SCP send completed: {} bytes", size);
        Ok(())
    }
    
    /// SCP接收目录（scp -r -t）
    fn handle_scp_receive_dir(&self, channel: &mut Channel, command: &str) -> Result<()> {
        warn!("SCP directory receive not implemented: {}", command);
        Err(anyhow!("SCP directory receive not implemented"))
    }
    
    /// SCP发送目录（scp -r -f）
    fn handle_scp_send_dir(&self, channel: &mut Channel, command: &str) -> Result<()> {
        warn!("SCP directory send not implemented: {}", command);
        Err(anyhow!("SCP directory send not implemented"))
    }
}