mirror of
https://github.com/val1813/kwcode.git
synced 2026-09-03 06:34:30 +08:00
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>
This commit is contained in:
@@ -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)
|
||||
|
||||
184
kaiwu/tools/ssh_session.py
Normal file
184
kaiwu/tools/ssh_session.py
Normal file
@@ -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})"
|
||||
Reference in New Issue
Block a user