475 lines
24 KiB
Python
475 lines
24 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""newapi_pool_provision.py -- 内测 new-api per-user ¥100 额度池离线预置脚本（WU2 · 接线点 B）
|
||
|
||
权威设计：docs/agent-specs/2026-07-07-内测-WU2-newapi额度接入-设计.md §3.3（离线 ops 预置脚本）
|
||
+ §3.5（池表 newapi_quota_pool 契约）。
|
||
|
||
它做什么
|
||
--------
|
||
在 mini-infra 上一次性预建 N 个专属 (new-api user + token + ¥100 额度) 账户，
|
||
把每个条目 emit 成可导入 game-cloud 池表 newapi_quota_pool 的 JSON + SQL，供 WU2 落地后灌数。
|
||
provision 全部离线做（S0 实测坐实：new-api admin 令牌不能替他人建 token、access_token
|
||
不可经 API 设，per-user token 必须直连其 postgres），运行时零 new-api admin 耦合。
|
||
|
||
每个池条目四步建成（S0 坐实的唯一可行路径）
|
||
------------------------------------------------
|
||
1. POST /api/user/（root 令牌）建用户，username = 确定性 `<前缀><序号>`；从 postgres 查 uid。
|
||
2. postgres 直写 UPDATE users SET access_token=<确定性 32 位串> WHERE id=uid
|
||
—— access_token 无法经 API 设，又是下一步「以该用户身份建 token」的前提。
|
||
3. POST /api/token/（以该用户 access_token + New-Api-User:uid 建 token），
|
||
remain_quota=¥100 折算、unlimited_quota=false、expired_time=-1（不过期）；从 postgres 查 token_id/key。
|
||
4. PUT /api/user/（root 令牌）设 user.quota=¥100 折算（账户自洽；权威余额闸仍是 token.remain_quota）。
|
||
|
||
幂等（可重跑补池）
|
||
------------------
|
||
按确定性 username 去重：重跑时若用户已存在则复用其 uid，逐步「补齐」缺失的 access_token /
|
||
token / quota（ensure 语义），不重复建号。补池 = 提高 --count 重跑。单条失败只记该条、不污染后续。
|
||
|
||
外部交互红线
|
||
------------
|
||
- HTTP：连接/读超时（urlopen timeout=读10s，内部连粒度由 socket 兜），非 2xx / success=false 归错，
|
||
5xx/超时重试 2 次（指数退避），4xx 不重试；全程绕系统代理（ProxyHandler({})，内网直连 100.64.x）。
|
||
- postgres：经 `docker exec infra-postgres psql` 执行（S0 验证路径），失败即抛、该条标记失败。
|
||
- 日志：token / access_token 一律脱敏（前后各留 4 位）；可追溯（时间戳 + 步骤 + 条目序号）。
|
||
|
||
凭据（不硬编码）
|
||
----------------
|
||
- postgres 密码：环境变量 NEWAPI_PG_PASSWORD（必填；见 docs/内网凭据与端点.md）。
|
||
- root 管理令牌：环境变量 NEWAPI_ROOT_TOKEN（选填）；未设则自动从 postgres users.id=1 读。
|
||
|
||
用法（在 mini-infra 上跑，docker 本地可用）
|
||
------------------------------------------
|
||
export NEWAPI_PG_PASSWORD=<见凭据文档>
|
||
# 预置 2 条测试条目（验证用，前缀明确可辨），dry-run 先看：
|
||
python3 newapi_pool_provision.py --count 2 --prefix neice_s0test_ --dry-run
|
||
python3 newapi_pool_provision.py --count 2 --prefix neice_s0test_
|
||
# 清理测试条目（删 tokens + users 行）：
|
||
python3 newapi_pool_provision.py --cleanup 'neice_s0test_%'
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
import os
|
||
import re
|
||
import secrets
|
||
import socket
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
# ─── 常量 ────────────────────────────────────────────────────────────────
|
||
NEWAPI_BASE = os.environ.get("NEWAPI_BASE", "http://localhost:3000") # 脚本宿主 mini-infra，本地直连
|
||
PG_CONTAINER = os.environ.get("NEWAPI_PG_CONTAINER", "infra-postgres")
|
||
PG_DB = os.environ.get("NEWAPI_PG_DB", "new-api")
|
||
PG_USER = os.environ.get("NEWAPI_PG_USER", "root")
|
||
ROOT_UID = 1 # new-api root 用户 id（qingse），持系统管理 access_token
|
||
HTTP_CONNECT_TIMEOUT = 5.0
|
||
HTTP_READ_TIMEOUT = 10.0
|
||
HTTP_RETRIES = 2 # 5xx / 超时最多重试 2 次
|
||
USERNAME_RE = re.compile(r"^[A-Za-z0-9_]+$") # 只允许安全 username，杜绝 SQL/命令注入
|
||
POOL_TABLE = "newapi_quota_pool"
|
||
|
||
log = logging.getLogger("newapi_pool")
|
||
|
||
|
||
def _mask(token: str | None) -> str:
|
||
"""脱敏敏感串：前后各留 4 位，中间打码。用于日志（红线：token/access_token 不整条打日志）。"""
|
||
if not token:
|
||
return "<empty>"
|
||
if len(token) <= 10:
|
||
return token[:2] + "***"
|
||
return f"{token[:4]}…{token[-4:]}(len={len(token)})"
|
||
|
||
|
||
# ─── postgres 直连（经 docker exec psql，S0 验证路径） ──────────────────────
|
||
class Postgres:
|
||
"""封装对 new-api postgres 的只读/写操作。所有写走事务语义（单条 SQL 天然原子）。"""
|
||
|
||
def __init__(self, password: str):
|
||
if not password:
|
||
raise RuntimeError(
|
||
"缺少 NEWAPI_PG_PASSWORD 环境变量——无法连接 new-api postgres。"
|
||
"见 docs/内网凭据与端点.md new-api 段。"
|
||
)
|
||
self._password = password
|
||
|
||
def _run(self, sql: str, operation: str = "query") -> str:
|
||
"""执行一条 SQL，返回 tuples-only、unaligned、以 | 分隔的原始 stdout。失败即抛。"""
|
||
cmd = [
|
||
# SQL 只经 stdin 传输；docker exec 必须带 -i，否则会静默执行空输入并返回成功。
|
||
"docker", "exec", "-i",
|
||
"-e", f"PGPASSWORD={self._password}",
|
||
PG_CONTAINER,
|
||
"psql", "-U", PG_USER, "-d", PG_DB,
|
||
"-t", "-A", "-F", "|", "-v", "ON_ERROR_STOP=1",
|
||
]
|
||
try:
|
||
# SQL 经 stdin 传给 psql，避免 access_token 出现在进程参数或错误日志。
|
||
proc = subprocess.run(cmd, input=sql, capture_output=True, text=True, timeout=30)
|
||
except subprocess.TimeoutExpired as e:
|
||
raise RuntimeError(f"postgres 执行超时 operation={operation}") from e
|
||
if proc.returncode != 0:
|
||
# psql 可能在 stderr 回显失败 SQL，因此这里只记录操作名与返回码，不复制原文。
|
||
raise RuntimeError(f"postgres 执行失败 operation={operation} rc={proc.returncode}")
|
||
return proc.stdout.strip()
|
||
|
||
def query_root_token(self) -> str:
|
||
out = self._run(f"SELECT access_token FROM users WHERE id={ROOT_UID};")
|
||
if not out:
|
||
raise RuntimeError("postgres 未查到 root(id=1) access_token")
|
||
return out.splitlines()[0].strip()
|
||
|
||
def query_user_id(self, username: str) -> int | None:
|
||
out = self._run(f"SELECT id FROM users WHERE username='{username}';")
|
||
line = out.splitlines()[0].strip() if out else ""
|
||
return int(line) if line else None
|
||
|
||
def query_access_token(self, uid: int) -> str:
|
||
out = self._run(f"SELECT access_token FROM users WHERE id={uid};")
|
||
return out.splitlines()[0].strip() if out else ""
|
||
|
||
def query_user_quota(self, uid: int) -> int | None:
|
||
out = self._run(f"SELECT quota FROM users WHERE id={uid};")
|
||
line = out.splitlines()[0].strip() if out else ""
|
||
return int(line) if line else None
|
||
|
||
def set_access_token(self, uid: int, access_token: str) -> None:
|
||
# access_token 无法经 API 设，必须 DB 直写（S0 坐实的关键步）
|
||
self._run(f"UPDATE users SET access_token='{access_token}' WHERE id={uid};", "set-access-token")
|
||
|
||
def query_token(self, uid: int, name: str) -> tuple[int, str] | None:
|
||
"""按 user_id + token name 查 token，返回 (token_id, token_key)。"""
|
||
out = self._run(
|
||
f"SELECT id, key FROM tokens WHERE user_id={uid} AND name='{name}' ORDER BY id LIMIT 1;"
|
||
)
|
||
line = out.splitlines()[0].strip() if out else ""
|
||
if not line or "|" not in line:
|
||
return None
|
||
tid, key = line.split("|", 1)
|
||
return int(tid), key.strip()
|
||
|
||
def delete_user_cascade(self, username_like: str) -> tuple[int, int]:
|
||
"""清理：按 username LIKE 删 users + 其 tokens。返回 (删 tokens 数, 删 users 数)。
|
||
安全护栏：拒绝任何可能命中 root(id<=1) 的模式；只允许显式 % 通配。"""
|
||
# 先查将被删的 uid 列表，护栏校验
|
||
uids_out = self._run(f"SELECT id FROM users WHERE username LIKE '{username_like}';")
|
||
uids = [int(x) for x in uids_out.splitlines() if x.strip()]
|
||
if not uids:
|
||
return (0, 0)
|
||
if any(u <= ROOT_UID for u in uids):
|
||
raise RuntimeError(f"清理护栏：模式 {username_like} 命中受保护用户 id<={ROOT_UID}，拒绝执行")
|
||
uid_csv = ",".join(str(u) for u in uids)
|
||
tok_out = self._run(f"WITH d AS (DELETE FROM tokens WHERE user_id IN ({uid_csv}) RETURNING 1) SELECT count(*) FROM d;")
|
||
usr_out = self._run(f"WITH d AS (DELETE FROM users WHERE id IN ({uid_csv}) RETURNING 1) SELECT count(*) FROM d;")
|
||
return (int(tok_out.strip() or 0), int(usr_out.strip() or 0))
|
||
|
||
|
||
# ─── new-api HTTP（stdlib urllib，绕系统代理，带超时+重试） ──────────────────
|
||
_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({})) # 空 ProxyHandler = 禁用系统代理
|
||
|
||
|
||
def _http_json(method: str, path: str, token: str, newapi_user: int, body: dict | None = None) -> dict:
|
||
"""对 new-api 发一个 JSON 请求。红线：绕代理 / 超时 / 非 2xx 归错 / 5xx 重试。
|
||
|
||
path 必须带尾斜杠（new-api gin RedirectTrailingSlash，否则 307）。
|
||
鉴权头：Authorization: Bearer <token> + New-Api-User: <uid>（缺后者即 401）。
|
||
"""
|
||
url = NEWAPI_BASE.rstrip("/") + path
|
||
data = json.dumps(body).encode("utf-8") if body is not None else None
|
||
headers = {
|
||
"Authorization": f"Bearer {token}",
|
||
"New-Api-User": str(newapi_user),
|
||
"Content-Type": "application/json",
|
||
}
|
||
last_err: Exception | None = None
|
||
for attempt in range(HTTP_RETRIES + 1):
|
||
req = urllib.request.Request(url, data=data, headers=headers, method=method)
|
||
try:
|
||
with _OPENER.open(req, timeout=HTTP_READ_TIMEOUT) as resp:
|
||
raw = resp.read().decode("utf-8")
|
||
payload = json.loads(raw) if raw else {}
|
||
# new-api 标准返回体 {success,message,data}；success=false 视为业务错
|
||
if isinstance(payload, dict) and payload.get("success") is False:
|
||
raise RuntimeError(f"{method} {path} 业务失败: {payload.get('message')}")
|
||
return payload
|
||
except urllib.error.HTTPError as e:
|
||
body_txt = e.read().decode("utf-8", "ignore")[:200]
|
||
last_err = RuntimeError(f"{method} {path} HTTP {e.code}: {body_txt}")
|
||
if 500 <= e.code < 600 and attempt < HTTP_RETRIES:
|
||
time.sleep(0.5 * (2 ** attempt))
|
||
log.warning("[http] %s %s 5xx 重试 %d/%d", method, path, attempt + 1, HTTP_RETRIES)
|
||
continue
|
||
raise last_err # 4xx 不重试
|
||
except (urllib.error.URLError, TimeoutError, OSError) as e:
|
||
last_err = RuntimeError(f"{method} {path} 网络错误: {e}")
|
||
if attempt < HTTP_RETRIES:
|
||
time.sleep(0.5 * (2 ** attempt))
|
||
log.warning("[http] %s %s 网络错重试 %d/%d: %s", method, path, attempt + 1, HTTP_RETRIES, e)
|
||
continue
|
||
raise last_err
|
||
raise last_err or RuntimeError("unreachable")
|
||
|
||
|
||
# ─── ¥ ↔ quota 折算（口径与 cost.py 同源，读网关而非硬编码） ─────────────────
|
||
def fetch_conversion() -> tuple[int, float]:
|
||
"""GET /api/status 拿 quota_per_unit / usd_exchange_rate（无需鉴权）。缺 usd 用 7.3 默认。"""
|
||
url = NEWAPI_BASE.rstrip("/") + "/api/status"
|
||
req = urllib.request.Request(url, method="GET")
|
||
with _OPENER.open(req, timeout=HTTP_READ_TIMEOUT) as resp:
|
||
data = json.loads(resp.read().decode("utf-8")).get("data", {})
|
||
qpu = int(data.get("quota_per_unit") or 500000)
|
||
usd = data.get("usd_exchange_rate")
|
||
usd = float(usd) if usd else 7.3 # DB options 未覆盖时默认 7.3
|
||
return qpu, usd
|
||
|
||
|
||
def compute_grant(yuan: float, qpu: int, usd: float) -> int:
|
||
"""¥ → quota：round(yuan / usd × qpu)。¥100 @ 500000/7.3 ≈ 6,849,315。"""
|
||
return round(yuan / usd * qpu)
|
||
|
||
|
||
def wait_for_db_visibility(query, attempts: int = 20, interval_seconds: float = 0.1):
|
||
"""等待 new-api API 写入在 PostgreSQL 可见；超出固定窗口后返回 None。"""
|
||
for attempt in range(attempts):
|
||
value = query()
|
||
if value is not None:
|
||
return value
|
||
if attempt + 1 < attempts:
|
||
time.sleep(interval_seconds)
|
||
return None
|
||
|
||
|
||
def validate_provision_request(start: int, count: int, prefix: str, hostname: str) -> list[str]:
|
||
"""校验离线补池边界，并返回确定性的目标用户名。
|
||
|
||
写 new-api 的工具只能在 mini-infra 执行；范围必须为正，用户名必须符合 new-api
|
||
安全字符与 20 字符上限，防止空跑被误记为成功或把不安全文本带入 SQL。
|
||
"""
|
||
short_hostname = hostname.split(".", 1)[0].lower()
|
||
# mini-infra 是 SSH/业务别名，minione-ubuntu-infra 是机器实际 hostname；两者指向同一受控宿主。
|
||
if short_hostname not in {"mini-infra", "minione-ubuntu-infra"}:
|
||
raise ValueError(f"额度池预置只允许在 mini-infra 离线执行，actual={short_hostname}")
|
||
if start < 1:
|
||
raise ValueError("--start 必须大于等于 1")
|
||
if count < 1:
|
||
raise ValueError("--count 必须大于等于 1")
|
||
if count > 1000:
|
||
raise ValueError("--count 单次不得超过 1000")
|
||
if not prefix or not USERNAME_RE.fullmatch(prefix):
|
||
raise ValueError("--prefix 只允许字母、数字和下划线")
|
||
usernames = [f"{prefix}{index:03d}" for index in range(start, start + count)]
|
||
if any(len(username) > 20 or not USERNAME_RE.fullmatch(username) for username in usernames):
|
||
raise ValueError("生成的 username 不符合安全字符或 20 字符上限")
|
||
return usernames
|
||
|
||
|
||
# ─── 单条目预置（ensure 语义，幂等可补齐） ──────────────────────────────────
|
||
def provision_one(pg: Postgres, root_token: str, username: str, grant: int,
|
||
qpu: int, usd: float, dry_run: bool) -> dict:
|
||
"""预置/补齐一个池条目，返回条目 dict（对齐 §3.5 池表契约）。任一步失败即抛，由调用方记为失败。"""
|
||
if not USERNAME_RE.match(username):
|
||
raise RuntimeError(f"非法 username: {username}")
|
||
|
||
if dry_run:
|
||
log.info("[dry-run] 计划预置 %s（grant=%d, unlimited=false, expired=-1）", username, grant)
|
||
return {"username": username, "dry_run": True, "grant_quota": grant}
|
||
|
||
# 步骤 1：建用户（幂等——已存在则复用）
|
||
uid = pg.query_user_id(username)
|
||
if uid is None:
|
||
password = secrets.token_urlsafe(12) # 随机密码 16 位（new-api Password max=20）；仅占位，从不密码登录
|
||
_http_json("POST", "/api/user/", root_token, ROOT_UID, {
|
||
"username": username,
|
||
"password": password,
|
||
"display_name": username[:20], # new-api DisplayName max=20，直接用 username（截断兜底）
|
||
})
|
||
uid = wait_for_db_visibility(lambda: pg.query_user_id(username))
|
||
if uid is None:
|
||
raise RuntimeError(f"建用户 {username} 后 postgres 未查到 uid")
|
||
log.info("[1/4] 建用户 %s → uid=%d", username, uid)
|
||
else:
|
||
log.info("[1/4] 用户 %s 已存在 uid=%d（复用）", username, uid)
|
||
|
||
# 步骤 2：DB 直写 access_token（确定性，可重复 UPDATE 幂等）
|
||
access_token = hashlib.sha256(f"neice-quota-pool::v1::{username}".encode()).hexdigest()[:32]
|
||
existing_at = pg.query_access_token(uid)
|
||
if existing_at != access_token:
|
||
pg.set_access_token(uid, access_token)
|
||
log.info("[2/4] DB 写 access_token uid=%d → %s", uid, _mask(access_token))
|
||
else:
|
||
log.info("[2/4] access_token uid=%d 已就绪 %s（复用）", uid, _mask(access_token))
|
||
|
||
# 步骤 3：以该用户身份建 token（幂等——已存在同名 token 则复用）
|
||
tok = pg.query_token(uid, username)
|
||
if tok is None:
|
||
_http_json("POST", "/api/token/", access_token, uid, {
|
||
"name": username,
|
||
"remain_quota": grant,
|
||
"unlimited_quota": False,
|
||
"expired_time": -1, # 不过期，避免自伤式过期（设计 §3.9）
|
||
})
|
||
tok = wait_for_db_visibility(lambda: pg.query_token(uid, username))
|
||
if tok is None:
|
||
raise RuntimeError(f"建 token 后 postgres 未查到 user={uid} name={username}")
|
||
log.info("[3/4] 建 token uid=%d → token_id=%d key=%s", uid, tok[0], _mask(tok[1]))
|
||
else:
|
||
log.info("[3/4] token uid=%d name=%s 已存在 token_id=%d（复用）", uid, username, tok[0])
|
||
token_id, token_key = tok
|
||
|
||
# 步骤 4：设 user.quota（root 令牌；账户自洽，权威余额闸仍是 token.remain_quota）
|
||
cur_quota = pg.query_user_quota(uid)
|
||
if cur_quota != grant:
|
||
_http_json("PUT", "/api/user/", root_token, ROOT_UID, {
|
||
"id": uid,
|
||
"username": username,
|
||
"display_name": username[:20],
|
||
"quota": grant,
|
||
"group": "default",
|
||
})
|
||
log.info("[4/4] 设 user.quota uid=%d → %d", uid, grant)
|
||
else:
|
||
log.info("[4/4] user.quota uid=%d 已=%d（复用）", uid, grant)
|
||
|
||
# 组装池表条目（§3.5 契约字段）
|
||
return {
|
||
"newapi_user_id": uid,
|
||
"newapi_token_id": token_id,
|
||
"newapi_token_key": token_key,
|
||
"grant_quota": grant,
|
||
"quota_per_unit_snapshot": qpu,
|
||
"usd_rate_snapshot": usd,
|
||
"status": "FREE",
|
||
"username": username, # 仅审计参考，非池表列
|
||
}
|
||
|
||
|
||
# ─── emit：JSON + SQL（对齐 §3.5，供 WU2 落地后导入 game-cloud） ─────────────
|
||
def emit_outputs(entries: list[dict], out_dir: Path, prefix: str) -> tuple[Path, Path]:
|
||
ts = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
stem = f"newapi_pool_{prefix.rstrip('_')}_{ts}"
|
||
json_path = out_dir / f"{stem}.json"
|
||
sql_path = out_dir / f"{stem}.sql"
|
||
|
||
# 只保留池表列（去掉 username 审计字段）
|
||
pool_cols = ["newapi_user_id", "newapi_token_id", "newapi_token_key",
|
||
"grant_quota", "quota_per_unit_snapshot", "usd_rate_snapshot", "status"]
|
||
json_rows = [{k: e[k] for k in pool_cols} for e in entries]
|
||
json_path.write_text(json.dumps(json_rows, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
||
# SQL：INSERT ... ON DUPLICATE KEY UPDATE（去重键 uk(newapi_user_id)，§3.3(a)）
|
||
# 刻意不在 ON DUPLICATE 里改 status/claimed_*——重跑不把已 CLAIMED 条目刷回 FREE（防免费续杯 / 防误 unclaim）。
|
||
lines = [
|
||
f"-- 内测 new-api 额度池预置产物 {ts}，共 {len(json_rows)} 条。",
|
||
f"-- 对齐 docs/agent-specs/2026-07-07-内测-WU2-newapi额度接入-设计.md §3.5 池表 {POOL_TABLE}（V32 待落）。",
|
||
f"-- 导入前提：game-cloud 已建 {POOL_TABLE} 表（含 uk(newapi_user_id)、审计列默认值）。",
|
||
f"-- ON DUPLICATE 只更 token/额度快照，不动 status/claimed_*（幂等补池不 unclaim 已绑条目）。",
|
||
]
|
||
for entry, r in zip(entries, json_rows):
|
||
# username 不含密钥，用于让受控导入器证明本批确为 neice_051..054。
|
||
lines.append(
|
||
f"-- pool_entry username={entry['username']} newapi_user_id={r['newapi_user_id']}"
|
||
)
|
||
key_sql = r["newapi_token_key"].replace("'", "''")
|
||
lines.append(
|
||
f"INSERT INTO {POOL_TABLE} "
|
||
f"(newapi_user_id, newapi_token_id, newapi_token_key, grant_quota, "
|
||
f"quota_per_unit_snapshot, usd_rate_snapshot, status) VALUES "
|
||
f"({r['newapi_user_id']}, {r['newapi_token_id']}, '{key_sql}', {r['grant_quota']}, "
|
||
f"{r['quota_per_unit_snapshot']}, {r['usd_rate_snapshot']}, '{r['status']}') "
|
||
f"ON DUPLICATE KEY UPDATE "
|
||
f"newapi_token_id=VALUES(newapi_token_id), newapi_token_key=VALUES(newapi_token_key), "
|
||
f"grant_quota=VALUES(grant_quota), quota_per_unit_snapshot=VALUES(quota_per_unit_snapshot), "
|
||
f"usd_rate_snapshot=VALUES(usd_rate_snapshot);"
|
||
)
|
||
sql_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||
# JSON/SQL 都含 token transport material，显式收紧权限，不依赖调用者 umask。
|
||
json_path.chmod(0o600)
|
||
sql_path.chmod(0o600)
|
||
return json_path, sql_path
|
||
|
||
|
||
# ─── 主流程 ──────────────────────────────────────────────────────────────
|
||
def cmd_provision(args) -> int:
|
||
pg = Postgres(os.environ.get("NEWAPI_PG_PASSWORD", ""))
|
||
root_token = os.environ.get("NEWAPI_ROOT_TOKEN") or pg.query_root_token()
|
||
log.info("root 管理令牌就绪 %s", _mask(root_token))
|
||
|
||
qpu, usd = fetch_conversion()
|
||
grant = compute_grant(args.yuan, qpu, usd)
|
||
log.info("折算：¥%.0f @ quota_per_unit=%d usd=%.4f → grant_quota=%d", args.yuan, qpu, usd, grant)
|
||
|
||
entries: list[dict] = []
|
||
failures: list[dict] = []
|
||
for username in args.usernames:
|
||
try:
|
||
entries.append(provision_one(pg, root_token, username, grant, qpu, usd, args.dry_run))
|
||
except Exception as e: # 单条失败不污染后续（红线）
|
||
log.error("[条目 %s] 预置失败：%s", username, e)
|
||
failures.append({"username": username, "error": str(e)})
|
||
|
||
log.info("完成：成功 %d / 失败 %d", len(entries), len(failures))
|
||
if args.dry_run:
|
||
log.info("dry-run 结束，未写 new-api、未 emit 文件。")
|
||
return 0 if not failures else 1
|
||
|
||
if entries:
|
||
json_path, sql_path = emit_outputs(entries, Path(args.out_dir), args.prefix)
|
||
log.info("emit JSON → %s", json_path)
|
||
log.info("emit SQL → %s", sql_path)
|
||
if failures:
|
||
log.warning("失败条目：%s", json.dumps(failures, ensure_ascii=False))
|
||
return 1
|
||
return 0
|
||
|
||
|
||
def cmd_cleanup(args) -> int:
|
||
"""清理：按 username LIKE 模式删 tokens + users（护栏拒绝命中 root）。"""
|
||
pg = Postgres(os.environ.get("NEWAPI_PG_PASSWORD", ""))
|
||
ntok, nusr = pg.delete_user_cascade(args.cleanup)
|
||
log.info("清理模式 %s：删 tokens %d 行 / users %d 行", args.cleanup, ntok, nusr)
|
||
return 0
|
||
|
||
|
||
def main() -> int:
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s %(levelname)s %(message)s",
|
||
stream=sys.stderr,
|
||
)
|
||
p = argparse.ArgumentParser(description="内测 new-api per-user ¥100 额度池离线预置")
|
||
p.add_argument("--count", type=int, default=1, help="目标条目数 N（本次预置到 start..start+count）")
|
||
p.add_argument("--start", type=int, default=1, help="起始序号（默认 1）")
|
||
p.add_argument("--prefix", default="neice_", help="username 前缀（默认 neice_；测试用 neice_s0test_）")
|
||
p.add_argument("--yuan", type=float, default=100.0, help="每条目额度（元，默认 100）")
|
||
p.add_argument("--out-dir", default="./newapi-pool-out", help="JSON/SQL 产物目录")
|
||
p.add_argument("--dry-run", action="store_true", help="只打印计划，不写 new-api、不 emit")
|
||
p.add_argument("--cleanup", metavar="LIKE_PATTERN",
|
||
help="清理模式：删匹配 username LIKE 的 users+tokens（如 'neice_s0test_%%'）")
|
||
args = p.parse_args()
|
||
|
||
hostname = socket.gethostname()
|
||
try:
|
||
if args.cleanup:
|
||
# cleanup 同样会写 new-api，不能从其它宿主绕过离线边界。
|
||
validate_provision_request(1, 1, "host_check_", hostname)
|
||
else:
|
||
args.usernames = validate_provision_request(args.start, args.count, args.prefix, hostname)
|
||
except ValueError as error:
|
||
p.error(str(error))
|
||
if args.cleanup:
|
||
return cmd_cleanup(args)
|
||
return cmd_provision(args)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|