group_fqcd_jr/tools/rotate_api_keys.py

"""轮换 `.env` 里的模型密钥（DashScope/Qwen 与 DeepSeek）。

## 为什么需要它

这两把 key **在聊天/文档里出现过明文**，必须轮换。手工改 `.env` 有三个坑：

1. **三个 Qwen 变量必须同一个值** —— `QWEN_API_KEY` / `QWEN_EMBEDDING_API_KEY` /
   `DASHSCOPE_API_KEY`。漏改一个，入库与检索就用了两把不同的 key，
   症状是「检索能跑、入库失败」或反过来（很难一眼看出是 key 的问题）。
2. **两个 DeepSeek 变量也必须同一个值** —— `DEEPSEEK_API_KEY` 与
   `OFFSITE_DEEPSEEK_API_KEY`（场外文档识别复用同一把）。
3. 写文件时若**带上 BOM**或**改变编码/换行**，`.env` 第一项就可能读不出来。

## 用法

    .venv\\Scripts\\python.exe tools\\rotate_api_keys.py            # 交互输入（不回显）
    .venv\\Scripts\\python.exe tools\\rotate_api_keys.py --check    # 只体检当前 5 个变量的取值关系

## 安全约定

- 输入**不回显**（`getpass`），**不写入任何日志**，**不回显 key 值**（只打印掩码与长度）。
- 改前自动备份为 `.env.bak-<时间戳>`（该模式已被 `.gitignore` 覆盖）。
- 任一校验不通过就**整体不写入**（先全量校验、后一次性落盘）。
"""

from __future__ import annotations

import argparse
import getpass
import shutil
import sys
from datetime import datetime
from pathlib import Path

ENV_PATH = Path(__file__).resolve().parent.parent / ".env"

#: 变量名 -> 归属的 key（同一个 key 必须写成同一个值）
QWEN_VARS = ("QWEN_API_KEY", "QWEN_EMBEDDING_API_KEY", "DASHSCOPE_API_KEY")
DEEPSEEK_VARS = ("DEEPSEEK_API_KEY", "OFFSITE_DEEPSEEK_API_KEY")
TARGETS = {name: "qwen" for name in QWEN_VARS} | {name: "deepseek" for name in DEEPSEEK_VARS}

MIN_LEN = 20


def mask(value: str) -> str:
    """只暴露前 6 位，其余按长度打码 —— 够判断"换没换"，不够还原。"""
    if not value:
        return "(空)"
    return f"{value[:6]}…({len(value)} 位)"


def parse_env(text: str) -> list[tuple[str, str]]:
    rows: list[tuple[str, str]] = []
    for line in text.split("\n"):
        if "=" in line and not line.lstrip().startswith("#"):
            key, _, value = line.partition("=")
            rows.append((key.strip(), value))
    return rows


def check(rows: list[tuple[str, str]]) -> int:
    """体检：三个 Qwen 是否同值、两个 DeepSeek 是否同值、是否为空/像占位符。"""
    values = dict(rows)
    problems = 0
    for group, names in (("Qwen（DashScope）", QWEN_VARS), ("DeepSeek", DEEPSEEK_VARS)):
        present = [values.get(n, "") for n in names]
        print(f"· {group}：")
        for name in names:
            print(f"    {name:<26} {mask(values.get(name, ''))}")
        uniq = {v for v in present if v}
        if not uniq:
            print("    ❌ 全部为空")
            problems += 1
        elif len(uniq) > 1:
            print("    ❌ **取值不一致** —— 这会让「入库/检索」或「主链路/场外链路」用到不同的 key")
            problems += 1
        elif len(present) != len([v for v in present if v]):
            print("    ❌ 有变量为空（必须同一个值）")
            problems += 1
        else:
            print("    ✅ 同值且非空")
    return problems


def rewrite(text: str, values: dict[str, str]) -> tuple[str, int]:
    """只替换目标行的值，其它行（含缩进、注释、顺序）原样保留。"""
    out: list[str] = []
    replaced = 0
    for line in text.split("\n"):
        stripped = line.lstrip()
        if stripped.startswith("#") or "=" not in line:
            out.append(line)
            continue
        key, sep, _old = line.partition("=")
        name = key.strip()
        if name in values and name in TARGETS:
            out.append(f"{name}{sep}{values[name]}")
            replaced += 1
        else:
            out.append(line)
    return "\n".join(out), replaced


def main() -> int:
    parser = argparse.ArgumentParser(description="轮换 .env 里的 Qwen / DeepSeek 密钥")
    parser.add_argument("--check", action="store_true", help="只体检，不修改")
    args = parser.parse_args()

    if not ENV_PATH.exists():
        print(f"❌ 找不到 {ENV_PATH}")
        return 2
    raw = ENV_PATH.read_bytes()
    if raw.startswith(b"\xef\xbb\xbf"):
        print("⚠️  .env 带 UTF-8 BOM：本次会去掉它（BOM 会让第一个变量名读不出来）")
    newline = "\r\n" if b"\r\n" in raw else "\n"
    text = raw.decode("utf-8-sig").replace("\r\n", "\n")
    rows = parse_env(text)

    print("=== .env 密钥体检（只显示掩码）===")
    problems = check(rows)
    if args.check:
        print()
        print("体检完成。" + ("存在不一致，建议先修。" if problems else "一致。"))
        return 1 if problems else 0

    missing = [n for n in TARGETS if n not in dict(rows)]
    if missing:
        print(f"❌ .env 里缺少这些变量：{missing}")
        return 2

    print()
    print("=== 开始轮换 ===")
    print("直接回车 = 跳过该项（保留旧值）。输入不会回显，也不会被记录。")
    print()
    new_values: dict[str, str] = {}
    for group, names in (("DashScope / Qwen", QWEN_VARS), ("DeepSeek", DEEPSEEK_VARS)):
        entered = getpass.getpass(f"{group} 的新 key（回车跳过）：").strip()
        if not entered:
            print(f"  · {group}：跳过，保留旧值")
            continue
        if not entered.startswith("sk-") or len(entered) < MIN_LEN:
            print(
                f"  ❌ {group}：格式不像 key（应以 'sk-' 开头且长度 ≥ {MIN_LEN}）—— **本次不做任何写入**"
            )
            return 3
        for name in names:
            new_values[name] = entered
        print(f"  · {group}：将更新 {len(names)} 个变量（{mask(entered)}）")

    if not new_values:
        print()
        print("没有输入任何 key，未做修改。")
        return 0

    backup = ENV_PATH.with_name(f".env.bak-{datetime.now():%Y%m%d-%H%M%S}")
    shutil.copy2(ENV_PATH, backup)

    updated, replaced = rewrite(text, new_values)
    if replaced != len(new_values):
        print(f"❌ 只匹配到 {replaced}/{len(new_values)} 行，已放弃写入（备份在 {backup.name}）")
        return 3
    ENV_PATH.write_bytes(updated.replace("\n", newline).encode("utf-8"))
    print()
    print(f"✅ 已更新 {replaced} 行；备份：{backup.name}")
    print(
        "   下一步：重启 API 与 Worker 让新 key 生效，然后跑一次链路复核（见仓库根 demo.ps1 的自检）。"
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())