opsxshell/src-tauri/src/services/ssh/sftp.rs

use std::{
    fs::{self, OpenOptions},
    io::{Read, Seek, SeekFrom, Write},
    path::{Path, PathBuf},
};

use ssh2::{FileStat, OpenFlags, OpenType, RenameFlags};
use tauri::{AppHandle, Emitter};
use uuid::Uuid;

use crate::domain::ssh::{SftpDirEntry, SftpFileAttrs, SftpProgressEvent};

use super::{lock_error, path::expand_home, SshService};

impl SshService {
    pub fn create_sftp(&self, connection_id: &str) -> Result<String, String> {
        let connection = self.connection(connection_id)?;
        connection
            .session
            .sftp()
            .map_err(|error| format!("初始化 SFTP 失败: {error}"))?;
        let id = format!("sftp-{}", Uuid::new_v4());
        self.sftps
            .lock()
            .map_err(lock_error)?
            .insert(id.clone(), connection_id.to_string());
        Ok(id)
    }

    pub fn sftp_list(&self, sftp_id: &str, remote_path: &str) -> Result<Vec<SftpDirEntry>, String> {
        let sftp = self.sftp(sftp_id)?;
        let mut entries = sftp
            .readdir(Path::new(remote_path))
            .map_err(|error| format!("读取远程目录失败: {error}"))?
            .into_iter()
            .filter_map(|(path, stat)| {
                let filename = path.file_name()?.to_string_lossy().to_string();
                if filename == "." || filename == ".." {
                    return None;
                }
                Some(dir_entry(filename, stat))
            })
            .collect::<Vec<_>>();
        entries.sort_by_key(|entry| entry.filename.to_lowercase());
        Ok(entries)
    }

    pub fn sftp_upload(
        &self,
        app: &AppHandle,
        sftp_id: &str,
        local_path: &str,
        remote_path: &str,
    ) -> Result<(), String> {
        let sftp = self.sftp(sftp_id)?;
        upload_file(
            &sftp,
            &expand_home(local_path),
            remote_path,
            |transferred, total| emit_progress(app, sftp_id, transferred, total),
        )
    }

    pub fn sftp_download(
        &self,
        app: &AppHandle,
        sftp_id: &str,
        remote_path: &str,
        local_path: &str,
    ) -> Result<(), String> {
        let sftp = self.sftp(sftp_id)?;
        let local_path = expand_home(local_path);
        download_file(&sftp, remote_path, &local_path, |transferred, total| {
            emit_progress(app, sftp_id, transferred, total)
        })
    }

    pub fn sftp_upload_dir(
        &self,
        app: &AppHandle,
        sftp_id: &str,
        local_dir: &str,
        remote_dir: &str,
    ) -> Result<(), String> {
        let sftp = self.sftp(sftp_id)?;
        upload_dir(
            &sftp,
            &expand_home(local_dir),
            remote_dir,
            &mut |transferred, total| emit_progress(app, sftp_id, transferred, total),
        )
    }

    pub fn sftp_download_dir(
        &self,
        app: &AppHandle,
        sftp_id: &str,
        remote_dir: &str,
        local_dir: &str,
    ) -> Result<(), String> {
        let sftp = self.sftp(sftp_id)?;
        let local_dir = expand_home(local_dir);
        download_dir(&sftp, remote_dir, &local_dir, &mut |transferred, total| {
            emit_progress(app, sftp_id, transferred, total)
        })
    }

    pub fn sftp_rename(
        &self,
        sftp_id: &str,
        old_remote_path: &str,
        new_remote_path: &str,
    ) -> Result<(), String> {
        self.sftp(sftp_id)?
            .rename(Path::new(old_remote_path), Path::new(new_remote_path), None)
            .map_err(|error| format!("远程重命名失败: {error}"))
    }

    pub fn sftp_delete(
        &self,
        sftp_id: &str,
        remote_path: &str,
        directory: bool,
    ) -> Result<(), String> {
        let sftp = self.sftp(sftp_id)?;
        if directory {
            sftp.rmdir(Path::new(remote_path))
                .map_err(|error| format!("删除远程目录失败: {error}"))
        } else {
            sftp.unlink(Path::new(remote_path))
                .map_err(|error| format!("删除远程文件失败: {error}"))
        }
    }

    pub fn sftp_mkdir(&self, sftp_id: &str, remote_path: &str) -> Result<(), String> {
        self.sftp(sftp_id)?
            .mkdir(Path::new(remote_path), 0o755)
            .map_err(|error| format!("创建远程目录失败: {error}"))
    }

