From a809367b86a659dfca95ab59b8f625a8ab2dffd1 Mon Sep 17 00:00:00 2001
From: Val-sss <154882199@qq.com>
Date: Thu, 30 Apr 2026 02:14:49 +0800
Subject: [PATCH] feat: persistent SSH session via paramiko

- Add kaiwu/tools/ssh_session.py: SSHSession class with connect/exec/upload/download/close
- Integrate into ToolExecutor: ssh_connect/ssh_exec/ssh_upload/ssh_download/ssh_close
- Guardrails apply to remote commands too (rm -rf blocked on SSH)
- Persistent connection: connect once, exec multiple commands without reconnecting
- Supports password auth and SSH key auth

Usage flow:
  executor.ssh_connect("183.222.230.89", port=22102, username="linux", password="xxx")
  executor.ssh_exec("systemctl status nginx")
  executor.ssh_exec("cd /app && cat config.yml")
  executor.ssh_close()

311 tests passing.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 kaiwu/tools/executor.py    |  69 +++++++++++++-
 kaiwu/tools/ssh_session.py | 184 +++++++++++++++++++++++++++++++++++++
 2 files changed, 251 insertions(+), 2 deletions(-)
 create mode 100644 kaiwu/tools/ssh_session.py

diff --git a/kaiwu/tools/executor.py b/kaiwu/tools/executor.py
index 78799da..5a49e91 100644
--- a/kaiwu/tools/executor.py
+++ b/kaiwu/tools/executor.py
@@ -1,11 +1,11 @@
 """
 Tool executor: self-implemented per FLEX-1 fallback.
-Provides read_file, write_file, run_bash, list_dir, git_commit.
+Provides read_file, write_file, run_bash, list_dir, git_commit, ssh_*.
 Interface is fixed (RED-4: transparent to user).
 
 Guardrails:
 - Dangerous commands blocked (rm -rf, git push --force, drop database, etc.)
-- Sensitive files protected (.env, credentials.json, id_rsa, etc.)
+- Sensitive files auto-backed up before overwrite (.env, credentials.json, etc.)
 - Write operations confined to project_root
 """
 
@@ -40,6 +40,7 @@ class ToolExecutor:
 
     def __init__(self, project_root: str = "."):
         self.project_root = os.path.abspath(project_root)
+        self._ssh_session = None  # Persistent SSH session
 
     def read_file(self, path: str) -> str:
         """Read file content. Path can be relative to project_root or absolute."""
@@ -200,3 +201,67 @@ class ToolExecutor:
             if protected in path_lower:
                 return True
         return False
