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 _DATA_ROOT = Path(os.getenv("QIP_DATA_DIR", str(PROJECT_ROOT / "data"))) SSH_SCRIPTS_DIR = _DATA_ROOT / "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