    pub fn sftp_stat(&self, sftp_id: &str, remote_path: &str) -> Result<FileStat, String> {
        self.sftp(sftp_id)?
            .lstat(Path::new(remote_path))
            .map_err(|error| format!("读取远程属性失败: {error}"))
    }

    pub fn sftp_chmod(&self, sftp_id: &str, remote_path: &str, mode: u32) -> Result<(), String> {
        self.sftp(sftp_id)?
            .setstat(
                Path::new(remote_path),
                FileStat {
                    size: None,
                    uid: None,
                    gid: None,
                    perm: Some(mode),
                    atime: None,
                    mtime: None,
                },
            )
            .map_err(|error| format!("修改远程权限失败: {error}"))
    }

    pub fn sftp_read_text(
        &self,
        sftp_id: &str,
        remote_path: &str,
    ) -> Result<(String, u64), String> {
        let sftp = self.sftp(sftp_id)?;
        let mut file = sftp
            .open(Path::new(remote_path))
            .map_err(|error| format!("打开远程文件失败: {error}"))?;
        let mut content = String::new();
        file.read_to_string(&mut content)
            .map_err(|error| format!("远程文件不是有效 UTF-8 文本: {error}"))?;
        let size = content.len() as u64;
        Ok((content, size))
    }

    pub fn sftp_write_text(
        &self,
        sftp_id: &str,
        remote_path: &str,
        content: &str,
    ) -> Result<(), String> {
        let sftp = self.sftp(sftp_id)?;
        let mut file = sftp
            .create(Path::new(remote_path))
            .map_err(|error| format!("创建远程文件失败: {error}"))?;
        file.write_all(content.as_bytes())
            .map_err(|error| format!("写入远程文件失败: {error}"))
    }

    pub(super) fn sftp(&self, sftp_id: &str) -> Result<ssh2::Sftp, String> {
        let connection_id = self
            .sftps
            .lock()
            .map_err(lock_error)?
            .get(sftp_id)
            .cloned()
            .ok_or_else(|| format!("SFTP 会话不存在: {sftp_id}"))?;
        self.connection(&connection_id)?
            .session
            .sftp()
            .map_err(|error| format!("创建 SFTP 通道失败: {error}"))
    }
}

fn dir_entry(filename: String, stat: FileStat) -> SftpDirEntry {
    let mode = stat.perm.unwrap_or(0);
    let kind = match mode & 0o170000 {
        0o040000 => 'd',
        0o120000 => 'l',
        _ => '-',
    };
    SftpDirEntry {
        longname: format!("{}{}", kind, permission_string(mode)),
        filename,
        attrs: SftpFileAttrs {
            mode,
            size: stat.size.unwrap_or(0),
            mtime: stat.mtime.unwrap_or(0),
            atime: stat.atime.unwrap_or(0),
        },
    }
}

fn permission_string(mode: u32) -> String {
    let mut output = String::with_capacity(9);
    for shift in [6, 3, 0] {
        let bits = (mode >> shift) & 7;
        output.push(if bits & 4 != 0 { 'r' } else { '-' });
        output.push(if bits & 2 != 0 { 'w' } else { '-' });
        output.push(if bits & 1 != 0 { 'x' } else { '-' });
    }
    output
}

pub(super) fn upload_file(
    sftp: &ssh2::Sftp,
    local_path: &Path,
    remote_path: &str,
    mut progress: impl FnMut(u64, u64),
) -> Result<(), String> {
    let mut source = fs::File::open(local_path)
        .map_err(|error| format!("打开本地文件失败 ({}): {error}", local_path.display()))?;
    let total = source
        .metadata()
        .map(|metadata| metadata.len())
        .unwrap_or(0);
    let partial_path = format!("{remote_path}.opsx-part");
    let mut offset = sftp
        .stat(Path::new(&partial_path))
        .ok()
        .and_then(|stat| stat.size)
        .unwrap_or(0);
    if offset > total {
        let _ = sftp.unlink(Path::new(&partial_path));
        offset = 0;
    }
    source
        .seek(SeekFrom::Start(offset))
        .map_err(|error| format!("定位本地断点失败: {error}"))?;
    let mut target = sftp
        .open_mode(
            Path::new(&partial_path),
            OpenFlags::WRITE | OpenFlags::CREATE,
            0o644,
            OpenType::File,
        )
        .map_err(|error| format!("打开远程临时文件失败 ({partial_path}): {error}"))?;
    target
        .seek(SeekFrom::Start(offset))
        .map_err(|error| format!("定位远程断点失败: {error}"))?;
    copy_with_progress(&mut source, &mut target, offset, total, |transferred| {
        progress(transferred, total)
    })?;
    drop(target);
    let uploaded = sftp
        .stat(Path::new(&partial_path))
        .ok()
        .and_then(|stat| stat.size)
        .unwrap_or(0);
    if uploaded != total {
        return Err(format!(
            "上传大小校验失败: 本地 {total} 字节，远程 {uploaded} 字节"
        ));
    }
    finalize_remote_transfer(sftp, &partial_path, remote_path)
}

