318 lines
10 KiB
Python
318 lines
10 KiB
Python
from __future__ import annotations
|
||
|
||
# 远端文本读写统一走 base64，避免 shell 转义和编码问题。
|
||
import base64
|
||
# 远端脚本里会拼接路径参数，需要显式转义。
|
||
import shlex
|
||
from pathlib import Path
|
||
# SSH / SCP 都直接复用系统 OpenSSH 客户端。
|
||
import subprocess
|
||
from typing import Any
|
||
|
||
from core.config import (
|
||
SSH_CONNECT_TIMEOUT_SECONDS,
|
||
SSH_KNOWN_HOSTS_PATH,
|
||
SSH_PORT,
|
||
SSH_PRIVATE_KEY_PATH,
|
||
SSH_STRICT_HOST_KEY_CHECKING,
|
||
SSH_USER,
|
||
SSH_USE_SUDO,
|
||
load_hosts,
|
||
utc_now,
|
||
)
|
||
|
||
|
||
# 远端文件缺失时统一打印这个标记，避免和普通命令输出混淆。
|
||
REMOTE_FILE_MISSING_MARKER = "__HY_APP_MONITOR_SSH_FILE_MISSING__"
|
||
|
||
|
||
def merged_host_record(host: dict[str, Any]) -> dict[str, Any]:
|
||
# 某些调用方只传 host 名，先尽量回填配置里的 IP / SSH 字段。
|
||
candidate = dict(host or {})
|
||
if candidate.get("ip") or candidate.get("sshHost"):
|
||
return candidate
|
||
|
||
host_name = str(candidate.get("host") or "").strip()
|
||
if not host_name:
|
||
return candidate
|
||
|
||
for item in load_hosts():
|
||
if item.get("host") != host_name:
|
||
continue
|
||
merged = dict(item)
|
||
merged.update({key: value for key, value in candidate.items() if value not in {"", None}})
|
||
return merged
|
||
|
||
return candidate
|
||
|
||
|
||
def ensure_ssh_runtime_files() -> None:
|
||
# known_hosts 放到项目 run 目录，首次连接时也能自动落盘。
|
||
SSH_KNOWN_HOSTS_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||
if not SSH_KNOWN_HOSTS_PATH.exists():
|
||
SSH_KNOWN_HOSTS_PATH.touch()
|
||
|
||
|
||
def require_private_key() -> Path:
|
||
# SSH 主链路必须拿到 deploy 机上的专用私钥。
|
||
key_path = Path(SSH_PRIVATE_KEY_PATH).expanduser()
|
||
if not key_path.exists():
|
||
raise RuntimeError(
|
||
f"missing ssh private key: {key_path} ; run ops/bootstrap_ssh_via_tat.py on the deploy chain first"
|
||
)
|
||
return key_path
|
||
|
||
|
||
def host_ssh_user(host: dict[str, Any]) -> str:
|
||
# 主机可单独覆盖 SSH 用户，未配置时走全局默认值。
|
||
merged = merged_host_record(host)
|
||
return str(merged.get("sshUser") or SSH_USER).strip() or SSH_USER
|
||
|
||
|
||
def host_ssh_host(host: dict[str, Any]) -> str:
|
||
# 优先取显式 sshHost，再回退到业务 IP。
|
||
merged = merged_host_record(host)
|
||
target = str(merged.get("sshHost") or merged.get("ip") or "").strip()
|
||
if not target:
|
||
raise RuntimeError(f"missing ssh host for {merged.get('host') or 'unknown host'}")
|
||
return target
|
||
|
||
|
||
def host_ssh_port(host: dict[str, Any]) -> int:
|
||
# 端口同样允许按主机覆写。
|
||
merged = merged_host_record(host)
|
||
return int(merged.get("sshPort") or SSH_PORT)
|
||
|
||
|
||
def host_ssh_use_sudo(host: dict[str, Any]) -> bool:
|
||
# root 登录时通常不需要 sudo，非 root 用户时可以通过配置统一启用 sudo -n。
|
||
merged = merged_host_record(host)
|
||
value = merged.get("sshUseSudo")
|
||
if isinstance(value, bool):
|
||
return value
|
||
if isinstance(value, str):
|
||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||
return SSH_USE_SUDO
|
||
|
||
|
||
def target_display(host: dict[str, Any]) -> str:
|
||
# 错误信息里统一带上 user@host:port，排查更直接。
|
||
return f"{host_ssh_user(host)}@{host_ssh_host(host)}:{host_ssh_port(host)}"
|
||
|
||
|
||
def ssh_base_args(host: dict[str, Any]) -> list[str]:
|
||
# 所有 SSH 命令统一带上专用 key、known_hosts 和超时配置。
|
||
ensure_ssh_runtime_files()
|
||
key_path = require_private_key()
|
||
return [
|
||
"ssh",
|
||
"-i",
|
||
str(key_path),
|
||
"-p",
|
||
str(host_ssh_port(host)),
|
||
"-o",
|
||
"BatchMode=yes",
|
||
"-o",
|
||
f"ConnectTimeout={SSH_CONNECT_TIMEOUT_SECONDS}",
|
||
"-o",
|
||
"ServerAliveInterval=15",
|
||
"-o",
|
||
"ServerAliveCountMax=3",
|
||
"-o",
|
||
f"StrictHostKeyChecking={SSH_STRICT_HOST_KEY_CHECKING}",
|
||
"-o",
|
||
f"UserKnownHostsFile={SSH_KNOWN_HOSTS_PATH}",
|
||
f"{host_ssh_user(host)}@{host_ssh_host(host)}",
|
||
]
|
||
|
||
|
||
def scp_base_args(host: dict[str, Any]) -> list[str]:
|
||
# SCP 也沿用同一套 key、known_hosts 和端口配置。
|
||
ensure_ssh_runtime_files()
|
||
key_path = require_private_key()
|
||
return [
|
||
"-i",
|
||
str(key_path),
|
||
"-P",
|
||
str(host_ssh_port(host)),
|
||
"-o",
|
||
"BatchMode=yes",
|
||
"-o",
|
||
f"ConnectTimeout={SSH_CONNECT_TIMEOUT_SECONDS}",
|
||
"-o",
|
||
f"StrictHostKeyChecking={SSH_STRICT_HOST_KEY_CHECKING}",
|
||
"-o",
|
||
f"UserKnownHostsFile={SSH_KNOWN_HOSTS_PATH}",
|
||
]
|
||
|
||
|
||
def run_host_script(
|
||
host: dict[str, Any],
|
||
script: str,
|
||
*,
|
||
command_name: str,
|
||
timeout_seconds: int,
|
||
) -> dict[str, Any]:
|
||
# 所有远端动作都通过 bash -se 执行，stdin 直接传脚本可以避开复杂转义。
|
||
remote_shell = ["bash", "-se"]
|
||
if host_ssh_use_sudo(host):
|
||
remote_shell = ["sudo", "-n", *remote_shell]
|
||
args = [*ssh_base_args(host), *remote_shell]
|
||
|
||
try:
|
||
result = subprocess.run(
|
||
args,
|
||
input=script,
|
||
text=True,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
timeout=timeout_seconds,
|
||
check=False,
|
||
)
|
||
except subprocess.TimeoutExpired as exc:
|
||
# 超时时仍把已有输出保留下来，页面和日志更容易定位卡在哪一步。
|
||
output = (exc.stdout or "") if isinstance(exc.stdout, str) else ""
|
||
return {
|
||
"status": "TIMEOUT",
|
||
"output": output,
|
||
"updatedAt": utc_now(),
|
||
"commandName": command_name,
|
||
"target": target_display(host),
|
||
}
|
||
except FileNotFoundError as exc:
|
||
raise RuntimeError("ssh command not found in PATH") from exc
|
||
|
||
output = str(result.stdout or "")
|
||
status = "SUCCESS" if result.returncode == 0 else "FAILED"
|
||
if status != "SUCCESS" and not output.strip():
|
||
output = f"{command_name} failed on {target_display(host)}: exit {result.returncode}"
|
||
|
||
return {
|
||
"status": status,
|
||
"output": output,
|
||
"updatedAt": utc_now(),
|
||
"commandName": command_name,
|
||
"target": target_display(host),
|
||
}
|
||
|
||
|
||
def run_host_script_checked(
|
||
host: dict[str, Any],
|
||
script: str,
|
||
*,
|
||
command_name: str,
|
||
timeout_seconds: int,
|
||
) -> dict[str, Any]:
|
||
# 需要强失败语义的调用方统一走这里，避免每个模块重复判断状态码。
|
||
result = run_host_script(host, script, command_name=command_name, timeout_seconds=timeout_seconds)
|
||
if result.get("status") != "SUCCESS":
|
||
detail = str(result.get("output") or result.get("status") or "remote command failed").strip()
|
||
raise RuntimeError(f"{command_name} failed on {target_display(host)}: {detail}")
|
||
return result
|
||
|
||
|
||
def read_remote_text(host: dict[str, Any], path: str, *, timeout_seconds: int = 60) -> str:
|
||
# 文本读取直接在远端整体 base64 后返回，本项目里的 env / compose 文件体积都足够小。
|
||
quoted_path = shlex.quote(path)
|
||
script = f"""set -eu
|
||
FILE={quoted_path}
|
||
if [ ! -f "$FILE" ]; then
|
||
printf '%s' {shlex.quote(REMOTE_FILE_MISSING_MARKER)}
|
||
exit 0
|
||
fi
|
||
python3 - "$FILE" <<'PY'
|
||
from pathlib import Path
|
||
import base64
|
||
import sys
|
||
|
||
print(base64.b64encode(Path(sys.argv[1]).read_bytes()).decode("ascii"), end="")
|
||
PY
|
||
"""
|
||
result = run_host_script_checked(host, script, command_name=f"read-{merged_host_record(host).get('host') or 'remote'}", timeout_seconds=timeout_seconds)
|
||
output = str(result.get("output") or "").strip()
|
||
if output == REMOTE_FILE_MISSING_MARKER:
|
||
raise RuntimeError(f"missing remote file on {target_display(host)}: {path}")
|
||
if not output:
|
||
return ""
|
||
return base64.b64decode(output).decode("utf-8")
|
||
|
||
|
||
def write_remote_text(host: dict[str, Any], path: str, content: str, *, timeout_seconds: int = 60) -> None:
|
||
# 写文件时先落临时文件再 mv，避免中途失败把线上文件写坏。
|
||
encoded = base64.b64encode(content.encode("utf-8")).decode("ascii")
|
||
quoted_path = shlex.quote(path)
|
||
script = f"""set -eu
|
||
FILE={quoted_path}
|
||
TMP="${{FILE}}.hyapp.$$"
|
||
mkdir -p "$(dirname "$FILE")"
|
||
printf '%s' {shlex.quote(encoded)} | base64 -d > "$TMP"
|
||
mv "$TMP" "$FILE"
|
||
"""
|
||
run_host_script_checked(
|
||
host,
|
||
script,
|
||
command_name=f"write-{merged_host_record(host).get('host') or 'remote'}",
|
||
timeout_seconds=timeout_seconds,
|
||
)
|
||
|
||
|
||
def copy_remote_file(host: dict[str, Any], source_path: str, target_path: str, *, timeout_seconds: int = 60) -> None:
|
||
# 线上 compose 备份直接在远端复制即可，不需要再把原文件拉回本机。
|
||
script = f"""set -eu
|
||
SOURCE={shlex.quote(source_path)}
|
||
TARGET={shlex.quote(target_path)}
|
||
if [ ! -f "$SOURCE" ]; then
|
||
printf '%s' {shlex.quote(REMOTE_FILE_MISSING_MARKER)}
|
||
exit 0
|
||
fi
|
||
mkdir -p "$(dirname "$TARGET")"
|
||
cp "$SOURCE" "$TARGET"
|
||
"""
|
||
result = run_host_script_checked(
|
||
host,
|
||
script,
|
||
command_name=f"copy-{merged_host_record(host).get('host') or 'remote'}",
|
||
timeout_seconds=timeout_seconds,
|
||
)
|
||
if str(result.get("output") or "").strip() == REMOTE_FILE_MISSING_MARKER:
|
||
raise RuntimeError(f"missing remote file on {target_display(host)}: {source_path}")
|
||
|
||
|
||
def copy_local_path_to_host(
|
||
host: dict[str, Any],
|
||
source_path: str | Path,
|
||
target_path: str,
|
||
*,
|
||
recursive: bool = False,
|
||
timeout_seconds: int = 120,
|
||
) -> None:
|
||
# Admin 前端发布会把本地 dist 目录直接传到远端临时目录。
|
||
args = ["scp", *scp_base_args(host)]
|
||
if recursive:
|
||
args.append("-r")
|
||
args.extend(
|
||
[
|
||
str(source_path),
|
||
f"{host_ssh_user(host)}@{host_ssh_host(host)}:{target_path}",
|
||
]
|
||
)
|
||
|
||
try:
|
||
result = subprocess.run(
|
||
args,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
text=True,
|
||
timeout=timeout_seconds,
|
||
check=False,
|
||
)
|
||
except subprocess.TimeoutExpired as exc:
|
||
output = (exc.stdout or "") if isinstance(exc.stdout, str) else ""
|
||
raise RuntimeError(f"scp timeout on {target_display(host)}: {output.strip() or source_path}") from exc
|
||
except FileNotFoundError as exc:
|
||
raise RuntimeError("scp command not found in PATH") from exc
|
||
|
||
if result.returncode != 0:
|
||
detail = str(result.stdout or "").strip()
|
||
raise RuntimeError(f"scp failed on {target_display(host)}: {detail or f'exit {result.returncode}'}")
|