补回 mint scope 中的 grok-cli:access,铸造后校验 JWT scope, 避免 cli-chat-proxy 返回 grok-cli-token-auth-required。 新增 remint_cli_scope.py,基于已有 sso 覆盖写回缺 scope 的 CPA auth; 同步更新 cpa/schema 文档说明 referrer 与 scope 双重要求。
293 lines
9.2 KiB
Python
293 lines
9.2 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""批量重铸 CPA auth:补回 grok-cli:access scope。
|
||
|
||
只覆盖写回,不删除原文件。默认仅处理缺 grok-cli:access 的 xai-*.json。
|
||
|
||
用法:
|
||
python remint_cli_scope.py --auth-dir /opt/cli-proxy-api/auths --dry-run
|
||
python remint_cli_scope.py --auth-dir /opt/cli-proxy-api/auths --limit 5
|
||
python remint_cli_scope.py --auth-dir /opt/cli-proxy-api/auths --concurrency 6
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import base64
|
||
import json
|
||
import os
|
||
import sys
|
||
import threading
|
||
import time
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
_ROOT = Path(__file__).resolve().parent
|
||
if str(_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(_ROOT))
|
||
|
||
|
||
def _jwt_payload(token: str) -> dict[str, Any]:
|
||
parts = (token or "").split(".")
|
||
if len(parts) < 2:
|
||
raise ValueError("invalid JWT")
|
||
seg = parts[1] + "=" * (-len(parts[1]) % 4)
|
||
return json.loads(base64.urlsafe_b64decode(seg.encode("ascii")))
|
||
|
||
|
||
def _scope_set(token: str) -> set[str]:
|
||
try:
|
||
return set(str(_jwt_payload(token).get("scope") or "").split())
|
||
except Exception:
|
||
return set()
|
||
|
||
|
||
def _needs_remint(data: dict[str, Any], *, only_missing_cli: bool) -> tuple[bool, str]:
|
||
at = str(data.get("access_token") or "").strip()
|
||
if not at:
|
||
return False, "no_access_token"
|
||
scopes = _scope_set(at)
|
||
has_cli = "grok-cli:access" in scopes
|
||
if only_missing_cli and has_cli:
|
||
return False, "already_has_cli_scope"
|
||
sso = str(data.get("sso") or "").strip()
|
||
if not sso:
|
||
return False, "no_sso"
|
||
return True, "missing_cli_scope" if not has_cli else "force"
|
||
|
||
|
||
def _atomic_write(path: Path, payload: dict[str, Any]) -> None:
|
||
tmp = path.with_suffix(path.suffix + f".tmp.{os.getpid()}.{threading.get_ident()}")
|
||
text = json.dumps(payload, ensure_ascii=False, indent=2) + "\n"
|
||
tmp.write_text(text, encoding="utf-8")
|
||
os.replace(tmp, path)
|
||
|
||
|
||
def remint_one(
|
||
path: Path,
|
||
*,
|
||
proxy: str | None,
|
||
dry_run: bool,
|
||
only_missing_cli: bool,
|
||
) -> dict[str, Any]:
|
||
import cpa
|
||
from oidc_mint.oauth_code import GROK_CLI_SCOPE, mint_from_sso, normalize_sso_cookie
|
||
|
||
try:
|
||
data = json.loads(path.read_text(encoding="utf-8", errors="replace"))
|
||
except Exception as exc: # noqa: BLE001
|
||
return {"ok": False, "path": str(path), "error": f"read: {exc}"}
|
||
|
||
need, reason = _needs_remint(data, only_missing_cli=only_missing_cli)
|
||
email = str(data.get("email") or "").strip()
|
||
if not need:
|
||
return {
|
||
"ok": True,
|
||
"skipped": True,
|
||
"path": str(path),
|
||
"email": email,
|
||
"reason": reason,
|
||
}
|
||
|
||
sso = normalize_sso_cookie(str(data.get("sso") or ""))
|
||
if dry_run:
|
||
return {
|
||
"ok": True,
|
||
"dry_run": True,
|
||
"path": str(path),
|
||
"email": email,
|
||
"reason": reason,
|
||
}
|
||
|
||
try:
|
||
tokens = mint_from_sso(
|
||
sso,
|
||
proxy=proxy,
|
||
require_referrer=True,
|
||
log=lambda m: None,
|
||
)
|
||
except Exception as exc: # noqa: BLE001
|
||
return {
|
||
"ok": False,
|
||
"path": str(path),
|
||
"email": email,
|
||
"error": f"mint: {exc}",
|
||
"reason": reason,
|
||
}
|
||
|
||
access = str(tokens.get("access_token") or "").strip()
|
||
refresh = str(tokens.get("refresh_token") or "").strip()
|
||
if GROK_CLI_SCOPE not in _scope_set(access):
|
||
return {
|
||
"ok": False,
|
||
"path": str(path),
|
||
"email": email,
|
||
"error": f"minted token still missing {GROK_CLI_SCOPE}",
|
||
"reason": reason,
|
||
}
|
||
|
||
try:
|
||
payload = cpa.build_cpa_xai_auth(
|
||
email=email or str(data.get("email") or ""),
|
||
access_token=access,
|
||
refresh_token=refresh,
|
||
id_token=tokens.get("id_token") or data.get("id_token"),
|
||
expires_in=tokens.get("expires_in"),
|
||
base_url=data.get("base_url") or cpa.CLI_BASE_URL,
|
||
sso=sso,
|
||
headers=data.get("headers") if isinstance(data.get("headers"), dict) else None,
|
||
)
|
||
# 保留原 disabled 状态,避免重铸把人工禁用号重新启用
|
||
if "disabled" in data:
|
||
payload["disabled"] = bool(data.get("disabled"))
|
||
_atomic_write(path, payload)
|
||
except Exception as exc: # noqa: BLE001
|
||
return {
|
||
"ok": False,
|
||
"path": str(path),
|
||
"email": email,
|
||
"error": f"write: {exc}",
|
||
"reason": reason,
|
||
}
|
||
|
||
return {
|
||
"ok": True,
|
||
"path": str(path),
|
||
"email": email or payload.get("email"),
|
||
"reason": reason,
|
||
"referrer": payload.get("referrer"),
|
||
"scope_ok": True,
|
||
}
|
||
|
||
|
||
def main(argv: list[str] | None = None) -> int:
|
||
p = argparse.ArgumentParser(description="Remint CPA auths missing grok-cli:access")
|
||
p.add_argument(
|
||
"--auth-dir",
|
||
default=os.environ.get("CPA_AUTH_DIR", str(_ROOT / "cpa_auths")),
|
||
help="CPA auth 目录(默认 ./cpa_auths 或 $CPA_AUTH_DIR)",
|
||
)
|
||
p.add_argument("--proxy", default=os.environ.get("https_proxy") or os.environ.get("http_proxy") or "")
|
||
p.add_argument("--concurrency", type=int, default=4)
|
||
p.add_argument("--limit", type=int, default=0, help="最多处理 N 个需重铸文件(0=不限)")
|
||
p.add_argument("--dry-run", action="store_true")
|
||
p.add_argument(
|
||
"--all",
|
||
action="store_true",
|
||
help="强制全部重铸(默认只处理缺 grok-cli:access 的)",
|
||
)
|
||
p.add_argument(
|
||
"--log-file",
|
||
default="",
|
||
help="结果日志路径(默认 auth-dir/remint_cli_scope.log)",
|
||
)
|
||
args = p.parse_args(argv)
|
||
|
||
auth_dir = Path(args.auth_dir).expanduser().resolve()
|
||
if not auth_dir.is_dir():
|
||
print(f"[!] auth-dir 不存在: {auth_dir}", flush=True)
|
||
return 2
|
||
|
||
proxy = (args.proxy or "").strip() or None
|
||
only_missing_cli = not args.all
|
||
files = sorted(auth_dir.glob("xai-*.json"))
|
||
targets: list[Path] = []
|
||
skipped = 0
|
||
for path in files:
|
||
try:
|
||
data = json.loads(path.read_text(encoding="utf-8", errors="replace"))
|
||
need, _ = _needs_remint(data, only_missing_cli=only_missing_cli)
|
||
if need:
|
||
targets.append(path)
|
||
else:
|
||
skipped += 1
|
||
except Exception:
|
||
skipped += 1
|
||
if args.limit and len(targets) >= args.limit:
|
||
break
|
||
|
||
log_path = Path(args.log_file).expanduser() if args.log_file else (auth_dir / "remint_cli_scope.log")
|
||
print(
|
||
f"[*] auth-dir={auth_dir} total={len(files)} targets={len(targets)} "
|
||
f"skipped_scan={skipped} concurrency={args.concurrency} "
|
||
f"proxy={proxy or '(direct)'} dry_run={args.dry_run}",
|
||
flush=True,
|
||
)
|
||
if not targets:
|
||
print("[*] 无需重铸", flush=True)
|
||
return 0
|
||
|
||
ok_n = 0
|
||
fail_n = 0
|
||
skip_n = 0
|
||
lock = threading.Lock()
|
||
started = time.time()
|
||
|
||
def _handle(path: Path) -> dict[str, Any]:
|
||
return remint_one(
|
||
path,
|
||
proxy=proxy,
|
||
dry_run=args.dry_run,
|
||
only_missing_cli=only_missing_cli,
|
||
)
|
||
|
||
def _record(res: dict[str, Any]) -> None:
|
||
nonlocal ok_n, fail_n, skip_n
|
||
line = json.dumps(res, ensure_ascii=False)
|
||
with lock:
|
||
with open(log_path, "a", encoding="utf-8") as f:
|
||
f.write(line + "\n")
|
||
if res.get("skipped") or res.get("dry_run"):
|
||
skip_n += 1
|
||
tag = "SKIP" if res.get("skipped") else "DRY"
|
||
elif res.get("ok"):
|
||
ok_n += 1
|
||
tag = "OK"
|
||
else:
|
||
fail_n += 1
|
||
tag = "FAIL"
|
||
email = res.get("email") or Path(str(res.get("path") or "")).name
|
||
extra = res.get("error") or res.get("reason") or ""
|
||
print(f"[{tag}] {email} {extra}", flush=True)
|
||
|
||
# 清空/追加日志头
|
||
with open(log_path, "a", encoding="utf-8") as f:
|
||
f.write(
|
||
json.dumps(
|
||
{
|
||
"event": "start",
|
||
"ts": int(time.time()),
|
||
"targets": len(targets),
|
||
"dry_run": bool(args.dry_run),
|
||
"proxy": proxy or "",
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
+ "\n"
|
||
)
|
||
|
||
if args.concurrency <= 1:
|
||
for path in targets:
|
||
_record(_handle(path))
|
||
else:
|
||
with ThreadPoolExecutor(max_workers=max(1, args.concurrency)) as ex:
|
||
futs = [ex.submit(_handle, path) for path in targets]
|
||
for fut in as_completed(futs):
|
||
try:
|
||
_record(fut.result())
|
||
except Exception as exc: # noqa: BLE001
|
||
_record({"ok": False, "error": f"worker: {exc}"})
|
||
|
||
elapsed = time.time() - started
|
||
print(
|
||
f"[*] done ok={ok_n} fail={fail_n} skip/dry={skip_n} "
|
||
f"elapsed={elapsed:.1f}s log={log_path}",
|
||
flush=True,
|
||
)
|
||
return 0 if fail_n == 0 else 1
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|