+
+    # ── SSH Session (persistent, paramiko) ──
+
+    def ssh_connect(
+        self,
+        host: str,
+        port: int = 22,
+        username: str = "root",
+        password: Optional[str] = None,
+        key_path: Optional[str] = None,
+    ) -> tuple[bool, str]:
+        """建立持久 SSH 连接。后续用 ssh_exec 执行命令。"""
+        from kaiwu.tools.ssh_session import SSHSession
+
+        # 关闭旧连接
+        if self._ssh_session and self._ssh_session.connected:
+            self._ssh_session.close()
+
+        self._ssh_session = SSHSession(
+            host=host, port=port, username=username,
+            password=password, key_path=key_path,
+        )
+        return self._ssh_session.connect()
+
+    def ssh_exec(self, command: str, timeout: float = 60.0) -> tuple[str, str, int]:
+        """在远程 SSH 会话中执行命令。返回 (stdout, stderr, returncode)。"""
+        if not self._ssh_session or not self._ssh_session.connected:
+            return "", "[ERROR] SSH未连接，请先用 ssh_connect 建立连接", -1
+
+        # Guardrail: 远程也拦截危险命令
+        blocked = self._check_dangerous(command)
+        if blocked:
+            logger.warning("[guardrail] Blocked dangerous SSH command: %s", command[:80])
+            return "", f"[BLOCKED] 远程危险操作被拦截: {blocked}", -2
+
+        result = self._ssh_session.exec(command, timeout=timeout)
+        return result["stdout"], result["stderr"], result["returncode"]
+
+    def ssh_upload(self, local_path: str, remote_path: str) -> tuple[bool, str]:
+        """上传本地文件到远程 SSH 服务器。"""
+        if not self._ssh_session or not self._ssh_session.connected:
+            return False, "SSH未连接"
+        full_local = self._resolve(local_path)
+        return self._ssh_session.upload(full_local, remote_path)
+
+    def ssh_download(self, remote_path: str, local_path: str) -> tuple[bool, str]:
+        """从远程 SSH 服务器下载文件到本地。"""
+        if not self._ssh_session or not self._ssh_session.connected:
+            return False, "SSH未连接"
+        full_local = self._resolve(local_path)
+        return self._ssh_session.download(remote_path, full_local)
+
+    def ssh_close(self) -> str:
+        """关闭 SSH 连接。"""
+        if self._ssh_session:
+            self._ssh_session.close()
+            self._ssh_session = None
+            return "SSH连接已关闭"
+        return "无活跃SSH连接"
+
+    @property
+    def ssh_connected(self) -> bool:
+        """检查 SSH 是否已连接。"""
+        return bool(self._ssh_session and self._ssh_session.connected)
diff --git a/kaiwu/tools/ssh_session.py b/kaiwu/tools/ssh_session.py
new file mode 100644
index 0000000..161dee9
--- /dev/null
+++ b/kaiwu/tools/ssh_session.py
@@ -0,0 +1,184 @@
+"""
+SSH Session Manager: paramiko 持久 SSH 连接。
+支持多轮命令交互，不需要每次重新连接。
+
+用法：
+  session = SSHSession("183.222.230.89", port=22102, username="linux", password="E5#ok")
+  session.connect()
+  stdout = session.exec("ls /app")
+  stdout = session.exec("cd /app && cat config.yml")
+  session.close()
+
+集成到 ToolExecutor：
+  executor.ssh_connect(host, port, username, password)
+  executor.ssh_exec("systemctl status nginx")
+  executor.ssh_close()
+"""
+
+import logging
+import time
+from typing import Optional
+
+import paramiko
+
+logger = logging.getLogger(__name__)
+
+
+class SSHSession:
+    """持久 SSH 会话，基于 paramiko。"""
+
+    def __init__(
+        self,
+        host: str,
+        port: int = 22,
+        username: str = "root",
+        password: Optional[str] = None,
+        key_path: Optional[str] = None,
+        timeout: float = 10.0,
+    ):
+        self.host = host
+        self.port = port
+        self.username = username
+        self.password = password
+        self.key_path = key_path
+        self.timeout = timeout
+        self._client: Optional[paramiko.SSHClient] = None
+        self._connected = False
+
+    @property
+    def connected(self) -> bool:
+        """检查连接是否存活。"""
+        if not self._client or not self._connected:
+            return False
+        try:
+            transport = self._client.get_transport()
+            return transport is not None and transport.is_active()
+        except Exception:
+            return False
+
+    def connect(self) -> tuple[bool, str]:
+        """
+        建立 SSH 连接。
+        返回 (success, message)。
+        """
+        try:
+            self._client = paramiko.SSHClient()
+            self._client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+            connect_kwargs = {
+                "hostname": self.host,
+                "port": self.port,
+                "username": self.username,
+                "timeout": self.timeout,
+            }
+
+            if self.key_path:
+                connect_kwargs["key_filename"] = self.key_path
+            elif self.password:
+                connect_kwargs["password"] = self.password
+
+            self._client.connect(**connect_kwargs)
+            self._connected = True
+            logger.info("[ssh] Connected to %s@%s:%d", self.username, self.host, self.port)
+            return True, f"已连接 {self.username}@{self.host}:{self.port}"
+
+        except paramiko.AuthenticationException:
+            self._connected = False
+            msg = f"认证失败：{self.username}@{self.host}:{self.port}"
+            logger.error("[ssh] %s", msg)
+            return False, msg
+        except paramiko.SSHException as e:
+            self._connected = False
+            msg = f"SSH错误：{e}"
+            logger.error("[ssh] %s", msg)
+            return False, msg
+        except Exception as e:
+            self._connected = False
+            msg = f"连接失败：{e}"
+            logger.error("[ssh] %s", msg)
+            return False, msg
+
+    def exec(self, command: str, timeout: float = 60.0) -> dict:
+        """
+        在远程执行命令。
+        返回 {"stdout": str, "stderr": str, "returncode": int, "elapsed": float}
+        """
+        if not self.connected:
+            return {
+                "stdout": "",
+                "stderr": "[ERROR] SSH未连接，请先执行 ssh_connect",
+                "returncode": -1,
+                "elapsed": 0.0,
+            }
+
+        t0 = time.time()
+        try:
+            stdin, stdout, stderr = self._client.exec_command(
+                command, timeout=timeout
+            )
+            # 等待命令完成
+            exit_code = stdout.channel.recv_exit_status()
+            out = stdout.read().decode("utf-8", errors="replace")
+            err = stderr.read().decode("utf-8", errors="replace")
+            elapsed = time.time() - t0
+
+            logger.info("[ssh] exec '%s' → rc=%d (%.1fs)", command[:60], exit_code, elapsed)
+            return {
+                "stdout": out,
+                "stderr": err,
+                "returncode": exit_code,
+                "elapsed": elapsed,
+            }
+
+        except Exception as e:
+            elapsed = time.time() - t0
+            logger.error("[ssh] exec failed: %s", e)
+            return {
+                "stdout": "",
+                "stderr": f"[ERROR] {e}",
+                "returncode": -1,
+                "elapsed": elapsed,
+            }
+
+    def upload(self, local_path: str, remote_path: str) -> tuple[bool, str]:
+        """上传文件到远程。"""
+        if not self.connected:
+            return False, "SSH未连接"
+        try:
+            sftp = self._client.open_sftp()
+            sftp.put(local_path, remote_path)
+            sftp.close()
+            logger.info("[ssh] Uploaded %s → %s", local_path, remote_path)
+            return True, f"已上传 {local_path} → {remote_path}"
+        except Exception as e:
+            return False, f"上传失败：{e}"
+
+    def download(self, remote_path: str, local_path: str) -> tuple[bool, str]:
+        """从远程下载文件。"""
+        if not self.connected:
+            return False, "SSH未连接"
+        try:
+            sftp = self._client.open_sftp()
+            sftp.get(remote_path, local_path)
+            sftp.close()
+            logger.info("[ssh] Downloaded %s → %s", remote_path, local_path)
+            return True, f"已下载 {remote_path} → {local_path}"
+        except Exception as e:
+            return False, f"下载失败：{e}"
+
+    def close(self):
+        """关闭连接。"""
+        if self._client:
+            try:
+                self._client.close()
+            except Exception:
+                pass
+            self._connected = False
+            logger.info("[ssh] Connection closed")
+
+    def __del__(self):
+        self.close()
+
+    def __repr__(self):
+        status = "connected" if self.connected else "disconnected"
+        return f"SSHSession({self.username}@{self.host}:{self.port}, {status})"