qihongkun
d810abdcee
纳入 FastAPI 后端、Vue 管理端、MCP 桥接与文档；通过 .gitignore 排除本地数据库与构建产物。 Co-authored-by: Cursor <cursoragent@cursor.com>
221 lines
8.2 KiB
Python
221 lines
8.2 KiB
Python
import asyncio
|
|
import contextlib
|
|
import os
|
|
import posixpath
|
|
import uuid
|
|
from collections.abc import AsyncIterator
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import paramiko
|
|
|
|
from ..schemas import SshProfileOut, SshScriptOut
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
|
|
SSH_SCRIPTS_DIR = PROJECT_ROOT / "data" / "ssh-scripts"
|
|
|
|
|
|
def resolve_script_body(script: SshScriptOut) -> str:
|
|
if script.source_type == "file":
|
|
if not script.storage_path:
|
|
raise ValueError("脚本文件路径为空")
|
|
file_path = (SSH_SCRIPTS_DIR / script.storage_path).resolve()
|
|
scripts_root = SSH_SCRIPTS_DIR.resolve()
|
|
if scripts_root not in file_path.parents and file_path != scripts_root:
|
|
raise ValueError("脚本文件路径非法")
|
|
if not file_path.is_file():
|
|
raise ValueError("脚本文件不存在")
|
|
return file_path.read_text(encoding="utf-8")
|
|
body = (script.content or "").strip()
|
|
if not body:
|
|
raise ValueError("脚本内容为空")
|
|
return body
|
|
|
|
|
|
def _build_connect_kwargs(profile: SshProfileOut, password: str) -> dict[str, Any]:
|
|
if not password.strip():
|
|
raise ValueError("执行前必须输入密码")
|
|
|
|
connect_kwargs: dict[str, Any] = {
|
|
"hostname": profile.host,
|
|
"port": int(profile.port or 22),
|
|
"username": profile.username,
|
|
"timeout": 15,
|
|
"banner_timeout": 15,
|
|
"auth_timeout": 15,
|
|
}
|
|
if profile.auth_type == "password":
|
|
connect_kwargs.update({"password": password, "look_for_keys": False, "allow_agent": False})
|
|
elif profile.auth_type == "key" and profile.key_path:
|
|
connect_kwargs.update({"key_filename": profile.key_path, "look_for_keys": False, "allow_agent": False})
|
|
else:
|
|
connect_kwargs.update({"look_for_keys": True, "allow_agent": True})
|
|
return connect_kwargs
|
|
|
|
|
|
def _connect_ssh(profile: SshProfileOut, password: str) -> paramiko.SSHClient:
|
|
if profile.auth_type == "key" and profile.key_path and not os.path.exists(profile.key_path):
|
|
raise ValueError("配置的 key_path 不存在")
|
|
|
|
client = paramiko.SSHClient()
|
|
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
client.connect(**_build_connect_kwargs(profile, password))
|
|
return client
|
|
|
|
|
|
def _upload_and_build_command(client: paramiko.SSHClient, script_body: str, script_name: str) -> str:
|
|
run_id = uuid.uuid4().hex[:12]
|
|
remote_dir = posixpath.join("/tmp", f"qip-script-{run_id}")
|
|
safe_name = "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in script_name) or "script.sh"
|
|
if not safe_name.endswith(".sh"):
|
|
safe_name = f"{safe_name}.sh"
|
|
remote_path = posixpath.join(remote_dir, safe_name)
|
|
|
|
stdin, stdout, stderr = client.exec_command(f"mkdir -p {remote_dir} && chmod 700 {remote_dir}")
|
|
exit_status = stdout.channel.recv_exit_status()
|
|
if exit_status != 0:
|
|
err = (stderr.read() or b"").decode("utf-8", errors="replace").strip()
|
|
raise RuntimeError(err or "无法在远端创建临时目录")
|
|
|
|
with client.open_sftp() as sftp:
|
|
with sftp.file(remote_path, "w") as remote_file:
|
|
remote_file.write(script_body)
|
|
sftp.chmod(remote_path, 0o700)
|
|
|
|
return f"bash {remote_path}"
|
|
|
|
|
|
def _cleanup_remote_session(stdout: Any, stderr: Any, client: paramiko.SSHClient | None) -> None:
|
|
for stream in (stdout, stderr):
|
|
channel = getattr(stream, "channel", None)
|
|
if channel is not None:
|
|
with contextlib.suppress(Exception):
|
|
channel.close()
|
|
if client is not None:
|
|
with contextlib.suppress(Exception):
|
|
client.close()
|
|
|
|
|
|
async def stream_script_run(
|
|
profile: SshProfileOut,
|
|
script: SshScriptOut,
|
|
password: str,
|
|
cancel_event: asyncio.Event | None = None,
|
|
) -> AsyncIterator[dict[str, Any]]:
|
|
cancel_event = cancel_event or asyncio.Event()
|
|
script_body = resolve_script_body(script)
|
|
client = await asyncio.to_thread(_connect_ssh, profile, password)
|
|
stdout = None
|
|
stderr = None
|
|
|
|
try:
|
|
remote_command = await asyncio.to_thread(_upload_and_build_command, client, script_body, script.name)
|
|
yield {"type": "meta", "command": remote_command, "profile": profile.name, "script": script.name}
|
|
|
|
stdin, stdout, stderr = await asyncio.to_thread(client.exec_command, remote_command)
|
|
if stdin:
|
|
with contextlib.suppress(Exception):
|
|
stdin.close()
|
|
|
|
while True:
|
|
if cancel_event.is_set():
|
|
yield {"type": "stopped", "detail": "用户手动停止", "ok": False, "exit_status": -1}
|
|
return
|
|
|
|
emitted = False
|
|
if stdout.channel.recv_ready():
|
|
chunk = await asyncio.to_thread(stdout.channel.recv, 4096)
|
|
if chunk:
|
|
emitted = True
|
|
yield {"type": "stdout", "data": chunk.decode("utf-8", errors="replace")}
|
|
if stderr.channel.recv_stderr_ready():
|
|
chunk = await asyncio.to_thread(stderr.channel.recv_stderr, 4096)
|
|
if chunk:
|
|
emitted = True
|
|
yield {"type": "stderr", "data": chunk.decode("utf-8", errors="replace")}
|
|
if (
|
|
stdout.channel.exit_status_ready()
|
|
and not stdout.channel.recv_ready()
|
|
and not stderr.channel.recv_stderr_ready()
|
|
):
|
|
break
|
|
if not emitted:
|
|
await asyncio.sleep(0.05)
|
|
|
|
if cancel_event.is_set():
|
|
yield {"type": "stopped", "detail": "用户手动停止", "ok": False, "exit_status": -1}
|
|
return
|
|
|
|
exit_status = await asyncio.to_thread(stdout.channel.recv_exit_status)
|
|
yield {"type": "done", "exit_status": int(exit_status), "ok": exit_status == 0}
|
|
finally:
|
|
_cleanup_remote_session(stdout, stderr, client)
|
|
|
|
|
|
async def run_script_collect(
|
|
profile: SshProfileOut,
|
|
script: SshScriptOut,
|
|
password: str,
|
|
timeout_seconds: float = 120.0,
|
|
max_output_chars: int = 200_000,
|
|
) -> dict[str, Any]:
|
|
timeout_value = max(5.0, min(float(timeout_seconds or 120.0), 600.0))
|
|
cancel_event = asyncio.Event()
|
|
stdout_parts: list[str] = []
|
|
stderr_parts: list[str] = []
|
|
meta: dict[str, Any] = {}
|
|
result: dict[str, Any] = {
|
|
"ok": False,
|
|
"exit_status": -1,
|
|
"stdout": "",
|
|
"stderr": "",
|
|
"command": "",
|
|
"timed_out": False,
|
|
"truncated": False,
|
|
}
|
|
|
|
async def consume() -> None:
|
|
nonlocal result
|
|
async for event in stream_script_run(profile, script, password, cancel_event):
|
|
event_type = str(event.get("type", ""))
|
|
if event_type == "meta":
|
|
meta.update(event)
|
|
result["command"] = str(event.get("command", ""))
|
|
continue
|
|
if event_type in {"stdout", "stderr"}:
|
|
chunk = str(event.get("data", ""))
|
|
bucket = stdout_parts if event_type == "stdout" else stderr_parts
|
|
bucket.append(chunk)
|
|
total_len = sum(len(item) for item in stdout_parts) + sum(len(item) for item in stderr_parts)
|
|
if total_len > max_output_chars:
|
|
result["truncated"] = True
|
|
cancel_event.set()
|
|
return
|
|
continue
|
|
if event_type == "done":
|
|
result["ok"] = bool(event.get("ok"))
|
|
result["exit_status"] = int(event.get("exit_status", 0))
|
|
return
|
|
if event_type == "stopped":
|
|
result["ok"] = False
|
|
result["exit_status"] = int(event.get("exit_status", -1))
|
|
result["stopped"] = True
|
|
return
|
|
if event_type == "error":
|
|
raise RuntimeError(str(event.get("detail", "ssh script run failed")))
|
|
|
|
try:
|
|
await asyncio.wait_for(consume(), timeout=timeout_value)
|
|
except asyncio.TimeoutError:
|
|
cancel_event.set()
|
|
result["timed_out"] = True
|
|
result["ok"] = False
|
|
result["exit_status"] = -1
|
|
finally:
|
|
result["stdout"] = "".join(stdout_parts)
|
|
result["stderr"] = "".join(stderr_parts)
|
|
if meta:
|
|
result["profile"] = meta.get("profile")
|
|
result["script"] = meta.get("script")
|
|
return result
|