358 lines
13 KiB
Python
358 lines
13 KiB
Python
"""sub2api 网关接入：系统用管理 API 为每位用户建号发 key。
|
||
|
||
设计（2026-10-27 独立部署 sub2api 后）：
|
||
- 系统侧只用管理 API 管用户：管理员登录 -> 建 sub2api 账号 ->
|
||
以该账号登录 -> 发 key -> key 存我方库，任务注入该用户的 key。
|
||
- key 创建是用户侧端点 POST /api/v1/keys（不是 /api/v1/api-keys，
|
||
后者只有管理员改组）；管理员不能代用户建/删 key，所以 refresh_token
|
||
也要存下来：以后轮换/吊销 key 时靠它换用户 access token。
|
||
- key 明文存 PG：与 github_app_config 的 App 私钥同一信任边界
|
||
（库只在宿主 loopback 可达，服务端 env 同机）；at-rest 加密是
|
||
后续项，不在本批假装已做。
|
||
- 零第三方依赖，只用标准库；HTTP 可注入 request 可调用对象，
|
||
便于单测（见 cloud/tests/test_sub2api.py）。
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import secrets
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
|
||
GATEWAY_TABLE_SQL = (
|
||
"CREATE TABLE IF NOT EXISTS user_gateway_keys("
|
||
"user_id TEXT PRIMARY KEY, "
|
||
"sub2api_user_id BIGINT NOT NULL, "
|
||
"api_key TEXT NOT NULL, "
|
||
"refresh_token TEXT, "
|
||
"created_at DOUBLE PRECISION, "
|
||
"updated_at DOUBLE PRECISION)")
|
||
|
||
|
||
class Sub2apiError(RuntimeError):
|
||
def __init__(self, message, status=None):
|
||
super().__init__(message)
|
||
self.status = status
|
||
|
||
|
||
def _default_request(method, url, body=None, token=None):
|
||
"""urllib 实现的 request：返回 (status, parsed_json_dict)。"""
|
||
headers = {"Content-Type": "application/json"}
|
||
if token:
|
||
headers["Authorization"] = "Bearer " + token
|
||
req = urllib.request.Request(
|
||
url, method=method,
|
||
data=json.dumps(body).encode() if body is not None else None,
|
||
headers=headers)
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
raw = resp.read() or b"{}"
|
||
return resp.status, json.loads(raw)
|
||
except urllib.error.HTTPError as e:
|
||
try:
|
||
data = json.loads(e.read() or b"{}")
|
||
except ValueError:
|
||
data = {}
|
||
return e.code, data
|
||
|
||
|
||
def gateway_origin():
|
||
"""sub2api origin（不带 /v1）：未配则返回 None。"""
|
||
raw = os.environ.get("SUB2API_BASE_URL", "").strip()
|
||
return raw.rstrip("/") or None
|
||
|
||
|
||
def gateway_name():
|
||
raw = os.environ.get("CLOUD_GATEWAY", "xcmapi").strip().lower()
|
||
return raw if raw in ("xcmapi", "sub2api") else "xcmapi"
|
||
|
||
|
||
def gateway_group_name():
|
||
return os.environ.get("SUB2API_GROUP_NAME", "codex-system").strip() \
|
||
or "codex-system"
|
||
|
||
|
||
def opening_balance():
|
||
"""开通时授的初始余额（美元）：默认 10，显式值优先。
|
||
|
||
新用户余额为 0 时网关鉴权层直接 403（余额门在任何上游路由
|
||
之前），所以开通必须同时授余额，否则发出的 key 立刻不可用。
|
||
显式配 0 表示不授（调用方自理）；非法/负值回落默认 10。
|
||
"""
|
||
raw = os.environ.get(
|
||
"CLOUD_SUB2API_OPENING_BALANCE_USD", "").strip()
|
||
if not raw:
|
||
return 10.0
|
||
try:
|
||
value = float(raw)
|
||
except ValueError:
|
||
return 10.0
|
||
return value if value >= 0 else 10.0
|
||
|
||
|
||
def ensure_gateway_keys_table(conn):
|
||
conn.execute(GATEWAY_TABLE_SQL)
|
||
|
||
|
||
def get_user_gateway_key(conn, user_id):
|
||
"""只读：返回该用户的网关行 dict 或 None，绝不触发开通。"""
|
||
ensure_gateway_keys_table(conn)
|
||
cur = conn.execute(
|
||
"SELECT user_id, sub2api_user_id, api_key, refresh_token "
|
||
"FROM user_gateway_keys WHERE user_id = ?",
|
||
(user_id,))
|
||
row = cur.fetchone()
|
||
if row is None:
|
||
return None
|
||
try:
|
||
keys = ("user_id", "sub2api_user_id", "api_key", "refresh_token")
|
||
return dict(zip(keys, row))
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
|
||
class Sub2apiAdmin:
|
||
"""管理员会话：登录、建/删用户、查组。token 自动刷新一次。"""
|
||
|
||
def __init__(self, origin, email, password, request=None):
|
||
if not origin:
|
||
raise Sub2apiError("SUB2API_BASE_URL 未配置")
|
||
if not email or not password:
|
||
raise Sub2apiError("SUB2API 管理员邮箱/密码未配置")
|
||
self.origin = origin
|
||
self.email = email
|
||
self.password = password
|
||
self.request = request or _default_request
|
||
self._token = None
|
||
self._refresh = None
|
||
self._expires_at = 0
|
||
|
||
@classmethod
|
||
def from_env(cls, request=None):
|
||
return cls(gateway_origin(),
|
||
os.environ.get("SUB2API_ADMIN_EMAIL", ""),
|
||
os.environ.get("SUB2API_ADMIN_PASSWORD", ""),
|
||
request=request)
|
||
|
||
def _login(self):
|
||
st, data = self.request(
|
||
"POST", self.origin + "/api/v1/auth/login",
|
||
{"email": self.email, "password": self.password})
|
||
payload = (data or {}).get("data") or {}
|
||
if st != 200 or not payload.get("access_token"):
|
||
raise Sub2apiError(
|
||
"sub2api 管理员登录失败: %s"
|
||
% ((data or {}).get("message") or st), status=st)
|
||
self._token = payload["access_token"]
|
||
self._refresh = payload.get("refresh_token")
|
||
try:
|
||
ttl = int(payload.get("expires_in") or 3600)
|
||
except (TypeError, ValueError):
|
||
ttl = 3600
|
||
self._expires_at = time.time() + max(60, ttl - 60)
|
||
|
||
def _refresh_token(self):
|
||
if not self._refresh:
|
||
self._login()
|
||
return
|
||
st, data = self.request(
|
||
"POST", self.origin + "/api/v1/auth/refresh",
|
||
{"refresh_token": self._refresh})
|
||
payload = (data or {}).get("data") or {}
|
||
if st == 200 and payload.get("access_token"):
|
||
self._token = payload["access_token"]
|
||
if payload.get("refresh_token"):
|
||
self._refresh = payload["refresh_token"]
|
||
try:
|
||
ttl = int(payload.get("expires_in") or 3600)
|
||
except (TypeError, ValueError):
|
||
ttl = 3600
|
||
self._expires_at = time.time() + max(60, ttl - 60)
|
||
return
|
||
self._login()
|
||
|
||
def _call(self, method, path, body=None):
|
||
if not self._token or time.time() >= self._expires_at:
|
||
self._login()
|
||
st, data = self.request(method, self.origin + path, body,
|
||
token=self._token)
|
||
if st == 401:
|
||
self._refresh_token()
|
||
st, data = self.request(method, self.origin + path, body,
|
||
token=self._token)
|
||
return st, data
|
||
|
||
def create_user(self, email, password, username):
|
||
st, data = self._call("POST", "/api/v1/admin/users", {
|
||
"email": email, "password": password,
|
||
"username": username, "role": "user"})
|
||
payload = (data or {}).get("data") or {}
|
||
if st not in (200, 201) or payload.get("id") is None:
|
||
raise Sub2apiError(
|
||
"sub2api 建用户失败: %s"
|
||
% ((data or {}).get("message") or st), status=st)
|
||
return payload["id"]
|
||
|
||
def delete_user(self, sub_user_id):
|
||
st, _data = self._call(
|
||
"DELETE", "/api/v1/admin/users/%s" % sub_user_id)
|
||
return st in (200, 204)
|
||
|
||
def set_balance(self, sub_user_id, amount):
|
||
st, data = self._call(
|
||
"POST", "/api/v1/admin/users/%s/balance" % sub_user_id,
|
||
{"balance": amount, "operation": "set",
|
||
"notes": "codex provisioning opening balance"})
|
||
if st != 200:
|
||
raise Sub2apiError(
|
||
"sub2api 授余额失败: %s"
|
||
% ((data or {}).get("message") or st), status=st)
|
||
|
||
def list_groups(self):
|
||
st, data = self._call("GET", "/api/v1/admin/groups")
|
||
payload = (data or {}).get("data")
|
||
if st != 200:
|
||
raise Sub2apiError("sub2api 查组失败: %s" % st, status=st)
|
||
if isinstance(payload, list):
|
||
return payload
|
||
if isinstance(payload, dict):
|
||
for key in ("items", "list", "groups"):
|
||
if isinstance(payload.get(key), list):
|
||
return payload[key]
|
||
return []
|
||
|
||
def resolve_group_id(self):
|
||
raw = os.environ.get("SUB2API_GROUP_ID", "").strip()
|
||
if raw:
|
||
try:
|
||
return int(raw)
|
||
except ValueError:
|
||
raise Sub2apiError("SUB2API_GROUP_ID 非法: %r" % raw)
|
||
want = gateway_group_name()
|
||
for grp in self.list_groups():
|
||
if isinstance(grp, dict) and grp.get("name") == want:
|
||
if grp.get("id") is not None:
|
||
return grp["id"]
|
||
raise Sub2apiError("sub2api 组不存在: %s" % want)
|
||
|
||
|
||
def user_login(origin, email, password, request=None):
|
||
"""以 sub2api 用户身份登录：返回 (access_token, refresh_token)。"""
|
||
req = request or _default_request
|
||
st, data = req("POST", origin + "/api/v1/auth/login",
|
||
{"email": email, "password": password})
|
||
payload = (data or {}).get("data") or {}
|
||
if st != 200 or not payload.get("access_token"):
|
||
raise Sub2apiError(
|
||
"sub2api 用户登录失败: %s"
|
||
% ((data or {}).get("message") or st), status=st)
|
||
return payload["access_token"], payload.get("refresh_token")
|
||
|
||
|
||
def user_create_key(origin, user_token, name, group_id, request=None):
|
||
"""以用户身份发 key：返回 (key_id, key_value)。"""
|
||
req = request or _default_request
|
||
body = {"name": name}
|
||
if group_id is not None:
|
||
body["group_id"] = group_id
|
||
st, data = req("POST", origin + "/api/v1/keys", body, token=user_token)
|
||
payload = (data or {}).get("data") or {}
|
||
value = payload.get("key") or payload.get("api_key") \
|
||
or payload.get("value")
|
||
if st not in (200, 201) or not value:
|
||
raise Sub2apiError(
|
||
"sub2api 发 key 失败: %s"
|
||
% ((data or {}).get("message") or st), status=st)
|
||
return payload.get("id"), value
|
||
|
||
|
||
def user_delete_key(origin, user_token, key_id, request=None):
|
||
req = request or _default_request
|
||
st, _data = req("DELETE", origin + "/api/v1/keys/%s" % key_id,
|
||
token=user_token)
|
||
return st in (200, 204)
|
||
|
||
|
||
def user_refresh_access(origin, refresh_token, request=None):
|
||
"""refresh_token 换 access_token：返回 (新 access, 新 refresh|None)。"""
|
||
req = request or _default_request
|
||
st, data = req("POST", origin + "/api/v1/auth/refresh",
|
||
{"refresh_token": refresh_token})
|
||
payload = (data or {}).get("data") or {}
|
||
if st != 200 or not payload.get("access_token"):
|
||
raise Sub2apiError("sub2api refresh 失败: %s" % st, status=st)
|
||
return payload["access_token"], payload.get("refresh_token")
|
||
|
||
|
||
def list_gateway_models(origin, api_key, request=None):
|
||
"""网关 /v1/models：返回模型 id 列表（key 无模型时为空列表）。"""
|
||
req = request or _default_request
|
||
st, data = req("GET", origin + "/v1/models", token=api_key)
|
||
if st != 200:
|
||
raise Sub2apiError("sub2api 查模型失败: %s" % st, status=st)
|
||
payload = (data or {}).get("data", data)
|
||
if isinstance(payload, dict):
|
||
payload = payload.get("data") or []
|
||
out = []
|
||
if isinstance(payload, list):
|
||
for item in payload:
|
||
if isinstance(item, dict) and item.get("id"):
|
||
out.append(item["id"])
|
||
elif isinstance(item, str):
|
||
out.append(item)
|
||
return out
|
||
|
||
|
||
def ensure_user_gateway_key(conn, user_id, admin=None, request=None):
|
||
"""幂等开通：已有行直接返回 key；否则建号->登录->发 key->授余额->入库。
|
||
|
||
入库失败时尽力清理刚建的 key 与账号（避免孤儿 key），再抛原错；
|
||
授余额失败同样走这条清理（余额为 0 的 key 在网关侧直接 403，
|
||
留着没有意义）。调用方负责 conn 的提交语义：本函数在入库后
|
||
commit（开通是独立事务，不跟调用方的业务事务绑在一起）。
|
||
"""
|
||
ensure_gateway_keys_table(conn)
|
||
row = get_user_gateway_key(conn, user_id)
|
||
if row is not None:
|
||
return row["api_key"]
|
||
origin = gateway_origin()
|
||
if not origin:
|
||
raise Sub2apiError("SUB2API_BASE_URL 未配置，无法开通网关账号")
|
||
admin = admin or Sub2apiAdmin.from_env(request=request)
|
||
req = request or _default_request
|
||
email = "codex-%s@codex.local" % user_id
|
||
username = "codex-" + str(user_id)[:12]
|
||
password = secrets.token_hex(16)
|
||
sub_id = admin.create_user(email, password, username)
|
||
key_id = None
|
||
try:
|
||
user_token, refresh = user_login(origin, email, password,
|
||
request=req)
|
||
group_id = admin.resolve_group_id()
|
||
key_id, key = user_create_key(origin, user_token,
|
||
"codex-task-key", group_id,
|
||
request=req)
|
||
amount = opening_balance()
|
||
if amount > 0:
|
||
admin.set_balance(sub_id, amount)
|
||
now = time.time()
|
||
conn.execute(
|
||
"INSERT INTO user_gateway_keys(user_id, sub2api_user_id, "
|
||
"api_key, refresh_token, created_at, updated_at) "
|
||
"VALUES(?, ?, ?, ?, ?, ?)",
|
||
(user_id, sub_id, key, refresh, now, now))
|
||
conn.commit()
|
||
except Exception:
|
||
if key_id is not None:
|
||
try:
|
||
user_delete_key(origin, user_token, key_id, request=req)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
admin.delete_user(sub_id)
|
||
except Exception:
|
||
pass
|
||
raise
|
||
return key
|