sshclient/core/ssh_client.py

"""
SSH 客户端核心模块
封装 paramiko，处理连接、命令执行、SFTP 文件传输。
所有 SSH 操作都通过此模块，与 UI 解耦。
"""
import os
import time
import threading
from pathlib import Path
from typing import Optional, Tuple, List

import paramiko
from paramiko import SSHClient, AutoAddPolicy, RSAKey, Ed25519Key
from paramiko.ssh_exception import AuthenticationException, SSHException


class SSHConnection:
    """单台主机的 SSH 连接管理"""

    def __init__(self, host: str, port: int = 22, username: str = "",
                 password: str = "", key_path: str = "", timeout: int = 10):
        self.host = host
        self.port = int(port) if port else 22
        self.username = username
        self.password = password
        self.key_path = key_path
        self.timeout = timeout
        self.client: Optional[SSHClient] = None
        self.sftp: Optional[paramiko.SFTPClient] = None
        self.connected = False
        self.last_error = ""

    def connect(self) -> Tuple[bool, str]:
        """建立连接；返回 (成功, 消息)"""
        try:
            self.client = SSHClient()
            self.client.set_missing_host_key_policy(AutoAddPolicy())
            connect_kwargs = {
                "hostname": self.host,
                "port": self.port,
                "username": self.username,
                "timeout": self.timeout,
                "allow_agent": False,
                "look_for_keys": False,
            }
            if self.key_path and os.path.isfile(self.key_path):
                pkey = self._load_key(self.key_path, self.password)
                connect_kwargs["pkey"] = pkey
                if self.password:
                    connect_kwargs["password"] = self.password
            else:
                connect_kwargs["password"] = self.password

            self.client.connect(**connect_kwargs)
            self.sftp = self.client.open_sftp()
            self.connected = True
            return True, f"已连接到 {self.username}@{self.host}:{self.port}"
        except AuthenticationException as e:
            self.last_error = f"认证失败: {e}"
        except SSHException as e:
            self.last_error = f"SSH 错误: {e}"
        except Exception as e:
            self.last_error = f"连接失败: {e}"
        self.connected = False
        return False, self.last_error

    def _load_key(self, path: str, passphrase: str = ""):
        """自动识别 RSA / Ed25519 私钥格式"""
        for loader in (Ed25519Key, RSAKey):
            try:
                return loader.from_private_key_file(path, password=passphrase or None)
            except paramiko.ssh_exception.PasswordRequiredException:
                raise
            except paramiko.ssh_exception.SSHException:
                continue
            except Exception:
                continue
        raise SSHException(f"无法加载私钥: {path}")

    def disconnect(self):
        """关闭 SFTP 和 SSH 连接"""
        for handle in (self.sftp, self.client):
            try:
                if handle:
                    handle.close()
            except Exception:
                pass
        self.sftp = None
        self.client = None
        self.connected = False

    def open_shell(self, term_type: str = "xterm", cols: int = 80, rows: int = 24):
        """打开一个交互式 shell 会话（用于真正的终端模拟器）"""
        if not self.connected or not self.client:
            return None
        try:
            chan = self.client.invoke_shell(term=term_type, width=cols, height=rows)
            chan.settimeout(0.0)  # 非阻塞
            return chan
        except Exception as e:
            self.last_error = f"打开 shell 失败: {e}"
            return None

    def resize_shell(self, chan, cols: int, rows: int):
        """终端大小变化时通知远端（vim/top 这类全屏程序需要）"""
        if not chan:
            return
        try:
            chan.resize_pty(width=cols, height=rows)
        except Exception:
            pass

    def exec_command(self, command: str, timeout: int = 30) -> Tuple[int, str, str]:
        """执行远程命令；返回 (退出码, stdout, stderr)"""
        if not self.connected or not self.client:
            return -1, "", "未连接"
        try:
            stdin, stdout, stderr = self.client.exec_command(command, timeout=timeout)
            out = stdout.read().decode("utf-8", errors="replace")
            err = stderr.read().decode("utf-8", errors="replace")
            code = stdout.channel.recv_exit_status()
            return code, out, err
        except Exception as e:
            return -1, "", f"执行错误: {e}"

    def list_dir(self, remote_path: str) -> List[dict]:
        """列出远程目录；返回 [{name, size, mtime, mode, is_dir}, ...]"""
        if not self.sftp:
            return []
        try:
            entries = []
            for attr in self.sftp.listdir_attr(remote_path):
                entries.append({
                    "name": attr.filename,
                    "size": attr.st_size or 0,
                    "mtime": attr.st_mtime or 0,
                    "mode": attr.st_mode or 0,
                    "is_dir": attr.st_mode is not None and (attr.st_mode & 0o170000) == 0o040000,
                })
            # 目录优先，再按名字排序
            entries.sort(key=lambda x: (not x["is_dir"], x["name"].lower()))
            return entries
        except Exception as e:
            self.last_error = f"列目录失败: {e}"
            return []

    def upload(self, local_path: str, remote_path: str, progress_cb=None) -> Tuple[bool, str]:
        """上传本地文件到远程；progress_cb(done, total) 回调"""
        if not self.sftp:
            return False, "SFTP 未就绪"
        try:
            total = os.path.getsize(local_path)
            done = [0]

            def _cb(transferred, _total):
                done[0] = transferred
                if progress_cb:
                    progress_cb(transferred, _total or total)

            self.sftp.put(local_path, remote_path, callback=_cb)
            return True, f"已上传 {os.path.basename(local_path)} ({total} bytes)"
        except Exception as e:
            return False, f"上传失败: {e}"

    def download(self, remote_path: str, local_path: str, progress_cb=None) -> Tuple[bool, str]:
        """下载远程文件到本地"""
        if not self.sftp:
            return False, "SFTP 未就绪"
        try:
            total = self.sftp.stat(remote_path).st_size
            done = [0]

            def _cb(transferred, _total):
                done[0] = transferred
                if progress_cb:
                    progress_cb(transferred, _total or total)

            self.sftp.get(remote_path, local_path, callback=_cb)
            return True, f"已下载到 {local_path}"
        except Exception as e:
            return False, f"下载失败: {e}"

    def mkdir(self, remote_path: str) -> Tuple[bool, str]:
        try:
            self.sftp.mkdir(remote_path)
            return True, f"已创建 {remote_path}"
        except Exception as e:
            return False, f"创建失败: {e}"

    def remove(self, remote_path: str) -> Tuple[bool, str]:
        try:
            try:
                self.sftp.remove(remote_path)
            except IOError:
                self.sftp.rmdir(remote_path)
            return True, f"已删除 {remote_path}"
        except Exception as e:
            return False, f"删除失败: {e}"

    def rename(self, old_path: str, new_path: str) -> Tuple[bool, str]:
        try:
            self.sftp.rename(old_path, new_path)
            return True, "已重命名"
        except Exception as e:
            return False, f"重命名失败: {e}"

    def stat(self, remote_path: str):
        try:
            return self.sftp.stat(remote_path)
        except Exception as e:
            self.last_error = f"stat 失败: {e}"
            return None