pub(super) fn download_file(
    sftp: &ssh2::Sftp,
    remote_path: &str,
    local_path: &Path,
    mut progress: impl FnMut(u64, u64),
) -> Result<(), String> {
    let mut source = sftp
        .open(Path::new(remote_path))
        .map_err(|error| format!("打开远程文件失败 ({remote_path}): {error}"))?;
    let total = source.stat().ok().and_then(|stat| stat.size).unwrap_or(0);
    let partial_path = local_partial_path(local_path);
    let mut offset = partial_path
        .metadata()
        .map(|metadata| metadata.len())
        .unwrap_or(0);
    if offset > total {
        fs::remove_file(&partial_path).map_err(|error| format!("清理无效下载断点失败: {error}"))?;
        offset = 0;
    }
    source
        .seek(SeekFrom::Start(offset))
        .map_err(|error| format!("定位远程断点失败: {error}"))?;
    let mut target = OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(false)
        .open(&partial_path)
        .map_err(|error| format!("打开本地临时文件失败 ({}): {error}", partial_path.display()))?;
    target
        .seek(SeekFrom::Start(offset))
        .map_err(|error| format!("定位本地断点失败: {error}"))?;
    copy_with_progress(&mut source, &mut target, offset, total, |transferred| {
        progress(transferred, total)
    })?;
    target
        .sync_all()
        .map_err(|error| format!("同步本地临时文件失败: {error}"))?;
    drop(target);
    let downloaded = partial_path
        .metadata()
        .map(|metadata| metadata.len())
        .unwrap_or(0);
    if downloaded != total {
        return Err(format!(
            "下载大小校验失败: 远程 {total} 字节，本地 {downloaded} 字节"
        ));
    }
    if local_path.is_file() {
        fs::remove_file(local_path).map_err(|error| format!("替换本地文件失败: {error}"))?;
    }
    fs::rename(&partial_path, local_path).map_err(|error| format!("提交下载文件失败: {error}"))
}

fn join_remote(parent: &str, name: &str) -> String {
    if parent.is_empty() || parent == "/" {
        format!("/{name}")
    } else {
        format!("{}/{name}", parent.trim_end_matches('/'))
    }
}

fn is_dir_mode(mode: u32) -> bool {
    mode & 0o170000 == 0o040000
}

fn ensure_remote_dir(sftp: &ssh2::Sftp, remote_dir: &str) -> Result<(), String> {
    let normalized = if remote_dir.is_empty() {
        "/".to_string()
    } else {
        remote_dir.trim_end_matches('/').to_string()
    };
    if normalized.is_empty() || normalized == "/" {
        return Ok(());
    }

    let mut parts = Vec::new();
    for part in normalized.trim_start_matches('/').split('/') {
        if part.is_empty() || part == "." {
            continue;
        }
        if part == ".." {
            return Err("远程目录路径不能包含 ..".into());
        }
        parts.push(part);
    }

    let mut current = String::new();
    for part in parts {
        current = if current.is_empty() {
            format!("/{part}")
        } else {
            format!("{current}/{part}")
        };
        match sftp.lstat(Path::new(&current)) {
            Ok(stat) if is_dir_mode(stat.perm.unwrap_or(0)) => continue,
            Ok(_) => return Err(format!("远程路径已存在且不是目录: {current}")),
            Err(_) => {
                sftp.mkdir(Path::new(&current), 0o755)
                    .map_err(|error| format!("创建远程目录失败 ({current}): {error}"))?;
            }
        }
    }
    Ok(())
}

