101 lines
3.7 KiB
Python
101 lines
3.7 KiB
Python
"""SSH 远程命令执行统一封装（STF 重启 / STF 设备管理共用）。
|
||
|
||
认证方式二选一：
|
||
1. `.env` 配置 STF_SSH_PASSWORD → 密码认证（paramiko 实现，跨平台，
|
||
无需 sshpass/ssh 等外部工具，Windows/Linux/macOS 通吃）
|
||
2. 未配置 → 退回系统 `ssh` 命令 + BatchMode=yes（免密密钥，与旧行为一致）
|
||
|
||
用法：
|
||
code, out, err = ssh_client.run("docker ps") # 返回 (退出码, stdout, stderr)
|
||
失败抛 SSHError（连接失败/超时），由调用方转成友好错误返回。
|
||
"""
|
||
import re
|
||
import socket
|
||
import subprocess
|
||
|
||
import paramiko
|
||
|
||
from config import STF_SSH_TARGET, STF_SSH_PASSWORD
|
||
from core.logger import get_logger
|
||
|
||
_log = get_logger("core.ssh")
|
||
|
||
|
||
class SSHError(Exception):
|
||
"""SSH 执行错误（连接失败/超时等）。"""
|
||
|
||
|
||
def _parse_target(target):
|
||
"""解析 STF_SSH_TARGET（user@host[:port] 或 host）→ (user, host, port)。"""
|
||
target = (target or "").strip()
|
||
user, host, port = "root", target, 22
|
||
m = re.match(r"^([^@]+)@(.+)$", target)
|
||
if m:
|
||
user, host = m.group(1), m.group(2)
|
||
if ":" in host:
|
||
host, p = host.rsplit(":", 1)
|
||
try:
|
||
port = int(p)
|
||
except ValueError:
|
||
pass
|
||
return user, host, port
|
||
|
||
|
||
def uses_password():
|
||
"""当前是否走密码认证。"""
|
||
return bool(STF_SSH_PASSWORD)
|
||
|
||
|
||
def run(cmd, timeout=25):
|
||
"""在部署机上执行一条命令。返回 (code, stdout, stderr)。"""
|
||
if uses_password():
|
||
return _run_paramiko(cmd, timeout)
|
||
return _run_system_ssh(cmd, timeout)
|
||
|
||
|
||
def _run_system_ssh(cmd, timeout):
|
||
"""免密密钥路径：系统 ssh + BatchMode=yes（不弹密码提示）。"""
|
||
ssh = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=5",
|
||
"-o", "StrictHostKeyChecking=no", STF_SSH_TARGET]
|
||
try:
|
||
r = subprocess.run(ssh + [cmd], capture_output=True, timeout=timeout)
|
||
except subprocess.TimeoutExpired:
|
||
raise SSHError(f"SSH 执行超时: {cmd[:80]}")
|
||
except FileNotFoundError:
|
||
raise SSHError("本机未安装 ssh，无法远程执行")
|
||
except Exception as e:
|
||
raise SSHError(f"SSH 连接失败: {e}")
|
||
return r.returncode, (r.stdout or b"").decode("utf-8", errors="replace"), \
|
||
(r.stderr or b"").decode("utf-8", errors="replace")
|
||
|
||
|
||
def _run_paramiko(cmd, timeout):
|
||
"""密码路径：paramiko SSHClient，**纯密码认证**。
|
||
|
||
look_for_keys=False + allow_agent=False：不尝试本地密钥、不询问 SSH agent
|
||
（如 Bitwarden），否则 agent 会弹授权框/用密钥绕过密码。
|
||
"""
|
||
user, host, port = _parse_target(STF_SSH_TARGET)
|
||
client = paramiko.SSHClient()
|
||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||
try:
|
||
client.connect(host, port=port, username=user, password=STF_SSH_PASSWORD,
|
||
timeout=8, banner_timeout=8, auth_timeout=8,
|
||
look_for_keys=False, allow_agent=False)
|
||
try:
|
||
stdin, stdout, stderr = client.exec_command(cmd, 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
|
||
finally:
|
||
client.close()
|
||
except paramiko.AuthenticationException:
|
||
raise SSHError(f"SSH 密码认证失败（请检查 .env 的 STF_SSH_PASSWORD）")
|
||
except socket.timeout:
|
||
raise SSHError(f"SSH 执行超时（{timeout}s 内未完成，命令可能耗时较长）")
|
||
except paramiko.SSHException as e:
|
||
raise SSHError(f"SSH 连接失败: {e}")
|
||
except Exception as e:
|
||
raise SSHError(f"SSH 连接失败: {e}")
|