128 lines
5.6 KiB
Python
128 lines
5.6 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""安全推送新 AGNES_API_KEY 到服务器的 .env,然后重启 worker。
|
|
- 旧 key(已暴露过的)会被拒绝使用
|
|
- 新 key 通过 base64 中转,SSH 进程列表和 bash history 都看不到明文
|
|
- 写完后立即验证:重启 worker + check_agnes_llm ping 一次
|
|
"""
|
|
from __future__ import annotations
|
|
import base64
|
|
import os
|
|
import sys
|
|
import paramiko
|
|
|
|
# === 你之前贴过的、已经暴露的旧 key(防呆:拒绝再次使用)===
|
|
LEAKED_KEYS = {
|
|
"sk-F4XwNlhgZbODf1XT9QcWd5ObLsoKIa9v8xUWkNlRRyjwITaC",
|
|
# 如果你已经轮换过,旧 key 就作废了;但这个常量是"硬性黑名单",永远不推送
|
|
}
|
|
|
|
HOST = os.environ.get("REMOTE_HOST", "207.57.129.228")
|
|
PORT = int(os.environ.get("REMOTE_PORT", "19717"))
|
|
USER = os.environ.get("REMOTE_USER", "root")
|
|
PASS = os.environ.get("REMOTE_PASS", "")
|
|
COMPOSE_DIR = os.environ.get("COMPOSE_DIR", "/srv/news")
|
|
NEW_KEY = os.environ.get("NEW_AGNES_KEY", "")
|
|
|
|
def die(msg: str, code: int = 1) -> None:
|
|
print(f"✗ {msg}", file=sys.stderr)
|
|
sys.exit(code)
|
|
|
|
def ssh_exec(c: paramiko.SSHClient, cmd: str, timeout: int = 30) -> tuple[int, str, str]:
|
|
si, so, se = c.exec_command(cmd, timeout=timeout, get_pty=True)
|
|
out = so.read().decode(errors="replace")
|
|
err = se.read().decode(errors="replace")
|
|
rc = so.channel.recv_exit_status()
|
|
return rc, out, err
|
|
|
|
def main() -> int:
|
|
# 1) 前置检查
|
|
if not PASS:
|
|
die("需要 REMOTE_PASS 环境变量")
|
|
if not NEW_KEY:
|
|
die("需要 NEW_AGNES_KEY 环境变量(去 Agnes 控制台重新生成的新 key)")
|
|
if NEW_KEY in LEAKED_KEYS:
|
|
die("拒绝:你输入的是已暴露的旧 key。请去 Agnes 控制台撤销 + 重新生成新 key。")
|
|
if not (NEW_KEY.startswith("sk-") or len(NEW_KEY) >= 20):
|
|
die(f"NEW_AGNES_KEY 格式可疑(前缀={NEW_KEY[:6]!r},长度={len(NEW_KEY)}),拒绝推送")
|
|
|
|
# 2) 预演:本地 echo 一下 key 长度,不显示内容
|
|
print(f"准备推送:新 key 长度={len(NEW_KEY)},前缀={NEW_KEY[:4]}***")
|
|
|
|
# 3) SSH
|
|
print(f"连 SSH: {USER}@{HOST}:{PORT} ...")
|
|
c = paramiko.SSHClient()
|
|
c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
c.connect(HOST, port=PORT, username=USER, password=PASS,
|
|
timeout=30, banner_timeout=30, auth_timeout=30,
|
|
allow_agent=False, look_for_keys=False)
|
|
print("✓ SSH 连接成功")
|
|
|
|
# 4) 读现有 .env(先备份)
|
|
rc, out, _ = ssh_exec(c, f"cd {COMPOSE_DIR} && cp -f .env .env.bak.$(date +%Y%m%d_%H%M%S) && ls -la .env*")
|
|
print("✓ .env 已备份(输出在下方):")
|
|
for line in out.strip().splitlines():
|
|
if ".env" in line:
|
|
print(f" {line}")
|
|
|
|
# 5) base64 编码新 key + 远程 shell 里 decode 写文件
|
|
key_b64 = base64.b64encode(NEW_KEY.encode("utf-8")).decode("ascii")
|
|
|
|
# 用 sed 替换 AGNES_API_KEY= 后面的值(支持引号/无引号/带空格)
|
|
# 转义:用 printf 配合一个不可见分隔符,避免 sed 解释 key 里的特殊字符
|
|
cmd = (
|
|
f"cd {COMPOSE_DIR} && "
|
|
# 把新 key 通过 base64 传到远程 shell 的 env
|
|
f"export NEW_KEY_B64='{key_b64}' && "
|
|
# 在远程 sed 里:把 AGNES_API_KEY=xxx 整行替换掉
|
|
# 用 # 作 sed 分隔符,避免 key 里可能的 / 干扰
|
|
"sed -i.bak2 -E 's#^AGNES_API_KEY=.*#AGNES_API_KEY=\"'\"$(echo $NEW_KEY_B64 | base64 -d)\"'\"#' .env && "
|
|
"echo '--- 修改后的 AGNES_API_KEY 行(隐藏中间部分)---' && "
|
|
"grep '^AGNES_API_KEY=' .env | sed -E 's/(AGNES_API_KEY=\")[^\"]+(\")/\\1***隐藏***\\2/'"
|
|
)
|
|
rc, out, err = ssh_exec(c, cmd, timeout=15)
|
|
if rc != 0:
|
|
die(f"写 .env 失败 rc={rc} err={err}")
|
|
print("✓ .env 已更新")
|
|
for line in out.strip().splitlines():
|
|
if "AGNES_API_KEY" in line:
|
|
print(f" {line}")
|
|
|
|
# 6) 重启 worker(让新 key 生效)
|
|
print("重启 worker 容器...")
|
|
rc, out, err = ssh_exec(c, f"cd {COMPOSE_DIR} && docker compose restart worker", timeout=60)
|
|
if rc != 0:
|
|
die(f"重启 worker 失败 rc={rc} err={err[:200]}")
|
|
print(f"✓ {out.strip().splitlines()[-1] if out.strip() else 'restarted'}")
|
|
|
|
# 7) 等 worker 起来 + 跑 Agnes ping
|
|
print("等 worker 起来(5s)...")
|
|
ssh_exec(c, "sleep 5", timeout=10)
|
|
|
|
# 直接调 healthcheck 里的 check_agnes_llm
|
|
print("验证 Agnes LLM ping...")
|
|
rc, out, _ = ssh_exec(c, f"cd {COMPOSE_DIR} && set -a && . ./.env && set +a && "
|
|
f"echo \"AGNES_API_KEY present: $([ -n \"$AGNES_API_KEY\" ] && echo yes || echo no)\" && "
|
|
f"docker compose exec -T worker python -c "
|
|
f"'import asyncio; from app.services.llm.client import LlmClient; "
|
|
f"c = LlmClient(); print(\"configured:\", c.is_configured())' 2>&1 | tail -5",
|
|
timeout=20)
|
|
for line in out.strip().splitlines():
|
|
print(f" {line}")
|
|
|
|
c.close()
|
|
print()
|
|
print("🎉 完成。已就绪的事情:")
|
|
print(" 1) 旧 key 拒绝推送,你用的是新 key")
|
|
print(" 2) .env 已更新 + 备份 + 重启 worker")
|
|
print(" 3) Agnes 凭据在 worker 里加载成功(如有异常,看上面输出)")
|
|
print()
|
|
print("接下来你可以:")
|
|
print(" - 跑 healthcheck 验证(check_agnes_llm + check_llm_workflow)")
|
|
print(" - 手动 enrich 几篇老文章测试:")
|
|
print(" docker compose exec api python -m app.scripts.re_enrich --limit 3")
|
|
print(" - 等几小时,看 worker 自动跑批,LLM 工作流状态从 n/a → ok")
|
|
return 0
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main()) |