fn upload_dir(
    sftp: &ssh2::Sftp,
    local_dir: &Path,
    remote_dir: &str,
    progress: &mut dyn FnMut(u64, u64),
) -> Result<(), String> {
    let meta = fs::metadata(local_dir)
        .map_err(|error| format!("读取本地目录失败 ({}): {error}", local_dir.display()))?;
    if !meta.is_dir() {
        return Err(format!("本地路径不是目录: {}", local_dir.display()));
    }
    ensure_remote_dir(sftp, remote_dir)?;

    let mut entries = fs::read_dir(local_dir)
        .map_err(|error| format!("遍历本地目录失败 ({}): {error}", local_dir.display()))?
        .collect::<Result<Vec<_>, _>>()
        .map_err(|error| format!("读取本地目录项失败: {error}"))?;
    entries.sort_by_key(|entry| entry.file_name());

    for entry in entries {
        let file_type = entry
            .file_type()
            .map_err(|error| format!("读取本地文件类型失败: {error}"))?;
        let name = entry.file_name().to_string_lossy().to_string();
        if name == "." || name == ".." {
            continue;
        }
        let local_child = entry.path();
        let remote_child = join_remote(remote_dir, &name);
        if file_type.is_dir() {
            upload_dir(sftp, &local_child, &remote_child, progress)?;
        } else if file_type.is_file() {
            upload_file(sftp, &local_child, &remote_child, |transferred, total| {
                progress(transferred, total)
            })?;
        }
        // skip symlinks / special files for safety
    }
    Ok(())
}

fn download_dir(
    sftp: &ssh2::Sftp,
    remote_dir: &str,
    local_dir: &Path,
    progress: &mut dyn FnMut(u64, u64),
) -> Result<(), String> {
    let stat = sftp
        .lstat(Path::new(remote_dir))
        .map_err(|error| format!("读取远程目录失败 ({remote_dir}): {error}"))?;
    if !is_dir_mode(stat.perm.unwrap_or(0)) {
        return Err(format!("远程路径不是目录: {remote_dir}"));
    }
    fs::create_dir_all(local_dir)
        .map_err(|error| format!("创建本地目录失败 ({}): {error}", local_dir.display()))?;

    let mut entries = sftp
        .readdir(Path::new(remote_dir))
        .map_err(|error| format!("遍历远程目录失败 ({remote_dir}): {error}"))?;
    entries.sort_by_key(|(path, _)| path.file_name().map(|name| name.to_os_string()));

    for (path, entry_stat) in entries {
        let name = match path.file_name() {
            Some(value) => value.to_string_lossy().to_string(),
            None => continue,
        };
        if name == "." || name == ".." {
            continue;
        }
        let remote_child = join_remote(remote_dir, &name);
        let local_child = local_dir.join(&name);
        if is_dir_mode(entry_stat.perm.unwrap_or(0)) {
            download_dir(sftp, &remote_child, &local_child, progress)?;
        } else {
            download_file(sftp, &remote_child, &local_child, |transferred, total| {
                progress(transferred, total)
            })?;
        }
    }
    Ok(())
}

pub(super) fn copy_with_progress(
    reader: &mut impl Read,
    writer: &mut impl Write,
    initial: u64,
    total: u64,
    mut progress: impl FnMut(u64),
) -> Result<(), String> {
    let mut buffer = [0_u8; 64 * 1024];
    let mut transferred = initial;
    progress(transferred);
    loop {
        let length = reader
            .read(&mut buffer)
            .map_err(|error| format!("读取传输数据失败: {error}"))?;
        if length == 0 {
            break;
        }
        writer
            .write_all(&buffer[..length])
            .map_err(|error| format!("写入传输数据失败: {error}"))?;
        transferred += length as u64;
        progress(transferred);
    }
    writer
        .flush()
        .map_err(|error| format!("刷新传输数据失败: {error}"))?;
    if transferred < total {
        progress(transferred);
    }
    Ok(())
}

fn finalize_remote_transfer(
    sftp: &ssh2::Sftp,
    partial_path: &str,
    remote_path: &str,
) -> Result<(), String> {
    let partial = Path::new(partial_path);
    let destination = Path::new(remote_path);
    if sftp
        .rename(
            partial,
            destination,
            Some(RenameFlags::OVERWRITE | RenameFlags::ATOMIC | RenameFlags::NATIVE),
        )
        .is_ok()
    {
        return Ok(());
    }
    let _ = sftp.unlink(destination);
    sftp.rename(partial, destination, Some(RenameFlags::OVERWRITE))
        .map_err(|error| format!("提交远程上传文件失败: {error}"))
}

pub(super) fn local_partial_path(destination: &Path) -> PathBuf {
    let mut value = destination.as_os_str().to_os_string();
    value.push(".opsx-part");
    PathBuf::from(value)
}

fn emit_progress(app: &AppHandle, sftp_id: &str, transferred: u64, total: u64) {
    let percentage = transferred
        .saturating_mul(100)
        .checked_div(total)
        .unwrap_or(100)
        .min(100) as u8;
    let _ = app.emit(
        "sftp:progress",
        SftpProgressEvent {
            sftp_id: sftp_id.to_string(),
            transferred,
            total,
            percentage,
        },
    );
}