Initial commit: grok-free-register-oss
Open-source Grok free registration CLI, xai_enroller auth pipeline, local auth service, tests and docs.
This commit is contained in:
chaos
124
.env.example
Normal file
124
.env.example
Normal file
@@ -0,0 +1,124 @@
|
||||
# grok-free-register 配置模板
|
||||
# 复制为 .env 后按需修改。未填写的项目会使用程序默认值。
|
||||
|
||||
# 邮箱模式
|
||||
# tempmail: 免费临时邮箱,适合快速试跑。
|
||||
# custom: 自建域名邮箱,需要 Cloudflare Email Routing 和本地收信服务。
|
||||
EMAIL_MODE=tempmail
|
||||
|
||||
# custom 模式专用
|
||||
# EMAIL_DOMAIN=example.com
|
||||
# EMAIL_API=http://127.0.0.1:8080
|
||||
|
||||
# 认证账号来源
|
||||
# auto: 配置 SSH 主机时读取远端，否则读取当前项目；一般保持默认即可。
|
||||
# XAI_AUTH_SERVICE_SOURCE=auto # auto | local | ssh
|
||||
# XAI_AUTH_SERVICE_REGISTER_ROOT=. # local 模式下的注册项目目录
|
||||
|
||||
# 注册与认证不在同一台机器时填写
|
||||
# XAI_AUTH_SERVICE_SSH_HOST=user@server.example
|
||||
# XAI_AUTH_SERVICE_SSH_IDENTITY=/path/to/key.pem
|
||||
# XAI_AUTH_SERVICE_REMOTE_ROOT=/opt/grok-free-register
|
||||
# XAI_AUTH_SERVICE_SYNC_SEC=30 # 同步新账号的间隔(秒)
|
||||
|
||||
# 认证节奏
|
||||
# XAI_AUTH_SERVICE_MIN_INTERVAL_SEC=10 # 两次认证请求的基础最小间隔(秒)
|
||||
# # 触发限流后会自动抬高(约18s起，最多约90s)，成功连串后再缓慢回落
|
||||
# XAI_AUTH_SERVICE_RETRY_SEC=60 # 被限流后的基础探测间隔(秒)；连续撞限会按 1.5x 增长，上限 300s
|
||||
# XAI_AUTH_SERVICE_LOG_MODE=user # user | debug
|
||||
|
||||
# 运行目标
|
||||
# TARGET=0 # 成功 N 个后停止;0=不限
|
||||
|
||||
# 日志显示
|
||||
# REGISTER_LOG_MODE=user # 默认:任务开始、结果、平均速度和限流等待
|
||||
# REGISTER_LOG_MODE=debug # 额外显示 T/Q/并发/阶段耗时调试面板
|
||||
|
||||
# 浏览器并发容量
|
||||
# PHYSICAL_CAP=0 # 0=按 CPU/内存自动估算;>0=手动指定
|
||||
# PHYSICAL_PER_CPU=2 # 自动估算时每个 CPU 核心对应的并发参考值
|
||||
# PHYSICAL_MEM_MB=512 # 自动估算时每个浏览器任务的内存预算(MB)
|
||||
# MIN_FREE_MEM_MB=500 # 自动估算时保留的内存(MB)
|
||||
# CAPACITY_PROFILE= # 可选:离线压测生成的静态 profile(JSON)
|
||||
|
||||
# 缓冲容量
|
||||
# T_SLOT_CAP=8
|
||||
# Q_SLOT_CAP=8
|
||||
# Q_PENDING_CAP=12
|
||||
|
||||
# Worker 数量
|
||||
# 通常不需要手动设置。0 表示根据容量自动派生。
|
||||
# S_WORKERS=0
|
||||
# P_WORKERS=0
|
||||
# C_WORKERS=0
|
||||
|
||||
# 资源有效期
|
||||
# T_MAX_AGE=300 # token 最大年龄(秒)
|
||||
# Q_MAX_AGE=120 # 验证码最大年龄(秒)
|
||||
|
||||
# 阶段超时
|
||||
# P_REQUEST_TIMEOUT=95 # 等待验证码返回超时(秒)
|
||||
# C_CONSUME_TIMEOUT=60 # 最终提交阶段超时(秒)
|
||||
|
||||
# 局部门控和批量发送
|
||||
# 这些项目默认会按容量派生。只有在压测定位后才建议覆盖。
|
||||
# T_TARGET=4 # token 缓冲目标
|
||||
# Q_TARGET=4 # 验证码缓冲目标
|
||||
# T_HIGH_WATER= # token 生产高水位
|
||||
# T_LOW_WATER= # token 生产恢复低水位
|
||||
# Q_HIGH_WATER= # 验证码生产高水位
|
||||
# Q_LOW_WATER= # 验证码生产恢复低水位
|
||||
# P_BATCH_MAX=4 # 单个发送页面最多发送的验证码请求数
|
||||
# P_SEND_CAP=0 # 0=不额外限制;>0=限制发送页面并发
|
||||
|
||||
# 页面和 token 获取
|
||||
# SOLVER_REUSE=1 # token 页面复用
|
||||
# MAX_SOLVER_REUSE=25 # 单个 token 页面最大复用次数
|
||||
# SOLVER_INITIAL_WAIT_MS=500 # token 注入后的首次等待
|
||||
# SOLVER_POLL_INTERVAL_MS=500 # token 轮询间隔
|
||||
# SOLVER_POLL_ATTEMPTS=100 # token 最大轮询次数
|
||||
# SOLVER_HARD_TIMEOUT=90 # 单次 solver 全链路硬截止(秒)
|
||||
# SOLVER_CLEANUP_TIMEOUT=5 # 异常页回收最长等待(秒)
|
||||
# SOLVER_FAST_CLICK=1 # 没有可见验证框时跳过慢点击
|
||||
# SOLVER_MOUSE_CLICK_RETRIES=3 # token 验证框中心点击次数;0=关闭
|
||||
# SOLVER_MOUSE_CLICK_INTERVAL_MS=600 # token 中心点击间隔(ms)
|
||||
# PAGE_GOTO_WAIT_UNTIL=domcontentloaded # 页面导航等待策略
|
||||
# PAGE_POST_WAIT_MS=500 # 页面导航后的短等待
|
||||
|
||||
# 可选性能开关
|
||||
# PAGE_BLOCK_STATIC_ASSETS=0 # 阻断部分静态资源,降低页面准备成本
|
||||
# C_HOT_PAGE_POOL=0 # 复用消费阶段页面,会增加常驻内存
|
||||
# C_HOT_PAGE_POOL_SIZE=0 # 0=按启动期并发派生;>0=手动指定
|
||||
# C_SET_COOKIE_VIA_REQUEST=0 # 用浏览器 request 访问 set-cookie URL;热页池开启时默认启用
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 代理（grok-free-register 本进程出口，给 Playwright + httpx）
|
||||
# 与下面 FlareSolverr 清障栈是两套东西，不要混。
|
||||
# Playwright 不会自动读环境变量，代码里会显式挂到 launch(proxy=...)。
|
||||
# ---------------------------------------------------------------------------
|
||||
# REGISTER_PROXY=http://127.0.0.1:40080 # 优先；指向本机 privoxy→WARP 时用这个
|
||||
# HTTP_PROXY=http://127.0.0.1:7890 # 兜底（REGISTER_PROXY 未设时）
|
||||
# HTTPS_PROXY=http://127.0.0.1:7890
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cloudflare 清障预热（调用外部 FlareSolverr，不在本仓起 FS）
|
||||
# 默认 FS: docker compose 暴露 127.0.0.1:8191
|
||||
# 预热根域名（无 path 后缀）: accounts.x.ai / x.ai / status.x.ai / console.x.ai / auth.x.ai
|
||||
# CLEARANCE_PROXY 是 FS *容器内* 浏览器的出口，必须是容器可达地址；
|
||||
# 不要把 host 的 127.0.0.1 塞进去。走 WARP 时用 docker 网内 privoxy。
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLEARANCE_ENABLED=1
|
||||
# FLARESOLVERR_URL=http://127.0.0.1:8191
|
||||
# CLEARANCE_URLS=https://accounts.x.ai,https://x.ai,https://status.x.ai,https://console.x.ai,https://auth.x.ai
|
||||
# CLEARANCE_PROXY=http://privoxy:8118
|
||||
# CLEARANCE_TIMEOUT_SEC=60
|
||||
# CLEARANCE_REFRESH_SEC=3000
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 本机注册 → 推 keys 到远端 auth（可选；scripts/push_keys_to_auth.sh）
|
||||
# 无内置默认主机，必须显式填写。
|
||||
# ---------------------------------------------------------------------------
|
||||
# AUTH_SSH_HOST=user@auth-server.example
|
||||
# AUTH_REMOTE_ROOT=/opt/grok-free-register
|
||||
# AUTH_SSH_IDENTITY=
|
||||
# AUTH_PUSH_INTERVAL=20
|
||||
19
.gitignore
vendored
Normal file
19
.gitignore
vendored
Normal file
@@ -0,0 +1,19 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
.env
|
||||
*.log
|
||||
logs/
|
||||
*.bak
|
||||
.venv/
|
||||
.setup.lock/
|
||||
.DS_Store
|
||||
.pytest_cache/
|
||||
xai-enroller-ledger.db*
|
||||
tmp_*.py
|
||||
|
||||
# runtime credential outputs under keys/ (keep product scripts tracked)
|
||||
keys/*
|
||||
!keys/acpa_watchdog.py
|
||||
!keys/sync_acc.py
|
||||
!keys/async_auth.sh
|
||||
!keys/README.md
|
||||
21
LICENSE
Normal file
21
LICENSE
Normal file
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 grok-free-register contributors
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
370
README.md
Normal file
370
README.md
Normal file
@@ -0,0 +1,370 @@
|
||||
# grok-free-register
|
||||
|
||||
`grok-free-register` 是一个命令行注册工具。程序会启动本机浏览器，完成页面操作、邮箱验证码处理和结果保存。
|
||||
|
||||
运行结果写入 `keys/` 目录。
|
||||
|
||||
## 快速开始
|
||||
|
||||
```bash
|
||||
git clone <your-fork-or-mirror>/grok-free-register.git
|
||||
cd grok-free-register
|
||||
bash start.sh
|
||||
```
|
||||
|
||||
首次运行会自动创建 `.venv`、安装依赖，并引导生成 `.env`。
|
||||
|
||||
需要完整说明时，按用途查看：
|
||||
|
||||
- [注册教程](docs/guides/registration.md)
|
||||
- [本地认证服务](docs/guides/auth-service.md)
|
||||
- [凭据库存与取用](docs/guides/credential-inventory.md)
|
||||
- [运行状态与排障](docs/guides/runtime-troubleshooting.md)
|
||||
|
||||
常用命令：
|
||||
|
||||
```bash
|
||||
bash start.sh # 按当前 .env 前台运行
|
||||
bash start.sh --target 100 # 成功 100 个后停止
|
||||
bash start.sh --max-mem 6G # 自动估算并发时最多使用 6G 内存
|
||||
bash start.sh --reconfig # 重新选择邮箱模式
|
||||
```
|
||||
|
||||
`start.sh` 直接在当前终端显示状态。按 `Ctrl-C` 停止，再次执行同一命令即可重启；不需要额外的会话管理或守护进程依赖。
|
||||
|
||||
代理与 CF 清障是两套配置，不要混：
|
||||
|
||||
```env
|
||||
# 本进程出口（Playwright + httpx 显式挂载；Playwright 不会自己读 env）
|
||||
REGISTER_PROXY=http://127.0.0.1:40080
|
||||
# 或兜底：
|
||||
# HTTP_PROXY=http://127.0.0.1:7890
|
||||
# HTTPS_PROXY=http://127.0.0.1:7890
|
||||
|
||||
# 清障预热：调用外部 FlareSolverr（本仓不跑 FS）
|
||||
CLEARANCE_ENABLED=1
|
||||
FLARESOLVERR_URL=http://127.0.0.1:8191
|
||||
CLEARANCE_URLS=https://accounts.x.ai,https://x.ai,https://status.x.ai,https://console.x.ai,https://auth.x.ai
|
||||
# FS 容器内出口（容器可达名）；走 WARP 时用 docker 网内 privoxy
|
||||
# CLEARANCE_PROXY=http://privoxy:8118
|
||||
```
|
||||
|
||||
## 邮箱模式
|
||||
|
||||
`tempmail` 是默认模式，不需要额外配置，适合快速试跑：
|
||||
|
||||
```env
|
||||
EMAIL_MODE=tempmail
|
||||
```
|
||||
|
||||
`custom` 是自建域名邮箱模式，适合长时间运行。需要一个已接入 Cloudflare Email Routing 的域名，并在运行机器上启动本项目的收信服务。
|
||||
|
||||
配置步骤：
|
||||
|
||||
1. 在 Cloudflare 为域名开启 Email Routing。
|
||||
2. 部署 `cloudflare/email-worker.js`。
|
||||
3. 在 Email Routing 中配置 catch-all，动作选择发送到该 Worker。
|
||||
4. 在运行机器上启动收信服务：
|
||||
|
||||
```bash
|
||||
bash start.sh --email-service
|
||||
```
|
||||
|
||||
5. 在 `.env` 中配置：
|
||||
|
||||
```env
|
||||
EMAIL_MODE=custom
|
||||
EMAIL_DOMAIN=example.com
|
||||
EMAIL_API=http://127.0.0.1:8080
|
||||
```
|
||||
|
||||
如果 Worker 需要回调本机服务，`WEBHOOK_URL` 应使用可访问的域名地址。
|
||||
|
||||
## 配置
|
||||
|
||||
完整模板见 `.env.example`。日常使用通常只需要配置邮箱模式、代理、目标数量和内存预算。
|
||||
|
||||
| 配置 | 默认值 | 说明 |
|
||||
|---|---:|---|
|
||||
| `EMAIL_MODE` | `tempmail` | 邮箱模式，支持 `tempmail` 和 `custom` |
|
||||
| `EMAIL_DOMAIN` | 空 | `custom` 模式使用的域名 |
|
||||
| `EMAIL_API` | `http://127.0.0.1:8080` | 本地收信服务地址 |
|
||||
| `TARGET` | `0` | 成功数量目标，`0` 表示不限 |
|
||||
| `PHYSICAL_CAP` | `0` | 浏览器并发上限，`0` 表示启动时自动估算 |
|
||||
| `PHYSICAL_PER_CPU` | `2` | 自动估算时每个 CPU 核心对应的并发参考值 |
|
||||
| `PHYSICAL_MEM_MB` | `512` | 自动估算时每个浏览器任务的内存预算 |
|
||||
| `MIN_FREE_MEM_MB` | `500` | 自动估算时保留的内存 |
|
||||
| `T_SLOT_CAP` | `8` | token 缓冲容量 |
|
||||
| `Q_SLOT_CAP` | `8` | 验证码缓冲容量 |
|
||||
| `Q_PENDING_CAP` | `12` | 等待验证码返回的请求上限 |
|
||||
| `SOLVER_MOUSE_CLICK_RETRIES` | `3` | token 验证框中心点击次数，`0` 表示关闭 |
|
||||
| `PAGE_BLOCK_STATIC_ASSETS` | `0` | 可选：阻断部分静态资源，降低页面准备成本 |
|
||||
| `C_HOT_PAGE_POOL` | `0` | 可选：复用消费阶段页面，减少页面重建开销 |
|
||||
|
||||
不确定怎么设置时，先保持默认值。性能压测时优先观察 `PHYSICAL_CAP` 和内存，不建议先改 Worker 数量。
|
||||
|
||||
## 运行日志
|
||||
|
||||
直接运行 `bash start.sh` 时，终端只输出任务开始、成功或失败、本次运行平均速度、累计数量和限流等待：
|
||||
|
||||
```text
|
||||
[→] 开始注册 #38
|
||||
[✓] 注册成功 #38 | 运行平均 47/分 | 累计 38
|
||||
[⏸] 触发限流 | 60秒后恢复探测
|
||||
[▶] 限流解除 | 实际等待 61秒
|
||||
```
|
||||
|
||||
需要调试并发、库存和阶段耗时时，使用：
|
||||
|
||||
```bash
|
||||
bash start.sh --debug
|
||||
```
|
||||
|
||||
它会在上述任务事件之外，每 8 秒输出一次完整的 T/Q、物理并发和阶段耗时面板。已有自动化环境也可继续使用 `REGISTER_LOG_MODE=debug`。
|
||||
|
||||
常用字段：
|
||||
|
||||
| 字段 | 含义 |
|
||||
|---|---|
|
||||
| `T` | 当前可用 token 数量 |
|
||||
| `Q` | 当前可用验证码数量 |
|
||||
| `phys` | 空闲浏览器并发许可 |
|
||||
| `s_phys` / `p_phys` / `c_phys` | S/P/C 获取浏览器许可的平均等待秒数 / 平均持有秒数 |
|
||||
| `p_stage` | P 阶段平均耗时：建邮箱 / 准备页面 / 发送请求 |
|
||||
| `c_stage` | C 阶段平均耗时：拿页面 / 验证码校验 / 注册提交 |
|
||||
| `c_hot` | C 热页池命中 / 未命中次数 |
|
||||
| `t_solve_avg` | 平均 token 获取时间 |
|
||||
| `q_sent` / `q_ret` | 已发送 / 已收到的验证码数量 |
|
||||
| `pair` | 已配对消费次数 |
|
||||
| `ok` / `fail` | 成功 / 失败数量 |
|
||||
| `rate` | 当前累计成功速率 |
|
||||
|
||||
简单判断：
|
||||
|
||||
- `T` 长期为 `0` 且 `Q` 有库存，通常是 token 获取较慢。
|
||||
- `Q` 长期为 `0` 且 `T` 有库存，通常是邮箱或验证码链路较慢。
|
||||
- `phys` 长期为 `0`，说明浏览器并发已经用满。
|
||||
- `t_solve_avg` 明显升高，通常表示浏览器压力、网络质量或 token 服务响应变慢。
|
||||
|
||||
可以用日志分析工具解析已有日志：
|
||||
|
||||
```bash
|
||||
python3 - <<'PY'
|
||||
from pathlib import Path
|
||||
from tools.runtime_log_analyzer import analyze_text
|
||||
print(analyze_text(Path("run.log").read_text()))
|
||||
PY
|
||||
```
|
||||
|
||||
## 输出文件
|
||||
|
||||
成功结果写入：
|
||||
|
||||
```text
|
||||
keys/accounts.txt
|
||||
keys/grok.txt
|
||||
```
|
||||
|
||||
`accounts.txt` 每行格式：
|
||||
|
||||
```text
|
||||
email:password:sso_token
|
||||
```
|
||||
|
||||
`keys/` 目录包含运行结果，默认不会提交到 Git。
|
||||
|
||||
认证成功后的 OAuth 出货默认在：
|
||||
|
||||
```text
|
||||
~/Downloads/grok-free-register-auth/authenticated/
|
||||
```
|
||||
|
||||
可用 `keys/async_auth.sh` 把出货整备成 `keys/cpa_ready/` 并同步探活通过的 `access_token` 到 `keys/acc.md`。
|
||||
|
||||
## 清零流水线（删库跑路）
|
||||
|
||||
只删 `keys/acc.md` / `accounts.txt` **不够**：`acpa_watchdog` 会从
|
||||
`~/Downloads/grok-free-register-auth/authenticated/` 把旧号再整备进
|
||||
`keys/cpa_ready/`，`sync_acc` 再写回 `acc.md`。状态文件 `_state.tsv` 也会残留幽灵 alive。
|
||||
|
||||
**全量清零**用项目根的 `reset_pipeline.sh`（默认 dry-run，必须加 `--yes` 才真删）：
|
||||
|
||||
```bash
|
||||
# 建议先停 register / auth-service / async_auth
|
||||
bash reset_pipeline.sh # dry-run，只列将删内容
|
||||
bash reset_pipeline.sh --yes # 真删：注册源 + 认证出货 + cpa_ready + state
|
||||
bash reset_pipeline.sh --yes --keep-register
|
||||
# 只清认证/CPA，保留 keys 里新注册的 SSO 源
|
||||
bash reset_pipeline.sh --yes --keep-cpa
|
||||
# 只清认证账本/源快照，保留 cpa_ready（少用）
|
||||
bash reset_pipeline.sh --yes --wipe-salt
|
||||
# 连 enrollment ledger salt 一并清（极少需要）
|
||||
```
|
||||
|
||||
会清掉的大致范围：
|
||||
|
||||
| 范围 | 路径 |
|
||||
|------|------|
|
||||
| 注册源 | `keys/accounts.txt`、`keys/grok.txt`、`keys/auth-sessions.jsonl`、`keys/acc.md` |
|
||||
| CPA | `keys/cpa_ready/xai-*.json`、`_state.tsv`、`_discarded/` |
|
||||
| 认证出货 | `$XAI_AUTH_SERVICE_LOCAL_DIR`（默认 `~/Downloads/grok-free-register-auth`）下的 `authenticated/`、`claimed/`、`enrollment-ledger.db*`、`source-snapshot.jsonl` |
|
||||
|
||||
认证目录可用环境变量覆盖：
|
||||
|
||||
```bash
|
||||
export XAI_AUTH_SERVICE_LOCAL_DIR=~/Downloads/grok-free-register-auth
|
||||
# 或
|
||||
export XAI_ENROLLER_LOCAL_AUTH_DIR=~/Downloads/grok-free-register-auth
|
||||
```
|
||||
|
||||
**只清历史归档、不动现网号池**用 `scripts/clean_history.sh`：
|
||||
|
||||
```bash
|
||||
bash scripts/clean_history.sh # dry-run
|
||||
bash scripts/clean_history.sh --yes # 清 zip / _discarded / logs / 杂包
|
||||
bash scripts/clean_history.sh --yes --deep # 再清 source-snapshot 等更深一层
|
||||
```
|
||||
|
||||
清完后重开：
|
||||
|
||||
```bash
|
||||
bash start.sh # 或远端注册机继续跑
|
||||
bash auth-service.sh # device-flow → authenticated/
|
||||
bash keys/async_auth.sh # acpa_watchdog + sync_acc
|
||||
```
|
||||
|
||||
## 项目结构
|
||||
|
||||
```text
|
||||
grok_register/ 注册核心与 custom 邮箱服务
|
||||
xai_enroller/ OAuth 认证服务
|
||||
keys/ acpa_watchdog / sync_acc / async_auth（运行时产物 gitignore）
|
||||
cloudflare/email-worker.js Cloudflare Email Routing Worker 示例
|
||||
start.sh 首次配置和运行
|
||||
auth-service.sh 认证服务入口
|
||||
reset_pipeline.sh 全量清零注册→认证→CPA 运行时数据
|
||||
scripts/clean_history.sh 只清历史归档，不动现网号池
|
||||
scripts/push_keys_to_auth.sh 本机 SSO 源推远端 auth（需显式 AUTH_SSH_*）
|
||||
setup.sh 安装依赖
|
||||
.env.example 配置模板
|
||||
tools/ 日志分析 + GrokSession→CPA/sub2api 前端转换
|
||||
tests/ 自动化测试
|
||||
docs/architecture.md 并发架构说明
|
||||
```
|
||||
|
||||
## 测试
|
||||
|
||||
测试依赖与运行依赖分开安装：
|
||||
|
||||
```bash
|
||||
.venv/bin/pip install -r tests/requirements.txt
|
||||
```
|
||||
|
||||
快速检查：
|
||||
|
||||
```bash
|
||||
python3 -m unittest tests.test_admission_gate tests.test_register_runtime_unittest tests.test_inventory_unittest tests.test_runtime_log_analyzer -v
|
||||
```
|
||||
|
||||
完整测试：
|
||||
|
||||
```bash
|
||||
python3 -m pytest tests -q
|
||||
```
|
||||
|
||||
## xAI OAuth Enroller
|
||||
|
||||
`xai_enroller/` 用于把已有 xAI 账号的 SSO 会话转换为 OAuth 凭据，并导入 CPA
|
||||
凭据库。它独立于注册流程运行：注册机不需要启动，已有账号也不需要重新注册。
|
||||
|
||||
### 默认同机模式
|
||||
|
||||
注册和认证在同一个项目目录运行时，无需配置来源：
|
||||
|
||||
```bash
|
||||
bash auth-service.sh
|
||||
```
|
||||
|
||||
认证服务默认读取本项目的 `keys/auth-sessions.jsonl` 和历史
|
||||
`keys/accounts.txt`，每 30 秒生成一次经过校验的原子快照。注册仍可同时运行；认证服务
|
||||
只读取完整记录，不会读取正在追加的半行。
|
||||
|
||||
### 分离设备的 SSH 模式
|
||||
|
||||
注册机和认证服务分开运行时，服务器继续写入注册结果，本地认证服务
|
||||
每 30 秒通过一次性 SSH 导出全量 JSONL 快照。快照经逐行校验、`fsync` 后原子替换到
|
||||
`~/Downloads/grok-free-register-auth/source-snapshot.jsonl`；同步失败会继续使用上一份有效快照。
|
||||
认证使用本机 CloakBrowser Chromium，
|
||||
成功结果写入 `~/Downloads/grok-free-register-auth/authenticated/`，运行状态、同步快照与
|
||||
认证文件分开保存；认证文件格式可以直接供 CPA 使用。
|
||||
|
||||
先把导出器同步到注册机的 `scripts/` 目录：
|
||||
|
||||
```bash
|
||||
scp scripts/export_registered_sessions.py user@your-server:/opt/grok-free-register/scripts/
|
||||
```
|
||||
|
||||
然后在本地配置 SSH 连接：
|
||||
|
||||
下面这些变量既可在终端 `export`，也可写入由 `.env.example` 复制出的 `.env`；
|
||||
`auth-service.sh` 会自动读取该文件。
|
||||
|
||||
```bash
|
||||
export XAI_AUTH_SERVICE_SSH_HOST=user@your-server
|
||||
export XAI_AUTH_SERVICE_SSH_IDENTITY=/path/to/ssh-key.pem # 使用 ssh-agent 时可省略
|
||||
export XAI_AUTH_SERVICE_REMOTE_ROOT=/opt/grok-free-register
|
||||
export XAI_AUTH_SERVICE_SYNC_SEC=30
|
||||
```
|
||||
|
||||
存在 `XAI_AUTH_SERVICE_SSH_HOST` 时会自动选择 SSH。也可显式设置
|
||||
`XAI_AUTH_SERVICE_SOURCE=local|ssh`；默认值 `auto` 优先兼容已有 SSH 配置，否则使用同机来源。
|
||||
|
||||
启动认证服务：
|
||||
|
||||
```bash
|
||||
bash auth-service.sh
|
||||
```
|
||||
|
||||
需要查看队列、重试、节拍和冷却探针时使用：
|
||||
|
||||
```bash
|
||||
bash auth-service.sh --debug
|
||||
```
|
||||
|
||||
默认每次认证至少间隔 10 秒；实测该值比无间隔运行有更高的长期平均成功速率。限流后
|
||||
每 60 秒只进行一次恢复探测。需要覆盖时使用：
|
||||
|
||||
```bash
|
||||
export XAI_AUTH_SERVICE_MIN_INTERVAL_SEC=10
|
||||
export XAI_AUTH_SERVICE_RETRY_SEC=60
|
||||
```
|
||||
|
||||
终端只在账号开始、认证结果、限流状态或控制状态变化时输出，并在底部保持 `认证> ` 输入行。`s` 查看状态，`p` 暂停，
|
||||
`r` 恢复，`c` 取消当前账号，`q` 退出；`Ctrl-C` 同样会退出。在输入行键入 `take 100`
|
||||
会把最新的 100 个可用凭证登记为已取用，并移动到独立批次目录。认证记录仍保留为
|
||||
`imported`，不会因为凭证被取出而重新认证。
|
||||
|
||||
可用凭证保存在：
|
||||
|
||||
```text
|
||||
~/Downloads/grok-free-register-auth/authenticated/
|
||||
```
|
||||
|
||||
已取用批次保存在：
|
||||
|
||||
```text
|
||||
~/Downloads/grok-free-register-auth/claimed/<batch-id>/
|
||||
```
|
||||
|
||||
库存状态保存在 `enrollment-ledger.db` 的 `credential_inventory` 表中，状态为
|
||||
`available`、`claiming` 或 `claimed`。每条记录预留 `note` 字段，默认留空。
|
||||
|
||||
`auth-service.sh` 首次运行会自动安装项目依赖。正式用户流程只需要这个 Bash 入口；底层 Python 模块保留给开发和测试，不作为另一套使用方式。
|
||||
|
||||
## 开发文档
|
||||
|
||||
[docs/architecture.md](docs/architecture.md) 记录并发模型、资源生命周期和必须保持的不变量。
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
55
auth-service.sh
Executable file
55
auth-service.sh
Executable file
@@ -0,0 +1,55 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
. scripts/ensure_runtime.sh
|
||||
ensure_runtime
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 后台拉起 CPA 整备 + acc token 汇总
|
||||
# acpa_watchdog.py → authenticated/ → keys/cpa_ready/ (整备 cli-chat-proxy 凭据)
|
||||
# sync_acc.py → keys/cpa_ready/ → keys/acc.md (只抽 access_token, 一行一个)
|
||||
# 两者都常驻轮询。不用 exec: exec 会替换 shell, 后台子进程变孤儿且 trap 失效。
|
||||
# 脚本退出(正常 / Ctrl-C / SIGTERM)时一并回收, 避免孤儿。
|
||||
# ---------------------------------------------------------------------------
|
||||
WATCHDOG="keys/acpa_watchdog.py"
|
||||
SYNC_ACC="keys/sync_acc.py"
|
||||
WATCHDOG_PID=""
|
||||
SYNC_PID=""
|
||||
|
||||
cleanup() {
|
||||
for pid in "${SYNC_PID:-}" "${WATCHDOG_PID:-}"; do
|
||||
if [ -n "$pid" ] && kill -0 "$pid" 2>/dev/null; then
|
||||
kill "$pid" 2>/dev/null || true
|
||||
wait "$pid" 2>/dev/null || true
|
||||
fi
|
||||
done
|
||||
}
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
mkdir -p logs
|
||||
|
||||
if [ -f "$WATCHDOG" ]; then
|
||||
.venv/bin/python "$WATCHDOG" </dev/null >>logs/watchdog.log 2>&1 &
|
||||
WATCHDOG_PID=$!
|
||||
echo "[auth-service] watchdog 已拉起 (pid=$WATCHDOG_PID, 日志 logs/watchdog.log)"
|
||||
else
|
||||
echo "[auth-service] 未找到 $WATCHDOG, 跳过自动整备" >&2
|
||||
fi
|
||||
|
||||
if [ -f "$SYNC_ACC" ]; then
|
||||
.venv/bin/python "$SYNC_ACC" </dev/null >>logs/sync_acc.log 2>&1 &
|
||||
SYNC_PID=$!
|
||||
echo "[auth-service] sync_acc 已拉起 (pid=$SYNC_PID, 日志 logs/sync_acc.log → keys/acc.md)"
|
||||
else
|
||||
echo "[auth-service] 未找到 $SYNC_ACC, 跳过 acc.md 汇总" >&2
|
||||
fi
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 前台跑认证服务 (注册出货)。退出码透传; trap cleanup 在 exit 时兜底回收后台进程。
|
||||
# ---------------------------------------------------------------------------
|
||||
set +e
|
||||
.venv/bin/python -m xai_enroller.service "$@"
|
||||
RC=$?
|
||||
set -e
|
||||
exit "$RC"
|
||||
61
cloudflare/email-worker.js
Normal file
61
cloudflare/email-worker.js
Normal file
@@ -0,0 +1,61 @@
|
||||
// ============================================================
|
||||
// Cloudflare Email Worker — 把收到的邮件 POST 到你的 webhook
|
||||
// 用于 grok-free-register 的「自建邮箱模式」(EMAIL_MODE=custom)
|
||||
//
|
||||
// 链路: 发件方 → Cloudflare Email Routing → 本 Worker → 你的 webhook(email_server.py)
|
||||
//
|
||||
// ⚠️ 重要(踩过的坑):env.WEBHOOK_URL 必须是【域名】,不能用【裸 IP】!
|
||||
// Cloudflare Workers 的 fetch 不允许直接请求裸 IP,会返回
|
||||
// "error code: 1003 (Direct IP access not allowed)" 并被静默吞掉
|
||||
// —— Worker 看起来成功、你的服务器却收不到任何请求。
|
||||
// 做法:给服务器 IP 加一条 DNS A 记录(如 hook.example.com,灰云/DNS-only),
|
||||
// WEBHOOK_URL 填 http://hook.example.com:8080/webhook。
|
||||
//
|
||||
// 部署:
|
||||
// npm create cloudflare@latest cf-mail-webhook
|
||||
// cd cf-mail-webhook && npm i postal-mime
|
||||
// # 用本文件替换 src/index.js
|
||||
// npx wrangler secret put WEBHOOK_URL # 输入 http://hook.example.com:8080/webhook
|
||||
// npx wrangler secret put WEBHOOK_TOKEN # 可选,鉴权用
|
||||
// npx wrangler deploy
|
||||
// 然后在 Cloudflare 后台:Email → Email Routing → Routing rules →
|
||||
// Catch-all → 动作选「Send to a Worker」→ 选本 Worker。
|
||||
// (子域名收信需在 Email Routing 单独启用该子域并配它自己的 catch-all。)
|
||||
// ============================================================
|
||||
import PostalMime from "postal-mime";
|
||||
|
||||
function trim(s, max = 20000) {
|
||||
if (!s) return "";
|
||||
return s.length > max ? s.slice(0, max) + "\n...[truncated]" : s;
|
||||
}
|
||||
|
||||
export default {
|
||||
async email(message, env, ctx) {
|
||||
if (!env.WEBHOOK_URL) {
|
||||
console.error("WEBHOOK_URL 未配置 (npx wrangler secret put WEBHOOK_URL)");
|
||||
return;
|
||||
}
|
||||
const parsed = await PostalMime.parse(message.raw);
|
||||
const payload = {
|
||||
from: message.from,
|
||||
to: message.to,
|
||||
subject: parsed.subject || message.headers.get("subject") || "",
|
||||
text: trim(parsed.text || ""),
|
||||
html: trim(parsed.html || ""),
|
||||
};
|
||||
|
||||
const headers = { "content-type": "application/json" };
|
||||
if (env.WEBHOOK_TOKEN) headers["x-webhook-token"] = env.WEBHOOK_TOKEN;
|
||||
|
||||
const res = await fetch(env.WEBHOOK_URL, {
|
||||
method: "POST",
|
||||
headers,
|
||||
body: JSON.stringify(payload),
|
||||
});
|
||||
// 日志只保留状态码，避免把 webhook 地址或上游响应写入 Worker 日志。
|
||||
console.log(`webhook delivery -> ${res.status}`);
|
||||
if (!res.ok) {
|
||||
console.error(`webhook delivery failed: ${res.status}`);
|
||||
}
|
||||
},
|
||||
};
|
||||
180
docs/architecture.md
Normal file
180
docs/architecture.md
Normal file
@@ -0,0 +1,180 @@
|
||||
# CSP 架构说明
|
||||
|
||||
本文档记录当前运行时架构。README 面向使用者；本文面向维护者，重点是并发边界、资源生命周期和测试不变量。
|
||||
|
||||
## 目标
|
||||
|
||||
运行时采用 CSP 风格的异步流水线：
|
||||
|
||||
- `S_Worker` 生产 `T`。
|
||||
- `P_Worker` 发起外部请求，等待并生产 `Q`。
|
||||
- `C_Worker` 原子获取一组 `T + Q` 并执行最终消费。
|
||||
|
||||
架构目标是资源所有权闭合、背压边界清晰、取消语义可测试。当前设计不包含中心调度器、动态角色选择或运行时动态并发分配。
|
||||
|
||||
## 核心组件
|
||||
|
||||
`Physical_Sem` 限制本地浏览器重操作并发。`P_Worker` 发出请求后，在等待 `Q` 返回期间不得持有该许可。
|
||||
|
||||
`T_Slot_Sem` 限制已入库 `T` 的容量。`T` 生成成功后才允许获取 slot。
|
||||
|
||||
`Q_Slot_Sem` 限制已返回并入库 `Q` 的容量。`Q` 真正返回前不得获取 slot。
|
||||
|
||||
`Q_Pending_Sem` 限制已发出但尚未终态的外部 `Q` 请求数量。
|
||||
|
||||
`Inventory` 是唯一库存门面。worker 不直接访问底层队列，只能调用：
|
||||
|
||||
```text
|
||||
put_t(env)
|
||||
put_q(env)
|
||||
claim_pair()
|
||||
```
|
||||
|
||||
`ResourceEnvelope` 把资源实体和库存 slot 绑定在一起。slot 只能释放一次。
|
||||
|
||||
`PairLease` 是 `claim_pair()` 返回的异步上下文管理器。pair 一旦 claim 成功，两个 envelope 的所有权转移给 lease，直到 consumer 退出上下文。
|
||||
|
||||
`AdmissionGate` 是局部生产准入门控，只根据库存深度和静态水位决定是否允许继续生产。它不选择 worker 角色，不搬运资源，不调整并发容量。
|
||||
|
||||
## Worker 流程
|
||||
|
||||
### S_Worker
|
||||
|
||||
1. 等待 `AdmissionGate` 允许生产 `T`。
|
||||
2. 获取 `Physical_Sem`。
|
||||
3. 生产 `T`。
|
||||
4. 释放 `Physical_Sem`。
|
||||
5. 调用 `ResourceEnvelope.create_with_slot(...)`，创建 envelope 并获取 `T_Slot_Sem`。
|
||||
6. 调用 `Inventory.put_t(...)`，把所有权转移给 `Inventory`。
|
||||
7. 所有权转移前如果异常或取消，释放已获取的 slot。
|
||||
|
||||
slot 获取和 envelope 创建必须绑定，避免取消落在“已获取 slot、尚未创建 envelope”的窗口。
|
||||
|
||||
### P_Worker
|
||||
|
||||
1. 等待 `AdmissionGate` 允许生产 `Q`。
|
||||
2. 获取 `Q_Pending_Sem`。
|
||||
3. 获取 `Physical_Sem`。
|
||||
4. 创建请求并发送。
|
||||
5. 释放 `Physical_Sem`。
|
||||
6. 在不持有 `Physical_Sem` 的情况下等待 `Q` 返回。
|
||||
7. `Q` 返回后调用 `ResourceEnvelope.create_with_slot(...)`，创建 envelope 并获取 `Q_Slot_Sem`。
|
||||
8. 调用 `Inventory.put_q(...)`，把所有权转移给 `Inventory`。
|
||||
9. 请求进入终态后释放 `Q_Pending_Sem`。
|
||||
|
||||
`Q_Pending_Sem` 表达外部在途上限；`Q_Slot_Sem` 只表达已返回库存容量。两者不能合并。
|
||||
|
||||
### C_Worker
|
||||
|
||||
1. 进入 `async with inventory.claim_pair() as pair`。
|
||||
2. 获取 `Physical_Sem`。
|
||||
3. 消费 pair。
|
||||
4. 释放 `Physical_Sem`。
|
||||
5. 退出 `PairLease`，释放两个库存 slot。
|
||||
|
||||
`C_Worker` 不允许先取单边资源再等待另一边。对 consumer 来说，pair claim 必须是原子的。
|
||||
|
||||
## Inventory 语义
|
||||
|
||||
第一版 `Inventory` 使用一把 lock 和一个 condition。lock 保护等待、复查、过期清理和弹出操作。
|
||||
|
||||
必须保持以下语义：
|
||||
|
||||
- 等待中的 consumer 不移除资源。
|
||||
- claim 成功时同时移除一个有效 `T` 和一个有效 `Q`。
|
||||
- 等待 pair 时被取消，不影响库存。
|
||||
- claim 成功后被取消，由 `PairLease` 核销两个 envelope。
|
||||
- worker 永远不能直接操作底层 `T` / `Q` 队列。
|
||||
|
||||
只要锁内逻辑保持很小，除 lazy expiry cleanup 外基本是 O(1)，单锁在当前版本可以接受。只有 profile 证明锁竞争成为真实瓶颈时，才考虑拆锁或分片。
|
||||
|
||||
## 过期模型
|
||||
|
||||
`ResourceEnvelope` 可以携带 `created_at` 和 `expires_at`。`Inventory` 在配对前可以丢弃已过期资源。
|
||||
|
||||
当前采用 lazy cleanup：
|
||||
|
||||
- `put_t`、`put_q`、`claim_pair` 被触发时顺手清理。
|
||||
- 清理会释放它看到的过期 envelope 对应 slot。
|
||||
- 系统完全静默时不会主动扫库。
|
||||
- 单边长期故障和静默停摆由监控暴露，不靠后台 sweeper 修复。
|
||||
|
||||
当前不做 worker 级回队。消费失败后，已 claim 的 `T` 和 `Q` 都由 `PairLease` 核销。
|
||||
|
||||
## 容量策略
|
||||
|
||||
容量边界由 Semaphore 表达。启动期容量优先级：
|
||||
|
||||
```text
|
||||
显式 PHYSICAL_CAP > CAPACITY_PROFILE > CPU/内存自动派生
|
||||
```
|
||||
|
||||
`CAPACITY_PROFILE` 只在启动时读取，是静态 profile，不是运行时调度器。
|
||||
|
||||
默认 worker 数量由容量派生：
|
||||
|
||||
```text
|
||||
S_WORKERS = Physical_Sem + 2
|
||||
P_WORKERS = Q_Pending_Sem + 2
|
||||
C_WORKERS = Physical_Sem + 2
|
||||
```
|
||||
|
||||
worker 数量不是主要调参入口。主要并发边界是各类容量许可，而不是 coroutine 循环数量。
|
||||
|
||||
## 当前不做
|
||||
|
||||
当前版本明确不包含：
|
||||
|
||||
- 中心调度器；
|
||||
- 运行时角色选择；
|
||||
- 动态打分；
|
||||
- 动态并发控制；
|
||||
- worker 级回队；
|
||||
- 高价值资源抢救策略；
|
||||
- 后台过期清扫；
|
||||
- 自动切换高风险浏览器模式。
|
||||
|
||||
这些方向可以单独实验，但不能混入基础所有权模型。
|
||||
|
||||
## 必须保持的不变量
|
||||
|
||||
实现和测试必须维持以下不变量：
|
||||
|
||||
- 每个已获取的库存 slot 最终释放一次且只释放一次。
|
||||
- 每个已准入的 pending 请求最终释放一次 `Q_Pending_Sem`。
|
||||
- `P_Worker` 等待 `Q` 返回时不持有 `Physical_Sem`。
|
||||
- `Q_Slot_Sem` 只在 `Q` 返回后获取。
|
||||
- `C_Worker` 只能通过 `Inventory.claim_pair()` 获取 `T` 和 `Q`。
|
||||
- `claim_pair()` 要么返回一个受 `PairLease` 保护的完整 pair，要么不返回资源。
|
||||
- 等待 pair 时取消，不移除库存资源。
|
||||
- claim pair 后取消，两个 envelope 由 `PairLease` 释放。
|
||||
- 触发清理时，过期资源不能被配对。
|
||||
- 监控只读，不修改 Semaphore、队列或 worker 状态。
|
||||
|
||||
## 测试
|
||||
|
||||
```bash
|
||||
.venv/bin/pip install -r tests/requirements.txt
|
||||
```
|
||||
|
||||
快速检查：
|
||||
|
||||
```bash
|
||||
python3 -m unittest tests.test_admission_gate tests.test_register_runtime_unittest tests.test_inventory_unittest tests.test_runtime_log_analyzer -v
|
||||
```
|
||||
|
||||
完整测试：
|
||||
|
||||
```bash
|
||||
python3 -m pytest tests -q
|
||||
```
|
||||
|
||||
重点测试文件：
|
||||
|
||||
- `tests.test_inventory_unittest`：库存、过期和 lease 行为。
|
||||
- `tests.test_register_runtime_unittest`：worker 运行语义和监控行。
|
||||
- `tests.test_admission_gate`：局部门控水位。
|
||||
- `tests.test_runtime_log_analyzer`：日志解析兼容性。
|
||||
- `tests/test_cancel.py`：取消边界。
|
||||
- `tests/test_property.py`：随机化不变量检查。
|
||||
- `tests/test_stress.py`：更高并发 fake-service 压测。
|
||||
67
docs/guides/auth-service.md
Normal file
67
docs/guides/auth-service.md
Normal file
@@ -0,0 +1,67 @@
|
||||
# 本地认证服务
|
||||
|
||||
认证服务把已有的 SSO 会话转换为 CPA 可直接读取的 OAuth 凭据。注册与认证可以在同一台机器，也可以分开运行。
|
||||
|
||||
## 默认同机运行
|
||||
|
||||
同一个项目目录已经或正在运行注册服务时，直接启动认证：
|
||||
|
||||
```bash
|
||||
bash auth-service.sh
|
||||
```
|
||||
|
||||
未配置 SSH 主机时，认证服务自动读取本项目 `keys/` 中的完整会话与历史账号。注册可以继续追加，认证服务只安装经过校验的完整快照。
|
||||
|
||||
## 配置远端同步
|
||||
|
||||
先把无密码导出器放到服务器项目目录：
|
||||
|
||||
```bash
|
||||
scp scripts/export_registered_sessions.py user@server.example:/opt/grok-free-register/scripts/
|
||||
```
|
||||
|
||||
在本地终端设置连接信息：
|
||||
|
||||
可以直接 `export`，也可以把 `.env.example` 复制为 `.env` 后填写；认证入口会自动读取 `.env`。
|
||||
|
||||
```bash
|
||||
export XAI_AUTH_SERVICE_SSH_HOST=user@server.example
|
||||
export XAI_AUTH_SERVICE_SSH_IDENTITY=/path/to/key.pem
|
||||
export XAI_AUTH_SERVICE_REMOTE_ROOT=/opt/grok-free-register
|
||||
```
|
||||
|
||||
使用 `ssh-agent` 时可省略 `XAI_AUTH_SERVICE_SSH_IDENTITY`。
|
||||
|
||||
设置了 `XAI_AUTH_SERVICE_SSH_HOST` 后，默认的 `auto` 模式会选择 SSH。需要明确覆盖时使用：
|
||||
|
||||
```bash
|
||||
export XAI_AUTH_SERVICE_SOURCE=local # 强制读取同机注册结果
|
||||
export XAI_AUTH_SERVICE_SOURCE=ssh # 强制使用 SSH，必须配置主机
|
||||
```
|
||||
|
||||
## 运行
|
||||
|
||||
```bash
|
||||
bash auth-service.sh
|
||||
```
|
||||
|
||||
首次运行会自动安装项目依赖。该命令在当前终端持续运行并直接接受控制命令；输入 `q` 或按 `Ctrl-C` 停止，再次执行同一命令即可重启。不需要额外的会话管理工具。
|
||||
|
||||
普通模式只在来源连接、发现新账号、任务开始、认证结果、限流和控制状态变化时输出。查看队列、重试、节拍和冷却探针时使用：
|
||||
|
||||
```bash
|
||||
bash auth-service.sh --debug
|
||||
```
|
||||
|
||||
运行中终端底部会保持 `认证> ` 输入行。日志更新不会清掉尚未提交的内容；直接输入命令并回车：
|
||||
|
||||
```text
|
||||
s 查看状态
|
||||
take N 取用 N 个凭据
|
||||
p 暂停
|
||||
r 恢复
|
||||
c 取消当前任务
|
||||
q 安全退出
|
||||
```
|
||||
|
||||
快照默认每 30 秒更新一次；内容无变化时终端保持安静。有效快照和已生成凭据会在重启后继续使用。
|
||||
19
docs/guides/credential-inventory.md
Normal file
19
docs/guides/credential-inventory.md
Normal file
@@ -0,0 +1,19 @@
|
||||
# 凭据库存与取用
|
||||
|
||||
认证成功的凭据保存在本地认证目录，并在 SQLite 库存中维护三种状态：
|
||||
|
||||
- `available`：已经认证、尚未取用；
|
||||
- `claiming`：正在移动到取用批次；
|
||||
- `claimed`：已经取用。
|
||||
|
||||
运行认证服务后输入：
|
||||
|
||||
```text
|
||||
take 100
|
||||
```
|
||||
|
||||
服务会选择最新的 100 个可用凭据，移动到 `claimed/<batch-id>/`，再把对应库存标记为 `claimed`。认证 ledger 中的 `imported` 记录仍然保留，所以这些账号不会被重新认证。
|
||||
|
||||
每条库存记录预留 `note` 字段，默认为空。取用操作不要求填写用途；以后需要备注时可直接更新该字段。
|
||||
|
||||
若文件移动中断，服务下次启动会恢复 `claiming` 状态。库存不足或凭据文件缺失时，操作失败但认证服务继续运行。
|
||||
54
docs/guides/registration.md
Normal file
54
docs/guides/registration.md
Normal file
@@ -0,0 +1,54 @@
|
||||
# 注册教程
|
||||
|
||||
## 开始运行
|
||||
|
||||
```bash
|
||||
git clone <your-fork-or-mirror>/grok-free-register.git
|
||||
cd grok-free-register
|
||||
bash start.sh
|
||||
```
|
||||
|
||||
首次运行会安装 Python、CloakBrowser Chromium 及其系统依赖，然后引导选择邮箱模式。以后再次执行 `bash start.sh` 会直接使用已有配置。
|
||||
|
||||
普通模式只显示服务启动、任务开始、注册成功或失败、本次运行平均速率、累计数量和限流状态。查看完整并发、库存和阶段耗时时使用：
|
||||
|
||||
```bash
|
||||
bash start.sh --debug
|
||||
```
|
||||
|
||||
常用参数：
|
||||
|
||||
```bash
|
||||
bash start.sh --target 100
|
||||
bash start.sh --max-mem 6G
|
||||
bash start.sh --reconfig
|
||||
```
|
||||
|
||||
未设置 `--target` 时服务持续运行，按 `Ctrl-C` 安全停止。
|
||||
再次执行 `bash start.sh` 即可重启。程序直接使用当前终端，不需要额外的会话管理工具。
|
||||
|
||||
## 配置邮箱
|
||||
|
||||
临时邮箱无需额外配置：
|
||||
|
||||
```env
|
||||
EMAIL_MODE=tempmail
|
||||
```
|
||||
|
||||
自建邮箱需要可接收邮件的域名和本项目的收信服务：
|
||||
|
||||
```env
|
||||
EMAIL_MODE=custom
|
||||
EMAIL_DOMAIN=example.com
|
||||
EMAIL_API=http://127.0.0.1:8080
|
||||
```
|
||||
|
||||
自建模式还需运行：
|
||||
|
||||
```bash
|
||||
bash start.sh --email-service
|
||||
```
|
||||
|
||||
性能参数默认会根据 CPU 和可用内存估算。除非正在压测，否则保持 `.env.example` 中的默认值即可。
|
||||
|
||||
成功结果写入 `keys/accounts.txt`、`keys/grok.txt` 和 `keys/auth-sessions.jsonl`；这些文件默认不提交到 Git。
|
||||
38
docs/guides/runtime-troubleshooting.md
Normal file
38
docs/guides/runtime-troubleshooting.md
Normal file
@@ -0,0 +1,38 @@
|
||||
# 运行状态与排障
|
||||
|
||||
## 普通模式
|
||||
|
||||
普通模式的每一行都对应一次状态变化：
|
||||
|
||||
- `[→]`：新任务开始；
|
||||
- `[✓]`：任务成功或来源连接完成；
|
||||
- `[✗]`：当前任务失败或被跳过；
|
||||
- `[⏸]`：进入限流冷却或服务暂停；
|
||||
- `[▶]`：限流解除或服务恢复；
|
||||
- `[!]`：来源、配置或流水线出现需要关注的问题。
|
||||
|
||||
认证端输入 `s` 可查看运行状态、待处理数量、当前阶段、本次运行平均速率、累计成功、可用和已取用凭据，以及是否处于限流。
|
||||
|
||||
## Debug 模式
|
||||
|
||||
注册端：
|
||||
|
||||
```bash
|
||||
bash start.sh --debug
|
||||
```
|
||||
|
||||
认证端：
|
||||
|
||||
```bash
|
||||
bash auth-service.sh --debug
|
||||
```
|
||||
|
||||
注册 Debug 面板包含 T/Q 库存、物理并发、S/P/C 阶段耗时和 token 求解时间。认证 Debug 状态包含 source/prepared/completion 队列、重试、授权节拍、冷却、单次探针和近五分钟滚动速率。
|
||||
|
||||
## 常见状态
|
||||
|
||||
限流后不会持续重试。认证端默认等待 60 秒，再只放行一个恢复探针；探针仍被限流时重新等待。注册端同样通过全局冷却闸门阻止并发任务漏过等待周期。
|
||||
|
||||
远端来源暂时断开时，本地认证服务继续使用上一份完整有效快照。恢复连接后会自动同步，不需要重启。
|
||||
|
||||
配置错误会指出缺少或非法的配置名，不输出 traceback。按提示检查 [注册配置](registration.md#配置邮箱) 或 [认证同步配置](auth-service.md#配置远端同步)。
|
||||
1
grok_register/__init__.py
Normal file
1
grok_register/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Registration service package."""
|
||||
395
grok_register/clearance.py
Normal file
395
grok_register/clearance.py
Normal file
@@ -0,0 +1,395 @@
|
||||
"""Cloudflare clearance prewarm via external FlareSolverr.
|
||||
|
||||
This module does NOT run FlareSolverr. It only calls an already-deployed
|
||||
FlareSolverr instance (default http://127.0.0.1:8191).
|
||||
|
||||
Proxy roles (do not mix):
|
||||
- REGISTER_PROXY / HTTP_PROXY / HTTPS_PROXY
|
||||
→ grok-free-register Playwright + httpx egress (this process)
|
||||
- CLEARANCE_PROXY
|
||||
→ optional proxy *inside* FlareSolverr's browser (must be reachable
|
||||
from the FS *container*, e.g. http://privoxy:8118)
|
||||
- FLARESOLVERR_URL
|
||||
→ host-side URL to reach FS (http://127.0.0.1:8191)
|
||||
|
||||
Clearance cookies are only valid on the same egress path they were minted on.
|
||||
If REGISTER_PROXY points at host Privoxy→WARP, set CLEARANCE_PROXY to the
|
||||
docker-internal privoxy URL so FS solves on the same WARP exit.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
# CF-fronted x.ai family roots (no path suffix; path pages share host clearance)
|
||||
DEFAULT_CLEARANCE_URLS = (
|
||||
"https://accounts.x.ai",
|
||||
"https://x.ai",
|
||||
"https://status.x.ai",
|
||||
"https://console.x.ai",
|
||||
"https://auth.x.ai",
|
||||
)
|
||||
|
||||
_lock = threading.RLock()
|
||||
_cache: dict[str, Any] = {
|
||||
"cookies": [], # Playwright cookie dicts
|
||||
"user_agent": "",
|
||||
"fetched_at": 0.0,
|
||||
"hosts": [],
|
||||
"last_error": "",
|
||||
}
|
||||
|
||||
|
||||
def _env_bool(name: str, default: bool = False) -> bool:
|
||||
raw = str(os.environ.get(name, "")).strip().lower()
|
||||
if not raw:
|
||||
return default
|
||||
if raw in {"1", "true", "yes", "on"}:
|
||||
return True
|
||||
if raw in {"0", "false", "no", "off"}:
|
||||
return False
|
||||
return default
|
||||
|
||||
|
||||
def _env_int(name: str, default: int) -> int:
|
||||
try:
|
||||
return int(str(os.environ.get(name, "")).strip() or default)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def clearance_enabled() -> bool:
|
||||
return _env_bool("CLEARANCE_ENABLED", False)
|
||||
|
||||
|
||||
def flaresolverr_url() -> str:
|
||||
return (os.environ.get("FLARESOLVERR_URL") or "http://127.0.0.1:8191").strip().rstrip("/")
|
||||
|
||||
|
||||
def clearance_timeout_sec() -> int:
|
||||
return max(10, _env_int("CLEARANCE_TIMEOUT_SEC", 60))
|
||||
|
||||
|
||||
def clearance_refresh_sec() -> int:
|
||||
return max(60, _env_int("CLEARANCE_REFRESH_SEC", 3000))
|
||||
|
||||
|
||||
def clearance_urls() -> list[str]:
|
||||
raw = (os.environ.get("CLEARANCE_URLS") or "").strip()
|
||||
if not raw:
|
||||
return list(DEFAULT_CLEARANCE_URLS)
|
||||
urls = []
|
||||
for part in raw.split(","):
|
||||
item = part.strip()
|
||||
if not item:
|
||||
continue
|
||||
if "://" not in item:
|
||||
item = "https://" + item
|
||||
# strip path/query — host-level prewarm only
|
||||
parsed = urlparse(item)
|
||||
if parsed.scheme and parsed.netloc:
|
||||
urls.append(f"{parsed.scheme}://{parsed.netloc}")
|
||||
else:
|
||||
urls.append(item.rstrip("/"))
|
||||
# de-dupe preserve order
|
||||
seen = set()
|
||||
out = []
|
||||
for u in urls:
|
||||
if u not in seen:
|
||||
seen.add(u)
|
||||
out.append(u)
|
||||
return out or list(DEFAULT_CLEARANCE_URLS)
|
||||
|
||||
|
||||
def register_proxy_url() -> str:
|
||||
"""Egress proxy for *this* process (Playwright / httpx)."""
|
||||
for key in ("REGISTER_PROXY", "HTTPS_PROXY", "HTTP_PROXY", "ALL_PROXY"):
|
||||
value = (os.environ.get(key) or "").strip()
|
||||
if value:
|
||||
return value
|
||||
return ""
|
||||
|
||||
|
||||
def clearance_proxy_url() -> str:
|
||||
"""Proxy for FlareSolverr browser (container-reachable). Empty = FS direct egress."""
|
||||
explicit = (os.environ.get("CLEARANCE_PROXY") or "").strip()
|
||||
if explicit:
|
||||
return explicit
|
||||
# Do NOT auto-forward host REGISTER_PROXY into FS — 127.0.0.1 inside
|
||||
# the container is not the host. Operator must set CLEARANCE_PROXY when
|
||||
# minting on WARP (e.g. http://privoxy:8118).
|
||||
return ""
|
||||
|
||||
|
||||
def playwright_proxy_settings() -> dict[str, str] | None:
|
||||
"""Playwright launch/context proxy dict, or None for direct."""
|
||||
url = register_proxy_url()
|
||||
if not url:
|
||||
return None
|
||||
return {"server": url}
|
||||
|
||||
|
||||
def httpx_proxy_mounts() -> str | None:
|
||||
"""Single proxy URL for httpx.AsyncClient(proxy=...), or None."""
|
||||
return register_proxy_url() or None
|
||||
|
||||
|
||||
def _host_from_url(url: str) -> str:
|
||||
try:
|
||||
return (urlparse(url).hostname or "").lower()
|
||||
except Exception:
|
||||
return ""
|
||||
|
||||
|
||||
def _fs_post(payload: dict[str, Any], timeout: float) -> dict[str, Any]:
|
||||
endpoint = f"{flaresolverr_url()}/v1"
|
||||
body = json.dumps(payload).encode("utf-8")
|
||||
req = urllib.request.Request(
|
||||
endpoint,
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=timeout) as response:
|
||||
raw = response.read()
|
||||
return json.loads(raw.decode("utf-8"))
|
||||
|
||||
|
||||
def _cookie_to_playwright(item: dict[str, Any], fallback_host: str) -> dict[str, Any] | None:
|
||||
name = str(item.get("name") or "").strip()
|
||||
if not name:
|
||||
return None
|
||||
value = str(item.get("value") or "")
|
||||
domain = str(item.get("domain") or "").strip()
|
||||
if not domain:
|
||||
domain = fallback_host
|
||||
if domain and not domain.startswith(".") and domain.count(".") >= 1:
|
||||
# keep host-only cookies host-only; CF often uses .x.ai
|
||||
pass
|
||||
path = str(item.get("path") or "/") or "/"
|
||||
cookie: dict[str, Any] = {
|
||||
"name": name,
|
||||
"value": value,
|
||||
"domain": domain,
|
||||
"path": path,
|
||||
}
|
||||
if "httpOnly" in item:
|
||||
cookie["httpOnly"] = bool(item.get("httpOnly"))
|
||||
if "secure" in item:
|
||||
cookie["secure"] = bool(item.get("secure"))
|
||||
else:
|
||||
cookie["secure"] = True
|
||||
same_site = item.get("sameSite") or item.get("same_site")
|
||||
if same_site:
|
||||
# Playwright: Strict | Lax | None
|
||||
normalized = str(same_site).capitalize()
|
||||
if normalized.lower() == "none":
|
||||
normalized = "None"
|
||||
if normalized in {"Strict", "Lax", "None"}:
|
||||
cookie["sameSite"] = normalized
|
||||
expires = item.get("expires")
|
||||
if expires is not None:
|
||||
try:
|
||||
exp = float(expires)
|
||||
# FS may return ms or session -1
|
||||
if exp > 0:
|
||||
if exp > 1e12:
|
||||
exp = exp / 1000.0
|
||||
cookie["expires"] = exp
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
return cookie
|
||||
|
||||
|
||||
def _merge_cookies(existing: list[dict[str, Any]], incoming: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
index: dict[tuple[str, str, str], int] = {}
|
||||
merged: list[dict[str, Any]] = []
|
||||
for cookie in existing:
|
||||
key = (
|
||||
str(cookie.get("name") or ""),
|
||||
str(cookie.get("domain") or ""),
|
||||
str(cookie.get("path") or "/"),
|
||||
)
|
||||
index[key] = len(merged)
|
||||
merged.append(cookie)
|
||||
for cookie in incoming:
|
||||
key = (
|
||||
str(cookie.get("name") or ""),
|
||||
str(cookie.get("domain") or ""),
|
||||
str(cookie.get("path") or "/"),
|
||||
)
|
||||
if key in index:
|
||||
merged[index[key]] = cookie
|
||||
else:
|
||||
index[key] = len(merged)
|
||||
merged.append(cookie)
|
||||
return merged
|
||||
|
||||
|
||||
def prewarm_clearance(*, force: bool = False) -> dict[str, Any]:
|
||||
"""Hit FlareSolverr for each CLEARANCE_URLS host; cache Playwright cookies.
|
||||
|
||||
Safe to call when disabled — returns status without network.
|
||||
"""
|
||||
if not clearance_enabled():
|
||||
return {
|
||||
"enabled": False,
|
||||
"ok": True,
|
||||
"skipped": True,
|
||||
"message": "CLEARANCE_ENABLED=0",
|
||||
"cookies": 0,
|
||||
}
|
||||
|
||||
with _lock:
|
||||
age = time.time() - float(_cache.get("fetched_at") or 0)
|
||||
if (
|
||||
not force
|
||||
and _cache.get("cookies")
|
||||
and age < clearance_refresh_sec()
|
||||
):
|
||||
return {
|
||||
"enabled": True,
|
||||
"ok": True,
|
||||
"cached": True,
|
||||
"age_sec": round(age, 1),
|
||||
"cookies": len(_cache["cookies"]),
|
||||
"hosts": list(_cache.get("hosts") or []),
|
||||
"user_agent": _cache.get("user_agent") or "",
|
||||
}
|
||||
|
||||
urls = clearance_urls()
|
||||
timeout = float(clearance_timeout_sec())
|
||||
fs_proxy = clearance_proxy_url()
|
||||
max_timeout_ms = int(timeout * 1000)
|
||||
merged: list[dict[str, Any]] = []
|
||||
user_agent = ""
|
||||
hosts_ok: list[str] = []
|
||||
errors: list[str] = []
|
||||
per_host: list[dict[str, Any]] = []
|
||||
|
||||
for url in urls:
|
||||
host = _host_from_url(url)
|
||||
payload: dict[str, Any] = {
|
||||
"cmd": "request.get",
|
||||
"url": url,
|
||||
"maxTimeout": max_timeout_ms,
|
||||
}
|
||||
if fs_proxy:
|
||||
payload["proxy"] = {"url": fs_proxy}
|
||||
t0 = time.time()
|
||||
try:
|
||||
data = _fs_post(payload, timeout=timeout + 15)
|
||||
elapsed = round(time.time() - t0, 2)
|
||||
if str(data.get("status") or "").lower() != "ok":
|
||||
errors.append(f"{host}: status={data.get('status')} msg={data.get('message')}")
|
||||
per_host.append({"host": host, "ok": False, "elapsed": elapsed, "error": data.get("message")})
|
||||
continue
|
||||
solution = data.get("solution") if isinstance(data.get("solution"), dict) else {}
|
||||
raw_cookies = solution.get("cookies") or []
|
||||
pw_cookies = []
|
||||
if isinstance(raw_cookies, list):
|
||||
for item in raw_cookies:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
converted = _cookie_to_playwright(item, host)
|
||||
if converted:
|
||||
pw_cookies.append(converted)
|
||||
merged = _merge_cookies(merged, pw_cookies)
|
||||
ua = str(solution.get("userAgent") or "").strip()
|
||||
if ua:
|
||||
user_agent = ua
|
||||
has_cf = any(c.get("name") == "cf_clearance" for c in pw_cookies)
|
||||
hosts_ok.append(host)
|
||||
per_host.append(
|
||||
{
|
||||
"host": host,
|
||||
"ok": True,
|
||||
"elapsed": elapsed,
|
||||
"http": solution.get("status"),
|
||||
"cookies": len(pw_cookies),
|
||||
"cf_clearance": has_cf,
|
||||
}
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 — surface to operator log
|
||||
elapsed = round(time.time() - t0, 2)
|
||||
errors.append(f"{host}: {exc}")
|
||||
per_host.append({"host": host, "ok": False, "elapsed": elapsed, "error": str(exc)})
|
||||
|
||||
with _lock:
|
||||
if merged:
|
||||
_cache["cookies"] = merged
|
||||
_cache["user_agent"] = user_agent
|
||||
_cache["fetched_at"] = time.time()
|
||||
_cache["hosts"] = hosts_ok
|
||||
_cache["last_error"] = "; ".join(errors)
|
||||
elif errors:
|
||||
_cache["last_error"] = "; ".join(errors)
|
||||
|
||||
return {
|
||||
"enabled": True,
|
||||
"ok": bool(merged) or not errors,
|
||||
"cached": False,
|
||||
"cookies": len(merged),
|
||||
"hosts": hosts_ok,
|
||||
"user_agent": user_agent,
|
||||
"errors": errors,
|
||||
"detail": per_host,
|
||||
"flaresolverr": flaresolverr_url(),
|
||||
"clearance_proxy": fs_proxy or "(direct)",
|
||||
"register_proxy": register_proxy_url() or "(direct)",
|
||||
}
|
||||
|
||||
|
||||
def cached_playwright_cookies() -> list[dict[str, Any]]:
|
||||
with _lock:
|
||||
return [dict(c) for c in (_cache.get("cookies") or [])]
|
||||
|
||||
|
||||
def cached_user_agent() -> str:
|
||||
with _lock:
|
||||
return str(_cache.get("user_agent") or "")
|
||||
|
||||
|
||||
async def apply_clearance_to_context(context) -> int:
|
||||
"""Inject cached CF cookies into a Playwright BrowserContext. Returns count."""
|
||||
cookies = cached_playwright_cookies()
|
||||
if not cookies or context is None:
|
||||
return 0
|
||||
try:
|
||||
await context.add_cookies(cookies)
|
||||
return len(cookies)
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
|
||||
def format_prewarm_log(result: dict[str, Any]) -> str:
|
||||
if result.get("skipped"):
|
||||
return "[clearance] 未启用 (CLEARANCE_ENABLED=0)"
|
||||
if result.get("cached"):
|
||||
return (
|
||||
f"[clearance] 复用缓存 cookies={result.get('cookies')} "
|
||||
f"age={result.get('age_sec')}s hosts={','.join(result.get('hosts') or [])}"
|
||||
)
|
||||
detail = result.get("detail") or []
|
||||
parts = []
|
||||
for row in detail:
|
||||
host = row.get("host")
|
||||
if row.get("ok"):
|
||||
cf = "cf+" if row.get("cf_clearance") else "cf-"
|
||||
parts.append(f"{host}:{row.get('elapsed')}s/{cf}")
|
||||
else:
|
||||
parts.append(f"{host}:FAIL")
|
||||
err = result.get("errors") or []
|
||||
suffix = f" errors={len(err)}" if err else ""
|
||||
return (
|
||||
f"[clearance] 预热完成 cookies={result.get('cookies')} "
|
||||
f"fs={result.get('flaresolverr')} reg_proxy={result.get('register_proxy')} "
|
||||
f"fs_proxy={result.get('clearance_proxy')} | " + " ".join(parts) + suffix
|
||||
)
|
||||
0
grok_register/core/__init__.py
Normal file
0
grok_register/core/__init__.py
Normal file
129
grok_register/core/admission.py
Normal file
129
grok_register/core/admission.py
Normal file
@@ -0,0 +1,129 @@
|
||||
"""Local admission gates for CSP workers.
|
||||
|
||||
AdmissionGate is not a scheduler: it never chooses a role and never moves
|
||||
resources. It only lets producers reserve local production intent against
|
||||
inventory watermarks, so fixed S/P workers can self-throttle without a central
|
||||
dispatcher.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
|
||||
class _TProductionLease:
|
||||
__slots__ = ("_gate", "_released")
|
||||
|
||||
def __init__(self, gate):
|
||||
self._gate = gate
|
||||
self._released = False
|
||||
|
||||
async def release(self):
|
||||
if self._released:
|
||||
return
|
||||
self._released = True
|
||||
async with self._gate._cond:
|
||||
self._gate.t_in_progress -= 1
|
||||
self._gate._cond.notify_all()
|
||||
|
||||
|
||||
class _QBatchLease:
|
||||
__slots__ = ("_gate", "count", "_remaining")
|
||||
|
||||
def __init__(self, gate, count):
|
||||
self._gate = gate
|
||||
self.count = count
|
||||
self._remaining = count
|
||||
|
||||
async def release_one(self):
|
||||
if self._remaining <= 0:
|
||||
return
|
||||
async with self._gate._cond:
|
||||
if self._remaining <= 0:
|
||||
return
|
||||
self._remaining -= 1
|
||||
self._gate.q_inflight -= 1
|
||||
self._gate._cond.notify_all()
|
||||
|
||||
async def release_all(self):
|
||||
if self._remaining <= 0:
|
||||
return
|
||||
async with self._gate._cond:
|
||||
if self._remaining <= 0:
|
||||
return
|
||||
self._gate.q_inflight -= self._remaining
|
||||
self._remaining = 0
|
||||
self._gate._cond.notify_all()
|
||||
|
||||
|
||||
class AdmissionGate:
|
||||
"""Watermark-based producer admission for T and Q.
|
||||
|
||||
T production uses high/low hysteresis because local T generation can be
|
||||
heavy and should not keep filling an already sufficient inventory.
|
||||
|
||||
Q production reserves a batch by current demand:
|
||||
q_high - q_depth - q_inflight
|
||||
capped by the worker's max batch. Each reservation is released when the
|
||||
corresponding pending Q requests have reached terminal state.
|
||||
"""
|
||||
|
||||
def __init__(self, inventory, *, t_low, t_high, q_low, q_high):
|
||||
if t_low < 0 or q_low < 0:
|
||||
raise ValueError("low watermarks must be non-negative")
|
||||
if t_high <= 0 or q_high <= 0:
|
||||
raise ValueError("high watermarks must be positive")
|
||||
if t_low > t_high:
|
||||
raise ValueError("t_low must be <= t_high")
|
||||
if q_low > q_high:
|
||||
raise ValueError("q_low must be <= q_high")
|
||||
|
||||
self.inventory = inventory
|
||||
self.t_low = t_low
|
||||
self.t_high = t_high
|
||||
self.q_low = q_low
|
||||
self.q_high = q_high
|
||||
self.t_in_progress = 0
|
||||
self.q_inflight = 0
|
||||
self._t_paused = False
|
||||
self._q_paused = False
|
||||
self._cond = asyncio.Condition()
|
||||
|
||||
async def acquire_t_production(self):
|
||||
"""Reserve one T production attempt."""
|
||||
async with self._cond:
|
||||
while True:
|
||||
total = self.inventory.t_depth + self.t_in_progress
|
||||
if total >= self.t_high:
|
||||
self._t_paused = True
|
||||
if self._t_paused and total <= self.t_low:
|
||||
self._t_paused = False
|
||||
if not self._t_paused and total < self.t_high:
|
||||
self.t_in_progress += 1
|
||||
self._cond.notify_all()
|
||||
return _TProductionLease(self)
|
||||
await self._cond.wait()
|
||||
|
||||
async def acquire_q_batch(self, *, max_batch):
|
||||
"""Reserve up to max_batch pending Q requests based on current demand."""
|
||||
if max_batch <= 0:
|
||||
raise ValueError("max_batch must be positive")
|
||||
|
||||
async with self._cond:
|
||||
while True:
|
||||
total = self.inventory.q_depth + self.q_inflight
|
||||
if total >= self.q_high:
|
||||
self._q_paused = True
|
||||
if self._q_paused and total <= self.q_low:
|
||||
self._q_paused = False
|
||||
demand = self.q_high - total
|
||||
if not self._q_paused and demand > 0:
|
||||
count = min(max_batch, demand)
|
||||
self.q_inflight += count
|
||||
if self.inventory.q_depth + self.q_inflight >= self.q_high:
|
||||
self._q_paused = True
|
||||
self._cond.notify_all()
|
||||
return _QBatchLease(self, count)
|
||||
await self._cond.wait()
|
||||
|
||||
async def notify_changed(self):
|
||||
"""Wake waiters after inventory depth changes."""
|
||||
async with self._cond:
|
||||
self._cond.notify_all()
|
||||
61
grok_register/core/envelope.py
Normal file
61
grok_register/core/envelope.py
Normal file
@@ -0,0 +1,61 @@
|
||||
"""
|
||||
ResourceEnvelope — T/Q 资源实体与库存 slot 的生命周期绑定
|
||||
|
||||
slot 获取与 envelope 创建必须通过 create_with_slot() 完成。
|
||||
release_slot_once() 幂等:重复调用只释放一次。
|
||||
"""
|
||||
import time
|
||||
import asyncio
|
||||
|
||||
|
||||
class ResourceEnvelope:
|
||||
"""绑定一个 T/Q 资源值与其对应的库存 slot。"""
|
||||
|
||||
__slots__ = ('kind', 'value', 'slot_sem', '_released', 'created_at', 'expires_at', 'meta')
|
||||
|
||||
def __init__(self, kind, value, slot_sem, *, created_at, expires_at=None, meta=None):
|
||||
self.kind = kind # 'T' or 'Q'
|
||||
self.value = value # T: turnstile token str; Q: {'email','password','code'}
|
||||
self.slot_sem = slot_sem # T_Slot_Sem or Q_Slot_Sem
|
||||
self._released = False
|
||||
self.created_at = created_at
|
||||
self.expires_at = expires_at
|
||||
self.meta = meta or {}
|
||||
|
||||
@classmethod
|
||||
async def create_with_slot(cls, kind, value, slot_sem, *, expires_at=None, meta=None):
|
||||
"""先获取库存 slot,再构造 envelope。slot 获取失败则抛异常,不创建 envelope。"""
|
||||
await slot_sem.acquire()
|
||||
try:
|
||||
return cls(
|
||||
kind, value, slot_sem,
|
||||
created_at=time.time(),
|
||||
expires_at=expires_at,
|
||||
meta=meta,
|
||||
)
|
||||
except BaseException:
|
||||
slot_sem.release()
|
||||
raise
|
||||
|
||||
def release_slot_once(self, reason=''):
|
||||
"""幂等释放 slot。重复调用安全。"""
|
||||
if not self._released:
|
||||
self._released = True
|
||||
self.slot_sem.release()
|
||||
|
||||
def discard(self, reason=''):
|
||||
"""资源未被消费,释放库存占用。"""
|
||||
self.release_slot_once(reason=f'discard:{reason}')
|
||||
|
||||
def is_expired(self, now=None):
|
||||
"""检查资源是否过期。"""
|
||||
if self.expires_at is None:
|
||||
return False
|
||||
return (now or time.time()) > self.expires_at
|
||||
|
||||
@property
|
||||
def released(self):
|
||||
return self._released
|
||||
|
||||
def __repr__(self):
|
||||
return f'Envelope({self.kind}, released={self._released})'
|
||||
169
grok_register/core/inventory.py
Normal file
169
grok_register/core/inventory.py
Normal file
@@ -0,0 +1,169 @@
|
||||
"""
|
||||
Inventory — 唯一库存门面 + PairLease
|
||||
|
||||
put_t / put_q: 所有权从 Worker 转移到 Inventory
|
||||
claim_pair(): 等待完整 T/Q pair,通过 PairLease 原子转移所有权
|
||||
内部单锁 + Condition,不做分片,不做无锁。
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
from collections import deque
|
||||
|
||||
from .envelope import ResourceEnvelope
|
||||
|
||||
|
||||
class PairLease:
|
||||
"""claim_pair 成功后的 pair 所有权对象。
|
||||
|
||||
用法:
|
||||
async with inventory.claim_pair() as pair:
|
||||
t_val, q_val = pair.t.value, pair.q.value
|
||||
...
|
||||
离开上下文时无论成功/失败/取消,都核销 T/Q 并释放 slot。
|
||||
"""
|
||||
|
||||
__slots__ = ('t', 'q', '_inventory')
|
||||
|
||||
def __init__(self, t_env, q_env, inventory):
|
||||
self.t = t_env
|
||||
self.q = q_env
|
||||
self._inventory = inventory
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
# 第一版策略: 无论结果如何,均核销 T/Q
|
||||
self.t.release_slot_once(reason='pair_settle')
|
||||
self.q.release_slot_once(reason='pair_settle')
|
||||
# 通知可能在等待的 claim_pair
|
||||
async with self._inventory._lock:
|
||||
self._inventory._cond.notify_all()
|
||||
return False
|
||||
|
||||
|
||||
class Inventory:
|
||||
"""T/Q 库存门面。
|
||||
|
||||
内部维护两个 deque 和一个 Lock+Condition。
|
||||
put_t/put_q 在锁内 O(1) 操作;claim_pair 在锁外等待 Condition。
|
||||
"""
|
||||
|
||||
def __init__(self, metrics=None):
|
||||
self._t_buf = deque()
|
||||
self._q_buf = deque()
|
||||
self._lock = asyncio.Lock()
|
||||
self._cond = asyncio.Condition(self._lock)
|
||||
self._metrics = metrics
|
||||
|
||||
# ── 入库 ──
|
||||
|
||||
async def put_t(self, env: ResourceEnvelope):
|
||||
"""将 T envelope 入库。调用前 env 属于 Worker,成功后属于 Inventory。"""
|
||||
async with self._lock:
|
||||
now = time.time()
|
||||
if env.is_expired(now):
|
||||
env.discard(reason='expired_on_put')
|
||||
if self._metrics:
|
||||
self._metrics.t_expired += 1
|
||||
return
|
||||
self._cleanup_expired(now)
|
||||
self._t_buf.append(env)
|
||||
if self._metrics:
|
||||
self._metrics.t_admitted += 1
|
||||
self._cond.notify_all()
|
||||
|
||||
async def put_q(self, env: ResourceEnvelope):
|
||||
"""将 Q envelope 入库。调用前 env 属于 Worker,成功后属于 Inventory。"""
|
||||
async with self._lock:
|
||||
now = time.time()
|
||||
if env.is_expired(now):
|
||||
env.discard(reason='expired_on_put')
|
||||
if self._metrics:
|
||||
self._metrics.q_expired += 1
|
||||
return
|
||||
self._cleanup_expired(now)
|
||||
self._q_buf.append(env)
|
||||
if self._metrics:
|
||||
self._metrics.q_admitted += 1
|
||||
self._cond.notify_all()
|
||||
|
||||
# ── Claim ──
|
||||
|
||||
def claim_pair(self):
|
||||
"""返回 PairLease 的 async context manager。
|
||||
|
||||
等待期间可取消,取消不弹出资源。成功后 T/Q 原子弹出并转移给 PairLease。
|
||||
"""
|
||||
return _PairClaim(self, self._metrics)
|
||||
|
||||
# ── 内部 ──
|
||||
|
||||
def _try_pop_pair(self):
|
||||
"""在锁内尝试弹出一对未过期的 T/Q。返回 (t_env, q_env) 或 None。"""
|
||||
now = time.time()
|
||||
# 清理过期
|
||||
self._cleanup_expired(now)
|
||||
if self._t_buf and self._q_buf:
|
||||
t_env = self._t_buf.popleft()
|
||||
q_env = self._q_buf.popleft()
|
||||
return t_env, q_env
|
||||
return None
|
||||
|
||||
def _cleanup_expired(self, now):
|
||||
"""lazy cleanup: 清理已过期的 T/Q 并释放 slot。锁内调用。"""
|
||||
self._cleanup_buf(self._t_buf, now, 't_expired')
|
||||
self._cleanup_buf(self._q_buf, now, 'q_expired')
|
||||
|
||||
def _cleanup_buf(self, buf, now, metric_name):
|
||||
"""扫描完整库存,保留未过期实体的原始 FIFO 顺序。"""
|
||||
kept = deque()
|
||||
while buf:
|
||||
env = buf.popleft()
|
||||
if env.is_expired(now):
|
||||
env.discard(reason='expired_in_inventory')
|
||||
if self._metrics:
|
||||
setattr(self._metrics, metric_name, getattr(self._metrics, metric_name) + 1)
|
||||
else:
|
||||
kept.append(env)
|
||||
buf.extend(kept)
|
||||
|
||||
@property
|
||||
def t_depth(self):
|
||||
return len(self._t_buf)
|
||||
|
||||
@property
|
||||
def q_depth(self):
|
||||
return len(self._q_buf)
|
||||
|
||||
|
||||
class _PairClaim:
|
||||
"""claim_pair() 返回的 async context manager。"""
|
||||
|
||||
__slots__ = ('_inventory', '_pair', '_metrics')
|
||||
|
||||
def __init__(self, inventory, metrics):
|
||||
self._inventory = inventory
|
||||
self._pair = None
|
||||
self._metrics = metrics
|
||||
|
||||
async def __aenter__(self):
|
||||
inv = self._inventory
|
||||
async with inv._lock:
|
||||
# 循环等待直到有 pair 或被取消
|
||||
while True:
|
||||
pair = inv._try_pop_pair()
|
||||
if pair is not None:
|
||||
t_env, q_env = pair
|
||||
self._pair = PairLease(t_env, q_env, inv)
|
||||
if self._metrics:
|
||||
self._metrics.pair_claimed += 1
|
||||
return self._pair
|
||||
# 没有 pair,等待 notify
|
||||
await inv._cond.wait()
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
if self._pair is not None:
|
||||
# PairLease.__aexit__ 处理核销
|
||||
return await self._pair.__aexit__(exc_type, exc, tb)
|
||||
return False
|
||||
215
grok_register/core/observer.py
Normal file
215
grok_register/core/observer.py
Normal file
@@ -0,0 +1,215 @@
|
||||
"""
|
||||
Observer — 只读观测层
|
||||
|
||||
只记录指标、生成日志。不修改 Semaphore,不调度 Worker,不释放资源。
|
||||
"""
|
||||
import time
|
||||
from collections import deque
|
||||
|
||||
|
||||
class Metrics:
|
||||
"""全局指标收集器。所有字段只在 asyncio 单线程中写入,无需加锁。"""
|
||||
|
||||
__slots__ = (
|
||||
'start_time',
|
||||
't_produced', 't_admitted', 't_claimed', 't_expired', 't_discarded',
|
||||
't_solve_count', 't_solve_seconds', 't_solve_failed',
|
||||
'solver_goto_seconds', 'solver_inject_seconds', 'solver_initial_seconds',
|
||||
'solver_click_seconds', 'solver_wait_seconds', 'solver_reused_count',
|
||||
'solver_visible_frame_count',
|
||||
's_physical_count', 's_physical_wait_seconds', 's_physical_hold_seconds',
|
||||
'p_physical_count', 'p_physical_wait_seconds', 'p_physical_hold_seconds',
|
||||
'c_physical_count', 'c_physical_wait_seconds', 'c_physical_hold_seconds',
|
||||
'p_email_create_count', 'p_email_create_seconds',
|
||||
'p_page_prepare_count', 'p_page_prepare_seconds',
|
||||
'p_send_count', 'p_send_seconds',
|
||||
'c_page_acquire_count', 'c_page_acquire_seconds',
|
||||
'c_verify_count', 'c_verify_seconds',
|
||||
'c_register_count', 'c_register_seconds',
|
||||
'c_hot_page_hits', 'c_hot_page_misses',
|
||||
'q_sent', 'q_returned', 'q_admitted', 'q_claimed', 'q_expired', 'q_discarded',
|
||||
'q_send_batches', 'q_send_batch_items',
|
||||
'pair_claimed', 'pair_consumed_ok', 'pair_consumed_fail',
|
||||
'success_count',
|
||||
'registration_starts',
|
||||
'_clock', 'started_monotonic', 'recent_success_times',
|
||||
)
|
||||
|
||||
def __init__(self, clock=time.monotonic):
|
||||
self._clock = clock
|
||||
self.started_monotonic = clock()
|
||||
self.recent_success_times = deque()
|
||||
self.start_time = time.time()
|
||||
# T 生命周期
|
||||
self.t_produced = 0
|
||||
self.t_admitted = 0
|
||||
self.t_claimed = 0
|
||||
self.t_expired = 0
|
||||
self.t_discarded = 0
|
||||
self.t_solve_count = 0
|
||||
self.t_solve_seconds = 0.0
|
||||
self.t_solve_failed = 0
|
||||
self.solver_goto_seconds = 0.0
|
||||
self.solver_inject_seconds = 0.0
|
||||
self.solver_initial_seconds = 0.0
|
||||
self.solver_click_seconds = 0.0
|
||||
self.solver_wait_seconds = 0.0
|
||||
self.solver_reused_count = 0
|
||||
self.solver_visible_frame_count = 0
|
||||
self.s_physical_count = 0
|
||||
self.s_physical_wait_seconds = 0.0
|
||||
self.s_physical_hold_seconds = 0.0
|
||||
self.p_physical_count = 0
|
||||
self.p_physical_wait_seconds = 0.0
|
||||
self.p_physical_hold_seconds = 0.0
|
||||
self.c_physical_count = 0
|
||||
self.c_physical_wait_seconds = 0.0
|
||||
self.c_physical_hold_seconds = 0.0
|
||||
self.p_email_create_count = 0
|
||||
self.p_email_create_seconds = 0.0
|
||||
self.p_page_prepare_count = 0
|
||||
self.p_page_prepare_seconds = 0.0
|
||||
self.p_send_count = 0
|
||||
self.p_send_seconds = 0.0
|
||||
self.c_page_acquire_count = 0
|
||||
self.c_page_acquire_seconds = 0.0
|
||||
self.c_verify_count = 0
|
||||
self.c_verify_seconds = 0.0
|
||||
self.c_register_count = 0
|
||||
self.c_register_seconds = 0.0
|
||||
self.c_hot_page_hits = 0
|
||||
self.c_hot_page_misses = 0
|
||||
# Q 生命周期
|
||||
self.q_sent = 0
|
||||
self.q_returned = 0
|
||||
self.q_admitted = 0
|
||||
self.q_claimed = 0
|
||||
self.q_expired = 0
|
||||
self.q_discarded = 0
|
||||
self.q_send_batches = 0
|
||||
self.q_send_batch_items = 0
|
||||
# Pair
|
||||
self.pair_claimed = 0
|
||||
self.pair_consumed_ok = 0
|
||||
self.pair_consumed_fail = 0
|
||||
# 成功数
|
||||
self.success_count = 0
|
||||
self.registration_starts = 0
|
||||
|
||||
def next_registration_task(self):
|
||||
self.registration_starts += 1
|
||||
return self.registration_starts
|
||||
|
||||
def record_success(self):
|
||||
self.success_count += 1
|
||||
self.recent_success_times.append(self._clock())
|
||||
|
||||
def five_minute_success_rate(self):
|
||||
now = self._clock()
|
||||
cutoff = now - 300.0
|
||||
while self.recent_success_times and self.recent_success_times[0] < cutoff:
|
||||
self.recent_success_times.popleft()
|
||||
if not self.recent_success_times:
|
||||
return None if self.success_count == 0 else 0.0
|
||||
elapsed = max(1.0, min(300.0, now - self.started_monotonic))
|
||||
return len(self.recent_success_times) * 60.0 / elapsed
|
||||
|
||||
def runtime_average_success_rate(self):
|
||||
if self.success_count == 0:
|
||||
return None
|
||||
elapsed = max(1.0, self._clock() - self.started_monotonic)
|
||||
return self.success_count * 60.0 / elapsed
|
||||
|
||||
def snapshot(self, inventory, sems):
|
||||
"""生成一行监控日志。"""
|
||||
elapsed = time.time() - self.start_time
|
||||
rate = self.success_count / (elapsed / 60) if elapsed > 60 else 0
|
||||
p_batch_avg = (
|
||||
self.q_send_batch_items / self.q_send_batches
|
||||
if self.q_send_batches else 0
|
||||
)
|
||||
t_solve_avg = (
|
||||
self.t_solve_seconds / self.t_solve_count
|
||||
if self.t_solve_count else 0
|
||||
)
|
||||
solver_goto_avg = (
|
||||
self.solver_goto_seconds / self.t_solve_count
|
||||
if self.t_solve_count else 0
|
||||
)
|
||||
solver_inject_avg = (
|
||||
self.solver_inject_seconds / self.t_solve_count
|
||||
if self.t_solve_count else 0
|
||||
)
|
||||
solver_initial_avg = (
|
||||
self.solver_initial_seconds / self.t_solve_count
|
||||
if self.t_solve_count else 0
|
||||
)
|
||||
solver_click_avg = (
|
||||
self.solver_click_seconds / self.t_solve_count
|
||||
if self.t_solve_count else 0
|
||||
)
|
||||
solver_wait_avg = (
|
||||
self.solver_wait_seconds / self.t_solve_count
|
||||
if self.t_solve_count else 0
|
||||
)
|
||||
solver_reuse_ratio = (
|
||||
self.solver_reused_count / self.t_solve_count
|
||||
if self.t_solve_count else 0
|
||||
)
|
||||
solver_visible_ratio = (
|
||||
self.solver_visible_frame_count / self.t_solve_count
|
||||
if self.t_solve_count else 0
|
||||
)
|
||||
s_phys_wait, s_phys_hold = self._avg_pair(
|
||||
self.s_physical_wait_seconds, self.s_physical_hold_seconds, self.s_physical_count
|
||||
)
|
||||
p_phys_wait, p_phys_hold = self._avg_pair(
|
||||
self.p_physical_wait_seconds, self.p_physical_hold_seconds, self.p_physical_count
|
||||
)
|
||||
c_phys_wait, c_phys_hold = self._avg_pair(
|
||||
self.c_physical_wait_seconds, self.c_physical_hold_seconds, self.c_physical_count
|
||||
)
|
||||
p_email_create = self._avg(self.p_email_create_seconds, self.p_email_create_count)
|
||||
p_page_prepare = self._avg(self.p_page_prepare_seconds, self.p_page_prepare_count)
|
||||
p_send_stage = self._avg(self.p_send_seconds, self.p_send_count)
|
||||
c_page_acquire = self._avg(self.c_page_acquire_seconds, self.c_page_acquire_count)
|
||||
c_verify = self._avg(self.c_verify_seconds, self.c_verify_count)
|
||||
c_register = self._avg(self.c_register_seconds, self.c_register_count)
|
||||
p_send_sem = sems.get("p_send")
|
||||
admission = sems.get("admission")
|
||||
p_send_part = f' p_send:{p_send_sem._value}' if p_send_sem is not None else ''
|
||||
admission_part = (
|
||||
f' t_prog:{admission.t_in_progress} q_inflight:{admission.q_inflight}'
|
||||
if admission is not None else ''
|
||||
)
|
||||
return (
|
||||
f'[*] T:{inventory.t_depth} Q:{inventory.q_depth} '
|
||||
f'phys:{sems["physical"]._value}{p_send_part} t_slot:{sems["t_slot"]._value} '
|
||||
f'q_slot:{sems["q_slot"]._value} q_pend:{sems["q_pending"]._value} '
|
||||
f'p_batch:{p_batch_avg:.1f}{admission_part} '
|
||||
f's_phys:{s_phys_wait:.2f}/{s_phys_hold:.2f} '
|
||||
f'p_phys:{p_phys_wait:.2f}/{p_phys_hold:.2f} '
|
||||
f'c_phys:{c_phys_wait:.2f}/{c_phys_hold:.2f} '
|
||||
f'p_stage:{p_email_create:.2f}/{p_page_prepare:.2f}/{p_send_stage:.2f} '
|
||||
f'c_stage:{c_page_acquire:.2f}/{c_verify:.2f}/{c_register:.2f} '
|
||||
f'c_hot:{self.c_hot_page_hits}/{self.c_hot_page_misses} '
|
||||
f't_solve_avg:{t_solve_avg:.1f} t_solve_fail:{self.t_solve_failed} '
|
||||
f'solver_goto:{solver_goto_avg:.2f} solver_inject:{solver_inject_avg:.2f} '
|
||||
f'solver_initial:{solver_initial_avg:.2f} solver_click:{solver_click_avg:.2f} '
|
||||
f'solver_wait:{solver_wait_avg:.2f} solver_reuse:{solver_reuse_ratio:.2f} '
|
||||
f'solver_visible:{solver_visible_ratio:.2f} '
|
||||
f't_prod:{self.t_produced} t_adm:{self.t_admitted} t_exp:{self.t_expired} '
|
||||
f'q_sent:{self.q_sent} q_ret:{self.q_returned} q_adm:{self.q_admitted} q_exp:{self.q_expired} '
|
||||
f'pair:{self.pair_claimed} ok:{self.pair_consumed_ok} fail:{self.pair_consumed_fail} '
|
||||
f'rate:{rate:.1f}/min #{self.success_count}'
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _avg(total, count):
|
||||
return total / count if count else 0
|
||||
|
||||
@staticmethod
|
||||
def _avg_pair(wait_total, hold_total, count):
|
||||
if not count:
|
||||
return 0, 0
|
||||
return wait_total / count, hold_total / count
|
||||
171
grok_register/email_server.py
Normal file
171
grok_register/email_server.py
Normal file
@@ -0,0 +1,171 @@
|
||||
"""
|
||||
邮件接收 API 服务器
|
||||
====================
|
||||
接收 Cloudflare Email Routing 转发的邮件，供注册服务查询。
|
||||
|
||||
端点:
|
||||
POST /webhook — Cloudflare 转发邮件到这里
|
||||
GET /check/<email> — 查询某邮箱的验证码
|
||||
GET /domains — 返回可用域名
|
||||
GET /health — 健康检查
|
||||
|
||||
用法:
|
||||
bash start.sh --email-service
|
||||
EMAIL_DOMAIN=your.domain bash start.sh --email-service --port 8080
|
||||
"""
|
||||
import os, re, json, time, sys
|
||||
try:
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
except ImportError:
|
||||
pass
|
||||
from http.server import HTTPServer, ThreadingHTTPServer, BaseHTTPRequestHandler
|
||||
from urllib.parse import urlparse, parse_qs
|
||||
from threading import Lock
|
||||
|
||||
# 配置
|
||||
DEFAULT_DOMAIN = os.environ.get("EMAIL_DOMAIN", "")
|
||||
DEFAULT_PORT = 8080
|
||||
|
||||
# 存储
|
||||
emails = {} # {email_address: [{"code": "ABC123", "time": timestamp, "raw": "..."}]}
|
||||
emails_lock = Lock()
|
||||
|
||||
# 清理过期邮件（5 分钟）
|
||||
def cleanup_old():
|
||||
now = time.time()
|
||||
with emails_lock:
|
||||
for addr in list(emails.keys()):
|
||||
emails[addr] = [e for e in emails[addr] if now - e['time'] < 300]
|
||||
if not emails[addr]:
|
||||
del emails[addr]
|
||||
|
||||
|
||||
def extract_code(text):
|
||||
"""从邮件内容提取验证码（ABC-DEF 或 ABCDEF 格式）"""
|
||||
# 格式1: ABC-DEF
|
||||
m = re.search(r'>([A-Z0-9]{3}-[A-Z0-9]{3})<', text)
|
||||
if m:
|
||||
return m.group(1).replace('-', '')
|
||||
# 格式2: 直接 6 位
|
||||
m = re.search(r'>([A-Z0-9]{6})<', text)
|
||||
if m:
|
||||
return m.group(1)
|
||||
# 格式3: 正文中的 6 位
|
||||
m = re.search(r'\b([A-Z0-9]{3}-?[A-Z0-9]{3})\b', text)
|
||||
if m:
|
||||
return m.group(1).replace('-', '')
|
||||
return None
|
||||
|
||||
|
||||
class EmailHandler(BaseHTTPRequestHandler):
|
||||
def do_POST(self):
|
||||
if self.path == '/webhook':
|
||||
content_length = int(self.headers.get('Content-Length', 0))
|
||||
body = self.rfile.read(content_length).decode('utf-8')
|
||||
|
||||
try:
|
||||
data = json.loads(body)
|
||||
except:
|
||||
data = {}
|
||||
|
||||
# Cloudflare Email Routing 格式
|
||||
to_addr = data.get('to', data.get('recipient', ''))
|
||||
from_addr = data.get('from', data.get('sender', ''))
|
||||
subject = data.get('subject', '')
|
||||
text = data.get('text', '')
|
||||
html = data.get('html', '')
|
||||
|
||||
# 提取验证码
|
||||
content = f"{subject}\n{text}\n{html}"
|
||||
code = extract_code(content)
|
||||
|
||||
if to_addr and code:
|
||||
with emails_lock:
|
||||
if to_addr not in emails:
|
||||
emails[to_addr] = []
|
||||
emails[to_addr].append({
|
||||
'code': code,
|
||||
'time': time.time(),
|
||||
'from': from_addr,
|
||||
'subject': subject
|
||||
})
|
||||
print(f'[+] {to_addr} code={code}', flush=True)
|
||||
|
||||
self.send_response(200)
|
||||
self.send_header('Content-Type', 'application/json')
|
||||
self.end_headers()
|
||||
self.wfile.write(json.dumps({"ok": True, "code": code}).encode())
|
||||
else:
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
|
||||
def do_GET(self):
|
||||
parsed = urlparse(self.path)
|
||||
path = parsed.path
|
||||
|
||||
if path == '/health':
|
||||
self._json({"status": "ok", "emails": len(emails)})
|
||||
|
||||
elif path == '/domains':
|
||||
self._json({"domains": [DEFAULT_DOMAIN]})
|
||||
|
||||
elif path.startswith('/check/'):
|
||||
addr = path[7:] # 去掉 /check/
|
||||
cleanup_old()
|
||||
with emails_lock:
|
||||
items = emails.get(addr, [])
|
||||
if items:
|
||||
# 返回最新的验证码
|
||||
latest = items[-1]
|
||||
self._json({"code": latest['code'], "from": latest['from']})
|
||||
else:
|
||||
self._json({"code": None})
|
||||
|
||||
elif path == '/list':
|
||||
# 列出所有有邮件的地址（调试用）
|
||||
cleanup_old()
|
||||
with emails_lock:
|
||||
result = {addr: len(msgs) for addr, msgs in emails.items()}
|
||||
self._json(result)
|
||||
|
||||
else:
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
|
||||
def _json(self, data):
|
||||
body = json.dumps(data).encode()
|
||||
self.send_response(200)
|
||||
self.send_header('Content-Type', 'application/json')
|
||||
self.send_header('Content-Length', str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, format, *args):
|
||||
print(f"[HTTP] {args[0] if args else format}", flush=True)
|
||||
|
||||
|
||||
def main():
|
||||
global DEFAULT_DOMAIN
|
||||
port = DEFAULT_PORT
|
||||
domain = DEFAULT_DOMAIN
|
||||
|
||||
for i, arg in enumerate(sys.argv[1:]):
|
||||
if arg == '--port' and i + 2 <= len(sys.argv):
|
||||
port = int(sys.argv[i + 2])
|
||||
if arg == '--domain' and i + 2 <= len(sys.argv):
|
||||
domain = sys.argv[i + 2]
|
||||
|
||||
DEFAULT_DOMAIN = domain
|
||||
|
||||
print(f"[*] Email server starting on :{port}", flush=True)
|
||||
print(f"[*] Domain: {domain}", flush=True)
|
||||
print(f"[*] Webhook: http://0.0.0.0:{port}/webhook", flush=True)
|
||||
print(f"[*] Check: http://localhost:{port}/check/<email>", flush=True)
|
||||
|
||||
server = ThreadingHTTPServer(('0.0.0.0', port), EmailHandler)
|
||||
server.serve_forever()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
2308
grok_register/register.py
Normal file
2308
grok_register/register.py
Normal file
File diff suppressed because it is too large
Load Diff
944
keys/acpa_watchdog.py
Executable file
944
keys/acpa_watchdog.py
Executable file
@@ -0,0 +1,944 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
acpa_watchdog.py — 自动 CPA 凭据整备驻守进程
|
||||
|
||||
职责:
|
||||
1. 盯 xai_enroller (auth-service.sh) 的出货目录:
|
||||
$XAI_ENROLLER_LOCAL_AUTH_DIR/authenticated/xai-*.json
|
||||
默认 ~/Downloads/grok-free-register-auth/authenticated/
|
||||
xai_enroller 产出的 CPA 文件 base_url=api.x.ai/v1 (付费口, 免费号必 402)
|
||||
+ headers=None (缺 grok-cli 身份头) + email=None。
|
||||
2. 每个新号原样不动地「整备」一份到 PROJECT/keys/cpa_ready/xai-*.json:
|
||||
- base_url → https://cli-chat-proxy.grok.com/v1 (免费 Grok 4.5 口)
|
||||
- headers → grok-cli 身份头(x-grok-client-version/identifier 等)
|
||||
- email → 从 id_token /access_token JWT 里挖(便于辨识, 不影响通信用)
|
||||
- 权限 0600 (合规: 含 refresh_token)
|
||||
3. 探活: 打 cli-chat-proxy /v1/responses model=grok-4.5
|
||||
请求对齐 new-api grok-cli 完整指纹（session 族头 + input 数组 + max_output_tokens 极小）。
|
||||
首探 ≤MAX_PROBES(默认 1，省 free 2M 额度)；200 → alive。
|
||||
新号落盘后 ACPA_PROBE_WARMUP_SEC(默认 3s) 再探，避免 mint 后瞬时 permission-denied。
|
||||
单次探针内 403 当场短重试 ACPA_403_IMMEDIATE_RETRIES×SLEEP（默认 2×4s）。
|
||||
仍 403 → chat_denied：不丢，RETEST_AFTER_SEC(默认 180s) 后回测 1 次；
|
||||
回测活 → alive；仍 403 → chat_dead 丢。
|
||||
429 free-usage-exhausted / rate_limited / 401 → 立刻丢。
|
||||
alive 永不复探。
|
||||
4. 状态: keys/cpa_ready/_state.tsv
|
||||
name\tsub\temail\tstatus\tts\tprobes
|
||||
每轮 prune：json 在 cpa_ready/_discarded 都不存在的行直接从 TSV 删掉
|
||||
（删库跑路后不会残留幽灵 alive；盘上有 json 但无 state 会重新进探活）。
|
||||
|
||||
只读 AUTH_DIR, 只写 keys/cpa_ready/。源 CPA 文件不动; 不重铸 device flow;
|
||||
不做 HTTP 灌库 (入库步骤留给灌库脚本, 此处只把号整备到可直接喂 cliproxy / new-api)。
|
||||
|
||||
CLI:
|
||||
python3 acpa_watchdog.py # 常驻轮询 (默认 2s), Ctrl-C 退
|
||||
python3 acpa_watchdog.py --once # 单次扫一轮现存号就退 (CLI 自检)
|
||||
python3 acpa_watchdog.py --interval 4
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# paths — portable: follow project root of this script; AUTH overridable via env
|
||||
# ---------------------------------------------------------------------------
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent.parent # .../grok-free-register
|
||||
# Align with xai_enroller.service DEFAULT_LOCAL_AUTH_DIR
|
||||
# (~/Downloads/grok-free-register-auth/authenticated by default)
|
||||
_AUTH_BASE = Path(
|
||||
os.environ.get(
|
||||
"XAI_ENROLLER_LOCAL_AUTH_DIR",
|
||||
os.environ.get(
|
||||
"XAI_AUTH_SERVICE_LOCAL_DIR",
|
||||
str(Path.home() / "Downloads" / "grok-free-register-auth"),
|
||||
),
|
||||
)
|
||||
).expanduser()
|
||||
AUTH_DIR = _AUTH_BASE / "authenticated"
|
||||
OUT_DIR = PROJECT_ROOT / "keys" / "cpa_ready"
|
||||
DISCARD_DIR = OUT_DIR / "_discarded"
|
||||
STATE_FILE = OUT_DIR / "_state.tsv"
|
||||
|
||||
# grok-cli 静态身份头（对齐 new-api common.GrokCLI* / SetupRequestHeader）
|
||||
# session/req/agent 等 per-request 头只在探针里生成，不写进落盘 json。
|
||||
CLIPROXY_BASE_URL = "https://cli-chat-proxy.grok.com/v1"
|
||||
CLIPROXY_HEADERS = {
|
||||
"x-grok-client-version": "0.2.93",
|
||||
"x-xai-token-auth": "xai-grok-cli",
|
||||
"X-XAI-Token-Auth": "xai-grok-cli",
|
||||
"x-authenticateresponse": "authenticate-response",
|
||||
"x-grok-client-identifier": "grok-shell",
|
||||
"x-compaction-at": "400000",
|
||||
"User-Agent": "grok-shell/0.2.93 (linux; x86_64)",
|
||||
}
|
||||
|
||||
RUN = True
|
||||
|
||||
|
||||
def _sig(signum, _frame):
|
||||
global RUN
|
||||
RUN = False
|
||||
|
||||
|
||||
signal.signal(signal.SIGINT, _sig)
|
||||
signal.signal(signal.SIGTERM, _sig)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
def log(msg: str) -> None:
|
||||
print(f"[acpa {time.strftime('%H:%M:%S')}] {msg}", flush=True)
|
||||
|
||||
|
||||
def b64url_decode(seg: str) -> bytes:
|
||||
seg += "=" * (-len(seg) % 4)
|
||||
return base64.urlsafe_b64decode(seg)
|
||||
|
||||
|
||||
def jwt_payload(token: str) -> dict:
|
||||
if not token or token.count(".") != 2:
|
||||
return {}
|
||||
try:
|
||||
return json.loads(b64url_decode(token.split(".")[1]))
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
|
||||
def load_source(path: Path) -> dict | None:
|
||||
try:
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
except Exception as exc:
|
||||
log(f" ✗ 读 {path.name} 失败: {exc}")
|
||||
return None
|
||||
|
||||
|
||||
def ensure_dir(p: Path) -> None:
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
os.chmod(p, 0o700)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def fallback_email(src: dict) -> str:
|
||||
"""xai_enroller 产的 file email 字段是 None。从 JWT 里挖一个标识。"""
|
||||
for key in ("id_token", "access_token"):
|
||||
pl = jwt_payload(src.get(key) or "")
|
||||
for f in ("email", "preferred_username", "sub"):
|
||||
v = pl.get(f)
|
||||
if v:
|
||||
return str(v)
|
||||
return ""
|
||||
|
||||
|
||||
def finalize(src: dict) -> tuple[str, dict]:
|
||||
"""把 xai_enroller 产的 raw CPA 整备成 cli-chat-proxy 可用格式。返回 (sub, dst)。"""
|
||||
dst = copy.deepcopy(src)
|
||||
dst["base_url"] = CLIPROXY_BASE_URL
|
||||
dst["headers"] = dict(CLIPROXY_HEADERS)
|
||||
# token_endpoint 不变 (auth.x.ai/oauth2/token, refresh 用)
|
||||
dst.setdefault("auth_kind", "oauth")
|
||||
dst.setdefault("type", "xai")
|
||||
if not dst.get("email"):
|
||||
dst["email"] = fallback_email(src)
|
||||
sub = dst.get("sub") or jwt_payload(dst.get("access_token", "")).get("sub", "")
|
||||
dst["sub"] = sub
|
||||
return sub, dst
|
||||
|
||||
|
||||
def write_out(name: str, entry: dict) -> Path:
|
||||
ensure_dir(OUT_DIR)
|
||||
target = OUT_DIR / name
|
||||
tmp = target.with_suffix(target.suffix + ".tmp")
|
||||
tmp.write_text(json.dumps(entry, separators=(",", ":"), ensure_ascii=False),
|
||||
encoding="utf-8")
|
||||
os.replace(tmp, target)
|
||||
os.chmod(target, 0o600)
|
||||
return target
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# probe liveness — 使用项目自带 .venv(已装 curl_cffi)
|
||||
# 完整 grok-cli 指纹；默认每号只首探 1 次（free ~2M，禁止连 ping）
|
||||
# ---------------------------------------------------------------------------
|
||||
_VENV_PY = PROJECT_ROOT / ".venv" / "bin" / "python"
|
||||
_PROBE_SRC = r'''
|
||||
import json, sys, uuid
|
||||
try:
|
||||
from curl_cffi import requests
|
||||
except Exception as e:
|
||||
print("NO_CURL_CFFI", str(e)); sys.exit(3)
|
||||
|
||||
fp = sys.argv[1]
|
||||
d = json.load(open(fp))
|
||||
at = d.get("access_token") or ""
|
||||
if not at:
|
||||
print(json.dumps({"status": "token_dead", "http": 401, "snippet": "no access_token"}))
|
||||
sys.exit(0)
|
||||
|
||||
# 对齐 new-api-fusion relay/channel/grok_cli SetupRequestHeader + common.GrokCLI*
|
||||
sid = str(uuid.uuid4())
|
||||
rid = str(uuid.uuid4())
|
||||
hdrs = {
|
||||
"Authorization": "Bearer " + at,
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "grok-shell/0.2.93 (linux; x86_64)",
|
||||
"x-xai-token-auth": "xai-grok-cli",
|
||||
"X-XAI-Token-Auth": "xai-grok-cli",
|
||||
"x-grok-client-identifier": "grok-shell",
|
||||
"x-grok-client-version": "0.2.93",
|
||||
"x-authenticateresponse": "authenticate-response",
|
||||
"x-compaction-at": "400000",
|
||||
"Connection": "Keep-Alive",
|
||||
"x-grok-session-id": sid,
|
||||
"x-grok-conv-id": sid,
|
||||
"x-grok-req-id": rid,
|
||||
"x-grok-turn-idx": "1",
|
||||
"x-grok-agent-id": "agent-" + rid[:8],
|
||||
"x-grok-model-override": "grok-4.5",
|
||||
}
|
||||
email = d.get("email") or ""
|
||||
sub = d.get("sub") or ""
|
||||
if email:
|
||||
hdrs["x-email"] = str(email)
|
||||
if sub:
|
||||
hdrs["x-userid"] = str(sub)
|
||||
|
||||
base = (d.get("base_url") or "https://cli-chat-proxy.grok.com/v1").rstrip("/")
|
||||
# 兼容 base 已带 /v1 或只有 host
|
||||
if base.endswith("/responses"):
|
||||
url = base
|
||||
else:
|
||||
url = base + "/responses"
|
||||
|
||||
# input 必须是 items 数组；裸字符串会被上游拒 / 误 403
|
||||
# max_output_tokens 压到极小，free 号额度只花探针这一枪
|
||||
body = {
|
||||
"model": "grok-4.5",
|
||||
"store": False,
|
||||
"stream": False,
|
||||
"max_output_tokens": 16,
|
||||
"input": [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "ok"}],
|
||||
}
|
||||
],
|
||||
}
|
||||
try:
|
||||
r = requests.post(url, json=body, headers=hdrs, impersonate="chrome", timeout=45)
|
||||
txt = r.text or ""
|
||||
code = r.status_code
|
||||
except Exception as e:
|
||||
print("EXC", str(e)[:200]); sys.exit(2)
|
||||
|
||||
status = "unknown"
|
||||
low = txt.lower()
|
||||
if code == 200 and ("output" in txt or "completed" in low or '"status"' in low):
|
||||
# 200 即视为通路可用；不要求模型真回 "ok"（省解码/推理）
|
||||
status = "alive"
|
||||
elif (
|
||||
code == 429
|
||||
or "free-usage-exhausted" in low
|
||||
or "used all the include" in low
|
||||
or "rate limit" in low
|
||||
or "rate_limit" in low
|
||||
):
|
||||
status = "free_exhausted" if ("free-usage" in low or "used all" in low) else "rate_limited"
|
||||
elif code == 402 or "personal-team-blocked" in txt or "spending-limit" in txt:
|
||||
status = "no_quota_paid_end"
|
||||
elif code == 403 or "permission-denied" in txt or "chat endpoint is denied" in low:
|
||||
status = "chat_denied"
|
||||
elif code == 401:
|
||||
status = "token_dead"
|
||||
elif code == 422:
|
||||
# body 形态问题：标 unknown 便于发现探针回归，不直接当号死
|
||||
status = "unknown"
|
||||
elif "forbidden" in low:
|
||||
status = "chat_denied"
|
||||
print(json.dumps({"status": status, "http": code, "snippet": txt[:200]}))
|
||||
sys.exit(0)
|
||||
'''
|
||||
|
||||
# 省额度：默认首探 1 次；403 当场短重试后再标 soft，默认 3 分钟后延迟回测 1 次。
|
||||
# alive 永不复探。429 free-usage-exhausted / token_dead 等真坏号立刻丢。
|
||||
MAX_PROBES = int(
|
||||
os.environ.get(
|
||||
"ACPA_MAX_PROBES",
|
||||
os.environ.get("ACPA_PROBE_RETRIES", "1"),
|
||||
)
|
||||
)
|
||||
PROBE_RETRY_SLEEP = float(os.environ.get("ACPA_PROBE_RETRY_SLEEP", "1.0"))
|
||||
# 新号 finalize 后稍等再首探：上游常对刚 mint 的 access 回瞬时 permission-denied
|
||||
PROBE_WARMUP_SEC = float(os.environ.get("ACPA_PROBE_WARMUP_SEC", "3.0"))
|
||||
# 单次 probe_one 内 403 当场短重试（不算额外 MAX_PROBES 预算）
|
||||
IMMEDIATE_403_RETRIES = int(os.environ.get("ACPA_403_IMMEDIATE_RETRIES", "2"))
|
||||
IMMEDIATE_403_SLEEP = float(os.environ.get("ACPA_403_IMMEDIATE_SLEEP", "4.0"))
|
||||
RETEST_AFTER_SEC = int(os.environ.get("ACPA_RETEST_AFTER_SEC", "180")) # 3 分钟
|
||||
RETEST_MAX = int(os.environ.get("ACPA_RETEST_MAX", "1")) # 403 延迟回测次数
|
||||
|
||||
# 终态：不再探（alive 保留；下列坏号 discard）
|
||||
TERMINAL_OK = frozenset({"alive"})
|
||||
TERMINAL_BAD = frozenset({
|
||||
"free_exhausted", # 429 免费额度耗尽 → 丢
|
||||
"rate_limited", # 其它 429 → 丢
|
||||
"token_dead", # 401
|
||||
"no_quota_paid_end",
|
||||
"chat_dead", # 5 分钟回测后仍 403 → 真坏，丢
|
||||
"discarded",
|
||||
})
|
||||
TERMINAL_STATUSES = TERMINAL_OK | TERMINAL_BAD
|
||||
|
||||
# 仅这些立刻/确认后丢出 cpa_ready（403 首探不在此列）
|
||||
DISCARD_STATUSES = frozenset({
|
||||
"free_exhausted",
|
||||
"rate_limited",
|
||||
"token_dead",
|
||||
"no_quota_paid_end",
|
||||
"chat_dead",
|
||||
})
|
||||
|
||||
# 等 5 分钟回测的软状态（保留 json，不进 acc，不丢）
|
||||
SOFT_RETEST_STATUSES = frozenset({
|
||||
"chat_denied",
|
||||
"forbidden_domain_ban", # 历史误标，按 403 软状态回测
|
||||
"probe_cap", # 历史：首探 2 次 403 被钉死，捞回可回测
|
||||
})
|
||||
|
||||
|
||||
def probe_one(path: Path) -> dict:
|
||||
"""用子进程跑探活(隔离 curl_cffi 环境)。返回 {status, http, snippet}."""
|
||||
import subprocess
|
||||
cp = subprocess.run(
|
||||
[str(_VENV_PY), "-c", _PROBE_SRC, str(path)],
|
||||
capture_output=True, text=True, timeout=60,
|
||||
)
|
||||
out = cp.stdout.strip()
|
||||
if out.startswith("alive") or out.lstrip().startswith("{"):
|
||||
try:
|
||||
idx = out.find("{")
|
||||
if idx >= 0:
|
||||
return json.loads(out[idx:])
|
||||
except Exception:
|
||||
pass
|
||||
return {"status": "probe_err", "http": -1, "snippet": (out or cp.stderr[:200])[:200]}
|
||||
|
||||
|
||||
def probe_one_soft(path: Path) -> dict:
|
||||
"""
|
||||
一次「逻辑探针」：底层请求 + 403/瞬时错误当场短重试。
|
||||
仍计为 probe_budgeted 的 1 次 used（不因短重试烧额外额度记账）。
|
||||
"""
|
||||
last = {"status": "probe_err", "http": -1, "snippet": ""}
|
||||
attempts = 1 + max(0, IMMEDIATE_403_RETRIES)
|
||||
for i in range(attempts):
|
||||
last = probe_one(path)
|
||||
st = last.get("status")
|
||||
if st in (
|
||||
"alive", "free_exhausted", "rate_limited",
|
||||
"no_quota_paid_end", "token_dead",
|
||||
):
|
||||
return last
|
||||
if st in ("chat_denied", "forbidden_domain_ban", "probe_err", "unknown"):
|
||||
if i + 1 < attempts:
|
||||
time.sleep(IMMEDIATE_403_SLEEP)
|
||||
continue
|
||||
break
|
||||
if last.get("status") == "forbidden_domain_ban":
|
||||
last = {
|
||||
"status": "chat_denied",
|
||||
"http": last.get("http", 403),
|
||||
"snippet": last.get("snippet", "")[:200],
|
||||
}
|
||||
return last
|
||||
|
||||
|
||||
def probe_budgeted(path: Path, budget: int) -> tuple[dict, int]:
|
||||
"""
|
||||
在 budget 次内探活。返回 (最后结果, 实际探活次数)。
|
||||
alive / 额度终态立刻停；403 可在预算内连探；最终 403 保持 chat_denied（不改 probe_cap）。
|
||||
"""
|
||||
budget = max(0, int(budget))
|
||||
if budget <= 0:
|
||||
return {"status": "chat_denied", "http": -1, "snippet": "no probe budget"}, 0
|
||||
last = {"status": "probe_err", "http": -1, "snippet": ""}
|
||||
used = 0
|
||||
for i in range(budget):
|
||||
last = probe_one_soft(path)
|
||||
used += 1
|
||||
st = last.get("status")
|
||||
if st in (
|
||||
"alive", "free_exhausted", "rate_limited",
|
||||
"no_quota_paid_end", "token_dead",
|
||||
):
|
||||
return last, used
|
||||
if i + 1 < budget and st in (
|
||||
"chat_denied", "forbidden_domain_ban", "probe_err", "unknown",
|
||||
):
|
||||
time.sleep(PROBE_RETRY_SLEEP)
|
||||
continue
|
||||
break
|
||||
# 首探/回测结束仍是 403 → 保持 chat_denied，交给延迟回测或 chat_dead
|
||||
if last.get("status") in ("forbidden_domain_ban",):
|
||||
last = {
|
||||
"status": "chat_denied",
|
||||
"http": last.get("http", 403),
|
||||
"snippet": last.get("snippet", "")[:200],
|
||||
}
|
||||
return last, used
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# state — name \t sub \t email \t status \t ts [\t probes]
|
||||
# ---------------------------------------------------------------------------
|
||||
def load_state() -> dict:
|
||||
"""{name: {sub, email, status, ts, probes}}"""
|
||||
if not STATE_FILE.exists():
|
||||
return {}
|
||||
st = {}
|
||||
for ln in STATE_FILE.read_text(encoding="utf-8").splitlines():
|
||||
parts = ln.split("\t")
|
||||
if len(parts) >= 5:
|
||||
name, sub, email, status, ts = parts[:5]
|
||||
probes = 0
|
||||
if len(parts) >= 6:
|
||||
try:
|
||||
probes = int(parts[5] or 0)
|
||||
except ValueError:
|
||||
probes = 0
|
||||
# 历史行没有 probes：已有终态按已探满处理，避免重启后重烧额度
|
||||
if probes <= 0 and status in TERMINAL_STATUSES | {
|
||||
"chat_denied", "forbidden_domain_ban", "unknown", "probe_err",
|
||||
}:
|
||||
probes = MAX_PROBES if status != "alive" else 1
|
||||
if probes <= 0 and status == "alive":
|
||||
probes = 1
|
||||
st[name] = {
|
||||
"sub": sub, "email": email, "status": status,
|
||||
"ts": ts, "probes": probes,
|
||||
}
|
||||
return st
|
||||
|
||||
|
||||
def write_state(st: dict) -> None:
|
||||
ensure_dir(OUT_DIR)
|
||||
lines = []
|
||||
for name, v in sorted(st.items()):
|
||||
lines.append("\t".join([
|
||||
name, v.get("sub", ""), v.get("email", ""),
|
||||
v.get("status", ""), str(v.get("ts", "")),
|
||||
str(int(v.get("probes") or 0)),
|
||||
]))
|
||||
tmp = STATE_FILE.with_suffix(".tsv.tmp")
|
||||
tmp.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8")
|
||||
os.replace(tmp, STATE_FILE)
|
||||
os.chmod(STATE_FILE, 0o600)
|
||||
|
||||
|
||||
def json_exists_for(name: str) -> bool:
|
||||
"""cpa_ready 主目录或 _discarded 里是否还有该号 json。"""
|
||||
return (OUT_DIR / name).exists() or (DISCARD_DIR / name).exists()
|
||||
|
||||
|
||||
def prune_orphan_state(st: dict) -> int:
|
||||
"""
|
||||
删库跑路 / 手工清 json 后：TSV 里既无 cpa_ready 也无 _discarded 文件的行直接 drop。
|
||||
避免幽灵 alive 把 sync/账本钉死。返回删除行数。
|
||||
"""
|
||||
if not st:
|
||||
return 0
|
||||
drop = [name for name in list(st.keys()) if not json_exists_for(name)]
|
||||
if not drop:
|
||||
return 0
|
||||
for name in drop:
|
||||
del st[name]
|
||||
write_state(st)
|
||||
log(f" 🧹 prune orphan state {len(drop)} 行（盘上无 json）")
|
||||
return len(drop)
|
||||
|
||||
|
||||
def should_discard(status: str, probes: int = 0) -> bool:
|
||||
"""仅真坏号丢：429 额度耗尽 / 401 / 回测后仍 403(chat_dead)。首探 403 不丢。"""
|
||||
return status in DISCARD_STATUSES
|
||||
|
||||
|
||||
def discard_bad(name: str, st: dict, *, reason: str = "") -> bool:
|
||||
"""
|
||||
坏号移出 keys/cpa_ready/ → keys/cpa_ready/_discarded/
|
||||
state 标记 discarded，sync_acc 只认 alive 故 acc.md 自动剔除。
|
||||
返回是否执行了丢弃。
|
||||
"""
|
||||
entry = st.get(name) or {}
|
||||
status = entry.get("status", reason or "bad")
|
||||
probes = _probes(entry) if entry else MAX_PROBES
|
||||
if status == "discarded":
|
||||
# 已丢过：确保 json 不在主目录
|
||||
p = OUT_DIR / name
|
||||
if p.exists():
|
||||
ensure_dir(DISCARD_DIR)
|
||||
dest = DISCARD_DIR / name
|
||||
try:
|
||||
os.replace(p, dest)
|
||||
except OSError:
|
||||
try:
|
||||
p.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
return False
|
||||
if not should_discard(status, probes) and not reason:
|
||||
return False
|
||||
|
||||
ensure_dir(DISCARD_DIR)
|
||||
src = OUT_DIR / name
|
||||
moved = False
|
||||
if src.exists():
|
||||
dest = DISCARD_DIR / name
|
||||
try:
|
||||
if dest.exists():
|
||||
dest.unlink()
|
||||
os.replace(src, dest)
|
||||
os.chmod(dest, 0o600)
|
||||
moved = True
|
||||
except OSError as exc:
|
||||
log(f" ✗ 丢弃移动失败 {name}: {exc}")
|
||||
try:
|
||||
src.unlink()
|
||||
moved = True
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
st[name] = {
|
||||
"sub": entry.get("sub", ""),
|
||||
"email": entry.get("email", ""),
|
||||
"status": "discarded",
|
||||
"ts": int(time.time()),
|
||||
"probes": probes if probes else MAX_PROBES,
|
||||
}
|
||||
# 旁路账本：便于扫为何丢
|
||||
try:
|
||||
ledger = DISCARD_DIR / "_discarded.tsv"
|
||||
with ledger.open("a", encoding="utf-8") as fh:
|
||||
fh.write(
|
||||
f"{int(time.time())}\t{name}\t{entry.get('email','')}\t"
|
||||
f"{status}\t{reason or status}\n"
|
||||
)
|
||||
os.chmod(ledger, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
log(
|
||||
f" 🗑 discard {name} [{status}]"
|
||||
f"{' → _discarded/' if moved else ' (no json)'}"
|
||||
f" email={entry.get('email','')}"
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
def sweep_discard(st: dict) -> int:
|
||||
"""扫 state + cpa_ready，丢掉所有真坏号 json。返回丢弃个数。"""
|
||||
n = 0
|
||||
for name in list(st.keys()):
|
||||
if not RUN:
|
||||
break
|
||||
v = st[name]
|
||||
if v.get("status") == "discarded":
|
||||
discard_bad(name, st) # 清残留 json
|
||||
continue
|
||||
# 软 403 已用完延迟回测次数 → 升格 chat_dead 再丢
|
||||
if v.get("status") in SOFT_RETEST_STATUSES:
|
||||
if _retests_done(v) >= RETEST_MAX and _age_sec(v) >= RETEST_AFTER_SEC:
|
||||
v["status"] = "chat_dead"
|
||||
st[name] = v
|
||||
if should_discard(v.get("status", ""), _probes(v)):
|
||||
if discard_bad(name, st):
|
||||
n += 1
|
||||
write_state(st)
|
||||
return n
|
||||
|
||||
|
||||
def rescue_soft_discards(st: dict) -> int:
|
||||
"""
|
||||
把误丢的 403/probe_cap 从 _discarded 捞回主目录，标 chat_denied，
|
||||
立刻可做一次延迟回测（ts 回拨到已到期）。
|
||||
账本原因为 soft，或 state 已 discarded 但 json 仍在 _discarded 且无硬坏标记。
|
||||
"""
|
||||
ensure_dir(DISCARD_DIR)
|
||||
ensure_dir(OUT_DIR)
|
||||
soft_reasons = SOFT_RETEST_STATUSES | {"probe_cap", "chat_denied", "forbidden_domain_ban"}
|
||||
soft_names: set[str] = set()
|
||||
ledger = DISCARD_DIR / "_discarded.tsv"
|
||||
if ledger.exists():
|
||||
for ln in ledger.read_text(encoding="utf-8").splitlines():
|
||||
parts = ln.split("\t")
|
||||
if len(parts) >= 5:
|
||||
name, prev_st, reason = parts[1], parts[3], parts[4]
|
||||
if prev_st in soft_reasons or reason in soft_reasons:
|
||||
soft_names.add(name)
|
||||
elif len(parts) >= 4:
|
||||
name, prev_st = parts[1], parts[3]
|
||||
if prev_st in soft_reasons:
|
||||
soft_names.add(name)
|
||||
|
||||
n = 0
|
||||
for p in sorted(DISCARD_DIR.glob("xai-*.json")):
|
||||
if not RUN:
|
||||
break
|
||||
name = p.name
|
||||
# 账本标 soft，或 state 缺/已 discarded 时默认按 soft 捞（旧误丢）
|
||||
entry = st.get(name) or {}
|
||||
prev = entry.get("status", "discarded")
|
||||
if name not in soft_names and prev not in ("discarded", "") and prev not in soft_reasons:
|
||||
continue
|
||||
if name not in soft_names and prev == "discarded":
|
||||
# 无账本线索：仍捞（用户确认 403 常误杀）；真 429 已不在 soft 集合时靠账本
|
||||
soft_names.add(name)
|
||||
if name not in soft_names:
|
||||
continue
|
||||
dest = OUT_DIR / name
|
||||
try:
|
||||
if dest.exists():
|
||||
dest.unlink()
|
||||
os.replace(p, dest)
|
||||
os.chmod(dest, 0o600)
|
||||
except OSError as exc:
|
||||
log(f" ✗ rescue 移动失败 {name}: {exc}")
|
||||
continue
|
||||
# ts 回拨：立刻进入可回测窗口（不额外空等 RETEST_AFTER_SEC）
|
||||
st[name] = {
|
||||
"sub": entry.get("sub", ""),
|
||||
"email": entry.get("email", ""),
|
||||
"status": "chat_denied",
|
||||
"ts": int(time.time()) - RETEST_AFTER_SEC,
|
||||
"probes": MAX_PROBES, # 留给 1 次回测额度（与当前 MAX_PROBES 对齐）
|
||||
}
|
||||
n += 1
|
||||
log(f" ↩ rescue {name} → chat_denied (retest due) probes={MAX_PROBES}")
|
||||
if n:
|
||||
write_state(st)
|
||||
return n
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# main loop
|
||||
# ---------------------------------------------------------------------------
|
||||
def _probes(entry: dict) -> int:
|
||||
try:
|
||||
return int(entry.get("probes") or 0)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def _age_sec(entry: dict) -> int:
|
||||
try:
|
||||
return max(0, int(time.time()) - int(entry.get("ts") or 0))
|
||||
except (TypeError, ValueError):
|
||||
return RETEST_AFTER_SEC
|
||||
|
||||
|
||||
def _retests_done(entry: dict) -> int:
|
||||
"""首探预算用 MAX_PROBES；超出部分算延迟回测次数。"""
|
||||
return max(0, _probes(entry) - MAX_PROBES)
|
||||
|
||||
|
||||
def should_probe(name: str, st: dict) -> bool:
|
||||
"""
|
||||
新号：首探（≤MAX_PROBES 次当场连探）。
|
||||
403 软状态：满首探后等 RETEST_AFTER_SEC，再给 RETEST_MAX 次回测。
|
||||
429/真坏/终态：不探。
|
||||
"""
|
||||
if name not in st:
|
||||
return True
|
||||
cur = st[name].get("status", "")
|
||||
if cur == "discarded" or cur in TERMINAL_STATUSES:
|
||||
return False
|
||||
probes = _probes(st[name])
|
||||
|
||||
if cur in SOFT_RETEST_STATUSES:
|
||||
if _retests_done(st[name]) >= RETEST_MAX:
|
||||
return False
|
||||
# 首探未打满（异常中断）→ 先补完首探，不算延迟回测
|
||||
if probes < MAX_PROBES:
|
||||
return True
|
||||
return _age_sec(st[name]) >= RETEST_AFTER_SEC
|
||||
|
||||
# 新号 / pending / 临时错误：未满首探预算可探
|
||||
if probes < MAX_PROBES and cur in ("", "pending", "probe_err", "unknown"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def process_file(name: str, st: dict, *, force_probe: bool = False) -> None:
|
||||
src_path = AUTH_DIR / name
|
||||
# 回测/整备也可只读 cpa_ready 已有 json（rescue 后 AUTH 可能仍在）
|
||||
out_existing = OUT_DIR / name
|
||||
if not src_path.exists() and not out_existing.exists():
|
||||
return
|
||||
|
||||
src = None
|
||||
if src_path.exists():
|
||||
src = load_source(src_path)
|
||||
if (not src or not src.get("access_token")) and out_existing.exists():
|
||||
src = load_source(out_existing)
|
||||
if not src:
|
||||
return
|
||||
if not src.get("access_token"):
|
||||
log(f" ✗ {name} 无 access_token, 跳过")
|
||||
return
|
||||
|
||||
if name in st and st[name].get("status") == "discarded":
|
||||
return
|
||||
|
||||
sub, entry = finalize(src)
|
||||
out_path = write_out(name, entry)
|
||||
|
||||
already = _probes(st[name]) if name in st else 0
|
||||
prev_status = st[name].get("status", "") if name in st else ""
|
||||
is_soft_retest = (
|
||||
prev_status in SOFT_RETEST_STATUSES
|
||||
and already >= MAX_PROBES
|
||||
and _retests_done(st.get(name, {})) < RETEST_MAX
|
||||
)
|
||||
|
||||
if not force_probe and not should_probe(name, st):
|
||||
if name in st and should_discard(st[name].get("status", ""), already):
|
||||
discard_bad(name, st)
|
||||
write_state(st)
|
||||
return
|
||||
# 软状态未到点：静默整备
|
||||
if prev_status in SOFT_RETEST_STATUSES:
|
||||
left = max(0, RETEST_AFTER_SEC - _age_sec(st[name]))
|
||||
if already >= MAX_PROBES and _retests_done(st[name]) < RETEST_MAX:
|
||||
log(
|
||||
f" • {name} 403 待回测 "
|
||||
f"({left}s / retest={_retests_done(st[name])}/{RETEST_MAX})"
|
||||
)
|
||||
return
|
||||
log(
|
||||
f" • {name} 整备覆盖 "
|
||||
f"({prev_status or '-'}/probes={already}) → {out_path.name}"
|
||||
)
|
||||
return
|
||||
|
||||
if force_probe and already >= MAX_PROBES + RETEST_MAX and not os.environ.get("ACPA_FORCE_IGNORE_CAP"):
|
||||
log(f" • {name} 已 probes={already}，跳过强制重探（省额度）")
|
||||
if should_discard(st.get(name, {}).get("status", ""), already):
|
||||
discard_bad(name, st)
|
||||
write_state(st)
|
||||
return
|
||||
|
||||
if is_soft_retest:
|
||||
budget = 1
|
||||
phase = f"403回测×1 (after {RETEST_AFTER_SEC}s, done={_retests_done(st[name])}/{RETEST_MAX})"
|
||||
else:
|
||||
budget = max(0, MAX_PROBES - already)
|
||||
phase = f"首探×≤{budget} (used={already}/{MAX_PROBES})"
|
||||
|
||||
# 新号首探前 warmup：刚 mint 的 access 常被上游瞬时 403
|
||||
if not is_soft_retest and PROBE_WARMUP_SEC > 0:
|
||||
time.sleep(PROBE_WARMUP_SEC)
|
||||
|
||||
# 新号首探前 warmup：刚 mint 的 access 常被上游瞬时 403
|
||||
if not is_soft_retest and PROBE_WARMUP_SEC > 0:
|
||||
time.sleep(PROBE_WARMUP_SEC)
|
||||
|
||||
log(f" ▸ {name} 整备完成, {phase} ...")
|
||||
r, used = probe_budgeted(out_path, budget)
|
||||
status = r.get("status", "probe_err")
|
||||
snip = r.get("snippet", "")
|
||||
new_probes = already + used
|
||||
|
||||
# 延迟回测仍 403 → 真坏 chat_dead
|
||||
if is_soft_retest and status in (
|
||||
"chat_denied", "forbidden_domain_ban", "probe_cap", "probe_err", "unknown",
|
||||
):
|
||||
status = "chat_dead"
|
||||
log(
|
||||
f" {name}: [chat_dead] 回测仍 403/失败 http={r.get('http')} "
|
||||
f"probes={new_probes} {snip[:70]}"
|
||||
)
|
||||
else:
|
||||
log(
|
||||
f" {name}: [{status}] http={r.get('http')} "
|
||||
f"probes={new_probes} {snip[:70]}"
|
||||
)
|
||||
|
||||
st[name] = {
|
||||
"sub": sub,
|
||||
"email": entry.get("email", ""),
|
||||
"status": status,
|
||||
"ts": int(time.time()),
|
||||
"probes": new_probes,
|
||||
}
|
||||
if should_discard(status, new_probes):
|
||||
discard_bad(name, st)
|
||||
write_state(st)
|
||||
|
||||
|
||||
def scan_once(st: dict) -> int:
|
||||
"""扫一轮：新号首探；403 到期回测；429 等真坏丢弃；alive 只刷整备。"""
|
||||
if not AUTH_DIR.exists():
|
||||
log(f"⚠ 出货目录不存在: {AUTH_DIR} (auth-service.sh 启动后才会建)")
|
||||
# 仍处理 cpa_ready 里待回测/rescue 的号
|
||||
pruned = prune_orphan_state(st)
|
||||
if pruned:
|
||||
log(f"本轮 prune 幽灵 state {pruned} 行")
|
||||
discarded = sweep_discard(st)
|
||||
if discarded:
|
||||
log(f"本轮丢弃坏号 {discarded} 个 → {DISCARD_DIR.name}/")
|
||||
|
||||
names: set[str] = set()
|
||||
if AUTH_DIR.exists():
|
||||
names.update(p.name for p in AUTH_DIR.glob("xai-*.json"))
|
||||
# cpa_ready 里所有现存 json 都纳入扫描：
|
||||
# - soft 到期回测
|
||||
# - 无 state（删 TSV 后残留 / 手工丢入）→ 当新号首探
|
||||
# - alive 走 should_probe=False 分支只刷整备
|
||||
if OUT_DIR.exists():
|
||||
names.update(p.name for p in OUT_DIR.glob("xai-*.json"))
|
||||
|
||||
n = 0
|
||||
for name in sorted(names):
|
||||
if not RUN:
|
||||
break
|
||||
if name in st and st[name].get("status") == "discarded":
|
||||
continue
|
||||
if name in st and not should_probe(name, st):
|
||||
if should_discard(st[name].get("status", ""), _probes(st[name])):
|
||||
discard_bad(name, st)
|
||||
write_state(st)
|
||||
continue
|
||||
src_path = AUTH_DIR / name if (AUTH_DIR / name).exists() else OUT_DIR / name
|
||||
src = load_source(src_path) if src_path.exists() else None
|
||||
if src and src.get("access_token"):
|
||||
_, entry = finalize(src)
|
||||
write_out(name, entry)
|
||||
continue
|
||||
process_file(name, st)
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(description="auto CPA 整备驻守")
|
||||
ap.add_argument("--interval", type=float, default=2.0)
|
||||
ap.add_argument("--once", action="store_true", help="扫一轮就退")
|
||||
ap.add_argument(
|
||||
"--reprobe-soft",
|
||||
action="store_true",
|
||||
help="立刻对 chat_denied/probe_cap 等软 403 做回测（仍 1 次/号，受 RETEST_MAX）",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--no-rescue",
|
||||
action="store_true",
|
||||
help="启动时不从 _discarded 捞回误丢的 403",
|
||||
)
|
||||
args = ap.parse_args()
|
||||
|
||||
ensure_dir(OUT_DIR)
|
||||
ensure_dir(DISCARD_DIR)
|
||||
st = load_state()
|
||||
n_prune = prune_orphan_state(st)
|
||||
|
||||
# 历史 probe_cap 并入 soft，可走延迟回测
|
||||
for name, v in st.items():
|
||||
if v.get("status") == "probe_cap":
|
||||
v["status"] = "chat_denied"
|
||||
|
||||
soft_n = sum(1 for v in st.values() if v.get("status") in SOFT_RETEST_STATUSES)
|
||||
bad_n = sum(
|
||||
1 for v in st.values()
|
||||
if v.get("status") != "discarded"
|
||||
and should_discard(v.get("status", ""), _probes(v))
|
||||
)
|
||||
log(
|
||||
f"启动 | 出货={AUTH_DIR} | 落盘={OUT_DIR} | state={len(st)} 条 "
|
||||
f"(启动 prune {n_prune}) | 首探≤{MAX_PROBES}/号 "
|
||||
f"| warmup={PROBE_WARMUP_SEC}s 403即重试={IMMEDIATE_403_RETRIES}×{IMMEDIATE_403_SLEEP}s "
|
||||
f"| 403回测 after={RETEST_AFTER_SEC}s ×{RETEST_MAX} "
|
||||
f"| soft_403={soft_n} | pending_discard={bad_n}"
|
||||
)
|
||||
|
||||
if not args.no_rescue:
|
||||
n_res = rescue_soft_discards(st)
|
||||
if n_res:
|
||||
log(f"启动捞回误丢 403 {n_res} 个（将回测）")
|
||||
|
||||
n0 = sweep_discard(st)
|
||||
if n0:
|
||||
log(f"启动丢弃真坏号 {n0} 个 → {DISCARD_DIR}/")
|
||||
|
||||
if st and not args.reprobe_soft:
|
||||
log(
|
||||
f"(state 已有: 刷 alive 整备；403 满 {RETEST_AFTER_SEC}s 回测×{RETEST_MAX}；"
|
||||
f"429 free_exhausted 立刻丢)"
|
||||
)
|
||||
for name in list(st.keys()):
|
||||
if not RUN:
|
||||
break
|
||||
if st[name].get("status") != "alive":
|
||||
continue
|
||||
src_path = AUTH_DIR / name
|
||||
if src_path.exists():
|
||||
src = load_source(src_path)
|
||||
if src and src.get("access_token"):
|
||||
_, entry = finalize(src)
|
||||
write_out(name, entry)
|
||||
|
||||
if args.reprobe_soft:
|
||||
# soft 立刻回测一次：拨到期 + 把 probes 收成 MAX_PROBES，
|
||||
# 避免旧版 probes=2 在新默认 MAX_PROBES=1 下被算成 retest 已用尽。
|
||||
# 只动 soft；alive 绝不复探（省 free 额度）。
|
||||
now = int(time.time())
|
||||
targets = []
|
||||
for name, v in st.items():
|
||||
if v.get("status") in SOFT_RETEST_STATUSES:
|
||||
v["ts"] = now - RETEST_AFTER_SEC
|
||||
v["probes"] = MAX_PROBES
|
||||
targets.append(name)
|
||||
write_state(st)
|
||||
log(f"--reprobe-soft: {len(targets)} 个 403 软状态立刻回测×1（alive 不动）")
|
||||
for name in targets:
|
||||
if not RUN:
|
||||
break
|
||||
process_file(name, st, force_probe=True)
|
||||
write_state(st)
|
||||
alive = sum(1 for v in st.values() if v.get("status") == "alive")
|
||||
dead = sum(
|
||||
1 for v in st.values()
|
||||
if v.get("status") in DISCARD_STATUSES or v.get("status") == "discarded"
|
||||
)
|
||||
soft = sum(1 for v in st.values() if v.get("status") in SOFT_RETEST_STATUSES)
|
||||
log(f"回测结束 | alive={alive} soft_403={soft} dead/discard={dead} state={len(st)}")
|
||||
return 0
|
||||
|
||||
while RUN:
|
||||
try:
|
||||
n = scan_once(st)
|
||||
if n:
|
||||
log(f"本轮处理 {n} 个号（新号/回测）")
|
||||
except Exception as exc:
|
||||
log(f"轮询异常: {type(exc).__name__} {exc}")
|
||||
if args.once:
|
||||
break
|
||||
slept = 0.0
|
||||
while slept < args.interval and RUN:
|
||||
time.sleep(0.25)
|
||||
slept += 0.25
|
||||
|
||||
log(f"退出 | state={len(st)} 条 写 {STATE_FILE}")
|
||||
write_state(st)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
188
keys/sync_acc.py
Executable file
188
keys/sync_acc.py
Executable file
@@ -0,0 +1,188 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
sync_acc.py — 把 keys/cpa_ready/ 里「探活通过」号的 access_token 同步进 keys/acc.md
|
||||
|
||||
默认只入库 acpa_watchdog 记为 alive 的号（读 keys/cpa_ready/_state.tsv）。
|
||||
403 forbidden / token_dead / probe_err 等不进 acc.md，避免坏号灌库。
|
||||
|
||||
用法:
|
||||
python3 sync_acc.py # 常驻轮询(默认 5s), Ctrl-C 退
|
||||
python3 sync_acc.py --once # 同步一次就退
|
||||
python3 sync_acc.py --all # 忽略探活状态，全量抽 token（排障用）
|
||||
python3 sync_acc.py --interval 10
|
||||
|
||||
acc.md 每轮全量重写(去重)，0600。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
PROJECT = Path(__file__).resolve().parent.parent # grok-free-register/
|
||||
SRC_DIR = PROJECT / "keys" / "cpa_ready"
|
||||
STATE_FILE = SRC_DIR / "_state.tsv"
|
||||
OUT_FILE = PROJECT / "keys" / "acc.md"
|
||||
|
||||
# 默认可入库的探活状态（discarded/坏号不在 cpa_ready 主目录，也不在此集合）
|
||||
ALIVE_STATUSES = frozenset({"alive"})
|
||||
SKIP_STATUSES = frozenset({
|
||||
"discarded", "free_exhausted", "rate_limited", "token_dead",
|
||||
"no_quota_paid_end", "probe_cap", "chat_denied", "chat_dead",
|
||||
"forbidden_domain_ban", "probe_err", "unknown",
|
||||
})
|
||||
|
||||
RUN = True
|
||||
|
||||
|
||||
def _sig(_s, _f):
|
||||
global RUN
|
||||
RUN = False
|
||||
|
||||
|
||||
signal.signal(signal.SIGINT, _sig)
|
||||
signal.signal(signal.SIGTERM, _sig)
|
||||
|
||||
|
||||
def log(msg: str) -> None:
|
||||
print(f"[sync_acc {time.strftime('%H:%M:%S')}] {msg}", flush=True)
|
||||
|
||||
|
||||
def load_alive_names() -> set[str] | None:
|
||||
"""
|
||||
读 _state.tsv → 状态为 alive 且盘上仍有 json 的文件名集合。
|
||||
无 state 文件返回 None。幽灵 alive（TSV 有、json 无）不计。
|
||||
"""
|
||||
if not STATE_FILE.exists():
|
||||
return None
|
||||
alive: set[str] = set()
|
||||
for ln in STATE_FILE.read_text(encoding="utf-8").splitlines():
|
||||
parts = ln.split("\t")
|
||||
if len(parts) < 4:
|
||||
continue
|
||||
name, status = parts[0], parts[3]
|
||||
if status in ALIVE_STATUSES:
|
||||
# 只认盘上真实存在的号，避免删库后 state 幽灵把计数钉死
|
||||
if (SRC_DIR / name).exists():
|
||||
alive.add(name)
|
||||
# discarded / 坏状态明确不进（双保险；主目录 json 也应已被 watchdog 挪走）
|
||||
elif status in SKIP_STATUSES:
|
||||
continue
|
||||
return alive
|
||||
|
||||
|
||||
def collect(*, require_alive: bool) -> tuple[list[str], dict[str, int]]:
|
||||
"""
|
||||
返回 (去重 token 列表, 计数)。
|
||||
require_alive=True 时只收 _state.tsv 标为 alive 且文件仍在的 xai-*.json。
|
||||
"""
|
||||
files = sorted(glob.glob(str(SRC_DIR / "xai-*.json")))
|
||||
alive_names = load_alive_names() if require_alive else None
|
||||
|
||||
seen: set[str] = set()
|
||||
out: list[str] = []
|
||||
stats = {
|
||||
"files": len(files),
|
||||
"alive_state": len(alive_names) if alive_names is not None else -1,
|
||||
"skipped_not_alive": 0,
|
||||
"no_tok": 0,
|
||||
"bad": 0,
|
||||
"no_state_skip_all": 0,
|
||||
}
|
||||
|
||||
if require_alive and alive_names is None:
|
||||
# 有 cpa json 但还没 state：宁可不灌，避免 403 进库
|
||||
if files:
|
||||
log(f"⚠ 无 {STATE_FILE.name}，拒绝全量入库（等 watchdog 探活）")
|
||||
stats["no_state_skip_all"] = len(files)
|
||||
return [], stats
|
||||
|
||||
for f in files:
|
||||
name = Path(f).name
|
||||
if require_alive and alive_names is not None and name not in alive_names:
|
||||
stats["skipped_not_alive"] += 1
|
||||
continue
|
||||
try:
|
||||
d = json.loads(Path(f).read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
stats["bad"] += 1
|
||||
continue
|
||||
tok = d.get("access_token")
|
||||
if not isinstance(tok, str) or not tok:
|
||||
stats["no_tok"] += 1
|
||||
continue
|
||||
if tok in seen:
|
||||
continue
|
||||
seen.add(tok)
|
||||
out.append(tok)
|
||||
return out, stats
|
||||
|
||||
|
||||
def write_acc(tokens: list[str]) -> None:
|
||||
OUT_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
# 固定同目录临时文件，避免 with_suffix 边界情况 + 保证 replace 原子
|
||||
tmp = OUT_FILE.parent / (OUT_FILE.name + ".tmp")
|
||||
tmp.write_text("\n".join(tokens) + ("\n" if tokens else ""), encoding="utf-8")
|
||||
os.replace(tmp, OUT_FILE)
|
||||
try:
|
||||
os.chmod(OUT_FILE, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def sync_once(*, require_alive: bool) -> int:
|
||||
toks, st = collect(require_alive=require_alive)
|
||||
write_acc(toks)
|
||||
mode = "alive-only" if require_alive else "all"
|
||||
log(
|
||||
f"sync → {OUT_FILE.name}: {len(toks)} tokens [{mode}] "
|
||||
f"(files={st['files']} state_alive={st['alive_state']} "
|
||||
f"skip_dead={st['skipped_not_alive']} no_tok={st['no_tok']} bad={st['bad']})"
|
||||
)
|
||||
return len(toks)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(
|
||||
description="cpa_ready access_token → keys/acc.md (default: alive only)"
|
||||
)
|
||||
ap.add_argument("--interval", type=float, default=5.0)
|
||||
ap.add_argument("--once", action="store_true", help="刷一次就退")
|
||||
ap.add_argument(
|
||||
"--all",
|
||||
action="store_true",
|
||||
help="忽略探活，全量抽 token（排障；会把 403 也写入 acc.md）",
|
||||
)
|
||||
args = ap.parse_args()
|
||||
require_alive = not args.all
|
||||
|
||||
if not SRC_DIR.exists():
|
||||
log(f"⚠ 目录不存在: {SRC_DIR} (watchdog 整备后才会建)")
|
||||
if args.once:
|
||||
return 0
|
||||
|
||||
sync_once(require_alive=require_alive)
|
||||
if args.once:
|
||||
return 0
|
||||
|
||||
log(f"常驻轮询 interval={args.interval}s alive_only={require_alive} (Ctrl-C 退)")
|
||||
while RUN:
|
||||
slept = 0.0
|
||||
while slept < args.interval and RUN:
|
||||
time.sleep(0.25)
|
||||
slept += 0.25
|
||||
if not RUN:
|
||||
break
|
||||
sync_once(require_alive=require_alive)
|
||||
|
||||
log("退出")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
7
requirements.txt
Normal file
7
requirements.txt
Normal file
@@ -0,0 +1,7 @@
|
||||
cloakbrowser>=0.3.0
|
||||
requests>=2.31.0
|
||||
python-dotenv>=1.0.0
|
||||
httpx>=0.28
|
||||
playwright>=1.55
|
||||
curl_cffi>=0.7.0
|
||||
curl_cffi>=0.7.0
|
||||
222
reset_pipeline.sh
Executable file
222
reset_pipeline.sh
Executable file
@@ -0,0 +1,222 @@
|
||||
#!/usr/bin/env bash
|
||||
# 一键清零注册 → 认证 → CPA 整条流水线的运行时数据，方便重开。
|
||||
#
|
||||
# 你只清 keys/acc.md / accounts.txt 不够：
|
||||
# acpa_watchdog 会从 ~/Downloads/grok-free-register-auth/authenticated/
|
||||
# 把旧号再整备进 keys/cpa_ready/，sync_acc 再写回 acc.md。
|
||||
#
|
||||
# 用法（本文件在项目根）:
|
||||
# bash reset_pipeline.sh # dry-run，只列将删内容
|
||||
# bash reset_pipeline.sh --yes # 真正删除
|
||||
# bash reset_pipeline.sh --yes --keep-register
|
||||
# # 只清认证/CPA，保留 keys 里新注册的 SSO 源
|
||||
# bash reset_pipeline.sh --yes --keep-cpa
|
||||
# # 只清认证账本/源快照，保留 cpa_ready（一般不需要）
|
||||
#
|
||||
# 建议先停 auth-service / register，再跑本脚本。
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
ROOT="$(cd "$(dirname "$0")" && pwd)"
|
||||
KEYS="$ROOT/keys"
|
||||
AUTH_DIR="${XAI_AUTH_SERVICE_LOCAL_DIR:-$HOME/Downloads/grok-free-register-auth}"
|
||||
|
||||
YES=0
|
||||
KEEP_REGISTER=0
|
||||
KEEP_CPA=0
|
||||
KEEP_AUTH_SALT=1
|
||||
|
||||
for arg in "$@"; do
|
||||
case "$arg" in
|
||||
--yes|-y) YES=1 ;;
|
||||
--keep-register) KEEP_REGISTER=1 ;;
|
||||
--keep-cpa) KEEP_CPA=1 ;;
|
||||
--wipe-salt) KEEP_AUTH_SALT=0 ;;
|
||||
-h|--help)
|
||||
sed -n '2,20p' "$0"
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "未知参数: $arg" >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# 运行时产物清单（脚本本身永不删）
|
||||
REGISTER_FILES=(
|
||||
"$KEYS/accounts.txt"
|
||||
"$KEYS/grok.txt"
|
||||
"$KEYS/auth-sessions.jsonl"
|
||||
"$KEYS/acc.md"
|
||||
)
|
||||
CPA_GLOBS=(
|
||||
"$KEYS/cpa_ready/xai-"*.json
|
||||
"$KEYS/cpa_ready/_state.tsv"
|
||||
)
|
||||
AUTH_PATHS=(
|
||||
"$AUTH_DIR/authenticated"
|
||||
"$AUTH_DIR/claimed"
|
||||
"$AUTH_DIR/enrollment-ledger.db"
|
||||
"$AUTH_DIR/source-snapshot.jsonl"
|
||||
)
|
||||
AUTH_SALT="$AUTH_DIR/.ledger-salt"
|
||||
|
||||
count_matches() {
|
||||
local n=0
|
||||
local p
|
||||
for p in "$@"; do
|
||||
# shellcheck disable=SC2086
|
||||
if compgen -G "$p" > /dev/null 2>&1; then
|
||||
while IFS= read -r -d '' f; do
|
||||
n=$((n + 1))
|
||||
done < <(find $p -type f -print0 2>/dev/null || true)
|
||||
# also count plain files matched by glob/file
|
||||
if [[ -f "$p" ]]; then
|
||||
n=$((n + 1))
|
||||
fi
|
||||
elif [[ -f "$p" ]]; then
|
||||
n=$((n + 1))
|
||||
elif [[ -d "$p" ]]; then
|
||||
while IFS= read -r -d '' f; do
|
||||
n=$((n + 1))
|
||||
done < <(find "$p" -type f -print0 2>/dev/null || true)
|
||||
fi
|
||||
done
|
||||
echo "$n"
|
||||
}
|
||||
|
||||
size_of() {
|
||||
local p="$1"
|
||||
if [[ -e "$p" ]]; then
|
||||
du -sh "$p" 2>/dev/null | awk '{print $1}'
|
||||
else
|
||||
echo "-"
|
||||
fi
|
||||
}
|
||||
|
||||
echo "=== grok-free-register 流水线清零 ==="
|
||||
echo "ROOT = $ROOT"
|
||||
echo "KEYS = $KEYS"
|
||||
echo "AUTH_DIR = $AUTH_DIR"
|
||||
echo
|
||||
|
||||
# 检测可能还在跑的服务
|
||||
if pgrep -af 'xai_enroller|acpa_watchdog|sync_acc|grok_register.register' 2>/dev/null | grep -v "reset_pipeline\|pgrep\|grep" >/dev/null; then
|
||||
echo "[!] 检测到相关进程仍在运行:"
|
||||
pgrep -af 'xai_enroller|acpa_watchdog|sync_acc|grok_register.register' 2>/dev/null | grep -v "reset_pipeline\|pgrep" || true
|
||||
echo " 建议先 Ctrl-C 停掉 auth-service / register，再 --yes"
|
||||
echo
|
||||
fi
|
||||
|
||||
echo "-- 将处理的目标 --"
|
||||
if [[ "$KEEP_REGISTER" -eq 0 ]]; then
|
||||
for f in "${REGISTER_FILES[@]}"; do
|
||||
if [[ -e "$f" ]]; then
|
||||
echo " [register] $(size_of "$f") $f"
|
||||
fi
|
||||
done
|
||||
else
|
||||
echo " [register] 跳过 (--keep-register)"
|
||||
fi
|
||||
|
||||
if [[ "$KEEP_CPA" -eq 0 ]]; then
|
||||
cpa_n=$(find "$KEYS/cpa_ready" -name 'xai-*.json' -type f 2>/dev/null | wc -l | tr -d ' ')
|
||||
echo " [cpa] ${cpa_n} 个 xai-*.json @ $KEYS/cpa_ready/"
|
||||
[[ -f "$KEYS/cpa_ready/_state.tsv" ]] && echo " [cpa] _state.tsv"
|
||||
else
|
||||
echo " [cpa] 跳过 (--keep-cpa)"
|
||||
fi
|
||||
|
||||
auth_n=$(find "$AUTH_DIR/authenticated" -type f 2>/dev/null | wc -l | tr -d ' ')
|
||||
echo " [auth] authenticated/ ${auth_n} 文件 ($(size_of "$AUTH_DIR/authenticated"))"
|
||||
[[ -d "$AUTH_DIR/claimed" ]] && echo " [auth] claimed/ $(find "$AUTH_DIR/claimed" -type f 2>/dev/null | wc -l | tr -d ' ') 文件"
|
||||
[[ -f "$AUTH_DIR/enrollment-ledger.db" ]] && echo " [auth] enrollment-ledger.db $(size_of "$AUTH_DIR/enrollment-ledger.db")"
|
||||
[[ -f "$AUTH_DIR/source-snapshot.jsonl" ]] && echo " [auth] source-snapshot.jsonl $(size_of "$AUTH_DIR/source-snapshot.jsonl") lines=$(wc -l < "$AUTH_DIR/source-snapshot.jsonl" 2>/dev/null | tr -d ' ')"
|
||||
if [[ "$KEEP_AUTH_SALT" -eq 0 && -f "$AUTH_SALT" ]]; then
|
||||
echo " [auth] .ledger-salt (将删，指纹会重算)"
|
||||
else
|
||||
echo " [auth] .ledger-salt 保留"
|
||||
fi
|
||||
echo
|
||||
|
||||
if [[ "$YES" -ne 1 ]]; then
|
||||
echo "dry-run 模式，未删除任何东西。"
|
||||
echo "确认后执行: bash reset_pipeline.sh --yes"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
rm_file() {
|
||||
local f="$1"
|
||||
if [[ -f "$f" || -L "$f" ]]; then
|
||||
rm -f -- "$f"
|
||||
echo " removed file $f"
|
||||
fi
|
||||
}
|
||||
|
||||
wipe_dir_contents() {
|
||||
local d="$1"
|
||||
if [[ -d "$d" ]]; then
|
||||
# 只清内容，保留目录本身
|
||||
find "$d" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} +
|
||||
echo " wiped dir $d/"
|
||||
fi
|
||||
}
|
||||
|
||||
echo "-- 执行删除 --"
|
||||
|
||||
if [[ "$KEEP_REGISTER" -eq 0 ]]; then
|
||||
for f in "${REGISTER_FILES[@]}"; do
|
||||
rm_file "$f"
|
||||
done
|
||||
# 重建空占位，避免下游 open 报错
|
||||
: > "$KEYS/accounts.txt"
|
||||
: > "$KEYS/grok.txt"
|
||||
: > "$KEYS/auth-sessions.jsonl"
|
||||
: > "$KEYS/acc.md"
|
||||
chmod 600 "$KEYS/acc.md" "$KEYS/auth-sessions.jsonl" 2>/dev/null || true
|
||||
echo " recreated empty accounts.txt / grok.txt / auth-sessions.jsonl / acc.md"
|
||||
fi
|
||||
|
||||
if [[ "$KEEP_CPA" -eq 0 ]]; then
|
||||
mkdir -p "$KEYS/cpa_ready"
|
||||
find "$KEYS/cpa_ready" -type f \( -name 'xai-*.json' -o -name '_state.tsv' \) -delete 2>/dev/null || true
|
||||
echo " wiped $KEYS/cpa_ready/xai-*.json + _state.tsv"
|
||||
fi
|
||||
|
||||
# 认证侧：这是旧号回灌的真正源头
|
||||
mkdir -p "$AUTH_DIR/authenticated" "$AUTH_DIR/claimed"
|
||||
wipe_dir_contents "$AUTH_DIR/authenticated"
|
||||
wipe_dir_contents "$AUTH_DIR/claimed"
|
||||
rm_file "$AUTH_DIR/enrollment-ledger.db"
|
||||
rm_file "$AUTH_DIR/source-snapshot.jsonl"
|
||||
# 偶发的 sqlite 旁路文件
|
||||
rm_file "$AUTH_DIR/enrollment-ledger.db-wal"
|
||||
rm_file "$AUTH_DIR/enrollment-ledger.db-shm"
|
||||
rm_file "$AUTH_DIR/enrollment-ledger.db-journal"
|
||||
|
||||
if [[ "$KEEP_AUTH_SALT" -eq 0 ]]; then
|
||||
rm_file "$AUTH_SALT"
|
||||
fi
|
||||
|
||||
# 兜底：项目内若误放了同名目录也清
|
||||
for extra in "$KEYS/authenticated" "$KEYS/claimed"; do
|
||||
if [[ -d "$extra" ]]; then
|
||||
wipe_dir_contents "$extra"
|
||||
fi
|
||||
done
|
||||
|
||||
echo
|
||||
echo "-- 清零后状态 --"
|
||||
echo " accounts.txt lines = $(wc -l < "$KEYS/accounts.txt" 2>/dev/null | tr -d ' ')"
|
||||
echo " auth-sessions = $(wc -l < "$KEYS/auth-sessions.jsonl" 2>/dev/null | tr -d ' ')"
|
||||
echo " acc.md tokens = $(wc -l < "$KEYS/acc.md" 2>/dev/null | tr -d ' ')"
|
||||
echo " cpa_ready json = $(find "$KEYS/cpa_ready" -name 'xai-*.json' -type f 2>/dev/null | wc -l | tr -d ' ')"
|
||||
echo " authenticated json = $(find "$AUTH_DIR/authenticated" -type f 2>/dev/null | wc -l | tr -d ' ')"
|
||||
echo " ledger.db = $([[ -f "$AUTH_DIR/enrollment-ledger.db" ]] && echo present || echo gone)"
|
||||
echo " snapshot = $([[ -f "$AUTH_DIR/source-snapshot.jsonl" ]] && echo present || echo gone)"
|
||||
echo
|
||||
echo "完成。重开顺序建议:"
|
||||
echo " 1) bash start.sh # 注册出 SSO → keys/auth-sessions.jsonl + accounts.txt"
|
||||
echo " 2) bash auth-service.sh # device-flow → AUTH_DIR/authenticated + acpa→cpa_ready + sync_acc"
|
||||
echo " 若只想清 CPA/认证、保留刚注的 SSO: 下次加 --keep-register"
|
||||
140
scripts/clean_history.sh
Executable file
140
scripts/clean_history.sh
Executable file
@@ -0,0 +1,140 @@
|
||||
#!/usr/bin/env bash
|
||||
# 快速清历史归档（不碰现网号池）。
|
||||
#
|
||||
# 只删这些「历史」：
|
||||
# keys/cpa_ready/cpa_ready_*.zip 导出包
|
||||
# keys/cpa_ready/_discarded/** 已丢弃号
|
||||
# keys/*.bak* / keys/*~ / keys/__pycache__
|
||||
# keys/async_auth.log / logs/*
|
||||
# keys/212.zip 等杂包
|
||||
#
|
||||
# 默认保留：
|
||||
# keys/cpa_ready/xai-*.json
|
||||
# keys/cpa_ready/_state.tsv
|
||||
# keys/acc.md / accounts.txt / auth-sessions.jsonl
|
||||
# AUTH_DIR/authenticated 现网出货
|
||||
#
|
||||
# 用法（项目根）:
|
||||
# bash scripts/clean_history.sh # dry-run
|
||||
# bash scripts/clean_history.sh --yes # 真删
|
||||
# bash scripts/clean_history.sh --yes --deep # 再清 source-snapshot / 旧 .bak 脚本
|
||||
#
|
||||
# 全量清零号池请用: bash reset_pipeline.sh --yes
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
|
||||
KEYS="$ROOT/keys"
|
||||
CPA="$KEYS/cpa_ready"
|
||||
AUTH_DIR="${XAI_AUTH_SERVICE_LOCAL_DIR:-${XAI_ENROLLER_LOCAL_AUTH_DIR:-$HOME/Downloads/grok-free-register-auth}}"
|
||||
LOGS="$ROOT/logs"
|
||||
|
||||
YES=0
|
||||
DEEP=0
|
||||
|
||||
for arg in "$@"; do
|
||||
case "$arg" in
|
||||
--yes|-y) YES=1 ;;
|
||||
--deep) DEEP=1 ;;
|
||||
-h|--help)
|
||||
sed -n '2,24p' "$0"
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "未知参数: $arg" >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
size_of() {
|
||||
local p="$1"
|
||||
if [[ -e "$p" ]]; then
|
||||
du -sh "$p" 2>/dev/null | awk '{print $1}'
|
||||
else
|
||||
echo "-"
|
||||
fi
|
||||
}
|
||||
|
||||
list_targets() {
|
||||
# zips / archives under cpa_ready and keys
|
||||
find "$CPA" -maxdepth 1 -type f \( -name 'cpa_ready_*.zip' -o -name '*.zip' -o -name '*.tar' -o -name '*.tar.gz' -o -name '*.tgz' \) 2>/dev/null || true
|
||||
find "$KEYS" -maxdepth 1 -type f \( -name '*.zip' -o -name '*.tar' -o -name '*.tar.gz' -o -name '*.tgz' \) 2>/dev/null || true
|
||||
# discarded
|
||||
if [[ -d "$CPA/_discarded" ]]; then
|
||||
find "$CPA/_discarded" -mindepth 1 2>/dev/null || true
|
||||
fi
|
||||
# logs / pyc / bak
|
||||
find "$KEYS" -maxdepth 1 -type f \( -name '*.log' -o -name '*.bak' -o -name '*.bak.*' -o -name '*~' \) 2>/dev/null || true
|
||||
find "$KEYS" -maxdepth 1 -type d -name '__pycache__' 2>/dev/null || true
|
||||
if [[ -d "$LOGS" ]]; then
|
||||
find "$LOGS" -type f 2>/dev/null || true
|
||||
fi
|
||||
if [[ "$DEEP" -eq 1 ]]; then
|
||||
# deep: old snapshot growth + script backups; still keep live pool
|
||||
[[ -f "$AUTH_DIR/source-snapshot.jsonl" ]] && echo "$AUTH_DIR/source-snapshot.jsonl"
|
||||
find "$KEYS" -maxdepth 1 -type f -name 'acpa_watchdog.py.bak*' 2>/dev/null || true
|
||||
fi
|
||||
}
|
||||
|
||||
mapfile -t TARGETS < <(list_targets | awk 'NF' | sort -u)
|
||||
|
||||
echo "=== grok-free-register 快速清历史 ==="
|
||||
echo "ROOT = $ROOT"
|
||||
echo "KEYS = $KEYS"
|
||||
echo "AUTH_DIR = $AUTH_DIR"
|
||||
echo "MODE = $([[ $YES -eq 1 ]] && echo APPLY || echo dry-run)$([[ $DEEP -eq 1 ]] && echo ' +deep' || true)"
|
||||
echo
|
||||
|
||||
if [[ "${#TARGETS[@]}" -eq 0 ]]; then
|
||||
echo "没有可清的历史文件。"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "-- 将处理 --"
|
||||
total=0
|
||||
for f in "${TARGETS[@]}"; do
|
||||
if [[ -e "$f" ]]; then
|
||||
echo " $(size_of "$f") $f"
|
||||
total=$((total + 1))
|
||||
fi
|
||||
done
|
||||
echo " 共 $total 项"
|
||||
echo
|
||||
|
||||
# live pool summary (never touched)
|
||||
echo "-- 现网号池（保留）--"
|
||||
echo " cpa xai-*.json = $(find "$CPA" -maxdepth 1 -name 'xai-*.json' -type f 2>/dev/null | wc -l | tr -d ' ')"
|
||||
echo " acc.md lines = $(wc -l < "$KEYS/acc.md" 2>/dev/null | tr -d ' ' || echo 0)"
|
||||
echo " sessions = $(wc -l < "$KEYS/auth-sessions.jsonl" 2>/dev/null | tr -d ' ' || echo 0)"
|
||||
echo " authenticated = $(find "$AUTH_DIR/authenticated" -type f 2>/dev/null | wc -l | tr -d ' ' || echo 0)"
|
||||
echo
|
||||
|
||||
if [[ "$YES" -ne 1 ]]; then
|
||||
echo "dry-run。确认后: bash scripts/clean_history.sh --yes"
|
||||
echo "更深一层(含 snapshot): bash scripts/clean_history.sh --yes --deep"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "-- 执行删除 --"
|
||||
for f in "${TARGETS[@]}"; do
|
||||
if [[ -d "$f" ]]; then
|
||||
rm -rf -- "$f"
|
||||
echo " removed dir $f"
|
||||
elif [[ -e "$f" ]]; then
|
||||
rm -f -- "$f"
|
||||
echo " removed file $f"
|
||||
fi
|
||||
done
|
||||
|
||||
# keep empty discarded dir
|
||||
mkdir -p "$CPA/_discarded" 2>/dev/null || true
|
||||
mkdir -p "$LOGS" 2>/dev/null || true
|
||||
|
||||
echo
|
||||
echo "-- 清后 --"
|
||||
echo " cpa zip left = $(find "$CPA" -maxdepth 1 -name '*.zip' -type f 2>/dev/null | wc -l | tr -d ' ')"
|
||||
echo " discarded files= $(find "$CPA/_discarded" -type f 2>/dev/null | wc -l | tr -d ' ')"
|
||||
echo " cpa xai-*.json = $(find "$CPA" -maxdepth 1 -name 'xai-*.json' -type f 2>/dev/null | wc -l | tr -d ' ')"
|
||||
echo " keys size = $(size_of "$KEYS")"
|
||||
echo "完成。"
|
||||
32
scripts/ensure_runtime.sh
Executable file
32
scripts/ensure_runtime.sh
Executable file
@@ -0,0 +1,32 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
ensure_runtime() {
|
||||
local lock_dir=".setup.lock"
|
||||
local acquired=0
|
||||
local attempt
|
||||
|
||||
for attempt in {1..300}; do
|
||||
if mkdir "$lock_dir" 2>/dev/null; then
|
||||
acquired=1
|
||||
break
|
||||
fi
|
||||
sleep 0.2
|
||||
done
|
||||
if [ "$acquired" -ne 1 ]; then
|
||||
echo "[!] 另一个安装进程长时间未结束，请稍后重试。" >&2
|
||||
return 1
|
||||
fi
|
||||
trap 'rmdir .setup.lock 2>/dev/null || true' EXIT INT TERM
|
||||
|
||||
if [ ! -d .venv ]; then
|
||||
echo "[*] 首次运行，安装依赖..."
|
||||
if ! bash setup.sh; then
|
||||
rmdir "$lock_dir"
|
||||
trap - EXIT INT TERM
|
||||
return 1
|
||||
fi
|
||||
fi
|
||||
|
||||
rmdir "$lock_dir"
|
||||
trap - EXIT INT TERM
|
||||
}
|
||||
206
scripts/export_registered_sessions.py
Executable file
206
scripts/export_registered_sessions.py
Executable file
@@ -0,0 +1,206 @@
|
||||
#!/usr/bin/env python3
|
||||
"""通过 SSH 导出不含密码的注册会话；兼容历史裸 SSO 记录。"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
COOKIE_FIELDS = frozenset(
|
||||
{
|
||||
"name",
|
||||
"value",
|
||||
"url",
|
||||
"domain",
|
||||
"path",
|
||||
"expires",
|
||||
"httpOnly",
|
||||
"secure",
|
||||
"sameSite",
|
||||
}
|
||||
)
|
||||
MAX_RECORD_BYTES = 256 * 1024
|
||||
LEGACY_COOKIE_SCOPE = {
|
||||
"name": "sso",
|
||||
"domain": "accounts.x.ai",
|
||||
"path": "/",
|
||||
"secure": True,
|
||||
"httpOnly": True,
|
||||
"sameSite": "Lax",
|
||||
}
|
||||
|
||||
|
||||
def _decode_json_line(raw, label):
|
||||
if len(raw) > MAX_RECORD_BYTES:
|
||||
raise ValueError(f"invalid {label} record")
|
||||
try:
|
||||
document = json.loads(raw.decode("utf-8"))
|
||||
email = document["email"]
|
||||
cookies = document["cookies"]
|
||||
except (UnicodeDecodeError, TypeError, ValueError, KeyError) as exc:
|
||||
raise ValueError(f"invalid {label} record") from exc
|
||||
if not isinstance(email, str) or not email or not isinstance(cookies, list) or not cookies:
|
||||
raise ValueError(f"invalid {label} record")
|
||||
try:
|
||||
email.encode("utf-8")
|
||||
except UnicodeEncodeError as exc:
|
||||
raise ValueError(f"invalid {label} record") from exc
|
||||
normalized = []
|
||||
for cookie in cookies:
|
||||
if not isinstance(cookie, dict):
|
||||
raise ValueError(f"invalid {label} record")
|
||||
filtered = {key: cookie[key] for key in COOKIE_FIELDS if key in cookie}
|
||||
if not all(
|
||||
isinstance(filtered.get(key), str) and filtered[key]
|
||||
for key in ("name", "value")
|
||||
):
|
||||
raise ValueError(f"invalid {label} record")
|
||||
scope = filtered.get("domain") or filtered.get("url")
|
||||
if not isinstance(scope, str) or not scope:
|
||||
raise ValueError(f"invalid {label} record")
|
||||
try:
|
||||
filtered["name"].encode("utf-8")
|
||||
filtered["value"].encode("utf-8")
|
||||
scope.encode("utf-8")
|
||||
except UnicodeEncodeError as exc:
|
||||
raise ValueError(f"invalid {label} record") from exc
|
||||
normalized.append(filtered)
|
||||
return {"email": email, "cookies": normalized}
|
||||
|
||||
|
||||
def _complete_lines(data):
|
||||
"""Return newline-terminated records and the unconsumed trailing bytes."""
|
||||
parts = data.split(b"\n")
|
||||
return parts[:-1], parts[-1]
|
||||
|
||||
|
||||
def _read_complete_file(path):
|
||||
if not path.exists():
|
||||
return []
|
||||
lines, _incomplete = _complete_lines(path.read_bytes())
|
||||
return [line for line in lines if line]
|
||||
|
||||
|
||||
def load_snapshots(path, *, raw_lines=None):
|
||||
snapshots = {}
|
||||
scopes = {}
|
||||
lines = _read_complete_file(path) if raw_lines is None else raw_lines
|
||||
for line in lines:
|
||||
document = _decode_json_line(line, "session")
|
||||
email = document["email"]
|
||||
cookies = document["cookies"]
|
||||
if email in snapshots:
|
||||
continue
|
||||
snapshots[email] = {"email": email, "cookies": cookies}
|
||||
for cookie in cookies:
|
||||
name = cookie.get("name")
|
||||
domain = cookie.get("domain")
|
||||
if name in {"sso", "sso-rw"} and domain:
|
||||
scopes[(name, domain)] = {
|
||||
"name": name,
|
||||
"domain": domain,
|
||||
"path": cookie.get("path", "/"),
|
||||
"secure": bool(cookie.get("secure", True)),
|
||||
"httpOnly": bool(cookie.get("httpOnly", True)),
|
||||
"sameSite": cookie.get("sameSite", "Lax"),
|
||||
}
|
||||
return snapshots, scopes
|
||||
|
||||
|
||||
def export_sessions(sessions_path, accounts_path, *, session_lines=None):
|
||||
snapshots, scopes = load_snapshots(sessions_path, raw_lines=session_lines)
|
||||
for document in snapshots.values():
|
||||
yield document
|
||||
if not accounts_path.exists():
|
||||
return
|
||||
legacy_scopes = list(scopes.values()) or [LEGACY_COOKIE_SCOPE]
|
||||
for raw in _read_complete_file(accounts_path):
|
||||
try:
|
||||
email, _password, sso = raw.decode("utf-8").rsplit(":", 2)
|
||||
except (UnicodeDecodeError, ValueError) as exc:
|
||||
raise ValueError("invalid account record") from exc
|
||||
if not email or not sso or email in snapshots:
|
||||
continue
|
||||
cookies = [{**scope, "value": sso} for scope in legacy_scopes]
|
||||
yield {"email": email, "cookies": cookies}
|
||||
|
||||
|
||||
def _write_document(document):
|
||||
payload = json.dumps(document, separators=(",", ":"))
|
||||
if len(payload.encode("utf-8")) > MAX_RECORD_BYTES:
|
||||
raise ValueError("invalid session record")
|
||||
print(payload, flush=False)
|
||||
|
||||
|
||||
def export_and_follow(sessions_path, accounts_path, *, poll_seconds=0.25):
|
||||
"""Emit a complete snapshot, then losslessly follow the same JSONL fd."""
|
||||
sessions_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not sessions_path.exists():
|
||||
fd = os.open(sessions_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
|
||||
os.close(fd)
|
||||
os.chmod(sessions_path, 0o600)
|
||||
stream = sessions_path.open("rb")
|
||||
|
||||
with stream:
|
||||
opened = os.fstat(stream.fileno())
|
||||
initial_lines, pending = _complete_lines(stream.read())
|
||||
if len(pending) > MAX_RECORD_BYTES:
|
||||
raise ValueError("invalid session record")
|
||||
for document in export_sessions(
|
||||
sessions_path,
|
||||
accounts_path,
|
||||
session_lines=[line for line in initial_lines if line],
|
||||
):
|
||||
_write_document(document)
|
||||
sys.stdout.flush()
|
||||
|
||||
while True:
|
||||
chunk = stream.read()
|
||||
if chunk:
|
||||
complete, pending = _complete_lines(pending + chunk)
|
||||
if len(pending) > MAX_RECORD_BYTES:
|
||||
raise ValueError("invalid session record")
|
||||
for raw in complete:
|
||||
if raw:
|
||||
_write_document(_decode_json_line(raw, "session"))
|
||||
sys.stdout.flush()
|
||||
continue
|
||||
|
||||
try:
|
||||
current = sessions_path.stat()
|
||||
except FileNotFoundError:
|
||||
return 3
|
||||
if (
|
||||
current.st_dev != opened.st_dev
|
||||
or current.st_ino != opened.st_ino
|
||||
or current.st_size < stream.tell()
|
||||
):
|
||||
return 3
|
||||
time.sleep(poll_seconds)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--follow", action="store_true")
|
||||
parser.add_argument("sessions_path", type=Path)
|
||||
parser.add_argument("accounts_path", type=Path)
|
||||
args = parser.parse_args()
|
||||
try:
|
||||
if args.follow:
|
||||
raise SystemExit(export_and_follow(args.sessions_path, args.accounts_path))
|
||||
for document in export_sessions(args.sessions_path, args.accounts_path):
|
||||
_write_document(document)
|
||||
except ValueError:
|
||||
raise SystemExit(4) from None
|
||||
except BrokenPipeError:
|
||||
try:
|
||||
sys.stdout.close()
|
||||
finally:
|
||||
raise SystemExit(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
153
scripts/push_keys_to_auth.sh
Executable file
153
scripts/push_keys_to_auth.sh
Executable file
@@ -0,0 +1,153 @@
|
||||
#!/usr/bin/env bash
|
||||
# 本机注册 → 推 SSO 源到远端 auth 机（local 模式吃 keys/）
|
||||
#
|
||||
# 动态公网 IP 不适合「远端 SSH 拉本机」；改由本机主动推：
|
||||
# keys/auth-sessions.jsonl
|
||||
# keys/accounts.txt
|
||||
# → AUTH_HOST:REMOTE_ROOT/keys/
|
||||
#
|
||||
# 用法（本机项目根或任意 cwd）:
|
||||
# bash scripts/push_keys_to_auth.sh # 推一次
|
||||
# bash scripts/push_keys_to_auth.sh --watch # 常驻，文件变了再推
|
||||
# bash scripts/push_keys_to_auth.sh --interval 15
|
||||
#
|
||||
# 环境变量（可写本机 .env，本脚本会 source）:
|
||||
# AUTH_SSH_HOST=user@auth-server.example
|
||||
# AUTH_REMOTE_ROOT=/opt/grok-free-register
|
||||
# AUTH_SSH_IDENTITY= # 可选 -i 路径
|
||||
# AUTH_PUSH_INTERVAL=20 # --watch 轮询秒
|
||||
#
|
||||
# 无内置默认主机/路径：必须显式设置 AUTH_SSH_HOST 与 AUTH_REMOTE_ROOT。
|
||||
|
||||
set -euo pipefail
|
||||
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
|
||||
cd "$ROOT"
|
||||
|
||||
# 调用方已 export 的 AUTH_* 优先（.env 不得盖掉），方便并行推多台 auth
|
||||
_PRE_AUTH_SSH_HOST="${AUTH_SSH_HOST-}"
|
||||
_PRE_AUTH_REMOTE_ROOT="${AUTH_REMOTE_ROOT-}"
|
||||
_PRE_AUTH_SSH_IDENTITY="${AUTH_SSH_IDENTITY-}"
|
||||
_PRE_AUTH_PUSH_INTERVAL="${AUTH_PUSH_INTERVAL-}"
|
||||
|
||||
if [[ -f "$ROOT/.env" ]]; then
|
||||
set -a
|
||||
# shellcheck disable=SC1091
|
||||
. "$ROOT/.env"
|
||||
set +a
|
||||
fi
|
||||
|
||||
AUTH_SSH_HOST="${_PRE_AUTH_SSH_HOST:-${AUTH_SSH_HOST:-}}"
|
||||
AUTH_REMOTE_ROOT="${_PRE_AUTH_REMOTE_ROOT:-${AUTH_REMOTE_ROOT:-}}"
|
||||
AUTH_SSH_IDENTITY="${_PRE_AUTH_SSH_IDENTITY:-${AUTH_SSH_IDENTITY:-}}"
|
||||
AUTH_PUSH_INTERVAL="${_PRE_AUTH_PUSH_INTERVAL:-${AUTH_PUSH_INTERVAL:-20}}"
|
||||
|
||||
if [[ -z "$AUTH_SSH_HOST" || -z "$AUTH_REMOTE_ROOT" ]]; then
|
||||
echo "push_keys_to_auth: set AUTH_SSH_HOST and AUTH_REMOTE_ROOT (env or .env)" >&2
|
||||
echo " e.g. AUTH_SSH_HOST=user@auth-server.example AUTH_REMOTE_ROOT=/opt/grok-free-register" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
WATCH=0
|
||||
INTERVAL="$AUTH_PUSH_INTERVAL"
|
||||
for a in "$@"; do
|
||||
case "$a" in
|
||||
--watch) WATCH=1 ;;
|
||||
--interval)
|
||||
# next arg handled below if present as --interval=N form preferred
|
||||
;;
|
||||
--interval=*) INTERVAL="${a#--interval=}" ;;
|
||||
-h|--help)
|
||||
sed -n '2,22p' "$0"
|
||||
exit 0
|
||||
;;
|
||||
esac
|
||||
done
|
||||
# support: --interval 15
|
||||
prev=""
|
||||
for a in "$@"; do
|
||||
if [[ "$prev" == "--interval" ]]; then
|
||||
INTERVAL="$a"
|
||||
fi
|
||||
prev="$a"
|
||||
done
|
||||
|
||||
SSH_OPTS=(-o BatchMode=yes -o ConnectTimeout=15 -o ServerAliveInterval=15)
|
||||
if [[ -n "$AUTH_SSH_IDENTITY" ]]; then
|
||||
SSH_OPTS+=(-i "$AUTH_SSH_IDENTITY")
|
||||
fi
|
||||
|
||||
KEYS_SRC="$ROOT/keys"
|
||||
REMOTE_KEYS="${AUTH_REMOTE_ROOT}/keys"
|
||||
# 只推 auth 源，不推 cpa_ready/acc（远端自己整备）
|
||||
FILES=(auth-sessions.jsonl accounts.txt)
|
||||
|
||||
log() { printf '[push_keys %s] %s\n' "$(date +%H:%M:%S)" "$*"; }
|
||||
|
||||
fingerprint() {
|
||||
local f
|
||||
for f in "${FILES[@]}"; do
|
||||
if [[ -f "$KEYS_SRC/$f" ]]; then
|
||||
# size + mtime + sha256 head — cheap change detect
|
||||
stat -c '%s %Y' "$KEYS_SRC/$f" 2>/dev/null || stat -f '%z %m' "$KEYS_SRC/$f"
|
||||
sha256sum "$KEYS_SRC/$f" 2>/dev/null | awk '{print $1}'
|
||||
else
|
||||
echo "missing:$f"
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
push_once() {
|
||||
local missing=0 f
|
||||
for f in "${FILES[@]}"; do
|
||||
if [[ ! -f "$KEYS_SRC/$f" ]]; then
|
||||
log "⚠ 缺 $KEYS_SRC/$f"
|
||||
missing=1
|
||||
fi
|
||||
done
|
||||
if [[ "$missing" -eq 1 ]]; then
|
||||
return 1
|
||||
fi
|
||||
|
||||
# 远端原子替换：先推 .push.tmp 再 mv
|
||||
ssh "${SSH_OPTS[@]}" "$AUTH_SSH_HOST" "mkdir -p $(printf %q "$REMOTE_KEYS") && chmod 700 $(printf %q "$REMOTE_KEYS") 2>/dev/null || true"
|
||||
|
||||
local remote_tmp remote_final
|
||||
for f in "${FILES[@]}"; do
|
||||
remote_tmp="${REMOTE_KEYS}/.${f}.push.tmp"
|
||||
remote_final="${REMOTE_KEYS}/${f}"
|
||||
# rsync over ssh when available; fallback scp
|
||||
if command -v rsync >/dev/null 2>&1; then
|
||||
rsync -az -e "ssh ${SSH_OPTS[*]}" \
|
||||
"$KEYS_SRC/$f" \
|
||||
"${AUTH_SSH_HOST}:${remote_tmp}"
|
||||
else
|
||||
scp "${SSH_OPTS[@]}" "$KEYS_SRC/$f" "${AUTH_SSH_HOST}:${remote_tmp}"
|
||||
fi
|
||||
ssh "${SSH_OPTS[@]}" "$AUTH_SSH_HOST" \
|
||||
"chmod 600 $(printf %q "$remote_tmp") && mv -f $(printf %q "$remote_tmp") $(printf %q "$remote_final")"
|
||||
done
|
||||
|
||||
local n_sess n_acc
|
||||
n_sess=$(wc -l < "$KEYS_SRC/auth-sessions.jsonl" | tr -d ' ')
|
||||
n_acc=$(wc -l < "$KEYS_SRC/accounts.txt" | tr -d ' ')
|
||||
log "→ ${AUTH_SSH_HOST}:${REMOTE_KEYS}/ sessions=${n_sess} accounts=${n_acc}"
|
||||
}
|
||||
|
||||
LAST_FP=""
|
||||
if [[ "$WATCH" -eq 0 ]]; then
|
||||
push_once
|
||||
exit $?
|
||||
fi
|
||||
|
||||
log "watch host=${AUTH_SSH_HOST} root=${AUTH_REMOTE_ROOT} interval=${INTERVAL}s"
|
||||
while true; do
|
||||
FP="$(fingerprint)"
|
||||
if [[ "$FP" != "$LAST_FP" ]]; then
|
||||
if push_once; then
|
||||
LAST_FP="$FP"
|
||||
else
|
||||
log "推送失败，${INTERVAL}s 后重试"
|
||||
fi
|
||||
fi
|
||||
sleep "$INTERVAL"
|
||||
done
|
||||
50
setup.sh
Executable file
50
setup.sh
Executable file
@@ -0,0 +1,50 @@
|
||||
#!/bin/bash
|
||||
# Grok Free Register — 一键安装脚本
|
||||
# 用法: bash setup.sh
|
||||
|
||||
set -e
|
||||
|
||||
echo "=== Grok Free Register 安装 ==="
|
||||
|
||||
# 检测系统
|
||||
if [ -f /etc/debian_version ]; then
|
||||
echo "[1/4] 安装系统依赖 (Debian/Ubuntu)..."
|
||||
sudo apt update -qq
|
||||
sudo apt install -y -qq \
|
||||
python3 python3-pip python3-venv \
|
||||
libatk1.0-0t64 libatk-bridge2.0-0t64 libcups2t64 \
|
||||
libdrm2 libxkbcommon0 libxcomposite1 libxdamage1 \
|
||||
libxfixes3 libxrandr2 libgbm1 libpango-1.0-0 \
|
||||
libcairo2 libasound2t64 libnspr4 libnss3 libxshmfence1 \
|
||||
2>/dev/null || true
|
||||
# 兼容旧版 Ubuntu
|
||||
sudo apt install -y -qq libatk1.0-0 libatk-bridge2.0-0 libcups2 libasound2 2>/dev/null || true
|
||||
elif [ -f /etc/redhat-release ]; then
|
||||
echo "[1/4] 安装系统依赖 (RHEL/CentOS)..."
|
||||
sudo yum install -y -q \
|
||||
python3 python3-pip \
|
||||
atk cups-libs libdrm libXcomposite libXdamage libXfixes libXrandr \
|
||||
mesa-libgbm pango cairo alsa-lib nspr nss libxshmfence \
|
||||
2>/dev/null || true
|
||||
else
|
||||
echo "[1/4] 未知系统，跳过系统依赖（如 Chrome 启动失败请手动安装）"
|
||||
fi
|
||||
|
||||
# Python 虚拟环境
|
||||
echo "[2/4] 创建 Python 环境..."
|
||||
python3 -m venv .venv
|
||||
.venv/bin/pip install -q --upgrade pip
|
||||
.venv/bin/pip install -q -r requirements.txt
|
||||
|
||||
# 本项目直接由 Playwright 启动二进制，因此显式准备 CloakBrowser Chromium。
|
||||
echo "[3/4] 下载 CloakBrowser Chromium..."
|
||||
.venv/bin/python -m cloakbrowser install
|
||||
|
||||
# 创建输出目录
|
||||
mkdir -p keys
|
||||
|
||||
echo "[4/4] 安装完成！"
|
||||
echo ""
|
||||
echo "运行注册服务: bash start.sh"
|
||||
echo "运行自建邮箱服务: bash start.sh --email-service"
|
||||
echo "运行认证服务: bash auth-service.sh"
|
||||
80
start.sh
Executable file
80
start.sh
Executable file
@@ -0,0 +1,80 @@
|
||||
#!/bin/bash
|
||||
# 一键启动:自动装依赖 → 引导配置 → 运行
|
||||
# 用法:
|
||||
# bash start.sh # 首次会引导选模式,之后直接启动
|
||||
# bash start.sh --reconfig # 重新选择邮箱模式
|
||||
# bash start.sh --debug # 保留完整调试面板
|
||||
set -e
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
. scripts/ensure_runtime.sh
|
||||
ensure_runtime
|
||||
|
||||
if [ "${1:-}" = "--email-service" ]; then
|
||||
shift
|
||||
if command -v flock >/dev/null 2>&1; then
|
||||
mkdir -p logs
|
||||
exec 8>logs/email-service.lock
|
||||
if ! flock -n 8; then
|
||||
echo "[!] 邮箱服务已经在运行。"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
echo "[*] 启动邮箱服务... (Ctrl-C 停止)"
|
||||
exec .venv/bin/python -m grok_register.email_server "$@"
|
||||
fi
|
||||
|
||||
reconfig=0
|
||||
register_args=()
|
||||
for arg in "$@"; do
|
||||
if [ "$arg" = "--reconfig" ]; then
|
||||
reconfig=1
|
||||
else
|
||||
register_args+=("$arg")
|
||||
fi
|
||||
done
|
||||
|
||||
# 同一工作目录只允许一个注册进程，避免重复启动同时写账号和日志。
|
||||
if command -v flock >/dev/null 2>&1; then
|
||||
mkdir -p logs
|
||||
exec 9>logs/register.lock
|
||||
if ! flock -n 9; then
|
||||
echo "[!] 注册机已经在运行。"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
# 1) 配置:无 .env 或显式 --reconfig 时进入引导
|
||||
if [ ! -f .env ] || [ "$reconfig" -eq 1 ]; then
|
||||
echo ""
|
||||
echo "选择邮箱模式:"
|
||||
echo " [1] 免费临时邮箱 (默认 · 零配置 · 直接回车 · 多 provider 自动 fallback)"
|
||||
echo " [2] 自建域名邮箱 (需 Cloudflare Email Routing + 本地 webhook)"
|
||||
read -rp "输入 1 或 2 [1]: " mode || mode=1
|
||||
if [ "$mode" = "2" ]; then
|
||||
read -rp " 你的域名 (如 example.com): " domain
|
||||
read -rp " webhook 地址 [http://127.0.0.1:8080]: " api
|
||||
api=${api:-http://127.0.0.1:8080}
|
||||
cat > .env <<ENV
|
||||
EMAIL_MODE=custom
|
||||
EMAIL_DOMAIN=${domain}
|
||||
EMAIL_API=${api}
|
||||
# CSP 容量(可选,0=按 CPU/内存启动期静态派生)
|
||||
# PHYSICAL_CAP=0
|
||||
# PHYSICAL_PER_CPU=2
|
||||
# PHYSICAL_MEM_MB=512
|
||||
# MIN_FREE_MEM_MB=500
|
||||
ENV
|
||||
echo ""
|
||||
echo "[!] custom 模式还需在另一终端运行收信服务:"
|
||||
echo " bash start.sh --email-service"
|
||||
echo " 并按 README「自建邮箱模式」配置 Cloudflare Email Worker。"
|
||||
else
|
||||
echo "EMAIL_MODE=tempmail" > .env
|
||||
fi
|
||||
echo "[*] 已写入 .env"
|
||||
fi
|
||||
|
||||
# 2) 运行
|
||||
echo "[*] 启动注册服务... (Ctrl-C 停止)"
|
||||
exec .venv/bin/python -m grok_register.register "${register_args[@]}"
|
||||
0
tests/__init__.py
Normal file
0
tests/__init__.py
Normal file
92
tests/conftest.py
Normal file
92
tests/conftest.py
Normal file
@@ -0,0 +1,92 @@
|
||||
"""
|
||||
共享 fixtures — 供四层测试共用
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
import asyncio
|
||||
import pytest
|
||||
|
||||
# 确保项目根目录在 sys.path 中
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from grok_register.core.inventory import Inventory
|
||||
from grok_register.core.observer import Metrics
|
||||
from tests.fakes import (
|
||||
FakeTurnstile, FakeEmailService, FakeRegisterAPI,
|
||||
Conservation, EventLog,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def metrics():
|
||||
return Metrics()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def inventory(metrics):
|
||||
return Inventory(metrics=metrics)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sems():
|
||||
"""标准 Semaphore 集合(小容量,便于测试)。"""
|
||||
return {
|
||||
'physical': asyncio.Semaphore(4),
|
||||
't_slot': asyncio.Semaphore(4),
|
||||
'q_slot': asyncio.Semaphore(4),
|
||||
'q_pending': asyncio.Semaphore(6),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def large_sems():
|
||||
"""大容量 Semaphore(压力测试用)。"""
|
||||
return {
|
||||
'physical': asyncio.Semaphore(32),
|
||||
't_slot': asyncio.Semaphore(64),
|
||||
'q_slot': asyncio.Semaphore(64),
|
||||
'q_pending': asyncio.Semaphore(48),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def conservation(sems, inventory):
|
||||
return Conservation(
|
||||
sems['t_slot'], sems['q_slot'], sems['q_pending'], inventory
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_turnstile():
|
||||
return FakeTurnstile()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_email():
|
||||
return FakeEmailService()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_register():
|
||||
return FakeRegisterAPI()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def event_log():
|
||||
return EventLog()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stop():
|
||||
return asyncio.Event()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def file_lock():
|
||||
return asyncio.Lock()
|
||||
|
||||
|
||||
# ── 跳过需要网络的测试(本地开发时) ──
|
||||
def pytest_configure(config):
|
||||
config.addinivalue_line("markers", "slow: 压力测试(较慢)")
|
||||
config.addinivalue_line("markers", "needs_network: 需要外部网络")
|
||||
317
tests/fakes.py
Normal file
317
tests/fakes.py
Normal file
@@ -0,0 +1,317 @@
|
||||
"""
|
||||
Fake 服务 + 测试工具
|
||||
|
||||
FakeTurnstile — 可控 T 生成器:可注入延迟、失败、取消点
|
||||
FakeEmailService — 可控 Q 提供者:可注入延迟、超时、验证码
|
||||
FakeConsumer — 可控 C 消费器:可注入延迟、失败
|
||||
Conservation — slot 守恒断言器
|
||||
CancelPoint — 在指定 await 边界注入取消
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
import dataclasses
|
||||
from typing import Optional, Callable, Any
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# Fake T 生成器 (替代 solve_one_turnstile)
|
||||
# ═══════════════════════════════════════════
|
||||
class FakeTurnstile:
|
||||
"""可控 token 生成器。
|
||||
|
||||
config:
|
||||
delay: 生成延迟(秒)
|
||||
fail_rate: 0.0~1.0,随机失败概率
|
||||
token_seq: 预设 token 序列,None 则自动递增
|
||||
"""
|
||||
|
||||
def __init__(self, delay=0.01, fail_rate=0.0, token_seq=None):
|
||||
self.delay = delay
|
||||
self._fail_rate = fail_rate
|
||||
self._seq = token_seq or []
|
||||
self._idx = 0
|
||||
self._call_count = 0
|
||||
self._fail_count = 0
|
||||
self._cancel_at: Optional[int] = None # 第 N 次调用时注入取消
|
||||
|
||||
def set_cancel_at(self, n: int):
|
||||
"""第 n 次调用时注入 CancelledError。"""
|
||||
self._cancel_at = n
|
||||
|
||||
async def generate(self) -> Optional[str]:
|
||||
"""生成一个 token。返回 None 表示失败,raise CancelledError 表示取消。"""
|
||||
self._call_count += 1
|
||||
|
||||
if self._cancel_at is not None and self._call_count == self._cancel_at:
|
||||
raise asyncio.CancelledError(f'FakeTurnstile cancel at call {self._call_count}')
|
||||
|
||||
await asyncio.sleep(self.delay)
|
||||
|
||||
if self._fail_count < self._fail_rate * self._call_count:
|
||||
return None
|
||||
# 简单随机: fail_rate 作为概率
|
||||
import random
|
||||
if random.random() < self._fail_rate:
|
||||
self._fail_count += 1
|
||||
return None
|
||||
|
||||
if self._idx < len(self._seq):
|
||||
tok = self._seq[self._idx]
|
||||
self._idx += 1
|
||||
return tok
|
||||
self._idx += 1
|
||||
return f'fake_token_{self._idx}'
|
||||
|
||||
@property
|
||||
def call_count(self):
|
||||
return self._call_count
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# Fake 邮箱服务 (替代 create_email + poll_code)
|
||||
# ═══════════════════════════════════════════
|
||||
@dataclasses.dataclass
|
||||
class FakeEmailResult:
|
||||
handle: str
|
||||
email: str
|
||||
password: str
|
||||
code: Optional[str] = None # None = 轮询超时
|
||||
create_delay: float = 0.01
|
||||
poll_delay: float = 0.01
|
||||
|
||||
|
||||
class FakeEmailService:
|
||||
"""可控邮箱 + 验证码服务。
|
||||
|
||||
results: 预设结果队列,每次 create+poll 消费一个。
|
||||
"""
|
||||
|
||||
def __init__(self, results=None):
|
||||
self._results = list(results or [])
|
||||
self._idx = 0
|
||||
self._create_count = 0
|
||||
self._poll_count = 0
|
||||
self._cancel_create_at: Optional[int] = None
|
||||
self._cancel_poll_at: Optional[int] = None
|
||||
|
||||
def set_cancel_create_at(self, n: int):
|
||||
self._cancel_create_at = n
|
||||
|
||||
def set_cancel_poll_at(self, n: int):
|
||||
self._cancel_poll_at = n
|
||||
|
||||
async def create_email(self):
|
||||
"""创建邮箱。返回 (handle, email, password)。"""
|
||||
self._create_count += 1
|
||||
if self._cancel_create_at and self._create_count == self._cancel_create_at:
|
||||
raise asyncio.CancelledError(f'FakeEmail.create cancel at {self._create_count}')
|
||||
|
||||
if self._idx < len(self._results):
|
||||
r = self._results[self._idx]
|
||||
await asyncio.sleep(r.create_delay)
|
||||
return r.handle, r.email, r.password
|
||||
self._idx += 1
|
||||
n = self._idx
|
||||
await asyncio.sleep(0.01)
|
||||
return f'handle_{n}', f'user_{n}@test.com', f'pass_{n}'
|
||||
|
||||
async def poll_code(self, handle: str, max_wait=90):
|
||||
"""轮询验证码。返回 code 或 None。"""
|
||||
self._poll_count += 1
|
||||
if self._cancel_poll_at and self._poll_count == self._cancel_poll_at:
|
||||
raise asyncio.CancelledError(f'FakeEmail.poll cancel at {self._poll_count}')
|
||||
|
||||
# 找匹配的预设结果
|
||||
for r in self._results:
|
||||
if r.handle == handle:
|
||||
await asyncio.sleep(r.poll_delay)
|
||||
return r.code
|
||||
await asyncio.sleep(0.01)
|
||||
return '000000'
|
||||
|
||||
@property
|
||||
def create_count(self):
|
||||
return self._create_count
|
||||
|
||||
@property
|
||||
def poll_count(self):
|
||||
return self._poll_count
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# Fake 注册 API (替代 grpc_verify + server_action_register)
|
||||
# ═══════════════════════════════════════════
|
||||
class FakeRegisterAPI:
|
||||
"""可控注册 API。"""
|
||||
|
||||
def __init__(self, success_rate=1.0, delay=0.01):
|
||||
self.success_rate = success_rate
|
||||
self.delay = delay
|
||||
self._register_count = 0
|
||||
self._cancel_at: Optional[int] = None
|
||||
self._fail_at: Optional[int] = None
|
||||
|
||||
def set_cancel_at(self, n: int):
|
||||
self._cancel_at = n
|
||||
|
||||
def set_fail_at(self, n: int):
|
||||
self._fail_at = n
|
||||
|
||||
async def register(self, email, password, code, token) -> Optional[str]:
|
||||
"""执行注册。返回 SSO 或 None。"""
|
||||
self._register_count += 1
|
||||
if self._cancel_at and self._register_count == self._cancel_at:
|
||||
raise asyncio.CancelledError(f'FakeRegister cancel at {self._register_count}')
|
||||
if self._fail_at and self._register_count == self._fail_at:
|
||||
return None
|
||||
await asyncio.sleep(self.delay)
|
||||
import random
|
||||
if random.random() < self.success_rate:
|
||||
return f'sso_{email}_{int(time.time())}'
|
||||
return None
|
||||
|
||||
@property
|
||||
def register_count(self):
|
||||
return self._register_count
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# Slot 守恒断言器
|
||||
# ═══════════════════════════════════════════
|
||||
class Conservation:
|
||||
"""检查 T/Q slot 守恒不变量。
|
||||
|
||||
守恒公式:
|
||||
slot_total = free_slots + inventory_depth + pairleased_held + worker_held
|
||||
"""
|
||||
|
||||
def __init__(self, t_slot_sem, q_slot_sem, q_pending_sem, inventory):
|
||||
self.t_slot_sem = t_slot_sem
|
||||
self.q_slot_sem = q_slot_sem
|
||||
self.q_pending_sem = q_pending_sem
|
||||
self.inventory = inventory
|
||||
# 追踪 Worker 持有的未发布资源
|
||||
self._worker_held_t = 0
|
||||
self._worker_held_q = 0
|
||||
|
||||
def worker_acquire_t(self):
|
||||
self._worker_held_t += 1
|
||||
|
||||
def worker_release_t(self):
|
||||
self._worker_held_t -= 1
|
||||
|
||||
def worker_acquire_q(self):
|
||||
self._worker_held_q += 1
|
||||
|
||||
def worker_release_q(self):
|
||||
self._worker_held_q -= 1
|
||||
|
||||
def check_t_conservation(self, t_slot_total: int, pairleased_t: int = 0):
|
||||
"""检查 T slot 守恒。
|
||||
|
||||
Args:
|
||||
t_slot_total: T_Slot_Sem 初始容量
|
||||
pairleased_t: 当前 PairLease 持有的 T 数量
|
||||
"""
|
||||
free = self.t_slot_sem._value
|
||||
inv = self.inventory.t_depth
|
||||
held = self._worker_held_t
|
||||
total = free + inv + pairleased_t + held
|
||||
assert total == t_slot_total, (
|
||||
f'T slot 守恒违反: free({free}) + inv({inv}) + pairleased({pairleased_t}) '
|
||||
f'+ worker_held({held}) = {total} != total({t_slot_total})'
|
||||
)
|
||||
|
||||
def check_q_conservation(self, q_slot_total: int, pairleased_q: int = 0):
|
||||
"""检查 Q slot 守恒。"""
|
||||
free = self.q_slot_sem._value
|
||||
inv = self.inventory.q_depth
|
||||
held = self._worker_held_q
|
||||
total = free + inv + pairleased_q + held
|
||||
assert total == q_slot_total, (
|
||||
f'Q slot 守恒违反: free({free}) + inv({inv}) + pairleased({pairleased_q}) '
|
||||
f'+ worker_held({held}) = {total} != total({q_slot_total})'
|
||||
)
|
||||
|
||||
def check_pending_conservation(self, pending_total: int, in_flight: int = 0):
|
||||
"""检查 Q pending 守恒。"""
|
||||
free = self.q_pending_sem._value
|
||||
total = free + in_flight
|
||||
assert total == pending_total, (
|
||||
f'Q pending 守恒违反: free({free}) + in_flight({in_flight}) = {total} '
|
||||
f'!= total({pending_total})'
|
||||
)
|
||||
|
||||
def check_all(self, t_total, q_total, pend_total, pairleased_t=0, pairleased_q=0, in_flight=0):
|
||||
"""一次性检查所有守恒。"""
|
||||
self.check_t_conservation(t_total, pairleased_t)
|
||||
self.check_q_conservation(q_total, pairleased_q)
|
||||
self.check_pending_conservation(pend_total, in_flight)
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# 取消注入器
|
||||
# ═══════════════════════════════════════════
|
||||
class CancelInjector:
|
||||
"""在指定 await 边界前后注入取消。"""
|
||||
|
||||
def __init__(self):
|
||||
self._points: dict[str, asyncio.Event] = {}
|
||||
self._cancel_task: Optional[asyncio.Task] = None
|
||||
|
||||
def register(self, name: str):
|
||||
"""注册一个取消点。"""
|
||||
self._points[name] = asyncio.Event()
|
||||
|
||||
async def wait_at(self, name: str):
|
||||
"""Worker 在此挂起,等待外部触发取消。"""
|
||||
if name in self._points:
|
||||
await self._points[name].wait()
|
||||
|
||||
def trigger(self, name: str, task: asyncio.Task):
|
||||
"""触发指定取消点,取消目标 task。"""
|
||||
if name in self._points:
|
||||
self._points[name].set()
|
||||
task.cancel()
|
||||
|
||||
async def cancel_after(self, delay: float, task: asyncio.Task):
|
||||
"""延迟后取消目标 task。"""
|
||||
await asyncio.sleep(delay)
|
||||
task.cancel()
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# 事件记录器(用于性质测试)
|
||||
# ═══════════════════════════════════════════
|
||||
@dataclasses.dataclass
|
||||
class Event:
|
||||
ts: float
|
||||
kind: str # 't_produced', 't_admitted', 't_claimed', 'q_sent', 'q_returned',
|
||||
# 'q_admitted', 'q_claimed', 'pair_claimed', 'pair_ok', 'pair_fail',
|
||||
# 't_expired', 'q_expired', 't_discarded', 'q_discarded',
|
||||
# 'sem_t_free', 'sem_q_free', 'sem_pend_free', 'inv_t', 'inv_q'
|
||||
detail: Any = None
|
||||
|
||||
|
||||
class EventLog:
|
||||
"""记录所有事件,用于事后分析。"""
|
||||
|
||||
def __init__(self):
|
||||
self.events: list[Event] = []
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def record(self, kind: str, detail=None):
|
||||
async with self._lock:
|
||||
self.events.append(Event(ts=time.time(), kind=kind, detail=detail))
|
||||
|
||||
def count(self, kind: str) -> int:
|
||||
return sum(1 for e in self.events if e.kind == kind)
|
||||
|
||||
def filter(self, kind: str) -> list[Event]:
|
||||
return [e for e in self.events if e.kind == kind]
|
||||
|
||||
def dump(self) -> str:
|
||||
lines = []
|
||||
for e in self.events:
|
||||
lines.append(f'{e.ts:.3f} {e.kind} {e.detail or ""}')
|
||||
return '\n'.join(lines)
|
||||
2
tests/requirements.txt
Normal file
2
tests/requirements.txt
Normal file
@@ -0,0 +1,2 @@
|
||||
pytest>=8.0.0
|
||||
pytest-asyncio>=0.24.0
|
||||
82
tests/test_admission_gate.py
Normal file
82
tests/test_admission_gate.py
Normal file
@@ -0,0 +1,82 @@
|
||||
import asyncio
|
||||
import unittest
|
||||
|
||||
from grok_register.core.admission import AdmissionGate
|
||||
|
||||
|
||||
class FakeInventory:
|
||||
def __init__(self):
|
||||
self.t_depth = 0
|
||||
self.q_depth = 0
|
||||
|
||||
|
||||
class AdmissionGateTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_t_production_blocks_at_high_water_and_wakes_at_low_water(self):
|
||||
inventory = FakeInventory()
|
||||
gate = AdmissionGate(inventory, t_low=1, t_high=3, q_low=1, q_high=3)
|
||||
|
||||
inventory.t_depth = 3
|
||||
blocked = asyncio.create_task(gate.acquire_t_production())
|
||||
await asyncio.sleep(0.02)
|
||||
self.assertFalse(blocked.done())
|
||||
|
||||
inventory.t_depth = 1
|
||||
await gate.notify_changed()
|
||||
lease = await asyncio.wait_for(blocked, timeout=1)
|
||||
|
||||
self.assertEqual(gate.t_in_progress, 1)
|
||||
await lease.release()
|
||||
self.assertEqual(gate.t_in_progress, 0)
|
||||
|
||||
async def test_q_batch_reserves_only_current_demand(self):
|
||||
inventory = FakeInventory()
|
||||
gate = AdmissionGate(inventory, t_low=1, t_high=3, q_low=2, q_high=5)
|
||||
|
||||
lease = await gate.acquire_q_batch(max_batch=4)
|
||||
|
||||
self.assertEqual(lease.count, 4)
|
||||
self.assertEqual(gate.q_inflight, 4)
|
||||
|
||||
second = await gate.acquire_q_batch(max_batch=4)
|
||||
self.assertEqual(second.count, 1)
|
||||
self.assertEqual(gate.q_inflight, 5)
|
||||
|
||||
blocked = asyncio.create_task(gate.acquire_q_batch(max_batch=4))
|
||||
await asyncio.sleep(0.02)
|
||||
self.assertFalse(blocked.done())
|
||||
|
||||
await lease.release_all()
|
||||
await second.release_all()
|
||||
inventory.q_depth = 1
|
||||
await gate.notify_changed()
|
||||
third = await asyncio.wait_for(blocked, timeout=1)
|
||||
self.assertEqual(third.count, 4)
|
||||
await third.release_all()
|
||||
|
||||
async def test_q_batch_waits_for_low_water_after_reaching_high_water(self):
|
||||
inventory = FakeInventory()
|
||||
gate = AdmissionGate(inventory, t_low=1, t_high=3, q_low=2, q_high=5)
|
||||
|
||||
lease = await gate.acquire_q_batch(max_batch=5)
|
||||
self.assertEqual(lease.count, 5)
|
||||
self.assertEqual(gate.q_inflight, 5)
|
||||
|
||||
await lease.release_one()
|
||||
self.assertEqual(gate.q_inflight, 4)
|
||||
|
||||
blocked = asyncio.create_task(gate.acquire_q_batch(max_batch=5))
|
||||
await asyncio.sleep(0.02)
|
||||
self.assertFalse(blocked.done())
|
||||
|
||||
await lease.release_one()
|
||||
await lease.release_one()
|
||||
await lease.release_one()
|
||||
|
||||
resumed = await asyncio.wait_for(blocked, timeout=1)
|
||||
self.assertEqual(resumed.count, 4)
|
||||
await lease.release_all()
|
||||
await resumed.release_all()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
475
tests/test_cancel.py
Normal file
475
tests/test_cancel.py
Normal file
@@ -0,0 +1,475 @@
|
||||
"""
|
||||
Layer 2: 取消注入测试
|
||||
|
||||
在每个 await 边界前后强制取消 Worker,验证:
|
||||
- S 取消后 Physical_Sem 和 T_Slot_Sem 不泄漏
|
||||
- P 取消后 Physical_Sem、Q_Pending_Sem、Q_Slot_Sem 不泄漏
|
||||
- C 取消后 Physical_Sem 不泄漏,已 claim 的 T/Q 被核销
|
||||
- 库存守恒: 取消不导致 slot 无界堆积或耗尽
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
import pytest
|
||||
|
||||
from grok_register.core.envelope import ResourceEnvelope
|
||||
from grok_register.core.inventory import Inventory
|
||||
from grok_register.core.observer import Metrics
|
||||
from tests.fakes import FakeTurnstile, FakeEmailService, FakeRegisterAPI, Conservation
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# S_Worker 取消测试
|
||||
# ═══════════════════════════════════════════
|
||||
|
||||
async def _s_worker_impl(browser, inventory, physical_sem, t_slot_sem, metrics, stop):
|
||||
"""S_Worker 简化实现,用于测试。"""
|
||||
while not stop.is_set():
|
||||
await physical_sem.acquire()
|
||||
try:
|
||||
token = await browser.generate()
|
||||
finally:
|
||||
physical_sem.release()
|
||||
|
||||
if token is None:
|
||||
metrics.t_discarded += 1
|
||||
await asyncio.sleep(0.05)
|
||||
continue
|
||||
|
||||
metrics.t_produced += 1
|
||||
now = time.time()
|
||||
try:
|
||||
t_env = await ResourceEnvelope.create_with_slot(
|
||||
'T', token, t_slot_sem, expires_at=now + 300
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
metrics.t_discarded += 1
|
||||
raise
|
||||
|
||||
try:
|
||||
await inventory.put_t(t_env)
|
||||
except Exception:
|
||||
t_env.discard(reason='put_failed')
|
||||
metrics.t_discarded += 1
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
|
||||
class TestSCancel:
|
||||
"""S_Worker 取消语义测试。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_generate_no_leak(self):
|
||||
"""S 在 generate() 期间被取消,sem 不泄漏。"""
|
||||
ft = FakeTurnstile(delay=0.5) # 慢生成
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
task = asyncio.create_task(
|
||||
_s_worker_impl(ft, inv, phys, t_slot, metrics, stop)
|
||||
)
|
||||
await asyncio.sleep(0.05) # 等它进入 generate
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# Physical_Sem 应完全释放
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
# T_Slot_Sem 不变(没生成成功)
|
||||
assert t_slot._value == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_put_t_no_leak(self):
|
||||
"""S 在 put_t 期间被取消(极端: put 内部取消),envelope 被 discard。"""
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
# 快速生成
|
||||
ft = FakeTurnstile(delay=0.001)
|
||||
|
||||
task = asyncio.create_task(
|
||||
_s_worker_impl(ft, inv, phys, t_slot, metrics, stop)
|
||||
)
|
||||
# 在 put_t 之前的极短窗口取消
|
||||
await asyncio.sleep(0.002)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert phys._value == 2
|
||||
# t_slot 可能是 4 或 3,但不应是负数
|
||||
assert t_slot._value >= 0
|
||||
assert t_slot._value <= 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s_repeated_cancel_no_accumulation(self):
|
||||
"""反复取消 S 多次,sem 计数不累积。"""
|
||||
ft = FakeTurnstile(delay=0.01)
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
for _ in range(20):
|
||||
task = asyncio.create_task(
|
||||
_s_worker_impl(ft, inv, phys, t_slot, metrics, stop)
|
||||
)
|
||||
await asyncio.sleep(0.02)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
# 守恒: free + inventory == total (成功入库的 envelope 占着 slot)
|
||||
assert t_slot._value + inv.t_depth == 4, f'T_Slot 守恒违反: free={t_slot._value} inv={inv.t_depth}'
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# P_Worker 取消测试
|
||||
# ═══════════════════════════════════════════
|
||||
|
||||
async def _p_worker_impl(email_svc, inventory, physical_sem, q_pending_sem,
|
||||
q_slot_sem, metrics, stop):
|
||||
"""P_Worker 简化实现,用于测试。"""
|
||||
loop = asyncio.get_event_loop()
|
||||
while not stop.is_set():
|
||||
await q_pending_sem.acquire()
|
||||
_pending_released = False
|
||||
try:
|
||||
await physical_sem.acquire()
|
||||
try:
|
||||
handle, email, password = await email_svc.create_email()
|
||||
sent = True # fake 总是成功
|
||||
finally:
|
||||
physical_sem.release()
|
||||
|
||||
if not sent:
|
||||
q_pending_sem.release()
|
||||
_pending_released = True
|
||||
continue
|
||||
|
||||
metrics.q_sent += 1
|
||||
|
||||
try:
|
||||
code = await asyncio.wait_for(
|
||||
email_svc.poll_code(handle), timeout=95
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
code = None
|
||||
|
||||
if code is None:
|
||||
metrics.q_discarded += 1
|
||||
q_pending_sem.release()
|
||||
_pending_released = True
|
||||
continue
|
||||
|
||||
metrics.q_returned += 1
|
||||
now = time.time()
|
||||
q_env = await ResourceEnvelope.create_with_slot(
|
||||
'Q', {'email': email, 'password': password, 'code': code},
|
||||
q_slot_sem, expires_at=now + 120,
|
||||
)
|
||||
|
||||
try:
|
||||
await inventory.put_q(q_env)
|
||||
except Exception:
|
||||
q_env.discard(reason='put_failed')
|
||||
metrics.q_discarded += 1
|
||||
|
||||
except asyncio.CancelledError:
|
||||
if not _pending_released:
|
||||
q_pending_sem.release()
|
||||
_pending_released = True
|
||||
raise
|
||||
except Exception:
|
||||
metrics.q_discarded += 1
|
||||
finally:
|
||||
if not _pending_released:
|
||||
q_pending_sem.release()
|
||||
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
|
||||
class TestPCancel:
|
||||
"""P_Worker 取消语义测试。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_create_email(self):
|
||||
"""P 在 create_email 期间被取消,所有 sem 正确释放。"""
|
||||
es = FakeEmailService()
|
||||
es.set_cancel_create_at(1)
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
q_pend = asyncio.Semaphore(4)
|
||||
q_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
task = asyncio.create_task(
|
||||
_p_worker_impl(es, inv, phys, q_pend, q_slot, metrics, stop)
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
assert q_pend._value == 4, f'Q_Pending 泄漏: {q_pend._value}'
|
||||
assert q_slot._value == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_poll_code(self):
|
||||
"""P 在 poll_code 期间被取消,不持有 Physical_Sem。"""
|
||||
es = FakeEmailService()
|
||||
es.set_cancel_poll_at(1)
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
q_pend = asyncio.Semaphore(4)
|
||||
q_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
task = asyncio.create_task(
|
||||
_p_worker_impl(es, inv, phys, q_pend, q_slot, metrics, stop)
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# P 等 Q 时不持有 Physical_Sem
|
||||
assert phys._value == 2, f'Physical_Sem 应完全释放: {phys._value}'
|
||||
assert q_pend._value == 4, f'Q_Pending 泄漏: {q_pend._value}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_p_repeated_cancel_no_accumulation(self):
|
||||
"""反复取消 P 多次,sem 计数不累积。"""
|
||||
es = FakeEmailService()
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
q_pend = asyncio.Semaphore(4)
|
||||
q_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
|
||||
for _ in range(20):
|
||||
task = asyncio.create_task(
|
||||
_p_worker_impl(es, inv, phys, q_pend, q_slot, metrics, stop)
|
||||
)
|
||||
await asyncio.sleep(0.03)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
assert q_pend._value == 4, f'Q_Pending 泄漏: {q_pend._value}'
|
||||
# 守恒: free + inventory == total
|
||||
assert q_slot._value + inv.q_depth == 4, f'Q_Slot 守恒违反: free={q_slot._value} inv={inv.q_depth}'
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# C_Worker 取消测试
|
||||
# ═══════════════════════════════════════════
|
||||
|
||||
async def _c_worker_impl(register_api, browser_unused, inventory, physical_sem,
|
||||
metrics, stop, file_lock):
|
||||
"""C_Worker 简化实现,用于测试。"""
|
||||
while not stop.is_set():
|
||||
try:
|
||||
async with inventory.claim_pair() as pair:
|
||||
t_val = pair.t.value
|
||||
q_val = pair.q.value
|
||||
await physical_sem.acquire()
|
||||
try:
|
||||
sso = await register_api.register(
|
||||
q_val['email'], q_val['password'], q_val['code'], t_val
|
||||
)
|
||||
if sso:
|
||||
metrics.pair_consumed_ok += 1
|
||||
metrics.success_count += 1
|
||||
else:
|
||||
metrics.pair_consumed_fail += 1
|
||||
finally:
|
||||
physical_sem.release()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
metrics.pair_consumed_fail += 1
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
|
||||
class TestCCancel:
|
||||
"""C_Worker 取消语义测试。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_register_pair_settled(self):
|
||||
"""C 在 register 期间被取消,pair 已 claim,必须核销。"""
|
||||
fra = FakeRegisterAPI(delay=0.5)
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(4)
|
||||
q_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
# 用带 metrics 的 inventory
|
||||
inv_with_metrics = Inventory(metrics=Metrics())
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_slot)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'email': 'a@b.com', 'password': 'p', 'code': '123'}, q_slot)
|
||||
await inv_with_metrics.put_t(t_env)
|
||||
await inv_with_metrics.put_q(q_env)
|
||||
|
||||
task = asyncio.create_task(
|
||||
_c_worker_impl(fra, None, inv_with_metrics, phys, metrics, stop, fl)
|
||||
)
|
||||
await asyncio.sleep(0.1) # 等它 claim 进入 register
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# pair 已核销: slot 应释放
|
||||
assert t_slot._value == 4, f'T slot 未核销: {t_slot._value}'
|
||||
assert q_slot._value == 4, f'Q slot 未核销: {q_slot._value}'
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_claim_wait_no_pop(self):
|
||||
"""C 在 claim_pair 等待时取消,不弹出资源。"""
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(4)
|
||||
q_slot = asyncio.Semaphore(4)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
fra = FakeRegisterAPI()
|
||||
|
||||
# 放一个 T 但不放 Q,让 C 等待
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_slot)
|
||||
await inv.put_t(t_env)
|
||||
|
||||
task = asyncio.create_task(
|
||||
_c_worker_impl(fra, None, inv, phys, metrics, stop, fl)
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# T 仍在 inventory 中
|
||||
assert inv.t_depth == 1, 'claim 等待时取消不应弹出 T'
|
||||
assert phys._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_c_repeated_cancel_no_accumulation(self):
|
||||
"""反复取消 C 多次,sem 计数不累积。"""
|
||||
fra = FakeRegisterAPI(delay=0.01)
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(2)
|
||||
t_slot = asyncio.Semaphore(10)
|
||||
q_slot = asyncio.Semaphore(10)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
for _ in range(10):
|
||||
# 每次放一对
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', f'tok_{_}', t_slot)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'email': f'e@{_}', 'password': 'p', 'code': '123'}, q_slot)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
task = asyncio.create_task(
|
||||
_c_worker_impl(fra, None, inv, phys, metrics, stop, fl)
|
||||
)
|
||||
await asyncio.sleep(0.03)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert phys._value == 2, f'Physical_Sem 泄漏: {phys._value}'
|
||||
# 所有 slot 应已核销
|
||||
assert t_slot._value == 10, f'T_Slot 泄漏: {t_slot._value}'
|
||||
assert q_slot._value == 10, f'Q_Slot 泄漏: {q_slot._value}'
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# 端到端取消: S+P+C 同时运行时取消
|
||||
# ═══════════════════════════════════════════
|
||||
class TestE2ECancel:
|
||||
"""端到端取消测试: S/P/C 同时运行时取消,验证全局守恒。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_all_workers_sem_reset(self):
|
||||
"""同时启动 S/P/C,然后全部取消,所有 sem 恢复初始值。"""
|
||||
ft = FakeTurnstile(delay=0.05)
|
||||
es = FakeEmailService()
|
||||
fra = FakeRegisterAPI(delay=0.05)
|
||||
|
||||
inv = Inventory()
|
||||
phys = asyncio.Semaphore(4)
|
||||
t_slot = asyncio.Semaphore(8)
|
||||
q_slot = asyncio.Semaphore(8)
|
||||
q_pend = asyncio.Semaphore(12)
|
||||
metrics = Metrics()
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
tasks = []
|
||||
for _ in range(3):
|
||||
tasks.append(asyncio.create_task(
|
||||
_s_worker_impl(ft, inv, phys, t_slot, metrics, stop)
|
||||
))
|
||||
tasks.append(asyncio.create_task(
|
||||
_p_worker_impl(es, inv, phys, q_pend, q_slot, metrics, stop)
|
||||
))
|
||||
tasks.append(asyncio.create_task(
|
||||
_c_worker_impl(fra, None, inv, phys, metrics, stop, fl)
|
||||
))
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# 全部停止后,守恒检查: free + inventory == total
|
||||
assert phys._value >= 0 and phys._value <= 4
|
||||
assert t_slot._value + inv.t_depth == 8, f'T 守恒违反: free={t_slot._value} inv={inv.t_depth}'
|
||||
assert q_slot._value + inv.q_depth == 8, f'Q 守恒违反: free={q_slot._value} inv={inv.q_depth}'
|
||||
assert q_pend._value >= 0 and q_pend._value <= 12
|
||||
|
||||
# 守恒检查
|
||||
cons = Conservation(t_slot, q_slot, q_pend, inv)
|
||||
cons.check_t_conservation(8)
|
||||
cons.check_q_conservation(8)
|
||||
cons.check_pending_conservation(12)
|
||||
491
tests/test_core.py
Normal file
491
tests/test_core.py
Normal file
@@ -0,0 +1,491 @@
|
||||
"""
|
||||
Layer 1: 单元测试
|
||||
|
||||
测 ResourceEnvelope、PairLease、Inventory 的局部语义:
|
||||
- release_slot_once() 幂等
|
||||
- create_with_slot() slot 获取失败不创建 envelope
|
||||
- is_expired() 正确判断
|
||||
- put_t/put_q 成功前所有权属于 Worker,成功后属于 Inventory
|
||||
- claim_pair 等待时取消不弹资源
|
||||
- claim_pair 成功后 pair 属于 PairLease
|
||||
- PairLease 退出时核销 T/Q
|
||||
- 过期资源不会交付给 C
|
||||
- lazy cleanup 只在事件触发时发生
|
||||
- discard() 幂等释放
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
import pytest
|
||||
|
||||
from grok_register.core.envelope import ResourceEnvelope
|
||||
from grok_register.core.inventory import Inventory, PairLease
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# ResourceEnvelope 测试
|
||||
# ═══════════════════════════════════════════
|
||||
class TestResourceEnvelope:
|
||||
"""ResourceEnvelope 语义测试。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_with_slot_acquires_sem(self):
|
||||
"""create_with_slot 成功后 semaphore 值减 1。"""
|
||||
sem = asyncio.Semaphore(3)
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok1', sem)
|
||||
assert sem._value == 2
|
||||
assert env.value == 'tok1'
|
||||
assert env.kind == 'T'
|
||||
assert not env.released
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_with_slot_fails_no_envelope(self):
|
||||
"""slot 获取失败(取消)时不创建 envelope,sem 不变。"""
|
||||
sem = asyncio.Semaphore(0)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
# 手动取消
|
||||
task = asyncio.current_task()
|
||||
task.cancel()
|
||||
try:
|
||||
await ResourceEnvelope.create_with_slot('T', 'tok1', sem)
|
||||
except asyncio.CancelledError:
|
||||
# sem 不应改变(因为 acquire 被取消了)
|
||||
raise
|
||||
# sem 仍为 0
|
||||
assert sem._value == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_slot_once_idempotent(self):
|
||||
"""release_slot_once() 幂等:重复调用只释放一次。"""
|
||||
sem = asyncio.Semaphore(3)
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok1', sem)
|
||||
assert sem._value == 2
|
||||
|
||||
env.release_slot_once(reason='test1')
|
||||
assert sem._value == 3
|
||||
assert env.released
|
||||
|
||||
# 第二次调用不应再次释放
|
||||
env.release_slot_once(reason='test2')
|
||||
assert sem._value == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discard_releases_slot(self):
|
||||
"""discard() 释放 slot。"""
|
||||
sem = asyncio.Semaphore(2)
|
||||
env = await ResourceEnvelope.create_with_slot('Q', {'code': '123'}, sem)
|
||||
assert sem._value == 1
|
||||
env.discard(reason='test')
|
||||
assert sem._value == 2
|
||||
assert env.released
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discard_idempotent(self):
|
||||
"""discard() 幂等。"""
|
||||
sem = asyncio.Semaphore(2)
|
||||
env = await ResourceEnvelope.create_with_slot('Q', 'val', sem)
|
||||
env.discard(reason='first')
|
||||
env.discard(reason='second')
|
||||
assert sem._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_expired_with_expires_at(self):
|
||||
"""有 expires_at 时正确判断过期。"""
|
||||
sem = asyncio.Semaphore(2)
|
||||
now = time.time()
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok', sem, expires_at=now - 1)
|
||||
assert env.is_expired(now=now)
|
||||
|
||||
env2 = await ResourceEnvelope.create_with_slot('T', 'tok2', sem, expires_at=now + 100)
|
||||
assert not env2.is_expired(now=now)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_expired_without_expires_at(self):
|
||||
"""没有 expires_at 时永远不过期。"""
|
||||
sem = asyncio.Semaphore(2)
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok', sem)
|
||||
assert not env.is_expired(now=time.time() + 99999)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_with_slot_meta(self):
|
||||
"""meta 字段正确传递。"""
|
||||
sem = asyncio.Semaphore(1)
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok', sem, meta={'source': 'test'})
|
||||
assert env.meta == {'source': 'test'}
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# Inventory 测试
|
||||
# ═══════════════════════════════════════════
|
||||
class TestInventory:
|
||||
"""Inventory 语义测试。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_put_t_then_claim(self):
|
||||
"""put_t 后 claim_pair 能拿到 pair。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'token1', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'abc'}, q_sem)
|
||||
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
assert inv.t_depth == 1
|
||||
assert inv.q_depth == 1
|
||||
|
||||
# claim
|
||||
async with inv.claim_pair() as pair:
|
||||
assert pair.t.value == 'token1'
|
||||
assert pair.q.value == {'code': 'abc'}
|
||||
# claim 后 inventory 应为空
|
||||
assert inv.t_depth == 0
|
||||
assert inv.q_depth == 0
|
||||
|
||||
# PairLease 退出后 slot 应被释放
|
||||
assert t_sem._value == 2
|
||||
assert q_sem._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claim_waits_for_both(self):
|
||||
"""只有 T 没有 Q 时,claim_pair 应等待。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
await inv.put_t(t_env)
|
||||
|
||||
claimed = False
|
||||
|
||||
async def claimer():
|
||||
nonlocal claimed
|
||||
async with inv.claim_pair() as pair:
|
||||
claimed = True
|
||||
|
||||
task = asyncio.create_task(claimer())
|
||||
await asyncio.sleep(0.1)
|
||||
assert not claimed, 'claim 不应在缺少 Q 时成功'
|
||||
|
||||
# 现在放入 Q
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
await asyncio.wait_for(task, timeout=2)
|
||||
assert claimed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claim_cancel_no_pop(self):
|
||||
"""claim_pair 等待时取消,不应弹出任何资源。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
await inv.put_t(t_env)
|
||||
assert inv.t_depth == 1
|
||||
|
||||
async def claimer():
|
||||
async with inv.claim_pair() as pair:
|
||||
pass # 不应到达
|
||||
|
||||
task = asyncio.create_task(claimer())
|
||||
await asyncio.sleep(0.1)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# T 仍在 inventory 中
|
||||
assert inv.t_depth == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pair_lease_settles_on_success(self):
|
||||
"""PairLease 成功退出时核销 T/Q 并释放 slot。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
async with inv.claim_pair() as pair:
|
||||
# slot 还在 pair 中
|
||||
assert t_sem._value == 1
|
||||
assert q_sem._value == 1
|
||||
|
||||
# 退出后 slot 释放
|
||||
assert t_sem._value == 2
|
||||
assert q_sem._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pair_lease_settles_on_exception(self):
|
||||
"""PairLease 异常退出时也核销 T/Q。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
async with inv.claim_pair() as pair:
|
||||
raise ValueError('boom')
|
||||
|
||||
# slot 仍被核销
|
||||
assert t_sem._value == 2
|
||||
assert q_sem._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pair_lease_settles_on_cancel(self):
|
||||
"""PairLease 取消退出时也核销 T/Q。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
async def consumer():
|
||||
async with inv.claim_pair() as pair:
|
||||
# 在 pair 内部被取消
|
||||
await asyncio.sleep(999)
|
||||
|
||||
task = asyncio.create_task(consumer())
|
||||
await asyncio.sleep(0.1)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# slot 仍被核销
|
||||
assert t_sem._value == 2
|
||||
assert q_sem._value == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_expired_t_not_delivered(self):
|
||||
"""过期的 T 不会交付给 C。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
now = time.time()
|
||||
# T 已过期
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'expired_tok', t_sem, expires_at=now - 10)
|
||||
# Q 有效
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem, expires_at=now + 100)
|
||||
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
# claim 应该把过期的 T 清理掉,然后等待新的 T
|
||||
claimed = False
|
||||
|
||||
async def claimer():
|
||||
nonlocal claimed
|
||||
async with inv.claim_pair() as pair:
|
||||
claimed = True
|
||||
|
||||
task = asyncio.create_task(claimer())
|
||||
await asyncio.sleep(0.2)
|
||||
assert not claimed, '过期 T 不应交付'
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# 过期 T 被清理,slot 被释放
|
||||
assert t_sem._value == 2
|
||||
# Q 仍在 inventory 中
|
||||
assert inv.q_depth == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_expired_q_not_delivered(self):
|
||||
"""过期的 Q 不会交付给 C。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
q_sem = asyncio.Semaphore(2)
|
||||
|
||||
now = time.time()
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem, expires_at=now + 100)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem, expires_at=now - 10)
|
||||
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
claimed = False
|
||||
|
||||
async def claimer():
|
||||
nonlocal claimed
|
||||
async with inv.claim_pair() as pair:
|
||||
claimed = True
|
||||
|
||||
task = asyncio.create_task(claimer())
|
||||
await asyncio.sleep(0.2)
|
||||
assert not claimed, '过期 Q 不应交付'
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert q_sem._value == 2
|
||||
assert inv.t_depth == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazy_cleanup_on_put(self):
|
||||
"""lazy cleanup: put 时清理已过期的库存。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(4)
|
||||
|
||||
now = time.time()
|
||||
# 放入 3 个过期的 T
|
||||
for i in range(3):
|
||||
env = await ResourceEnvelope.create_with_slot(
|
||||
'T', f'expired_{i}', t_sem, expires_at=now - 10
|
||||
)
|
||||
await inv.put_t(env)
|
||||
|
||||
# 由于 put 时 lazy cleanup 已清理过期项,库存应为空
|
||||
assert inv.t_depth == 0
|
||||
# slot 应已释放
|
||||
assert t_sem._value == 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazy_cleanup_silent_until_next_event(self):
|
||||
"""lazy cleanup 不在静默期主动释放,下一次 put 才触发清理。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(4)
|
||||
|
||||
env = await ResourceEnvelope.create_with_slot(
|
||||
'T', 'will_expire', t_sem, expires_at=time.time() + 0.05
|
||||
)
|
||||
await inv.put_t(env)
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
assert inv.t_depth == 1
|
||||
assert t_sem._value == 3
|
||||
|
||||
fresh = await ResourceEnvelope.create_with_slot(
|
||||
'T', 'fresh', t_sem, expires_at=time.time() + 100
|
||||
)
|
||||
await inv.put_t(fresh)
|
||||
|
||||
assert inv.t_depth == 1
|
||||
assert t_sem._value == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazy_cleanup_on_claim(self):
|
||||
"""lazy cleanup: claim_pair 时清理过期库存。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(4)
|
||||
q_sem = asyncio.Semaphore(4)
|
||||
|
||||
now = time.time()
|
||||
# 直接注入过期项到 inventory 内部(绕过 put 的清理)
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'expired', t_sem, expires_at=now - 10)
|
||||
inv._t_buf.append(env)
|
||||
|
||||
# 有效的 Q
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': 'x'}, q_sem, expires_at=now + 100)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
# claim 时应清理过期的 T
|
||||
claimed = False
|
||||
async def claimer():
|
||||
nonlocal claimed
|
||||
async with inv.claim_pair() as pair:
|
||||
claimed = True
|
||||
|
||||
task = asyncio.create_task(claimer())
|
||||
await asyncio.sleep(0.2)
|
||||
assert not claimed # T 被清了,没有 pair
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert inv.t_depth == 0
|
||||
assert t_sem._value == 4 # slot 已释放
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fifo_order(self):
|
||||
"""Inventory 按 FIFO 顺序配对。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(10)
|
||||
q_sem = asyncio.Semaphore(10)
|
||||
|
||||
for i in range(3):
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', f'tok_{i}', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': f'c_{i}'}, q_sem)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
pairs = []
|
||||
for _ in range(3):
|
||||
async with inv.claim_pair() as pair:
|
||||
pairs.append((pair.t.value, pair.q.value['code']))
|
||||
|
||||
assert pairs == [('tok_0', 'c_0'), ('tok_1', 'c_1'), ('tok_2', 'c_2')]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_concurrent_claims(self):
|
||||
"""多个 C 同时 claim,不应重复消费同一个 T/Q。"""
|
||||
inv = Inventory()
|
||||
t_sem = asyncio.Semaphore(10)
|
||||
q_sem = asyncio.Semaphore(10)
|
||||
|
||||
n = 5
|
||||
for i in range(n):
|
||||
t_env = await ResourceEnvelope.create_with_slot('T', f'tok_{i}', t_sem)
|
||||
q_env = await ResourceEnvelope.create_with_slot('Q', {'code': f'c_{i}'}, q_sem)
|
||||
await inv.put_t(t_env)
|
||||
await inv.put_q(q_env)
|
||||
|
||||
results = []
|
||||
|
||||
async def claimer(idx):
|
||||
async with inv.claim_pair() as pair:
|
||||
results.append((idx, pair.t.value))
|
||||
|
||||
tasks = [asyncio.create_task(claimer(i)) for i in range(n)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# 每个 T 只被消费一次
|
||||
t_values = [r[1] for r in results]
|
||||
assert len(set(t_values)) == n, f'重复消费: {t_values}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metrics_admitted(self):
|
||||
"""put 成功时 metrics.t_admitted / q_admitted 递增。"""
|
||||
from grok_register.core.observer import Metrics
|
||||
m = Metrics()
|
||||
inv = Inventory(metrics=m)
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem)
|
||||
await inv.put_t(env)
|
||||
assert m.t_admitted == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metrics_expired(self):
|
||||
"""过期资源 put 时 metrics.t_expired 递增。"""
|
||||
from grok_register.core.observer import Metrics
|
||||
m = Metrics()
|
||||
inv = Inventory(metrics=m)
|
||||
t_sem = asyncio.Semaphore(2)
|
||||
|
||||
env = await ResourceEnvelope.create_with_slot('T', 'tok', t_sem, expires_at=time.time() - 10)
|
||||
await inv.put_t(env)
|
||||
assert m.t_expired == 1
|
||||
assert m.t_admitted == 0
|
||||
58
tests/test_inventory_unittest.py
Normal file
58
tests/test_inventory_unittest.py
Normal file
@@ -0,0 +1,58 @@
|
||||
import asyncio
|
||||
import time
|
||||
import unittest
|
||||
|
||||
from grok_register.core.envelope import ResourceEnvelope
|
||||
from grok_register.core.inventory import Inventory
|
||||
from grok_register.core.observer import Metrics
|
||||
|
||||
|
||||
class InventoryUnitTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_create_with_slot_releases_slot_if_constructor_fails(self):
|
||||
class BrokenEnvelope(ResourceEnvelope):
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
sem = asyncio.Semaphore(1)
|
||||
|
||||
with self.assertRaises(RuntimeError):
|
||||
await BrokenEnvelope.create_with_slot("T", "tok", sem)
|
||||
|
||||
self.assertEqual(sem._value, 1)
|
||||
|
||||
async def test_lazy_cleanup_scans_past_fresh_head(self):
|
||||
metrics = Metrics()
|
||||
inventory = Inventory(metrics=metrics)
|
||||
t_sem = asyncio.Semaphore(3)
|
||||
q_sem = asyncio.Semaphore(1)
|
||||
now = time.time()
|
||||
|
||||
fresh_head = await ResourceEnvelope.create_with_slot(
|
||||
"T", "fresh_head", t_sem, expires_at=now + 100
|
||||
)
|
||||
stale_tail = await ResourceEnvelope.create_with_slot(
|
||||
"T", "stale_tail", t_sem, expires_at=now + 100
|
||||
)
|
||||
fresh_tail = await ResourceEnvelope.create_with_slot(
|
||||
"T", "fresh_tail", t_sem, expires_at=now + 100
|
||||
)
|
||||
|
||||
await inventory.put_t(fresh_head)
|
||||
await inventory.put_t(stale_tail)
|
||||
await inventory.put_t(fresh_tail)
|
||||
self.assertEqual(t_sem._value, 0)
|
||||
|
||||
stale_tail.expires_at = time.time() - 10
|
||||
|
||||
q_env = await ResourceEnvelope.create_with_slot(
|
||||
"Q", "q", q_sem, expires_at=now + 100
|
||||
)
|
||||
await inventory.put_q(q_env)
|
||||
|
||||
self.assertEqual(t_sem._value, 1)
|
||||
self.assertEqual(metrics.t_expired, 1)
|
||||
self.assertEqual(inventory.t_depth, 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
482
tests/test_property.py
Normal file
482
tests/test_property.py
Normal file
@@ -0,0 +1,482 @@
|
||||
"""
|
||||
Layer 3: 性质测试 (Property Tests)
|
||||
|
||||
随机生成 S/P/C 成功、失败、取消、超时、过期事件,持续检查:
|
||||
- slot 守恒 (T/Q 总量恒定)
|
||||
- 无重复 claim (同一个 T/Q 不被两个 PairLease 拿到)
|
||||
- 无单边持有等待 (C 不先拿 T 再等 Q)
|
||||
- 过期资源不交付
|
||||
- 取消后不泄漏
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
import random
|
||||
import pytest
|
||||
|
||||
from grok_register.core.envelope import ResourceEnvelope
|
||||
from grok_register.core.inventory import Inventory
|
||||
from grok_register.core.observer import Metrics
|
||||
from tests.fakes import FakeTurnstile, FakeEmailService, FakeRegisterAPI, Conservation, EventLog
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# 随机性质测试引擎
|
||||
# ═══════════════════════════════════════════
|
||||
|
||||
class PropertyEngine:
|
||||
"""驱动 S/P/C 随机行为,持续检查不变量。"""
|
||||
|
||||
def __init__(self, t_slot_cap=6, q_slot_cap=6, q_pending_cap=8, phys_cap=4,
|
||||
t_max_age=2.0, q_max_age=2.0, seed=None):
|
||||
self.rng = random.Random(seed)
|
||||
self.t_slot_sem = asyncio.Semaphore(t_slot_cap)
|
||||
self.q_slot_sem = asyncio.Semaphore(q_slot_cap)
|
||||
self.q_pending_sem = asyncio.Semaphore(q_pending_cap)
|
||||
self.phys_sem = asyncio.Semaphore(phys_cap)
|
||||
self.t_slot_cap = t_slot_cap
|
||||
self.q_slot_cap = q_slot_cap
|
||||
self.q_pending_cap = q_pending_cap
|
||||
self.t_max_age = t_max_age
|
||||
self.q_max_age = q_max_age
|
||||
|
||||
self.metrics = Metrics()
|
||||
self.inventory = Inventory(metrics=self.metrics)
|
||||
self.conservation = Conservation(
|
||||
self.t_slot_sem, self.q_slot_sem, self.q_pending_sem, self.inventory
|
||||
)
|
||||
self.events = EventLog()
|
||||
|
||||
self._claimed_tokens = set() # 检测重复 claim
|
||||
self._pairleased_t = 0 # 当前 PairLease 持有的 T
|
||||
self._pairleased_q = 0 # 当前 PairLease 持有的 Q
|
||||
|
||||
def check_invariants_soft(self):
|
||||
"""运行时软检查: 非负、不超限。
|
||||
|
||||
注意: Worker 运行期间 slot 可能在 create_with_slot 和 put_t 之间的
|
||||
envelope 中(worker 持有但无法从外部追踪)。严格守恒检查用
|
||||
check_invariants_final() 在所有 Worker 停止后执行。
|
||||
"""
|
||||
assert self.t_slot_sem._value >= 0, 'T slot 为负'
|
||||
assert self.q_slot_sem._value >= 0, 'Q slot 为负'
|
||||
assert self.q_pending_sem._value >= 0, 'Q pending 为负'
|
||||
assert self.q_pending_sem._value <= self.q_pending_cap, 'Q pending 超限'
|
||||
assert self.phys_sem._value >= 0, 'Physical 为负'
|
||||
# inventory 深度不超过 slot 总量
|
||||
assert self.inventory.t_depth <= self.t_slot_cap, 'T inventory 超限'
|
||||
assert self.inventory.q_depth <= self.q_slot_cap, 'Q inventory 超限'
|
||||
|
||||
def check_invariants_final(self):
|
||||
"""最终严格守恒检查: 所有 Worker 停止后执行。
|
||||
|
||||
此时没有 Worker 持有中间状态的 slot,也没有 PairLease。
|
||||
free + inventory == total
|
||||
"""
|
||||
self.conservation.check_t_conservation(self.t_slot_cap)
|
||||
self.conservation.check_q_conservation(self.q_slot_cap)
|
||||
self.conservation.check_pending_conservation(self.q_pending_cap)
|
||||
|
||||
async def s_produce(self):
|
||||
"""S 产生一个 T。返回 True/False/None(cancelled)。"""
|
||||
# 模拟物理资源使用
|
||||
await self.phys_sem.acquire()
|
||||
delay = self.rng.uniform(0.001, 0.02)
|
||||
try:
|
||||
await asyncio.sleep(delay)
|
||||
finally:
|
||||
self.phys_sem.release()
|
||||
|
||||
# 模拟随机失败
|
||||
if self.rng.random() < 0.1:
|
||||
self.metrics.t_discarded += 1
|
||||
await self.events.record('t_discarded', 'gen_fail')
|
||||
return False
|
||||
|
||||
token = f'tok_{self.metrics.t_produced}'
|
||||
self.metrics.t_produced += 1
|
||||
await self.events.record('t_produced', token)
|
||||
|
||||
# 创建 envelope + 入库
|
||||
# 关键: create_with_slot 成功后 slot 在 envelope 中,
|
||||
# 必须追踪 worker_held 直到 put_t 成功或 discard。
|
||||
now = time.time()
|
||||
expire = now + self.t_max_age if self.rng.random() > 0.15 else now - 1
|
||||
t_env = None
|
||||
try:
|
||||
t_env = await ResourceEnvelope.create_with_slot(
|
||||
'T', token, self.t_slot_sem, expires_at=expire
|
||||
)
|
||||
# slot 已获取,标记 worker 持有
|
||||
self.conservation.worker_acquire_t()
|
||||
await self.inventory.put_t(t_env)
|
||||
# put 成功,所有权转移到 inventory
|
||||
self.conservation.worker_release_t()
|
||||
except Exception:
|
||||
if t_env is not None and not t_env.released:
|
||||
t_env.discard(reason='error_cleanup')
|
||||
self.conservation.worker_release_t()
|
||||
self.metrics.t_discarded += 1
|
||||
return False
|
||||
except asyncio.CancelledError:
|
||||
if t_env is not None and not t_env.released:
|
||||
t_env.discard(reason='cancel_cleanup')
|
||||
self.conservation.worker_release_t()
|
||||
self.metrics.t_discarded += 1
|
||||
raise
|
||||
|
||||
self.check_invariants_soft()
|
||||
return True
|
||||
|
||||
async def p_produce(self):
|
||||
"""P 产生一个 Q。返回 True/False/None(cancelled)。"""
|
||||
await self.q_pending_sem.acquire()
|
||||
_pend_released = False
|
||||
|
||||
try:
|
||||
# 创建邮箱 + 发请求(持有 Physical)
|
||||
await self.phys_sem.acquire()
|
||||
try:
|
||||
delay = self.rng.uniform(0.001, 0.02)
|
||||
await asyncio.sleep(delay)
|
||||
finally:
|
||||
self.phys_sem.release()
|
||||
|
||||
# 模拟随机失败
|
||||
if self.rng.random() < 0.1:
|
||||
self.metrics.q_discarded += 1
|
||||
self.q_pending_sem.release()
|
||||
_pend_released = True
|
||||
await self.events.record('q_discarded', 'send_fail')
|
||||
return False
|
||||
|
||||
self.metrics.q_sent += 1
|
||||
|
||||
# 等待 Q 返回(不持有 Physical)
|
||||
poll_delay = self.rng.uniform(0.001, 0.05)
|
||||
await asyncio.sleep(poll_delay)
|
||||
|
||||
# 模拟超时/失败
|
||||
if self.rng.random() < 0.1:
|
||||
self.metrics.q_discarded += 1
|
||||
self.q_pending_sem.release()
|
||||
_pend_released = True
|
||||
await self.events.record('q_discarded', 'poll_timeout')
|
||||
return False
|
||||
|
||||
code = f'{self.rng.randint(100000, 999999)}'
|
||||
self.metrics.q_returned += 1
|
||||
await self.events.record('q_returned', code)
|
||||
|
||||
# 创建 envelope + 入库
|
||||
now = time.time()
|
||||
expire = now + self.q_max_age if self.rng.random() > 0.15 else now - 1
|
||||
q_env = None
|
||||
try:
|
||||
q_env = await ResourceEnvelope.create_with_slot(
|
||||
'Q', {'email': f'e@{code}', 'password': 'p', 'code': code},
|
||||
self.q_slot_sem, expires_at=expire,
|
||||
)
|
||||
self.conservation.worker_acquire_q()
|
||||
await self.inventory.put_q(q_env)
|
||||
self.conservation.worker_release_q()
|
||||
except Exception:
|
||||
if q_env is not None and not q_env.released:
|
||||
q_env.discard(reason='error_cleanup')
|
||||
self.conservation.worker_release_q()
|
||||
self.metrics.q_discarded += 1
|
||||
except asyncio.CancelledError:
|
||||
if q_env is not None and not q_env.released:
|
||||
q_env.discard(reason='cancel_cleanup')
|
||||
self.conservation.worker_release_q()
|
||||
self.metrics.q_discarded += 1
|
||||
raise
|
||||
|
||||
except asyncio.CancelledError:
|
||||
if not _pend_released:
|
||||
self.q_pending_sem.release()
|
||||
_pend_released = True
|
||||
raise
|
||||
except Exception:
|
||||
self.metrics.q_discarded += 1
|
||||
finally:
|
||||
if not _pend_released:
|
||||
self.q_pending_sem.release()
|
||||
|
||||
self.check_invariants_soft()
|
||||
return True
|
||||
|
||||
async def c_consume(self):
|
||||
"""C 消费一个 pair。返回 True/False/None(cancelled)。"""
|
||||
try:
|
||||
async with self.inventory.claim_pair() as pair:
|
||||
t_val = pair.t.value
|
||||
q_val = pair.q.value
|
||||
|
||||
# 检测重复 claim
|
||||
assert t_val not in self._claimed_tokens, f'重复 claim: {t_val}'
|
||||
self._claimed_tokens.add(t_val)
|
||||
|
||||
# pairleased 计数 — 必须在 claim 后立即标记,
|
||||
# 否则在下一个 await 前被取消会导致守恒违反
|
||||
self._pairleased_t += 1
|
||||
self._pairleased_q += 1
|
||||
|
||||
# 消费(持有 Physical)
|
||||
await self.phys_sem.acquire()
|
||||
try:
|
||||
delay = self.rng.uniform(0.001, 0.02)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
# 模拟随机失败
|
||||
if self.rng.random() < 0.15:
|
||||
self.metrics.pair_consumed_fail += 1
|
||||
await self.events.record('pair_fail', t_val)
|
||||
else:
|
||||
self.metrics.pair_consumed_ok += 1
|
||||
self.metrics.success_count += 1
|
||||
await self.events.record('pair_ok', t_val)
|
||||
finally:
|
||||
self.phys_sem.release()
|
||||
|
||||
# PairLease 退出后(正常路径)
|
||||
self._pairleased_t -= 1
|
||||
self._pairleased_q -= 1
|
||||
|
||||
except asyncio.CancelledError:
|
||||
# PairLease.__aexit__ 已核销 slot。
|
||||
# 如果 _pairleased 已 +1,则 -1;否则不操作。
|
||||
if self._pairleased_t > 0:
|
||||
self._pairleased_t -= 1
|
||||
if self._pairleased_q > 0:
|
||||
self._pairleased_q -= 1
|
||||
raise
|
||||
except Exception:
|
||||
self.metrics.pair_consumed_fail += 1
|
||||
if self._pairleased_t > 0:
|
||||
self._pairleased_t -= 1
|
||||
if self._pairleased_q > 0:
|
||||
self._pairleased_q -= 1
|
||||
|
||||
self.check_invariants_soft()
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
class TestProperty:
|
||||
"""性质测试: 随机事件流 + 持续不变量检查。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_random_events_10s(self):
|
||||
"""随机 S/P/C 事件流 10 秒,持续检查不变量。"""
|
||||
engine = PropertyEngine(seed=42)
|
||||
stop = asyncio.Event()
|
||||
errors = []
|
||||
|
||||
async def s_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.s_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'S: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
async def p_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.p_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'P: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
async def c_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.c_consume()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'C: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
tasks = []
|
||||
for _ in range(2):
|
||||
tasks.append(asyncio.create_task(s_loop()))
|
||||
tasks.append(asyncio.create_task(p_loop()))
|
||||
tasks.append(asyncio.create_task(c_loop()))
|
||||
|
||||
await asyncio.sleep(10)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# 最终不变量检查
|
||||
engine.check_invariants_final()
|
||||
|
||||
assert not errors, f'不变量违反:\n' + '\n'.join(errors)
|
||||
print(f'\n[Property] t_produced={engine.metrics.t_produced} '
|
||||
f'q_sent={engine.metrics.q_sent} q_returned={engine.metrics.q_returned} '
|
||||
f'pair_ok={engine.metrics.pair_consumed_ok} pair_fail={engine.metrics.pair_consumed_fail} '
|
||||
f'success={engine.metrics.success_count}')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_random_with_cancellations(self):
|
||||
"""随机事件 + 随机取消 Worker,不变量仍成立。"""
|
||||
engine = PropertyEngine(seed=123)
|
||||
stop = asyncio.Event()
|
||||
errors = []
|
||||
|
||||
async def s_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.s_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'S: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
async def p_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.p_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'P: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
async def c_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.c_consume()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'C: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.005, 0.03))
|
||||
|
||||
async def cancel_loop():
|
||||
"""随机取消并重启 Worker。"""
|
||||
while not stop.is_set():
|
||||
await asyncio.sleep(engine.rng.uniform(0.5, 2.0))
|
||||
if tasks:
|
||||
idx = engine.rng.randint(0, len(tasks) - 1)
|
||||
t = tasks[idx]
|
||||
if not t.done():
|
||||
t.cancel()
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
# 重启
|
||||
role = idx % 3
|
||||
if role == 0:
|
||||
tasks[idx] = asyncio.create_task(s_loop())
|
||||
elif role == 1:
|
||||
tasks[idx] = asyncio.create_task(p_loop())
|
||||
else:
|
||||
tasks[idx] = asyncio.create_task(c_loop())
|
||||
|
||||
tasks = []
|
||||
for _ in range(2):
|
||||
tasks.append(asyncio.create_task(s_loop()))
|
||||
tasks.append(asyncio.create_task(p_loop()))
|
||||
tasks.append(asyncio.create_task(c_loop()))
|
||||
tasks.append(asyncio.create_task(cancel_loop()))
|
||||
|
||||
await asyncio.sleep(10)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
engine.check_invariants_final()
|
||||
assert not errors, f'不变量违反:\n' + '\n'.join(errors)
|
||||
print(f'\n[Property+Cancel] t_produced={engine.metrics.t_produced} '
|
||||
f'pair_ok={engine.metrics.pair_consumed_ok} '
|
||||
f'success={engine.metrics.success_count}')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extreme_ttl_pressure(self):
|
||||
"""极端 TTL: 所有资源几乎立即过期,验证过期清理不破坏守恒。"""
|
||||
engine = PropertyEngine(
|
||||
t_slot_cap=4, q_slot_cap=4, q_pending_cap=6, phys_cap=3,
|
||||
t_max_age=0.05, q_max_age=0.05, # 50ms 过期
|
||||
seed=789,
|
||||
)
|
||||
stop = asyncio.Event()
|
||||
errors = []
|
||||
|
||||
async def s_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.s_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'S: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.01, 0.05))
|
||||
|
||||
async def p_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.p_produce()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'P: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.01, 0.05))
|
||||
|
||||
async def c_loop():
|
||||
while not stop.is_set():
|
||||
try:
|
||||
await engine.c_consume()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except AssertionError as e:
|
||||
errors.append(f'C: {e}')
|
||||
break
|
||||
await asyncio.sleep(engine.rng.uniform(0.01, 0.05))
|
||||
|
||||
tasks = []
|
||||
for _ in range(2):
|
||||
tasks.append(asyncio.create_task(s_loop()))
|
||||
tasks.append(asyncio.create_task(p_loop()))
|
||||
tasks.append(asyncio.create_task(c_loop()))
|
||||
|
||||
await asyncio.sleep(5)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
engine.check_invariants_final()
|
||||
assert not errors, f'不变量违反:\n' + '\n'.join(errors)
|
||||
print(f'\n[TTL Pressure] t_produced={engine.metrics.t_produced} '
|
||||
f't_expired={engine.metrics.t_expired} q_expired={engine.metrics.q_expired} '
|
||||
f'pair_ok={engine.metrics.pair_consumed_ok}')
|
||||
1325
tests/test_register_runtime_unittest.py
Normal file
1325
tests/test_register_runtime_unittest.py
Normal file
File diff suppressed because it is too large
Load Diff
132
tests/test_runtime_log_analyzer.py
Normal file
132
tests/test_runtime_log_analyzer.py
Normal file
@@ -0,0 +1,132 @@
|
||||
import unittest
|
||||
|
||||
from tools.runtime_log_analyzer import (
|
||||
parse_monitor_lines,
|
||||
parse_solver_timelines,
|
||||
summarize_monitor_rows,
|
||||
summarize_solver_timelines,
|
||||
)
|
||||
|
||||
|
||||
class RuntimeLogAnalyzerTests(unittest.TestCase):
|
||||
def test_parses_csp_monitor_rows_and_recent_rates(self):
|
||||
text = "\n".join(
|
||||
[
|
||||
"[*] T:0 Q:3 phys:0 t_slot:3 q_slot:0 q_pend:0 "
|
||||
"t_prod:76 t_adm:76 t_exp:0 q_sent:87 q_ret:87 q_adm:79 q_exp:0 "
|
||||
"pair:76 ok:71 fail:0 rate:8.6/min #71",
|
||||
"[*] T:0 Q:3 phys:0 t_slot:3 q_slot:0 q_pend:0 "
|
||||
"t_prod:91 t_adm:91 t_exp:0 q_sent:101 q_ret:101 q_adm:94 q_exp:0 "
|
||||
"pair:91 ok:86 fail:0 rate:8.8/min #86",
|
||||
]
|
||||
)
|
||||
|
||||
rows = parse_monitor_lines(text)
|
||||
summary = summarize_monitor_rows(rows)
|
||||
|
||||
self.assertEqual(rows[0].kind, "csp")
|
||||
self.assertEqual(summary["last_ok"], 86)
|
||||
self.assertAlmostEqual(summary["last_cumulative_rate"], 8.8)
|
||||
self.assertAlmostEqual(summary["recent_ok_per_min"], 9.89, places=2)
|
||||
self.assertAlmostEqual(summary["recent_t_prod_per_min"], 9.89, places=2)
|
||||
self.assertEqual(summary["last_q_minus_t"], 3)
|
||||
self.assertEqual(summary["last_q_return_minus_t_prod"], 10)
|
||||
|
||||
def test_parses_state_machine_monitor_rows(self):
|
||||
text = "\n".join(
|
||||
[
|
||||
"[*] slots:7/8 act:7 cpu:83% avg:76%/85 mem:4841M "
|
||||
"T:4 Q:0 sent:0 got:0(0%) rate:0.0/min #0",
|
||||
"[*] slots:8/8 act:8 cpu:91% avg:88%/85 mem:5012M "
|
||||
"T:6 Q:5 sent:61 got:48(79%) rate:7.8/min #40",
|
||||
"[*] slots:6/8 act:6 cpu:100% avg:90%/85 mem:5351M "
|
||||
"T:1 Q:6 sent:68 got:68(100%) rate:8.8/min #59",
|
||||
]
|
||||
)
|
||||
|
||||
rows = parse_monitor_lines(text)
|
||||
summary = summarize_monitor_rows(rows)
|
||||
|
||||
self.assertEqual(rows[-1].kind, "state_machine")
|
||||
self.assertEqual(summary["last_ok"], 59)
|
||||
self.assertAlmostEqual(summary["last_cumulative_rate"], 8.8)
|
||||
self.assertEqual(summary["last_q_minus_t"], 5)
|
||||
self.assertEqual(summary["last_slots"], 6)
|
||||
|
||||
def test_parses_csp_v2_monitor_rows_with_admission_fields(self):
|
||||
text = (
|
||||
"[*] T:1 Q:2 phys:3 p_send:1 t_slot:4 q_slot:5 q_pend:6 "
|
||||
"p_batch:3.5 t_prog:1 q_inflight:4 "
|
||||
"s_phys:0.50/3.00 p_phys:0.20/1.40 c_phys:0.30/2.50 "
|
||||
"p_stage:0.50/0.80/0.40 c_stage:0.30/0.40/1.50 c_hot:4/1 "
|
||||
"solver_goto:0.10 solver_inject:0.20 solver_initial:0.50 "
|
||||
"solver_click:0.01 solver_wait:11.80 solver_reuse:0.75 solver_visible:0.00 "
|
||||
"t_prod:20 t_adm:18 t_exp:1 q_sent:24 q_ret:22 q_adm:20 q_exp:0 "
|
||||
"pair:17 ok:16 fail:1 rate:9.1/min #16"
|
||||
)
|
||||
|
||||
rows = parse_monitor_lines(text)
|
||||
summary = summarize_monitor_rows(rows)
|
||||
|
||||
self.assertEqual(len(rows), 1)
|
||||
self.assertEqual(rows[0].kind, "csp")
|
||||
self.assertEqual(summary["last_ok"], 16)
|
||||
self.assertEqual(summary["last_q_return_minus_t_prod"], 2)
|
||||
self.assertEqual(summary["last_solver_wait"], 11.8)
|
||||
self.assertEqual(summary["last_solver_reuse"], 0.75)
|
||||
self.assertEqual(summary["last_phys"], 3)
|
||||
self.assertEqual(summary["last_p_batch"], 3.5)
|
||||
self.assertEqual(summary["last_t_prog"], 1)
|
||||
self.assertEqual(summary["last_q_inflight"], 4)
|
||||
self.assertEqual(summary["last_s_phys_hold"], 3.0)
|
||||
self.assertEqual(summary["last_p_phys_wait"], 0.2)
|
||||
self.assertEqual(summary["last_c_phys_hold"], 2.5)
|
||||
self.assertEqual(summary["last_physical_hold_leader"], "s")
|
||||
self.assertEqual(summary["last_physical_wait_leader"], "s")
|
||||
self.assertEqual(summary["last_p_email_create"], 0.5)
|
||||
self.assertEqual(summary["last_p_page_prepare"], 0.8)
|
||||
self.assertEqual(summary["last_p_send"], 0.4)
|
||||
self.assertEqual(summary["last_c_page_acquire"], 0.3)
|
||||
self.assertEqual(summary["last_c_verify"], 0.4)
|
||||
self.assertEqual(summary["last_c_register"], 1.5)
|
||||
self.assertEqual(summary["last_c_hot_hits"], 4)
|
||||
self.assertEqual(summary["last_c_hot_misses"], 1)
|
||||
|
||||
def test_summarizes_solver_timeline_events(self):
|
||||
text = (
|
||||
'[solver_timeline] [{"t":0.0,"event":"page_trace_after_inject",'
|
||||
'"page_trace":{"created_at":0.0,"render_called_at":100.0,"render_returned_at":120.0,'
|
||||
'"token_written_at":null}},'
|
||||
'{"t":0.5,"event":"click_before","attempt":1,"token_len":0,'
|
||||
'"dom":{"widget":{"present":true,"visible":true},'
|
||||
'"turnstile_iframe_count":0,'
|
||||
'"element_at_center":{"tag":"DIV","is_iframe":false}}},'
|
||||
'{"t":3.0,"event":"click_after","attempt":1,"token_len":0,'
|
||||
'"click_call_ms":2500.0,'
|
||||
'"click_trace":{"mouse_move1_ms":500.0,"mouse_move2_ms":1500.0,'
|
||||
'"mouse_down_ms":300.0,"mouse_up_ms":100.0}},'
|
||||
'{"t":3.0,"event":"poll_start"},'
|
||||
'{"t":4.0,"event":"poll_done","ok":true,"token_len":730,'
|
||||
'"poll_attempts":2,"first_token_attempt":2,"poll_read_ms_avg":12.0,'
|
||||
'"poll_read_ms_max":20.0,'
|
||||
'"page_trace":{"created_at":0.0,"render_called_at":100.0,"render_returned_at":120.0,'
|
||||
'"token_written_at":3900.0}}]'
|
||||
)
|
||||
|
||||
timelines = parse_solver_timelines(text)
|
||||
summary = summarize_solver_timelines(timelines)
|
||||
|
||||
self.assertEqual(len(timelines), 1)
|
||||
self.assertEqual(summary["solver_timeline_count"], 1)
|
||||
self.assertEqual(summary["ok_count"], 1)
|
||||
self.assertEqual(summary["avg_click_call_ms"], 2500.0)
|
||||
self.assertEqual(summary["avg_click_mouse_move_ms"], 2000.0)
|
||||
self.assertEqual(summary["avg_render_to_token_ms"], 3800.0)
|
||||
self.assertEqual(summary["avg_token_write_to_poll_done_ms"], 100.0)
|
||||
self.assertEqual(summary["avg_poll_attempts"], 2.0)
|
||||
self.assertEqual(summary["center_iframe_hit_ratio"], 0.0)
|
||||
self.assertEqual(summary["turnstile_iframe_seen_ratio"], 0.0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
64
tests/test_shell_entrypoints.py
Normal file
64
tests/test_shell_entrypoints.py
Normal file
@@ -0,0 +1,64 @@
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _entrypoint_workspace(tmp_path):
|
||||
for name in ("start.sh", "auth-service.sh", "setup.sh"):
|
||||
shutil.copy2(ROOT / name, tmp_path / name)
|
||||
scripts = tmp_path / "scripts"
|
||||
scripts.mkdir()
|
||||
shutil.copy2(ROOT / "scripts" / "ensure_runtime.sh", scripts / "ensure_runtime.sh")
|
||||
python = tmp_path / ".venv" / "bin" / "python"
|
||||
python.parent.mkdir(parents=True)
|
||||
python.write_text(
|
||||
"#!/bin/sh\nprintf '%s\\n' \"$@\" > \"$ENTRY_CAPTURE\"\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
python.chmod(0o755)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def _run_entry(workspace, script, *args):
|
||||
capture = workspace / "capture.txt"
|
||||
environment = {**os.environ, "ENTRY_CAPTURE": str(capture)}
|
||||
completed = subprocess.run(
|
||||
["bash", script, *args],
|
||||
cwd=workspace,
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert completed.returncode == 0, completed.stderr
|
||||
return capture.read_text(encoding="utf-8").splitlines()
|
||||
|
||||
|
||||
def test_start_dispatches_registration_and_email_as_independent_modules(tmp_path):
|
||||
workspace = _entrypoint_workspace(tmp_path)
|
||||
(workspace / ".env").write_text("EMAIL_MODE=tempmail\n", encoding="utf-8")
|
||||
|
||||
assert _run_entry(workspace, "start.sh", "--target", "3") == [
|
||||
"-m",
|
||||
"grok_register.register",
|
||||
"--target",
|
||||
"3",
|
||||
]
|
||||
assert _run_entry(workspace, "start.sh", "--email-service", "--port", "9090") == [
|
||||
"-m",
|
||||
"grok_register.email_server",
|
||||
"--port",
|
||||
"9090",
|
||||
]
|
||||
|
||||
|
||||
def test_auth_service_keeps_the_supported_python_module_internal(tmp_path):
|
||||
workspace = _entrypoint_workspace(tmp_path)
|
||||
assert _run_entry(workspace, "auth-service.sh", "--debug") == [
|
||||
"-m",
|
||||
"xai_enroller.service",
|
||||
"--debug",
|
||||
]
|
||||
411
tests/test_stress.py
Normal file
411
tests/test_stress.py
Normal file
@@ -0,0 +1,411 @@
|
||||
"""
|
||||
Layer 4: 压力测试
|
||||
|
||||
真实 asyncio 并发跑较长时间,检查:
|
||||
- pending 泄漏
|
||||
- 库存卡死
|
||||
- claim 饥饿
|
||||
- 等待时间周期性尖峰
|
||||
- slot 守恒在长时间运行后仍成立
|
||||
"""
|
||||
import asyncio
|
||||
import time
|
||||
import statistics
|
||||
import pytest
|
||||
|
||||
from grok_register.core.envelope import ResourceEnvelope
|
||||
from grok_register.core.inventory import Inventory
|
||||
from grok_register.core.observer import Metrics
|
||||
from tests.fakes import Conservation, EventLog
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════
|
||||
# S/P/C 真实并发 Worker (使用 fake 服务)
|
||||
# ═══════════════════════════════════════════
|
||||
|
||||
async def stress_s_worker(wid, t_gen, inventory, phys, t_slot, metrics, stop):
|
||||
"""压力测试 S_Worker。"""
|
||||
while not stop.is_set():
|
||||
await phys.acquire()
|
||||
try:
|
||||
token = await t_gen.generate()
|
||||
finally:
|
||||
phys.release()
|
||||
|
||||
if token is None:
|
||||
metrics.t_discarded += 1
|
||||
continue
|
||||
|
||||
metrics.t_produced += 1
|
||||
now = time.time()
|
||||
try:
|
||||
t_env = await ResourceEnvelope.create_with_slot(
|
||||
'T', token, t_slot, expires_at=now + 60
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
metrics.t_discarded += 1
|
||||
raise
|
||||
|
||||
try:
|
||||
await inventory.put_t(t_env)
|
||||
except Exception:
|
||||
t_env.discard(reason='put_fail')
|
||||
metrics.t_discarded += 1
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
|
||||
async def stress_p_worker(wid, email_svc, inventory, phys, q_pend, q_slot, metrics, stop):
|
||||
"""压力测试 P_Worker。"""
|
||||
loop = asyncio.get_event_loop()
|
||||
while not stop.is_set():
|
||||
await q_pend.acquire()
|
||||
_pend_released = False
|
||||
try:
|
||||
await phys.acquire()
|
||||
try:
|
||||
handle, email, password = await email_svc.create_email()
|
||||
sent = True
|
||||
finally:
|
||||
phys.release()
|
||||
|
||||
if not sent:
|
||||
q_pend.release()
|
||||
_pend_released = True
|
||||
continue
|
||||
|
||||
metrics.q_sent += 1
|
||||
|
||||
try:
|
||||
code = await asyncio.wait_for(email_svc.poll_code(handle), timeout=95)
|
||||
except (asyncio.TimeoutError, Exception):
|
||||
code = None
|
||||
|
||||
if code is None:
|
||||
metrics.q_discarded += 1
|
||||
q_pend.release()
|
||||
_pend_released = True
|
||||
continue
|
||||
|
||||
metrics.q_returned += 1
|
||||
now = time.time()
|
||||
q_env = await ResourceEnvelope.create_with_slot(
|
||||
'Q', {'email': email, 'password': password, 'code': code},
|
||||
q_slot, expires_at=now + 60,
|
||||
)
|
||||
|
||||
try:
|
||||
await inventory.put_q(q_env)
|
||||
except Exception:
|
||||
q_env.discard(reason='put_fail')
|
||||
metrics.q_discarded += 1
|
||||
|
||||
except asyncio.CancelledError:
|
||||
if not _pend_released:
|
||||
q_pend.release()
|
||||
_pend_released = True
|
||||
raise
|
||||
except Exception:
|
||||
metrics.q_discarded += 1
|
||||
finally:
|
||||
if not _pend_released:
|
||||
q_pend.release()
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
|
||||
async def stress_c_worker(wid, register_api, inventory, phys, metrics, stop, file_lock):
|
||||
"""压力测试 C_Worker。"""
|
||||
claim_wait_times = []
|
||||
while not stop.is_set():
|
||||
t0 = time.time()
|
||||
try:
|
||||
async with inventory.claim_pair() as pair:
|
||||
claim_wait_times.append(time.time() - t0)
|
||||
t_val = pair.t.value
|
||||
q_val = pair.q.value
|
||||
|
||||
await phys.acquire()
|
||||
try:
|
||||
sso = await register_api.register(
|
||||
q_val['email'], q_val['password'], q_val['code'], t_val
|
||||
)
|
||||
if sso:
|
||||
metrics.pair_consumed_ok += 1
|
||||
metrics.success_count += 1
|
||||
else:
|
||||
metrics.pair_consumed_fail += 1
|
||||
finally:
|
||||
phys.release()
|
||||
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
metrics.pair_consumed_fail += 1
|
||||
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
return claim_wait_times
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
class TestStress:
|
||||
"""压力测试: 长时间真实并发。"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_steady_state_30s(self):
|
||||
"""稳态运行 30 秒,检查守恒、无泄漏、无饥饿。"""
|
||||
from tests.fakes import FakeTurnstile, FakeEmailService, FakeRegisterAPI
|
||||
|
||||
t_cap, q_cap, pend_cap, phys_cap = 16, 16, 24, 8
|
||||
t_slot_sem = asyncio.Semaphore(t_cap)
|
||||
q_slot_sem = asyncio.Semaphore(q_cap)
|
||||
q_pend_sem = asyncio.Semaphore(pend_cap)
|
||||
phys_sem = asyncio.Semaphore(phys_cap)
|
||||
|
||||
ft = FakeTurnstile(delay=0.005, fail_rate=0.05)
|
||||
es = FakeEmailService()
|
||||
fra = FakeRegisterAPI(success_rate=0.9, delay=0.005)
|
||||
|
||||
metrics = Metrics()
|
||||
inv = Inventory(metrics=metrics)
|
||||
cons = Conservation(t_slot_sem, q_slot_sem, q_pend_sem, inv)
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
tasks = []
|
||||
all_claim_waits = []
|
||||
|
||||
# S workers
|
||||
for i in range(4):
|
||||
tasks.append(asyncio.create_task(
|
||||
stress_s_worker(i, ft, inv, phys_sem, t_slot_sem, metrics, stop)
|
||||
))
|
||||
# P workers
|
||||
for i in range(6):
|
||||
tasks.append(asyncio.create_task(
|
||||
stress_p_worker(i, es, inv, phys_sem, q_pend_sem, q_slot_sem, metrics, stop)
|
||||
))
|
||||
# C workers
|
||||
for i in range(4):
|
||||
tasks.append(asyncio.create_task(
|
||||
stress_c_worker(i, fra, inv, phys_sem, metrics, stop, fl)
|
||||
))
|
||||
|
||||
# 监控 + 守恒检查
|
||||
async def monitor():
|
||||
snapshots = []
|
||||
while not stop.is_set():
|
||||
await asyncio.sleep(2)
|
||||
try:
|
||||
cons.check_t_conservation(t_cap)
|
||||
cons.check_q_conservation(q_cap)
|
||||
# Q_Pending: 只做软检查(非负、不超限),不做守恒检查
|
||||
# 因为 P worker 正常 acquire/release 期间 free 值会波动
|
||||
assert q_pend_sem._value >= 0, 'Q pending 为负'
|
||||
assert q_pend_sem._value <= pend_cap, 'Q pending 超限'
|
||||
except AssertionError as e:
|
||||
snapshots.append(f'INVARIANT VIOLATION: {e}')
|
||||
snapshots.append(
|
||||
f't={metrics.t_produced} q_sent={metrics.q_sent} q_ret={metrics.q_returned} '
|
||||
f'pair_ok={metrics.pair_consumed_ok} inv_t={inv.t_depth} inv_q={inv.q_depth} '
|
||||
f'phys={phys_sem._value} t_slot={t_slot_sem._value} q_slot={q_slot_sem._value}'
|
||||
)
|
||||
return snapshots
|
||||
|
||||
mon_task = asyncio.create_task(monitor())
|
||||
|
||||
await asyncio.sleep(30)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
snapshots = await mon_task
|
||||
|
||||
# 最终守恒检查(所有 Worker 已停止)
|
||||
cons.check_t_conservation(t_cap)
|
||||
cons.check_q_conservation(q_cap)
|
||||
# Q_Pending: 软检查
|
||||
assert q_pend_sem._value >= 0 and q_pend_sem._value <= pend_cap
|
||||
|
||||
# 检查是否有 INVARIANT VIOLATION
|
||||
violations = [s for s in snapshots if 'INVARIANT' in s]
|
||||
assert not violations, f'守恒违反:\n' + '\n'.join(violations)
|
||||
|
||||
print(f'\n[Stress 30s] t_produced={metrics.t_produced} '
|
||||
f'q_sent={metrics.q_sent} q_returned={metrics.q_returned} '
|
||||
f'pair_ok={metrics.pair_consumed_ok} pair_fail={metrics.pair_consumed_fail} '
|
||||
f'success={metrics.success_count}')
|
||||
print(f'Snapshots ({len(snapshots)}):')
|
||||
for s in snapshots[-5:]:
|
||||
print(f' {s}')
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_claim_starvation(self):
|
||||
"""检查 C 不会永远等不到 pair (claim 饥饿)。"""
|
||||
from tests.fakes import FakeTurnstile, FakeEmailService, FakeRegisterAPI
|
||||
|
||||
t_cap, q_cap, pend_cap, phys_cap = 4, 4, 6, 2
|
||||
t_slot_sem = asyncio.Semaphore(t_cap)
|
||||
q_slot_sem = asyncio.Semaphore(q_cap)
|
||||
q_pend_sem = asyncio.Semaphore(pend_cap)
|
||||
phys_sem = asyncio.Semaphore(phys_cap)
|
||||
|
||||
ft = FakeTurnstile(delay=0.002, fail_rate=0.0)
|
||||
es = FakeEmailService()
|
||||
fra = FakeRegisterAPI(success_rate=1.0, delay=0.002)
|
||||
|
||||
metrics = Metrics()
|
||||
inv = Inventory(metrics=metrics)
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
claim_wait_samples = []
|
||||
last_success_time = time.time()
|
||||
|
||||
async def c_worker(wid):
|
||||
nonlocal last_success_time
|
||||
while not stop.is_set():
|
||||
t0 = time.time()
|
||||
try:
|
||||
async with inv.claim_pair() as pair:
|
||||
wait = time.time() - t0
|
||||
claim_wait_samples.append(wait)
|
||||
|
||||
await phys_sem.acquire()
|
||||
try:
|
||||
sso = await fra.register(
|
||||
pair.q.value['email'], pair.q.value['password'],
|
||||
pair.q.value['code'], pair.t.value
|
||||
)
|
||||
if sso:
|
||||
metrics.pair_consumed_ok += 1
|
||||
metrics.success_count += 1
|
||||
last_success_time = time.time()
|
||||
finally:
|
||||
phys_sem.release()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception:
|
||||
metrics.pair_consumed_fail += 1
|
||||
await asyncio.sleep(0.005)
|
||||
|
||||
tasks = []
|
||||
for i in range(3):
|
||||
tasks.append(asyncio.create_task(stress_s_worker(i, ft, inv, phys_sem, t_slot_sem, metrics, stop)))
|
||||
for i in range(4):
|
||||
tasks.append(asyncio.create_task(stress_p_worker(i, es, inv, phys_sem, q_pend_sem, q_slot_sem, metrics, stop)))
|
||||
for i in range(3):
|
||||
tasks.append(asyncio.create_task(c_worker(i)))
|
||||
|
||||
await asyncio.sleep(15)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# 检查: claim 等待时间不应有极端尖峰
|
||||
if claim_wait_samples:
|
||||
p50 = statistics.median(claim_wait_samples)
|
||||
p95 = sorted(claim_wait_samples)[int(len(claim_wait_samples) * 0.95)]
|
||||
p99 = sorted(claim_wait_samples)[int(len(claim_wait_samples) * 0.99)]
|
||||
print(f'\n[Starvation] claim_wait p50={p50:.3f}s p95={p95:.3f}s p99={p99:.3f}s')
|
||||
print(f' samples={len(claim_wait_samples)} pair_ok={metrics.pair_consumed_ok}')
|
||||
# p99 不应超过 5 秒(在 fake 服务下)
|
||||
assert p99 < 5.0, f'claim 饥饿: p99={p99:.3f}s 太高'
|
||||
|
||||
# 检查: 最近成功消费不应太久
|
||||
time_since_last = time.time() - last_success_time
|
||||
assert time_since_last < 5.0, f'太久没有成功消费: {time_since_last:.1f}s'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_high_cancellation_churn(self):
|
||||
"""高取消翻转: Worker 频繁取消重启,系统仍保持守恒。"""
|
||||
from tests.fakes import FakeTurnstile, FakeEmailService, FakeRegisterAPI
|
||||
|
||||
t_cap, q_cap, pend_cap, phys_cap = 8, 8, 12, 4
|
||||
t_slot_sem = asyncio.Semaphore(t_cap)
|
||||
q_slot_sem = asyncio.Semaphore(q_cap)
|
||||
q_pend_sem = asyncio.Semaphore(pend_cap)
|
||||
phys_sem = asyncio.Semaphore(phys_cap)
|
||||
|
||||
ft = FakeTurnstile(delay=0.003, fail_rate=0.1)
|
||||
es = FakeEmailService()
|
||||
fra = FakeRegisterAPI(success_rate=0.85, delay=0.003)
|
||||
|
||||
metrics = Metrics()
|
||||
inv = Inventory(metrics=metrics)
|
||||
cons = Conservation(t_slot_sem, q_slot_sem, q_pend_sem, inv)
|
||||
stop = asyncio.Event()
|
||||
fl = asyncio.Lock()
|
||||
|
||||
tasks = []
|
||||
violations = []
|
||||
|
||||
def spawn_worker(kind, idx):
|
||||
if kind == 'S':
|
||||
return asyncio.create_task(stress_s_worker(idx, ft, inv, phys_sem, t_slot_sem, metrics, stop))
|
||||
elif kind == 'P':
|
||||
return asyncio.create_task(stress_p_worker(idx, es, inv, phys_sem, q_pend_sem, q_slot_sem, metrics, stop))
|
||||
else:
|
||||
return asyncio.create_task(stress_c_worker(idx, fra, inv, phys_sem, metrics, stop, fl))
|
||||
|
||||
# 初始 workers
|
||||
for i in range(3):
|
||||
tasks.append(spawn_worker('S', i))
|
||||
tasks.append(spawn_worker('P', i))
|
||||
tasks.append(spawn_worker('C', i))
|
||||
|
||||
async def churn():
|
||||
import random
|
||||
while not stop.is_set():
|
||||
await asyncio.sleep(0.3)
|
||||
# 随机取消一个,重启一个
|
||||
if tasks:
|
||||
idx = random.randint(0, len(tasks) - 1)
|
||||
t = tasks[idx]
|
||||
if not t.done():
|
||||
t.cancel()
|
||||
tasks[idx] = spawn_worker(
|
||||
random.choice(['S', 'P', 'C']),
|
||||
random.randint(0, 9)
|
||||
)
|
||||
# 定期守恒检查(软检查)
|
||||
try:
|
||||
cons.check_t_conservation(t_cap)
|
||||
cons.check_q_conservation(q_cap)
|
||||
assert q_pend_sem._value >= 0 and q_pend_sem._value <= pend_cap
|
||||
except AssertionError as e:
|
||||
violations.append(str(e))
|
||||
|
||||
churn_task = asyncio.create_task(churn())
|
||||
tasks.append(churn_task)
|
||||
|
||||
await asyncio.sleep(15)
|
||||
stop.set()
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
try:
|
||||
await t
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# 最终守恒检查
|
||||
try:
|
||||
cons.check_t_conservation(t_cap)
|
||||
cons.check_q_conservation(q_cap)
|
||||
assert q_pend_sem._value >= 0 and q_pend_sem._value <= pend_cap
|
||||
except AssertionError as e:
|
||||
violations.append(f'FINAL: {e}')
|
||||
|
||||
assert not violations, f'守恒违反:\n' + '\n'.join(violations[:5])
|
||||
print(f'\n[High Churn] t_produced={metrics.t_produced} '
|
||||
f'pair_ok={metrics.pair_consumed_ok} success={metrics.success_count}')
|
||||
298
tests/test_xai_auth_pipeline.py
Normal file
298
tests/test_xai_auth_pipeline.py
Normal file
@@ -0,0 +1,298 @@
|
||||
import asyncio
|
||||
import threading
|
||||
|
||||
from xai_enroller.auth_pipeline import AuthPipeline
|
||||
from xai_enroller.ledger import Ledger
|
||||
from xai_enroller.models import (
|
||||
AuthorizationResult,
|
||||
AuthorizationStatus,
|
||||
DeviceFlow,
|
||||
JobStatus,
|
||||
OAuthCredential,
|
||||
PipelineState,
|
||||
SinkReceipt,
|
||||
SourceRecord,
|
||||
)
|
||||
|
||||
|
||||
class EmptySource:
|
||||
async def records(self):
|
||||
if False:
|
||||
yield None
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
class NoopExecutor:
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
def _pipeline(tmp_path):
|
||||
return AuthPipeline(
|
||||
source=EmptySource(),
|
||||
protocol=object(),
|
||||
executor=NoopExecutor(),
|
||||
sink=object(),
|
||||
ledger=Ledger(tmp_path / "ledger.db", b"salt"),
|
||||
)
|
||||
|
||||
|
||||
def test_cancel_active_is_idempotent_while_cleanup_is_in_progress(tmp_path):
|
||||
async def scenario():
|
||||
pipeline = _pipeline(tmp_path)
|
||||
|
||||
async def wait_forever():
|
||||
await asyncio.Future()
|
||||
|
||||
task = asyncio.create_task(wait_forever())
|
||||
pipeline._authorization_task = task
|
||||
pipeline._authorization_cancellable = True
|
||||
assert await pipeline.cancel_active()
|
||||
assert not await pipeline.cancel_active()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
class SingleSource:
|
||||
async def records(self):
|
||||
yield SourceRecord("stock", "opaque")
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
class ImmediateProtocol:
|
||||
async def start_device_flow(self):
|
||||
return DeviceFlow(
|
||||
"device",
|
||||
"code",
|
||||
"https://accounts.x.ai/oauth2/device",
|
||||
600,
|
||||
0,
|
||||
)
|
||||
|
||||
async def poll_token(self, **_kwargs):
|
||||
return OAuthCredential(
|
||||
"access",
|
||||
"refresh",
|
||||
None,
|
||||
"Bearer",
|
||||
3600,
|
||||
"later",
|
||||
"now",
|
||||
"subject",
|
||||
"https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
|
||||
class ImmediateExecutor:
|
||||
def __init__(self):
|
||||
self.calls = 0
|
||||
|
||||
async def start(self):
|
||||
return None
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
async def confirm(self, *_args):
|
||||
self.calls += 1
|
||||
return AuthorizationResult(AuthorizationStatus.AUTHORIZED, "confirmed")
|
||||
|
||||
|
||||
class BlockingThreadSink:
|
||||
def __init__(self):
|
||||
self.entered = threading.Event()
|
||||
self.release = threading.Event()
|
||||
self.pipeline = None
|
||||
self.calls = 0
|
||||
|
||||
def _store(self):
|
||||
self.calls += 1
|
||||
self.entered.set()
|
||||
self.release.wait(2)
|
||||
return SinkReceipt("opaque")
|
||||
|
||||
async def store(self, _credential):
|
||||
receipt = await asyncio.to_thread(self._store)
|
||||
self.pipeline.request_stop()
|
||||
return receipt
|
||||
|
||||
|
||||
def test_sink_and_ledger_commit_cannot_be_cancelled_mid_commit(tmp_path):
|
||||
async def scenario():
|
||||
executor = ImmediateExecutor()
|
||||
sink = BlockingThreadSink()
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
events = []
|
||||
pipeline = AuthPipeline(
|
||||
source=SingleSource(),
|
||||
protocol=ImmediateProtocol(),
|
||||
executor=executor,
|
||||
sink=sink,
|
||||
ledger=ledger,
|
||||
timeout=2,
|
||||
event_callback=lambda kind, data: events.append((kind, data)),
|
||||
)
|
||||
sink.pipeline = pipeline
|
||||
run = asyncio.create_task(pipeline.run())
|
||||
assert await asyncio.to_thread(sink.entered.wait, 1)
|
||||
assert not await pipeline.cancel_active()
|
||||
sink.release.set()
|
||||
await asyncio.wait_for(run, 2)
|
||||
assert executor.calls == 1
|
||||
assert sink.calls == 1
|
||||
assert ledger.aggregate_counts()["imported_unique"] == 1
|
||||
started = next(data for kind, data in events if kind == "authorization_started")
|
||||
imported = next(
|
||||
data
|
||||
for kind, data in events
|
||||
if kind == "result" and data["status"] == "imported"
|
||||
)
|
||||
assert imported["task_number"] == started["task_number"]
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_pending_total_uses_current_snapshot_set_difference(tmp_path):
|
||||
class SnapshotSource(EmptySource):
|
||||
snapshot_fingerprints = frozenset({"a", "b", "c"})
|
||||
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
imported = ledger.start_fingerprint("b")
|
||||
ledger.finish(imported, JobStatus.IMPORTED, "imported", "receipt-b")
|
||||
outside = ledger.start_fingerprint("outside")
|
||||
ledger.finish(outside, JobStatus.IMPORTED, "imported", "receipt-outside")
|
||||
pipeline = AuthPipeline(
|
||||
source=SnapshotSource(),
|
||||
protocol=object(),
|
||||
executor=NoopExecutor(),
|
||||
sink=object(),
|
||||
ledger=ledger,
|
||||
)
|
||||
|
||||
assert pipeline.status()["pending_total"] == 2
|
||||
assert pipeline.status()["five_minute_imports_per_minute"] is None
|
||||
assert pipeline.status()["lifetime_imports_per_minute"] is None
|
||||
|
||||
|
||||
def test_pending_total_is_unknown_before_first_valid_snapshot(tmp_path):
|
||||
pipeline = _pipeline(tmp_path)
|
||||
|
||||
assert pipeline.status()["pending_total"] is None
|
||||
|
||||
|
||||
def test_rate_limit_gate_grows_cooldown_on_consecutive_trips():
|
||||
class Clock:
|
||||
def __init__(self):
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self):
|
||||
return self.now
|
||||
|
||||
async def scenario():
|
||||
from xai_enroller.auth_pipeline import GlobalRateLimitGate
|
||||
|
||||
clock = Clock()
|
||||
gate = GlobalRateLimitGate(clock=clock, base_cooldown=60.0)
|
||||
assert await gate.rate_limited()
|
||||
assert gate.COOLDOWN_SECONDS == 60.0
|
||||
assert gate.snapshot()["cooldown_seconds"] == 60.0
|
||||
|
||||
resume = asyncio.Event()
|
||||
resume.set()
|
||||
clock.now = 60.0
|
||||
probe = await gate.wait_for_permission(resume)
|
||||
assert probe is not None
|
||||
assert await gate.rate_limited(probe)
|
||||
assert gate.COOLDOWN_SECONDS == 90.0
|
||||
assert gate.snapshot()["cooldown_remaining_seconds"] == 90.0
|
||||
|
||||
clock.now = 150.0
|
||||
probe = await gate.wait_for_permission(resume)
|
||||
elapsed = await gate.authorized(probe)
|
||||
assert elapsed == 150.0
|
||||
assert gate.COOLDOWN_SECONDS == 60.0
|
||||
assert gate.snapshot()["rate_limit_trips"] == 0
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_minimum_start_interval_raises_after_rate_limit_and_decays():
|
||||
from xai_enroller.auth_pipeline import MinimumStartInterval
|
||||
|
||||
def approx(value, rel=1e-6):
|
||||
class Approx:
|
||||
def __eq__(self, other):
|
||||
return abs(float(other) - float(value)) <= rel * max(1.0, abs(value))
|
||||
|
||||
def __repr__(self):
|
||||
return f"approx({value})"
|
||||
|
||||
return Approx()
|
||||
|
||||
interval = MinimumStartInterval(10.0)
|
||||
assert interval.seconds == 10.0
|
||||
raised = interval.note_rate_limited()
|
||||
assert raised == 18.0
|
||||
raised = interval.note_rate_limited()
|
||||
assert raised == approx(18.0 * 1.55)
|
||||
for _ in range(MinimumStartInterval.SUCCESS_STREAK_TO_DECAY):
|
||||
interval.note_success()
|
||||
assert interval.seconds < raised
|
||||
# Keep decaying until the base floor is restored.
|
||||
for _ in range(40):
|
||||
interval.note_success()
|
||||
assert interval.seconds == 10.0
|
||||
|
||||
|
||||
def test_queued_source_imported_while_paused_is_not_authorized(tmp_path):
|
||||
async def scenario():
|
||||
executor = ImmediateExecutor()
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
pipeline = AuthPipeline(
|
||||
source=EmptySource(),
|
||||
protocol=ImmediateProtocol(),
|
||||
executor=executor,
|
||||
sink=object(),
|
||||
ledger=ledger,
|
||||
timeout=2,
|
||||
)
|
||||
source = SourceRecord("stock", "opaque")
|
||||
fingerprint = ledger.fingerprint(source.source_id)
|
||||
|
||||
assert await pipeline.admit(source)
|
||||
pipeline.pause()
|
||||
await pipeline.rate_gate.rate_limited()
|
||||
run = asyncio.create_task(pipeline.run())
|
||||
try:
|
||||
async with asyncio.timeout(1):
|
||||
while pipeline._authorization_task is None:
|
||||
await asyncio.sleep(0)
|
||||
|
||||
external_job_id = ledger.start_fingerprint(fingerprint)
|
||||
ledger.finish(
|
||||
external_job_id,
|
||||
JobStatus.IMPORTED,
|
||||
"imported",
|
||||
"external-receipt",
|
||||
)
|
||||
pipeline.resume()
|
||||
|
||||
async with asyncio.timeout(1):
|
||||
while (
|
||||
executor.calls == 0
|
||||
and pipeline._states.get(fingerprint) is not PipelineState.IMPORTED
|
||||
):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert executor.calls == 0
|
||||
assert pipeline._states[fingerprint] is PipelineState.IMPORTED
|
||||
finally:
|
||||
pipeline.request_stop()
|
||||
await asyncio.wait_for(run, 2)
|
||||
|
||||
asyncio.run(scenario())
|
||||
391
tests/test_xai_auth_service.py
Normal file
391
tests/test_xai_auth_service.py
Normal file
@@ -0,0 +1,391 @@
|
||||
import asyncio
|
||||
import io
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from xai_enroller.models import JobStatus
|
||||
from xai_enroller.service import (
|
||||
AuthServiceSettings,
|
||||
AuthPipelineRunner,
|
||||
EventTerminal,
|
||||
InteractiveCommandPrompt,
|
||||
resolve_auth_log_mode,
|
||||
)
|
||||
from xai_enroller.ledger import Ledger
|
||||
from xai_enroller.inventory import InventoryError
|
||||
|
||||
|
||||
def test_auth_log_mode_prefers_cli_and_rejects_invalid_values():
|
||||
assert resolve_auth_log_mode([], {}) == "user"
|
||||
assert resolve_auth_log_mode(
|
||||
[], {"XAI_AUTH_SERVICE_LOG_MODE": "debug"}
|
||||
) == "debug"
|
||||
assert resolve_auth_log_mode(
|
||||
["--debug"], {"XAI_AUTH_SERVICE_LOG_MODE": "user"}
|
||||
) == "debug"
|
||||
with pytest.raises(ValueError, match="XAI_AUTH_SERVICE_LOG_MODE"):
|
||||
resolve_auth_log_mode([], {"XAI_AUTH_SERVICE_LOG_MODE": "verbose"})
|
||||
|
||||
|
||||
def test_user_terminal_reports_progress_without_internal_or_secret_fields():
|
||||
messages = []
|
||||
terminal = EventTerminal(mode="user", output=messages.append)
|
||||
terminal.emit(
|
||||
(
|
||||
"startup",
|
||||
{
|
||||
"available": 7,
|
||||
"claimed": 3,
|
||||
"destination": "authenticated/",
|
||||
"source_kind": "local",
|
||||
"ssh_host": "do-not-print.example",
|
||||
},
|
||||
)
|
||||
)
|
||||
terminal.emit(
|
||||
(
|
||||
"authorization_started",
|
||||
{
|
||||
"task_number": 3,
|
||||
"attempt_number": 1,
|
||||
"pending_total": 41,
|
||||
"source_queue": 64,
|
||||
"source_id": "secret@example.test",
|
||||
},
|
||||
)
|
||||
)
|
||||
terminal.emit(
|
||||
(
|
||||
"result",
|
||||
{
|
||||
"status": "imported",
|
||||
"reason": "imported",
|
||||
"task_number": 3,
|
||||
"five_minute_imports_per_minute": 4.25,
|
||||
"lifetime_imports_per_minute": 1.75,
|
||||
"imported_unique": 120,
|
||||
"available": 117,
|
||||
"source_id": "secret@example.test",
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
assert messages == [
|
||||
"[✓] 本地认证服务已启动 | 来源 本机 | 输出 authenticated/ | 待处理 — | 可用 7",
|
||||
"[→] 开始认证 #3 | 待处理 41",
|
||||
"[✓] 认证成功 #3 | 运行平均 1.75/分 | 累计 120 | 可用 117",
|
||||
]
|
||||
rendered = "\n".join(messages)
|
||||
assert "source_queue" not in rendered
|
||||
assert "secret@example.test" not in rendered
|
||||
assert "do-not-print.example" not in rendered
|
||||
|
||||
|
||||
def test_user_status_distinguishes_unknown_runtime_rate_from_zero_rate():
|
||||
base = {
|
||||
"state": "running",
|
||||
"pending_total": None,
|
||||
"active_stage": "idle",
|
||||
"imported_unique": 0,
|
||||
"available": 0,
|
||||
"claimed": 0,
|
||||
"cooldown": False,
|
||||
}
|
||||
messages = []
|
||||
terminal = EventTerminal(mode="user", output=messages.append)
|
||||
terminal.emit(("status", {**base, "lifetime_imports_per_minute": None}))
|
||||
terminal.emit(("status", {**base, "lifetime_imports_per_minute": 0.0}))
|
||||
|
||||
assert "运行平均 —" in messages[0]
|
||||
assert "运行平均 0.00/分" in messages[1]
|
||||
assert "source_queue" not in messages[0]
|
||||
|
||||
|
||||
def test_user_terminal_omits_missing_task_number_before_authorization_starts():
|
||||
messages = []
|
||||
terminal = EventTerminal(mode="user", output=messages.append)
|
||||
|
||||
terminal.emit(
|
||||
(
|
||||
"result",
|
||||
{
|
||||
"status": "failed",
|
||||
"reason": "device_flow_failed",
|
||||
"task_number": None,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
assert messages == ["[✗] 认证未完成 | 暂时失败，将自动重试"]
|
||||
assert "None" not in messages[0]
|
||||
|
||||
|
||||
def test_debug_terminal_keeps_aggregate_diagnostics_and_sanitizes_unknown_events():
|
||||
messages = []
|
||||
terminal = EventTerminal(mode="debug", output=messages.append)
|
||||
terminal.emit(
|
||||
(
|
||||
"status",
|
||||
{
|
||||
"state": "running",
|
||||
"source_queue": 2,
|
||||
"prepared_queue": 1,
|
||||
"completion_queue": 0,
|
||||
"active_stage": "authorization",
|
||||
"retry_waiting": 1,
|
||||
"next_retry_seconds": 5.0,
|
||||
"authorization_starts": 3,
|
||||
"cooldown": False,
|
||||
"cooldown_remaining_seconds": 0.0,
|
||||
"probe_in_flight": False,
|
||||
"min_authorization_interval_seconds": 10.0,
|
||||
"pacing_remaining_seconds": 2.0,
|
||||
"imported_unique": 2,
|
||||
"attempted_unique": 3,
|
||||
"rate_limited": 1,
|
||||
"five_minute_imports_per_minute": 2.0,
|
||||
"lifetime_imports_per_minute": 1.0,
|
||||
"available": 2,
|
||||
"claiming": 0,
|
||||
"claimed": 0,
|
||||
},
|
||||
)
|
||||
)
|
||||
terminal.emit(
|
||||
(
|
||||
"future_event",
|
||||
{"reason": "internal_error", "token": "do-not-print-token"},
|
||||
)
|
||||
)
|
||||
|
||||
assert "queues=2/1/0" in messages[0]
|
||||
assert messages[1] == "• debug event=future_event reason=internal_error"
|
||||
assert "do-not-print-token" not in "\n".join(messages)
|
||||
|
||||
|
||||
def test_terminal_output_failure_does_not_escape():
|
||||
def broken_output(_message):
|
||||
raise OSError("closed terminal")
|
||||
|
||||
terminal = EventTerminal(mode="user", output=broken_output)
|
||||
terminal.emit(("service_stopped", {}))
|
||||
|
||||
|
||||
def test_interactive_prompt_restores_partially_typed_take_command_after_events():
|
||||
output = io.StringIO()
|
||||
commands = []
|
||||
prompt = InteractiveCommandPrompt(
|
||||
output=output,
|
||||
interactive=True,
|
||||
prompt="认证> ",
|
||||
)
|
||||
prompt.start(commands.append)
|
||||
prompt.feed("take 12")
|
||||
|
||||
prompt.write_event("[↻] 发现新账号 3")
|
||||
|
||||
assert output.getvalue().endswith("[↻] 发现新账号 3\n认证> take 12")
|
||||
|
||||
prompt.feed("\r")
|
||||
|
||||
assert commands == ["take 12"]
|
||||
assert output.getvalue().endswith("\n认证> ")
|
||||
|
||||
|
||||
def test_auth_service_configuration_errors_are_actionable_without_traceback(tmp_path):
|
||||
environment = os.environ.copy()
|
||||
environment.pop("XAI_AUTH_SERVICE_SSH_HOST", None)
|
||||
environment.pop("XAI_AUTH_SERVICE_LOG_MODE", None)
|
||||
environment["XAI_AUTH_SERVICE_SOURCE"] = "ssh"
|
||||
missing = subprocess.run(
|
||||
[sys.executable, "-m", "xai_enroller.service"],
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert missing.returncode == 2
|
||||
assert "XAI_AUTH_SERVICE_SSH_HOST" in missing.stderr
|
||||
assert "docs/guides/auth-service.md" in missing.stderr
|
||||
assert "Traceback" not in missing.stderr
|
||||
|
||||
environment["XAI_AUTH_SERVICE_LOG_MODE"] = "verbose"
|
||||
invalid = subprocess.run(
|
||||
[sys.executable, "-m", "xai_enroller.service"],
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert invalid.returncode == 2
|
||||
assert "XAI_AUTH_SERVICE_LOG_MODE" in invalid.stderr
|
||||
assert "Traceback" not in invalid.stderr
|
||||
|
||||
environment["XAI_AUTH_SERVICE_LOG_MODE"] = "user"
|
||||
environment.pop("XAI_AUTH_SERVICE_SOURCE", None)
|
||||
environment["XAI_AUTH_SERVICE_SSH_HOST"] = "user@example.test"
|
||||
environment["XAI_ENROLLER_LOCAL_AUTH_DIR"] = str(tmp_path / "auth")
|
||||
environment["XAI_ENROLLER_TIMEOUT_SEC"] = "not-a-number"
|
||||
invalid_enroller_setting = subprocess.run(
|
||||
[sys.executable, "-m", "xai_enroller.service"],
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert invalid_enroller_setting.returncode == 2
|
||||
assert "XAI_ENROLLER_TIMEOUT_SEC" in invalid_enroller_setting.stderr
|
||||
assert "docs/guides/auth-service.md" in invalid_enroller_setting.stderr
|
||||
assert "Traceback" not in invalid_enroller_setting.stderr
|
||||
|
||||
|
||||
def test_auth_service_startup_failure_is_sanitized_without_traceback(tmp_path):
|
||||
blocked_destination = tmp_path / "private-output-path"
|
||||
blocked_destination.write_text("not a directory", encoding="utf-8")
|
||||
environment = os.environ.copy()
|
||||
environment["XAI_AUTH_SERVICE_SSH_HOST"] = "user@example.test"
|
||||
environment["XAI_ENROLLER_LOCAL_AUTH_DIR"] = str(blocked_destination)
|
||||
|
||||
failed = subprocess.run(
|
||||
[sys.executable, "-m", "xai_enroller.service"],
|
||||
env=environment,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
assert failed.returncode == 1
|
||||
assert "认证服务异常终止" in failed.stderr
|
||||
assert "bash auth-service.sh --debug" in failed.stderr
|
||||
assert "Traceback" not in failed.stderr
|
||||
assert str(blocked_destination) not in failed.stderr
|
||||
|
||||
|
||||
def test_ledger_recognizes_previously_imported_sources(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"test-salt")
|
||||
imported = ledger.start("done@example.test")
|
||||
ledger.finish(imported, JobStatus.IMPORTED, "imported")
|
||||
failed = ledger.start("retry@example.test")
|
||||
ledger.finish(failed, JobStatus.SINK_FAILED, "sink_failed")
|
||||
|
||||
assert ledger.has_imported("done@example.test") is True
|
||||
assert ledger.has_imported("retry@example.test") is False
|
||||
|
||||
|
||||
def test_auth_service_settings_defaults_local_and_preserves_ssh_auto_detection(tmp_path):
|
||||
local = AuthServiceSettings.from_environ(
|
||||
{"XAI_AUTH_SERVICE_REGISTER_ROOT": str(tmp_path)}
|
||||
)
|
||||
remote = AuthServiceSettings.from_environ(
|
||||
{
|
||||
"XAI_AUTH_SERVICE_SSH_HOST": "ubuntu@example.test",
|
||||
"XAI_AUTH_SERVICE_SYNC_SEC": "45",
|
||||
}
|
||||
)
|
||||
|
||||
assert local.source_kind == "local"
|
||||
assert local.ssh_host is None
|
||||
assert local.register_root == str(tmp_path)
|
||||
assert remote.source_kind == "ssh"
|
||||
assert remote.ssh_host == "ubuntu@example.test"
|
||||
assert remote.sync_seconds == 45
|
||||
assert remote.remote_root == "/opt/grok-free-register"
|
||||
|
||||
|
||||
def test_auth_service_settings_explicit_source_overrides_auto_detection():
|
||||
local = AuthServiceSettings.from_environ(
|
||||
{
|
||||
"XAI_AUTH_SERVICE_SOURCE": "local",
|
||||
"XAI_AUTH_SERVICE_SSH_HOST": "ignored@example.test",
|
||||
}
|
||||
)
|
||||
assert local.source_kind == "local"
|
||||
|
||||
with pytest.raises(ValueError, match="XAI_AUTH_SERVICE_SSH_HOST"):
|
||||
AuthServiceSettings.from_environ({"XAI_AUTH_SERVICE_SOURCE": "ssh"})
|
||||
with pytest.raises(ValueError, match="XAI_AUTH_SERVICE_SOURCE"):
|
||||
AuthServiceSettings.from_environ({"XAI_AUTH_SERVICE_SOURCE": "unknown"})
|
||||
|
||||
|
||||
def test_pipeline_runner_takes_a_credential_batch_and_reports_inventory():
|
||||
class InventoryLedger:
|
||||
def inventory_counts(self):
|
||||
return {"available": 7, "claiming": 0, "claimed": 3}
|
||||
|
||||
class Pipeline:
|
||||
ledger = InventoryLedger()
|
||||
|
||||
def status(self):
|
||||
return {"state": "running"}
|
||||
|
||||
class Inventory:
|
||||
def take(self, count):
|
||||
assert count == 3
|
||||
return SimpleNamespace(
|
||||
batch_id="batch-1",
|
||||
directory=Path("/tmp/claimed/batch-1"),
|
||||
moved=3,
|
||||
note="",
|
||||
)
|
||||
|
||||
async def scenario():
|
||||
events = []
|
||||
runner = AuthPipelineRunner(Pipeline(), events.append, inventory=Inventory())
|
||||
assert await runner.handle_command("take 3") is True
|
||||
assert await runner.handle_command("s") is True
|
||||
return events
|
||||
|
||||
assert asyncio.run(scenario()) == [
|
||||
(
|
||||
"inventory_taken",
|
||||
{
|
||||
"batch_id": "batch-1",
|
||||
"directory": "/tmp/claimed/batch-1",
|
||||
"moved": 3,
|
||||
"available": 7,
|
||||
"claiming": 0,
|
||||
"claimed": 3,
|
||||
},
|
||||
),
|
||||
(
|
||||
"status",
|
||||
{
|
||||
"state": "running",
|
||||
"available": 7,
|
||||
"claiming": 0,
|
||||
"claimed": 3,
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_pipeline_runner_reports_inventory_failure_without_stopping():
|
||||
class Ledger:
|
||||
def inventory_counts(self):
|
||||
return {"available": 1, "claiming": 1, "claimed": 0}
|
||||
|
||||
class Pipeline:
|
||||
ledger = Ledger()
|
||||
|
||||
class Inventory:
|
||||
def take(self, _count):
|
||||
raise InventoryError("credential file is missing")
|
||||
|
||||
async def scenario():
|
||||
events = []
|
||||
runner = AuthPipelineRunner(Pipeline(), events.append, inventory=Inventory())
|
||||
assert await runner.handle_command("take 1") is True
|
||||
return events
|
||||
|
||||
assert asyncio.run(scenario()) == [
|
||||
(
|
||||
"inventory_error",
|
||||
{
|
||||
"reason": "credential file is missing",
|
||||
"available": 1,
|
||||
"claiming": 1,
|
||||
"claimed": 0,
|
||||
},
|
||||
)
|
||||
]
|
||||
58
tests/test_xai_credential_inventory.py
Normal file
58
tests/test_xai_credential_inventory.py
Normal file
@@ -0,0 +1,58 @@
|
||||
from xai_enroller.inventory import CredentialInventory
|
||||
from xai_enroller.ledger import Ledger
|
||||
from xai_enroller.models import JobStatus
|
||||
|
||||
|
||||
def import_credential(ledger, source, receipt):
|
||||
job_id = ledger.start(source)
|
||||
ledger.finish(job_id, JobStatus.IMPORTED, "imported", receipt)
|
||||
return job_id
|
||||
|
||||
|
||||
def test_take_moves_available_credentials_without_changing_imported_jobs(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
available = tmp_path / "authenticated"
|
||||
claimed = tmp_path / "claimed"
|
||||
available.mkdir()
|
||||
first_job = import_credential(ledger, "first", "xai-first")
|
||||
second_job = import_credential(ledger, "second", "xai-second")
|
||||
(available / "xai-first.json").write_text("{}\n", encoding="utf-8")
|
||||
(available / "xai-second.json").write_text("{}\n", encoding="utf-8")
|
||||
|
||||
batch = CredentialInventory(ledger, available, claimed).take(2)
|
||||
|
||||
assert batch.moved == 2
|
||||
assert batch.note == ""
|
||||
assert not list(available.glob("*.json"))
|
||||
assert sorted(path.name for path in batch.directory.glob("*.json")) == [
|
||||
"xai-first.json",
|
||||
"xai-second.json",
|
||||
]
|
||||
assert ledger.inventory_counts() == {
|
||||
"available": 0,
|
||||
"claiming": 0,
|
||||
"claimed": 2,
|
||||
}
|
||||
assert ledger.get(first_job)["status"] == JobStatus.IMPORTED.value
|
||||
assert ledger.get(second_job)["status"] == JobStatus.IMPORTED.value
|
||||
|
||||
|
||||
def test_recover_finishes_a_batch_interrupted_before_file_move(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
available = tmp_path / "authenticated"
|
||||
claimed = tmp_path / "claimed"
|
||||
available.mkdir()
|
||||
import_credential(ledger, "source", "xai-recover")
|
||||
(available / "xai-recover.json").write_text("{}\n", encoding="utf-8")
|
||||
assert ledger.claim_available(1, "batch-recover") == ["xai-recover"]
|
||||
|
||||
recovered = CredentialInventory(ledger, available, claimed).recover()
|
||||
|
||||
assert recovered == 1
|
||||
assert not (available / "xai-recover.json").exists()
|
||||
assert (claimed / "batch-recover" / "xai-recover.json").exists()
|
||||
assert ledger.inventory_counts() == {
|
||||
"available": 0,
|
||||
"claiming": 0,
|
||||
"claimed": 1,
|
||||
}
|
||||
125
tests/test_xai_enroller_config.py
Normal file
125
tests/test_xai_enroller_config.py
Normal file
@@ -0,0 +1,125 @@
|
||||
import os
|
||||
import sqlite3
|
||||
import stat
|
||||
|
||||
import pytest
|
||||
|
||||
from xai_enroller.config import Settings
|
||||
from xai_enroller.sources import FileSourceAdapter, SQLiteSourceAdapter
|
||||
|
||||
|
||||
def base_env(tmp_path, **overrides):
|
||||
values = {
|
||||
"XAI_ENROLLER_SOURCE_KIND": "file",
|
||||
"XAI_ENROLLER_SOURCE_FILE": str(tmp_path / "sessions.tsv"),
|
||||
"XAI_ENROLLER_SOURCE_SALT": "test-salt",
|
||||
"XAI_ENROLLER_LEDGER_PATH": str(tmp_path / "ledger.db"),
|
||||
}
|
||||
values.update(overrides)
|
||||
return values
|
||||
|
||||
|
||||
def test_settings_rejects_public_http_cpa_sink(tmp_path):
|
||||
with pytest.raises(ValueError, match="HTTPS"):
|
||||
Settings.from_environ(
|
||||
base_env(
|
||||
tmp_path,
|
||||
XAI_ENROLLER_SINK="cpa",
|
||||
XAI_ENROLLER_CPA_BASE_URL="http://example.test",
|
||||
XAI_ENROLLER_CPA_MANAGEMENT_SECRET="secret",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_settings_defaults_are_bounded_and_redacted(tmp_path):
|
||||
settings = Settings.from_environ(base_env(tmp_path))
|
||||
assert settings.target == 1
|
||||
assert settings.concurrency == 1
|
||||
assert settings.retry_attempts == 0
|
||||
assert settings.executor == "http"
|
||||
assert settings.poll_interval == 5.0
|
||||
redacted = settings.redacted_dict()
|
||||
assert "test-salt" not in repr(redacted)
|
||||
assert "CPA_MANAGEMENT_SECRET" not in repr(redacted)
|
||||
assert "source_salt" not in redacted
|
||||
|
||||
|
||||
def test_settings_accepts_bounded_retry_attempts(tmp_path):
|
||||
settings = Settings.from_environ(
|
||||
base_env(tmp_path, XAI_ENROLLER_RETRY_ATTEMPTS="2")
|
||||
)
|
||||
|
||||
assert settings.retry_attempts == 2
|
||||
|
||||
|
||||
def test_settings_accepts_remote_source_without_a_local_source_file(tmp_path):
|
||||
settings = Settings.from_environ(
|
||||
base_env(
|
||||
tmp_path,
|
||||
XAI_ENROLLER_SOURCE_KIND="remote",
|
||||
XAI_ENROLLER_SOURCE_FILE="",
|
||||
)
|
||||
)
|
||||
|
||||
assert settings.source_kind == "remote"
|
||||
assert settings.source_file is None
|
||||
assert settings.source_db is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("key", "value", "message"),
|
||||
[
|
||||
("XAI_ENROLLER_SOURCE_KIND", "other", "source"),
|
||||
("XAI_ENROLLER_AUTH_EXECUTOR", "other", "executor"),
|
||||
("XAI_ENROLLER_CONCURRENCY", "0", "concurrency"),
|
||||
("XAI_ENROLLER_POLL_SEC", "0", "poll"),
|
||||
("XAI_ENROLLER_RETRY_ATTEMPTS", "-1", "retry"),
|
||||
("XAI_ENROLLER_RETRY_ATTEMPTS", "4", "retry"),
|
||||
("XAI_ENROLLER_TARGET", "0", "target"),
|
||||
("XAI_ENROLLER_TARGET", "101", "target"),
|
||||
],
|
||||
)
|
||||
def test_settings_reject_invalid_bounds(tmp_path, key, value, message):
|
||||
with pytest.raises(ValueError, match=message):
|
||||
Settings.from_environ(base_env(tmp_path, **{key: value}))
|
||||
|
||||
|
||||
def test_file_source_requires_private_regular_file(tmp_path):
|
||||
path = tmp_path / "sessions.tsv"
|
||||
path.write_text("one\tsecret\n", encoding="utf-8")
|
||||
path.chmod(stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP)
|
||||
with pytest.raises(ValueError, match="0600"):
|
||||
list(FileSourceAdapter(path).records())
|
||||
|
||||
|
||||
def test_file_source_suppresses_duplicates_without_leaking_tokens(tmp_path):
|
||||
path = tmp_path / "sessions.tsv"
|
||||
path.write_text("one\tsecret-a\none\tsecret-b\n", encoding="utf-8")
|
||||
path.chmod(0o600)
|
||||
records = list(FileSourceAdapter(path).records())
|
||||
assert len(records) == 1
|
||||
assert records[0].source_id == "one"
|
||||
assert "secret" not in repr(records)
|
||||
|
||||
|
||||
def test_sqlite_source_uses_fixed_read_only_query_and_hmac_ids(tmp_path):
|
||||
path = tmp_path / "accounts.db"
|
||||
connection = sqlite3.connect(path)
|
||||
connection.execute(
|
||||
"CREATE TABLE accounts (token TEXT, status TEXT, deleted_at TEXT, updated_at INTEGER)"
|
||||
)
|
||||
connection.executemany(
|
||||
"INSERT INTO accounts VALUES (?, ?, ?, ?)",
|
||||
[
|
||||
("old", "active", None, 1),
|
||||
("new", "active", None, 2),
|
||||
("deleted", "active", "2026-01-01", 3),
|
||||
("inactive", "inactive", None, 4),
|
||||
],
|
||||
)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
records = list(SQLiteSourceAdapter(path, b"salt").records())
|
||||
assert [record.sso_token for record in records] == ["new", "old"]
|
||||
assert records[0].source_id != records[1].source_id
|
||||
assert "new" not in repr(records)
|
||||
293
tests/test_xai_enroller_coordinator.py
Normal file
293
tests/test_xai_enroller_coordinator.py
Normal file
@@ -0,0 +1,293 @@
|
||||
import asyncio
|
||||
import sqlite3
|
||||
|
||||
import pytest
|
||||
|
||||
from xai_enroller.coordinator import EnrollmentCoordinator
|
||||
from xai_enroller.models import (
|
||||
AuthorizationResult,
|
||||
AuthorizationStatus,
|
||||
DeviceFlow,
|
||||
JobStatus,
|
||||
OAuthCredential,
|
||||
SinkReceipt,
|
||||
SourceRecord,
|
||||
)
|
||||
from xai_enroller.sources import FileSourceAdapter, SQLiteSourceAdapter
|
||||
|
||||
|
||||
class Source:
|
||||
async def records(self):
|
||||
for index in range(3):
|
||||
yield SourceRecord(f"source-{index}", f"token-{index}")
|
||||
|
||||
|
||||
class Protocol:
|
||||
def __init__(self):
|
||||
self.polls = 0
|
||||
|
||||
async def start_device_flow(self):
|
||||
return DeviceFlow("device", "code", "https://accounts.x.ai/device", 60, 0)
|
||||
|
||||
async def poll_token(self, **kwargs):
|
||||
self.polls += 1
|
||||
return OAuthCredential(
|
||||
"access",
|
||||
"refresh",
|
||||
None,
|
||||
"Bearer",
|
||||
3600,
|
||||
"2026-07-11T00:00:00Z",
|
||||
"2026-07-11T00:00:00Z",
|
||||
"subject",
|
||||
"https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
|
||||
class Executor:
|
||||
active = 0
|
||||
max_active = 0
|
||||
|
||||
async def confirm(self, source, flow):
|
||||
type(self).active += 1
|
||||
type(self).max_active = max(type(self).max_active, type(self).active)
|
||||
await asyncio.sleep(0)
|
||||
type(self).active -= 1
|
||||
if source.source_id == "source-1":
|
||||
return AuthorizationResult(AuthorizationStatus.NEEDS_BROWSER, "challenge")
|
||||
return AuthorizationResult(AuthorizationStatus.AUTHORIZED, "confirmed")
|
||||
|
||||
|
||||
class Sink:
|
||||
async def store(self, credential):
|
||||
if credential.subject == "subject":
|
||||
return SinkReceipt("receipt")
|
||||
|
||||
|
||||
class FlakyProtocol(Protocol):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.calls = 0
|
||||
|
||||
async def start_device_flow(self):
|
||||
self.calls += 1
|
||||
if self.calls == 1:
|
||||
raise OSError("temporary transport")
|
||||
return await super().start_device_flow()
|
||||
|
||||
|
||||
class InvalidProtocol(Protocol):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.calls = 0
|
||||
|
||||
async def start_device_flow(self):
|
||||
self.calls += 1
|
||||
raise ValueError("invalid device response")
|
||||
|
||||
|
||||
class FailingExecutor(Executor):
|
||||
async def confirm(self, source, flow):
|
||||
raise OSError("confirmation transport")
|
||||
|
||||
|
||||
class SlowDeviceFlowProtocol(Protocol):
|
||||
async def start_device_flow(self):
|
||||
await asyncio.sleep(0.04)
|
||||
return await super().start_device_flow()
|
||||
|
||||
|
||||
class SlowExecutor(Executor):
|
||||
async def confirm(self, source, flow):
|
||||
await asyncio.sleep(0.04)
|
||||
return await super().confirm(source, flow)
|
||||
|
||||
|
||||
def test_coordinator_limits_concurrency_and_maps_terminal_results(tmp_path):
|
||||
protocol = Protocol()
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=protocol,
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=2,
|
||||
timeout=5,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=3))
|
||||
assert [result.status for result in results] == [
|
||||
JobStatus.IMPORTED,
|
||||
JobStatus.NEEDS_BROWSER,
|
||||
JobStatus.IMPORTED,
|
||||
]
|
||||
assert Executor.max_active == 2
|
||||
assert protocol.polls == 2
|
||||
|
||||
|
||||
def test_coordinator_accepts_explicit_synced_records(tmp_path):
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=Protocol(),
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
results = asyncio.run(
|
||||
coordinator.run_records([SourceRecord("synced-account", "synced-sso")])
|
||||
)
|
||||
|
||||
assert [result.source_id for result in results] == ["synced-account"]
|
||||
assert results[0].status is JobStatus.IMPORTED
|
||||
|
||||
|
||||
def test_coordinator_emits_device_flow_after_the_flow_is_created(tmp_path):
|
||||
events = []
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=Protocol(),
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
event_callback=lambda kind, data: events.append((kind, data)),
|
||||
)
|
||||
|
||||
asyncio.run(coordinator.run_records([SourceRecord("synced-account", "synced-sso")]))
|
||||
|
||||
assert events == [
|
||||
(
|
||||
"device_flow",
|
||||
{
|
||||
"source_id": "synced-account",
|
||||
"user_code": "code",
|
||||
"verification_url": "https://accounts.x.ai/device",
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_coordinator_without_sink_never_reports_imported(tmp_path):
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=Protocol(),
|
||||
executor=Executor(),
|
||||
sink=None,
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
assert results[0].status is JobStatus.SINK_FAILED
|
||||
|
||||
|
||||
def test_coordinator_retries_only_before_device_flow_creation(tmp_path):
|
||||
protocol = FlakyProtocol()
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=protocol,
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
retry_attempts=1,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
assert results[0].status is JobStatus.IMPORTED
|
||||
assert protocol.calls == 2
|
||||
|
||||
|
||||
def test_coordinator_does_not_retry_after_device_flow_creation(tmp_path):
|
||||
protocol = Protocol()
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=protocol,
|
||||
executor=FailingExecutor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
retry_attempts=3,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
assert results[0].status is JobStatus.TRANSPORT_FAILED
|
||||
assert protocol.polls == 0
|
||||
|
||||
|
||||
def test_coordinator_does_not_retry_non_transport_device_flow_error(tmp_path):
|
||||
protocol = InvalidProtocol()
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=protocol,
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
retry_attempts=3,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
assert results[0].status is JobStatus.TRANSPORT_FAILED
|
||||
assert protocol.calls == 1
|
||||
|
||||
|
||||
def test_coordinator_applies_timeout_to_the_entire_attempt(tmp_path):
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=Source(),
|
||||
protocol=SlowDeviceFlowProtocol(),
|
||||
executor=SlowExecutor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"salt",
|
||||
concurrency=1,
|
||||
timeout=0.06,
|
||||
retry_attempts=3,
|
||||
)
|
||||
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
|
||||
assert results[0].status is JobStatus.TIMEOUT
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source_kind", ["file", "sqlite"])
|
||||
def test_coordinator_consumes_real_source_adapters(tmp_path, source_kind):
|
||||
if source_kind == "file":
|
||||
source_path = tmp_path / "sessions.tsv"
|
||||
source_path.write_text("file-source\tfile-token\n", encoding="utf-8")
|
||||
source_path.chmod(0o600)
|
||||
source = FileSourceAdapter(source_path)
|
||||
else:
|
||||
source_path = tmp_path / "accounts.db"
|
||||
connection = sqlite3.connect(source_path)
|
||||
connection.execute(
|
||||
"CREATE TABLE accounts (token TEXT, status TEXT, deleted_at TEXT, updated_at INTEGER)"
|
||||
)
|
||||
connection.execute("INSERT INTO accounts VALUES (?, ?, ?, ?)", ("db-token", "active", None, 1))
|
||||
connection.commit()
|
||||
connection.close()
|
||||
source = SQLiteSourceAdapter(source_path, b"source-salt")
|
||||
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=source,
|
||||
protocol=Protocol(),
|
||||
executor=Executor(),
|
||||
sink=Sink(),
|
||||
ledger_path=tmp_path / f"{source_kind}-ledger.db",
|
||||
ledger_salt=b"ledger-salt",
|
||||
concurrency=1,
|
||||
timeout=5,
|
||||
)
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
assert len(results) == 1
|
||||
assert results[0].status is JobStatus.IMPORTED
|
||||
411
tests/test_xai_enroller_executors.py
Normal file
411
tests/test_xai_enroller_executors.py
Normal file
@@ -0,0 +1,411 @@
|
||||
import asyncio
|
||||
from contextlib import suppress
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from xai_enroller.executors import HTTPProbeExecutor, PlaywrightExecutor
|
||||
from xai_enroller.models import AuthorizationStatus, DeviceFlow, SourceRecord
|
||||
|
||||
|
||||
def test_http_probe_classifies_403_challenge_without_submission():
|
||||
calls = []
|
||||
|
||||
def handler(request):
|
||||
calls.append(request)
|
||||
return httpx.Response(403, text="<html>challenge-platform</html>")
|
||||
|
||||
executor = HTTPProbeExecutor(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler), follow_redirects=False)
|
||||
)
|
||||
source = SourceRecord("source", "sso-token")
|
||||
flow = DeviceFlow("device", "ABCD", "https://accounts.x.ai/oauth2/device?user_code=ABCD", 60, 1)
|
||||
result = asyncio.run(executor.confirm(source, flow))
|
||||
assert result.status is AuthorizationStatus.NEEDS_BROWSER
|
||||
assert len(calls) == 1
|
||||
assert calls[0].method == "GET"
|
||||
assert "sso-token" not in repr(result)
|
||||
|
||||
|
||||
def test_http_probe_does_not_follow_redirects_or_submit_forms():
|
||||
def handler(request):
|
||||
return httpx.Response(302, headers={"location": "https://evil.example/"})
|
||||
|
||||
executor = HTTPProbeExecutor(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler), follow_redirects=False)
|
||||
)
|
||||
result = asyncio.run(
|
||||
executor.confirm(
|
||||
SourceRecord("source", "token"),
|
||||
DeviceFlow("device", "ABCD", "https://accounts.x.ai/oauth2/device", 60, 1),
|
||||
)
|
||||
)
|
||||
assert result.status is AuthorizationStatus.NEEDS_INTERACTION
|
||||
|
||||
|
||||
class FakeButton:
|
||||
def __init__(self, count):
|
||||
self._count = count
|
||||
|
||||
async def count(self):
|
||||
return self._count
|
||||
|
||||
@property
|
||||
def first(self):
|
||||
return self
|
||||
|
||||
async def click(self):
|
||||
return None
|
||||
|
||||
|
||||
class FakePage:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
async def goto(self, url, wait_until):
|
||||
self.url = url
|
||||
|
||||
def locator(self, selector):
|
||||
return self
|
||||
|
||||
async def inner_text(self):
|
||||
return "Authorize access"
|
||||
|
||||
def get_by_role(self, role, name):
|
||||
return FakeButton(1 if name == "Authorize" else 0)
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class FakeContext:
|
||||
def __init__(self):
|
||||
self.cookies = []
|
||||
self.page = FakePage()
|
||||
self.closed = False
|
||||
|
||||
async def add_cookies(self, cookies):
|
||||
self.cookies.extend(cookies)
|
||||
|
||||
async def new_page(self):
|
||||
return self.page
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class FakeBrowser:
|
||||
def __init__(self):
|
||||
self.launches = 0
|
||||
self.contexts = []
|
||||
self.closed = False
|
||||
|
||||
async def new_context(self):
|
||||
context = FakeContext()
|
||||
self.contexts.append(context)
|
||||
return context
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class FakeChromium:
|
||||
def __init__(self, browser):
|
||||
self.browser = browser
|
||||
self.launch_options = None
|
||||
|
||||
async def launch(self, **options):
|
||||
self.launch_options = options
|
||||
self.browser.launches += 1
|
||||
return self.browser
|
||||
|
||||
|
||||
class FakePlaywright:
|
||||
def __init__(self, browser):
|
||||
self.chromium = FakeChromium(browser)
|
||||
self.stopped = False
|
||||
|
||||
async def stop(self):
|
||||
self.stopped = True
|
||||
|
||||
|
||||
class FakePlaywrightFactory:
|
||||
def __init__(self):
|
||||
self.browser = FakeBrowser()
|
||||
self.playwright = FakePlaywright(self.browser)
|
||||
|
||||
def __call__(self):
|
||||
factory = self
|
||||
|
||||
class Runner:
|
||||
async def start(self):
|
||||
return factory.playwright
|
||||
|
||||
return Runner()
|
||||
|
||||
|
||||
class FailingPageContext(FakeContext):
|
||||
async def new_page(self):
|
||||
raise RuntimeError("page creation failed")
|
||||
|
||||
|
||||
class FailingPageBrowser(FakeBrowser):
|
||||
async def new_context(self):
|
||||
context = FailingPageContext()
|
||||
self.contexts.append(context)
|
||||
return context
|
||||
|
||||
|
||||
def test_playwright_executor_isolates_context_and_scopes_exact_cookie():
|
||||
factory = FakePlaywrightFactory()
|
||||
executor = PlaywrightExecutor(playwright_factory=factory)
|
||||
source = SourceRecord("source", "secret-token")
|
||||
flow = DeviceFlow("device", "ABCD", "https://accounts.x.ai/oauth2/device", 60, 1)
|
||||
result = asyncio.run(executor.confirm(source, flow))
|
||||
asyncio.run(executor.close())
|
||||
assert result.status is AuthorizationStatus.AUTHORIZED
|
||||
assert factory.browser.launches == 1
|
||||
context = factory.browser.contexts[0]
|
||||
assert context.cookies == [{
|
||||
"name": "sso",
|
||||
"value": "secret-token",
|
||||
"domain": "accounts.x.ai",
|
||||
"path": "/",
|
||||
"secure": True,
|
||||
"httpOnly": True,
|
||||
}]
|
||||
assert context.page.closed
|
||||
assert context.closed
|
||||
assert factory.browser.closed
|
||||
|
||||
|
||||
def test_playwright_executor_closes_context_when_page_creation_fails():
|
||||
factory = FakePlaywrightFactory()
|
||||
factory.browser = FailingPageBrowser()
|
||||
factory.playwright = FakePlaywright(factory.browser)
|
||||
executor = PlaywrightExecutor(playwright_factory=factory)
|
||||
|
||||
result = asyncio.run(
|
||||
executor.confirm(
|
||||
SourceRecord("source", "secret-token"),
|
||||
DeviceFlow("device", "ABCD", "https://accounts.x.ai/oauth2/device", 60, 1),
|
||||
)
|
||||
)
|
||||
asyncio.run(executor.close())
|
||||
|
||||
assert result.status is AuthorizationStatus.NEEDS_INTERACTION
|
||||
assert factory.browser.contexts[0].closed
|
||||
|
||||
|
||||
def test_playwright_executor_launches_configured_fingerprint_browser():
|
||||
factory = FakePlaywrightFactory()
|
||||
executor = PlaywrightExecutor(
|
||||
playwright_factory=factory,
|
||||
executable_path="/opt/cloakbrowser/chrome",
|
||||
)
|
||||
|
||||
asyncio.run(executor.start())
|
||||
asyncio.run(executor.close())
|
||||
|
||||
assert factory.playwright.chromium.launch_options == {
|
||||
"headless": True,
|
||||
"executable_path": "/opt/cloakbrowser/chrome",
|
||||
}
|
||||
|
||||
|
||||
def test_playwright_executor_uses_fingerprint_browser_path_from_environment(monkeypatch):
|
||||
monkeypatch.setenv("XAI_ENROLLER_BROWSER_EXECUTABLE", "/opt/cloakbrowser/chrome")
|
||||
factory = FakePlaywrightFactory()
|
||||
executor = PlaywrightExecutor(playwright_factory=factory)
|
||||
|
||||
asyncio.run(executor.start())
|
||||
asyncio.run(executor.close())
|
||||
|
||||
assert factory.playwright.chromium.launch_options == {
|
||||
"headless": True,
|
||||
"executable_path": "/opt/cloakbrowser/chrome",
|
||||
}
|
||||
|
||||
|
||||
def test_playwright_executor_finds_macos_cloakbrowser_binary(monkeypatch):
|
||||
macos_binary = (
|
||||
"/Users/test/.cloakbrowser/chromium-145/Chromium.app/Contents/MacOS/Chromium"
|
||||
)
|
||||
monkeypatch.delenv("XAI_ENROLLER_BROWSER_EXECUTABLE", raising=False)
|
||||
monkeypatch.setattr(
|
||||
"xai_enroller.executors.glob.glob",
|
||||
lambda pattern: [macos_binary] if "Chromium.app" in pattern else [],
|
||||
)
|
||||
|
||||
assert PlaywrightExecutor._find_executable_path() == macos_binary
|
||||
|
||||
|
||||
def test_playwright_attempt_cleanup_defers_repeated_cancellation():
|
||||
class Page:
|
||||
def __init__(self):
|
||||
self.entered = asyncio.Event()
|
||||
self.release = asyncio.Event()
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.entered.set()
|
||||
await self.release.wait()
|
||||
self.closed = True
|
||||
|
||||
class Context:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
async def scenario():
|
||||
page = Page()
|
||||
context = Context()
|
||||
task = asyncio.create_task(
|
||||
PlaywrightExecutor._close_attempt_resources(page, context)
|
||||
)
|
||||
await page.entered.wait()
|
||||
task.cancel()
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
page.release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert page.closed
|
||||
assert context.closed
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_playwright_attempt_cleanup_has_a_true_wall_clock_deadline(monkeypatch):
|
||||
class Page:
|
||||
def __init__(self):
|
||||
self.entered = asyncio.Event()
|
||||
self.release = asyncio.Event()
|
||||
|
||||
async def close(self):
|
||||
self.entered.set()
|
||||
while not self.release.is_set():
|
||||
try:
|
||||
await self.release.wait()
|
||||
except asyncio.CancelledError:
|
||||
# Playwright transports can remain stuck while cancellation
|
||||
# propagates. A hard deadline must not await that forever.
|
||||
continue
|
||||
|
||||
async def scenario():
|
||||
monkeypatch.setattr(PlaywrightExecutor, "CLOSE_TIMEOUT_SECONDS", 0.02)
|
||||
page = Page()
|
||||
task = asyncio.create_task(
|
||||
PlaywrightExecutor._close_attempt_resources(page, None)
|
||||
)
|
||||
await page.entered.wait()
|
||||
done, _pending = await asyncio.wait({task}, timeout=0.10)
|
||||
completed_within_deadline = task in done
|
||||
|
||||
page.release.set()
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
assert completed_within_deadline
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_playwright_confirmation_has_hard_deadline_and_restarts_browser(monkeypatch):
|
||||
class StubbornGotoPage(FakePage):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.entered = asyncio.Event()
|
||||
self.release = asyncio.Event()
|
||||
|
||||
async def goto(self, url, wait_until):
|
||||
self.url = url
|
||||
self.entered.set()
|
||||
while not self.release.is_set():
|
||||
try:
|
||||
await self.release.wait()
|
||||
except asyncio.CancelledError:
|
||||
# Model a wedged Playwright transport that does not finish
|
||||
# propagating task cancellation.
|
||||
continue
|
||||
|
||||
class StubbornGotoContext(FakeContext):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.page = StubbornGotoPage()
|
||||
|
||||
class StubbornGotoBrowser(FakeBrowser):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.close_entered = asyncio.Event()
|
||||
self.close_release = asyncio.Event()
|
||||
|
||||
async def new_context(self):
|
||||
context = StubbornGotoContext()
|
||||
self.contexts.append(context)
|
||||
return context
|
||||
|
||||
async def close(self):
|
||||
self.close_entered.set()
|
||||
while not self.close_release.is_set():
|
||||
try:
|
||||
await self.close_release.wait()
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
self.closed = True
|
||||
|
||||
class RotatingPlaywrightFactory:
|
||||
def __init__(self):
|
||||
self.browsers = [StubbornGotoBrowser(), FakeBrowser()]
|
||||
self.playwrights = [FakePlaywright(browser) for browser in self.browsers]
|
||||
self.starts = 0
|
||||
|
||||
def __call__(self):
|
||||
factory = self
|
||||
|
||||
class Runner:
|
||||
async def start(self):
|
||||
playwright = factory.playwrights[factory.starts]
|
||||
factory.starts += 1
|
||||
return playwright
|
||||
|
||||
return Runner()
|
||||
|
||||
async def scenario():
|
||||
monkeypatch.setattr(PlaywrightExecutor, "ATTEMPT_TIMEOUT_SECONDS", 0.02)
|
||||
monkeypatch.setattr(PlaywrightExecutor, "CLOSE_TIMEOUT_SECONDS", 0.02)
|
||||
factory = RotatingPlaywrightFactory()
|
||||
executor = PlaywrightExecutor(playwright_factory=factory)
|
||||
source = SourceRecord("source", "secret-token")
|
||||
flow = DeviceFlow(
|
||||
"device", "ABCD", "https://accounts.x.ai/oauth2/device", 60, 1
|
||||
)
|
||||
|
||||
first = asyncio.create_task(executor.confirm(source, flow))
|
||||
while not factory.browsers[0].contexts:
|
||||
await asyncio.sleep(0)
|
||||
stuck_page = factory.browsers[0].contexts[0].page
|
||||
await stuck_page.entered.wait()
|
||||
done, _pending = await asyncio.wait({first}, timeout=0.10)
|
||||
completed_within_deadline = first in done
|
||||
|
||||
stuck_page.release.set()
|
||||
factory.browsers[0].close_release.set()
|
||||
first_result = await first
|
||||
await asyncio.sleep(0)
|
||||
second_result = await executor.confirm(source, flow)
|
||||
await executor.close()
|
||||
|
||||
assert completed_within_deadline
|
||||
assert first_result.status is AuthorizationStatus.NEEDS_INTERACTION
|
||||
assert first_result.reason_code == "confirmation_timeout"
|
||||
assert second_result.status is AuthorizationStatus.AUTHORIZED
|
||||
assert factory.starts == 2
|
||||
assert factory.browsers[0].close_entered.is_set()
|
||||
assert factory.browsers[0].closed
|
||||
assert factory.playwrights[0].stopped
|
||||
|
||||
asyncio.run(scenario())
|
||||
188
tests/test_xai_enroller_integration.py
Normal file
188
tests/test_xai_enroller_integration.py
Normal file
@@ -0,0 +1,188 @@
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
from xai_enroller.coordinator import EnrollmentCoordinator
|
||||
from xai_enroller.executors import PlaywrightExecutor
|
||||
from xai_enroller.models import SourceRecord
|
||||
from xai_enroller.protocol import XAIProfile, XAIProtocol
|
||||
from xai_enroller.sinks import CPAAuthFileSink
|
||||
|
||||
|
||||
class SingleSource:
|
||||
async def records(self):
|
||||
yield SourceRecord("source-1", "sso-secret")
|
||||
|
||||
|
||||
class FakeButton:
|
||||
async def count(self):
|
||||
return 1
|
||||
|
||||
@property
|
||||
def first(self):
|
||||
return self
|
||||
|
||||
async def click(self):
|
||||
return None
|
||||
|
||||
|
||||
class FakePage:
|
||||
async def goto(self, url, wait_until):
|
||||
self.url = url
|
||||
|
||||
def locator(self, selector):
|
||||
return self
|
||||
|
||||
async def inner_text(self):
|
||||
return "Authorize access"
|
||||
|
||||
def get_by_role(self, role, name):
|
||||
return FakeButton() if name == "Authorize" else type(
|
||||
"NoButton",
|
||||
(),
|
||||
{"count": staticmethod(lambda: asyncio.sleep(0, result=0))},
|
||||
)()
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
class FakeContext:
|
||||
def __init__(self, browser):
|
||||
self.browser = browser
|
||||
self.cookies = []
|
||||
|
||||
async def add_cookies(self, cookies):
|
||||
self.cookies.extend(cookies)
|
||||
|
||||
async def new_page(self):
|
||||
return FakePage()
|
||||
|
||||
async def close(self):
|
||||
self.browser.live_contexts -= 1
|
||||
|
||||
|
||||
class FakeBrowser:
|
||||
def __init__(self):
|
||||
self.launches = 0
|
||||
self.live_contexts = 0
|
||||
self.max_live_contexts = 0
|
||||
|
||||
async def new_context(self):
|
||||
self.live_contexts += 1
|
||||
self.max_live_contexts = max(self.max_live_contexts, self.live_contexts)
|
||||
return FakeContext(self)
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
class FakeChromium:
|
||||
def __init__(self, browser):
|
||||
self.browser = browser
|
||||
|
||||
async def launch(self, **options):
|
||||
self.browser.launches += 1
|
||||
return self.browser
|
||||
|
||||
|
||||
class FakePlaywright:
|
||||
def __init__(self, browser):
|
||||
self.chromium = FakeChromium(browser)
|
||||
|
||||
async def stop(self):
|
||||
return None
|
||||
|
||||
|
||||
class FakePlaywrightFactory:
|
||||
def __init__(self):
|
||||
self.browser = FakeBrowser()
|
||||
self.playwright = FakePlaywright(self.browser)
|
||||
|
||||
def __call__(self):
|
||||
factory = self
|
||||
|
||||
class Runner:
|
||||
async def start(self):
|
||||
return factory.playwright
|
||||
|
||||
return Runner()
|
||||
|
||||
|
||||
def test_adapter_boundary_hands_xai_credential_to_cpa_without_ledger_secrets(tmp_path):
|
||||
def xai_handler(request):
|
||||
if request.url.path == "/.well-known/openid-configuration":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_authorization_endpoint": "https://auth.x.ai/oauth2/device/code",
|
||||
"token_endpoint": "https://auth.x.ai/oauth2/token",
|
||||
},
|
||||
)
|
||||
if request.url.path == "/oauth2/device/code":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_code": "device-secret",
|
||||
"user_code": "ABCD",
|
||||
"verification_uri": "https://accounts.x.ai/oauth2/device",
|
||||
"verification_uri_complete": (
|
||||
"https://accounts.x.ai/oauth2/device?user_code=ABCD"
|
||||
),
|
||||
"expires_in": 60,
|
||||
"interval": 0,
|
||||
},
|
||||
)
|
||||
if request.url.path == "/oauth2/token":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"access_token": "access-secret",
|
||||
"refresh_token": "refresh-secret",
|
||||
"id_token": "id-secret",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
"sub": "subject-1",
|
||||
},
|
||||
)
|
||||
raise AssertionError(f"unexpected xAI request: {request.url}")
|
||||
|
||||
cpa_requests = []
|
||||
|
||||
def cpa_handler(request):
|
||||
cpa_requests.append(request)
|
||||
return httpx.Response(201)
|
||||
|
||||
xai_client = httpx.AsyncClient(transport=httpx.MockTransport(xai_handler))
|
||||
cpa_client = httpx.AsyncClient(transport=httpx.MockTransport(cpa_handler))
|
||||
factory = FakePlaywrightFactory()
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=SingleSource(),
|
||||
protocol=XAIProtocol(xai_client, XAIProfile.default()),
|
||||
executor=PlaywrightExecutor(playwright_factory=factory),
|
||||
sink=CPAAuthFileSink(
|
||||
"https://cpa.example",
|
||||
"management-secret",
|
||||
cpa_client,
|
||||
name_secret=b"name-secret",
|
||||
),
|
||||
ledger_path=tmp_path / "ledger.db",
|
||||
ledger_salt=b"ledger-salt",
|
||||
)
|
||||
|
||||
try:
|
||||
results = asyncio.run(coordinator.run(target=1))
|
||||
finally:
|
||||
asyncio.run(xai_client.aclose())
|
||||
asyncio.run(cpa_client.aclose())
|
||||
|
||||
assert results[0].status.value == "imported"
|
||||
assert factory.browser.launches == 1
|
||||
assert factory.browser.max_live_contexts == 1
|
||||
assert cpa_requests[0].url.path == "/v0/management/auth-files"
|
||||
assert cpa_requests[0].url.params["name"].startswith("xai-")
|
||||
assert json.loads(cpa_requests[0].content)["refresh_token"] == "refresh-secret"
|
||||
ledger_bytes = (tmp_path / "ledger.db").read_bytes()
|
||||
for secret in ("sso-secret", "device-secret", "access-secret", "refresh-secret", "id-secret"):
|
||||
assert secret.encode() not in ledger_bytes
|
||||
163
tests/test_xai_enroller_ledger.py
Normal file
163
tests/test_xai_enroller_ledger.py
Normal file
@@ -0,0 +1,163 @@
|
||||
import sqlite3
|
||||
|
||||
import pytest
|
||||
|
||||
from xai_enroller.ledger import Ledger
|
||||
from xai_enroller.models import JobStatus
|
||||
|
||||
|
||||
def test_ledger_persists_only_redacted_terminal_fields(tmp_path):
|
||||
path = tmp_path / "ledger.db"
|
||||
ledger = Ledger(path, b"salt")
|
||||
job_id = ledger.start("source", attempt=1)
|
||||
ledger.finish(job_id, JobStatus.SINK_FAILED, "sink_failed", "receipt")
|
||||
raw = path.read_bytes()
|
||||
for secret in [
|
||||
"sso-token",
|
||||
"device-code",
|
||||
"https://accounts.x.ai",
|
||||
"access-token",
|
||||
"refresh-token",
|
||||
"id-token",
|
||||
"person@example.com",
|
||||
]:
|
||||
assert secret.encode() not in raw
|
||||
row = ledger.get(job_id)
|
||||
assert row["status"] == "sink_failed"
|
||||
assert row["reason_code"] == "sink_failed"
|
||||
assert row["sink_receipt_fingerprint"] == "receipt"
|
||||
assert "source" not in repr(row["source_fingerprint"])
|
||||
|
||||
|
||||
def test_ledger_recovers_pending_jobs(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
job_id = ledger.start("source", attempt=1)
|
||||
ledger.recover_pending()
|
||||
assert ledger.get(job_id)["status"] == JobStatus.CANCELLED.value
|
||||
|
||||
|
||||
def test_ledger_backfills_imported_receipts_into_available_inventory(tmp_path):
|
||||
path = tmp_path / "ledger.db"
|
||||
with sqlite3.connect(path) as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TABLE jobs (
|
||||
job_id INTEGER PRIMARY KEY,
|
||||
source_fingerprint TEXT NOT NULL,
|
||||
attempt_number INTEGER NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
started_at TEXT NOT NULL,
|
||||
finished_at TEXT,
|
||||
reason_code TEXT,
|
||||
sink_receipt_fingerprint TEXT
|
||||
)
|
||||
"""
|
||||
)
|
||||
connection.executemany(
|
||||
"INSERT INTO jobs(source_fingerprint, attempt_number, status, started_at, "
|
||||
"finished_at, reason_code, sink_receipt_fingerprint) VALUES (?, 1, ?, ?, ?, ?, ?)",
|
||||
[
|
||||
("source-old", "imported", "2026-01-01", "2026-01-02", "imported", "receipt-old"),
|
||||
("source-new", "imported", "2026-01-03", "2026-01-04", "imported", "receipt-new"),
|
||||
("source-failed", "sink_failed", "2026-01-05", "2026-01-06", "sink_failed", "receipt-failed"),
|
||||
],
|
||||
)
|
||||
|
||||
ledger = Ledger(path, b"salt")
|
||||
|
||||
assert ledger.inventory_counts() == {
|
||||
"available": 2,
|
||||
"claiming": 0,
|
||||
"claimed": 0,
|
||||
}
|
||||
assert ledger.claim_available(2, "batch-backfill") == [
|
||||
"receipt-new",
|
||||
"receipt-old",
|
||||
]
|
||||
|
||||
|
||||
def test_imported_finish_and_inventory_insert_are_atomic(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
job_id = ledger.start("source", attempt=1)
|
||||
with sqlite3.connect(ledger.path) as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TRIGGER reject_inventory_insert
|
||||
BEFORE INSERT ON credential_inventory
|
||||
BEGIN
|
||||
SELECT RAISE(ABORT, 'inventory insert rejected');
|
||||
END
|
||||
"""
|
||||
)
|
||||
|
||||
with pytest.raises(sqlite3.IntegrityError, match="inventory insert rejected"):
|
||||
ledger.finish(job_id, JobStatus.IMPORTED, "imported", "receipt")
|
||||
|
||||
assert ledger.get(job_id)["status"] == "pending"
|
||||
assert ledger.inventory_counts()["available"] == 0
|
||||
|
||||
|
||||
def test_claim_available_uses_latest_finish_and_tracks_batch(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
jobs = []
|
||||
for source, receipt in [
|
||||
("source-a", "receipt-a"),
|
||||
("source-b", "receipt-b"),
|
||||
("source-c", "receipt-c"),
|
||||
]:
|
||||
job_id = ledger.start(source)
|
||||
ledger.finish(job_id, JobStatus.IMPORTED, "imported", receipt)
|
||||
jobs.append(job_id)
|
||||
with sqlite3.connect(ledger.path) as connection:
|
||||
connection.executemany(
|
||||
"UPDATE jobs SET finished_at=? WHERE job_id=?",
|
||||
[
|
||||
("2026-01-01T00:00:00+00:00", jobs[0]),
|
||||
("2026-01-03T00:00:00+00:00", jobs[1]),
|
||||
("2026-01-02T00:00:00+00:00", jobs[2]),
|
||||
],
|
||||
)
|
||||
|
||||
assert ledger.claim_available(2, "batch-1") == ["receipt-b", "receipt-c"]
|
||||
assert ledger.claim_available(2, "batch-2") == ["receipt-a"]
|
||||
assert ledger.pending_claims("batch-1") == [
|
||||
{
|
||||
"sink_receipt_fingerprint": "receipt-b",
|
||||
"batch_id": "batch-1",
|
||||
"note": "",
|
||||
},
|
||||
{
|
||||
"sink_receipt_fingerprint": "receipt-c",
|
||||
"batch_id": "batch-1",
|
||||
"note": "",
|
||||
},
|
||||
]
|
||||
assert ledger.inventory_counts() == {
|
||||
"available": 0,
|
||||
"claiming": 3,
|
||||
"claimed": 0,
|
||||
}
|
||||
|
||||
|
||||
def test_claim_completion_and_recovery_are_batch_scoped(tmp_path):
|
||||
ledger = Ledger(tmp_path / "ledger.db", b"salt")
|
||||
for source, receipt in [
|
||||
("source-a", "receipt-a"),
|
||||
("source-b", "receipt-b"),
|
||||
("source-c", "receipt-c"),
|
||||
]:
|
||||
job_id = ledger.start(source)
|
||||
ledger.finish(job_id, JobStatus.IMPORTED, "imported", receipt)
|
||||
|
||||
assert ledger.claim_available(2, "batch-complete")
|
||||
assert ledger.claim_available(1, "batch-recover")
|
||||
assert ledger.mark_claimed("batch-complete", "delivered") == 2
|
||||
assert ledger.recover_claiming("batch-recover", note="worker_restart") == 1
|
||||
|
||||
assert ledger.pending_claims() == []
|
||||
assert ledger.inventory_counts() == {
|
||||
"available": 1,
|
||||
"claiming": 0,
|
||||
"claimed": 2,
|
||||
}
|
||||
assert ledger.claim_available(1, "batch-retry") == ["receipt-a"]
|
||||
245
tests/test_xai_enroller_protocol.py
Normal file
245
tests/test_xai_enroller_protocol.py
Normal file
@@ -0,0 +1,245 @@
|
||||
import base64
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from xai_enroller.models import AuthorizationStatus
|
||||
from xai_enroller.protocol import XAIProtocol, XAIProfile
|
||||
|
||||
|
||||
def _unsigned_jwt(payload):
|
||||
encoded = base64.urlsafe_b64encode(
|
||||
json.dumps(payload, separators=(",", ":")).encode()
|
||||
).decode().rstrip("=")
|
||||
return f"header.{encoded}.signature"
|
||||
|
||||
|
||||
def test_default_profile_matches_observed_grok_cli_device_contract():
|
||||
profile = XAIProfile.default()
|
||||
|
||||
assert profile.client_id == "b1a00492-073a-47ea-816f-4c329264a828"
|
||||
assert profile.scope == "openid profile email offline_access grok-cli:access api:access"
|
||||
assert profile.version == "grok-cli-device-v1"
|
||||
|
||||
|
||||
def test_device_flow_uses_discovered_endpoints_and_profile_scope():
|
||||
requests = []
|
||||
|
||||
def handler(request):
|
||||
requests.append(request)
|
||||
if request.url.path == "/.well-known/openid-configuration":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_authorization_endpoint": "https://auth.x.ai/oauth2/device/code",
|
||||
"token_endpoint": "https://auth.x.ai/oauth2/token",
|
||||
},
|
||||
)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_code": "device",
|
||||
"user_code": "ABCD",
|
||||
"verification_uri": "https://accounts.x.ai/oauth2/device",
|
||||
"expires_in": 60,
|
||||
"interval": 1,
|
||||
},
|
||||
)
|
||||
|
||||
profile = XAIProfile(client_id="client", scope="openid offline_access")
|
||||
protocol = XAIProtocol(httpx.AsyncClient(transport=httpx.MockTransport(handler)), profile)
|
||||
|
||||
async def run():
|
||||
flow = await protocol.start_device_flow()
|
||||
return flow
|
||||
|
||||
flow = __import__("asyncio").run(run())
|
||||
assert flow.verification_url == "https://accounts.x.ai/oauth2/device?user_code=ABCD"
|
||||
assert dict(__import__("urllib.parse", fromlist=["parse_qsl"]).parse_qsl(requests[1].content.decode())) == {
|
||||
"client_id": "client",
|
||||
"scope": "openid offline_access",
|
||||
}
|
||||
|
||||
|
||||
def test_device_flow_uses_configured_poll_interval_when_issuer_omits_it():
|
||||
def handler(request):
|
||||
if request.url.path == "/.well-known/openid-configuration":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_authorization_endpoint": "https://auth.x.ai/oauth2/device/code",
|
||||
"token_endpoint": "https://auth.x.ai/oauth2/token",
|
||||
},
|
||||
)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_code": "device",
|
||||
"user_code": "ABCD",
|
||||
"verification_uri": "https://accounts.x.ai/oauth2/device",
|
||||
"expires_in": 60,
|
||||
},
|
||||
)
|
||||
|
||||
protocol = XAIProtocol(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
XAIProfile(client_id="client", scope="openid"),
|
||||
default_poll_interval=7,
|
||||
)
|
||||
|
||||
flow = __import__("asyncio").run(protocol.start_device_flow())
|
||||
|
||||
assert flow.interval == 7
|
||||
|
||||
|
||||
def test_device_flow_preserves_complete_verification_url():
|
||||
def handler(request):
|
||||
if request.url.path == "/.well-known/openid-configuration":
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_authorization_endpoint": "https://auth.x.ai/oauth2/device/code",
|
||||
"token_endpoint": "https://auth.x.ai/oauth2/token",
|
||||
},
|
||||
)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_code": "device",
|
||||
"user_code": "ABCD",
|
||||
"verification_uri": "https://accounts.x.ai/oauth2/device",
|
||||
"verification_uri_complete": "https://accounts.x.ai/oauth2/device?user_code=ABCD&state=issuer",
|
||||
"expires_in": 60,
|
||||
},
|
||||
)
|
||||
|
||||
protocol = XAIProtocol(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
XAIProfile(client_id="client", scope="openid"),
|
||||
)
|
||||
|
||||
flow = __import__("asyncio").run(protocol.start_device_flow())
|
||||
|
||||
assert flow.verification_url == (
|
||||
"https://accounts.x.ai/oauth2/device?user_code=ABCD&state=issuer"
|
||||
)
|
||||
|
||||
|
||||
def test_xai_url_allowlist_rejects_missing_hostname():
|
||||
assert XAIProtocol._allowed_url("https:///oauth2/device") is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"discovery",
|
||||
[
|
||||
{
|
||||
"device_authorization_endpoint": "http://auth.x.ai/device",
|
||||
"token_endpoint": "https://auth.x.ai/token",
|
||||
},
|
||||
{
|
||||
"device_authorization_endpoint": "https://evil.example/device",
|
||||
"token_endpoint": "https://auth.x.ai/token",
|
||||
},
|
||||
],
|
||||
)
|
||||
def test_discovery_rejects_non_xai_or_non_https(discovery):
|
||||
def handler(request):
|
||||
return httpx.Response(200, json=discovery)
|
||||
|
||||
protocol = XAIProtocol(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
XAIProfile(client_id="client", scope="openid"),
|
||||
)
|
||||
with pytest.raises(ValueError, match="allowlist"):
|
||||
__import__("asyncio").run(protocol.discover())
|
||||
|
||||
|
||||
def test_token_polling_maps_pending_slow_down_and_denial():
|
||||
responses = iter(
|
||||
[
|
||||
httpx.Response(400, json={"error": "authorization_pending"}),
|
||||
httpx.Response(400, json={"error": "slow_down"}),
|
||||
httpx.Response(400, json={"error": "access_denied"}),
|
||||
]
|
||||
)
|
||||
|
||||
def handler(request):
|
||||
return next(responses)
|
||||
|
||||
protocol = XAIProtocol(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
XAIProfile(client_id="client", scope="openid"),
|
||||
sleep=lambda _: __import__("asyncio").sleep(0),
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="oauth_denied"):
|
||||
__import__("asyncio").run(
|
||||
protocol.poll_token(
|
||||
endpoint="https://auth.x.ai/oauth2/token",
|
||||
flow=type("Flow", (), {"device_code": "device", "interval": 0})(),
|
||||
timeout=1,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_token_response_requires_refresh_token_and_never_exposes_body():
|
||||
def handler(request):
|
||||
return httpx.Response(200, json={"access_token": "access", "token_type": "Bearer"})
|
||||
|
||||
protocol = XAIProtocol(
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
XAIProfile(client_id="client", scope="openid"),
|
||||
)
|
||||
with pytest.raises(RuntimeError) as error:
|
||||
__import__("asyncio").run(
|
||||
protocol.poll_token(
|
||||
endpoint="https://auth.x.ai/oauth2/token",
|
||||
flow=type("Flow", (), {"device_code": "device", "interval": 0})(),
|
||||
timeout=1,
|
||||
)
|
||||
)
|
||||
assert "access" not in str(error.value)
|
||||
|
||||
|
||||
def test_credential_subject_prefers_id_token_subject():
|
||||
credential = XAIProtocol._credential(
|
||||
{
|
||||
"access_token": _unsigned_jwt({"principal_id": "access-subject"}),
|
||||
"refresh_token": "refresh",
|
||||
"id_token": _unsigned_jwt(
|
||||
{"sub": "stable-subject", "email": "ignored@example.test"}
|
||||
),
|
||||
"sub": "response-subject",
|
||||
},
|
||||
"https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
assert credential.subject == "stable-subject"
|
||||
|
||||
|
||||
def test_credential_subject_falls_back_to_access_token_principal_id():
|
||||
credential = XAIProtocol._credential(
|
||||
{
|
||||
"access_token": _unsigned_jwt({"principal_id": "stable-principal"}),
|
||||
"refresh_token": "refresh",
|
||||
"id_token": "malformed",
|
||||
"sub": "response-subject",
|
||||
},
|
||||
"https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
assert credential.subject == "stable-principal"
|
||||
|
||||
|
||||
def test_credential_subject_ignores_email_and_falls_back_to_response_subject():
|
||||
credential = XAIProtocol._credential(
|
||||
{
|
||||
"access_token": "not-a-jwt",
|
||||
"refresh_token": "refresh",
|
||||
"id_token": _unsigned_jwt({"email": "ignored@example.test"}),
|
||||
"sub": "response-subject",
|
||||
},
|
||||
"https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
assert credential.subject == "response-subject"
|
||||
79
tests/test_xai_enroller_sinks.py
Normal file
79
tests/test_xai_enroller_sinks.py
Normal file
@@ -0,0 +1,79 @@
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
from xai_enroller.models import OAuthCredential
|
||||
from xai_enroller.sinks import CPAAuthFileSink, SinkError
|
||||
|
||||
|
||||
def credential():
|
||||
return OAuthCredential(
|
||||
access_token="access",
|
||||
refresh_token="refresh",
|
||||
id_token="id-token",
|
||||
token_type="Bearer",
|
||||
expires_in=3600,
|
||||
expires_at="2026-07-11T00:00:00Z",
|
||||
last_refresh="2026-07-11T00:00:00Z",
|
||||
subject="subject",
|
||||
token_endpoint="https://auth.x.ai/oauth2/token",
|
||||
)
|
||||
|
||||
|
||||
def test_cpa_sink_builds_pinned_document_and_safe_name():
|
||||
requests = []
|
||||
|
||||
def handler(request):
|
||||
requests.append(request)
|
||||
return httpx.Response(201)
|
||||
|
||||
sink = CPAAuthFileSink(
|
||||
"https://cpa.example",
|
||||
"management-secret",
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
name_secret=b"name-secret",
|
||||
)
|
||||
receipt = asyncio.run(sink.store(credential()))
|
||||
assert requests[0].url.path == "/v0/management/auth-files"
|
||||
assert requests[0].headers["authorization"] == "Bearer management-secret"
|
||||
assert requests[0].url.params["name"].startswith("xai-")
|
||||
assert requests[0].url.params["name"].endswith(".json")
|
||||
assert json.loads(requests[0].content) == {
|
||||
"type": "xai",
|
||||
"access_token": "access",
|
||||
"refresh_token": "refresh",
|
||||
"id_token": "id-token",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
"expired": "2026-07-11T00:00:00Z",
|
||||
"last_refresh": "2026-07-11T00:00:00Z",
|
||||
"sub": "subject",
|
||||
"base_url": "https://api.x.ai/v1",
|
||||
"token_endpoint": "https://auth.x.ai/oauth2/token",
|
||||
"auth_kind": "oauth",
|
||||
}
|
||||
assert "access" not in repr(receipt)
|
||||
|
||||
|
||||
def test_cpa_sink_maps_non_2xx_without_retry_or_secret_leak():
|
||||
count = 0
|
||||
|
||||
def handler(request):
|
||||
nonlocal count
|
||||
count += 1
|
||||
return httpx.Response(500, text="access")
|
||||
|
||||
sink = CPAAuthFileSink(
|
||||
"https://cpa.example",
|
||||
"management-secret",
|
||||
httpx.AsyncClient(transport=httpx.MockTransport(handler)),
|
||||
name_secret=b"name-secret",
|
||||
)
|
||||
try:
|
||||
asyncio.run(sink.store(credential()))
|
||||
except SinkError as error:
|
||||
assert "access" not in str(error)
|
||||
else:
|
||||
raise AssertionError("expected SinkError")
|
||||
assert count == 1
|
||||
240
tests/test_xai_remote_snapshot.py
Normal file
240
tests/test_xai_remote_snapshot.py
Normal file
@@ -0,0 +1,240 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
|
||||
from xai_enroller.models import SourceRecord
|
||||
from xai_enroller.remote_stream import (
|
||||
DiskSnapshotSource,
|
||||
LocalSnapshotSynchronizer,
|
||||
SSHSnapshotSynchronizer,
|
||||
)
|
||||
|
||||
|
||||
def _document(source_id):
|
||||
return json.dumps(
|
||||
{
|
||||
"email": source_id,
|
||||
"cookies": [
|
||||
{
|
||||
"name": "sso",
|
||||
"value": "opaque",
|
||||
"domain": "accounts.x.ai",
|
||||
"path": "/",
|
||||
}
|
||||
],
|
||||
},
|
||||
separators=(",", ":"),
|
||||
).encode()
|
||||
|
||||
|
||||
class FakeStream:
|
||||
def __init__(self, lines):
|
||||
self._lines = iter(lines)
|
||||
|
||||
async def readline(self):
|
||||
return next(self._lines, b"")
|
||||
|
||||
async def read(self, _size=-1):
|
||||
return b""
|
||||
|
||||
|
||||
class FakeProcess:
|
||||
def __init__(self, lines, returncode=0):
|
||||
self.stdout = FakeStream(lines)
|
||||
self.stderr = FakeStream([])
|
||||
self.returncode = None
|
||||
self._final_returncode = returncode
|
||||
|
||||
async def wait(self):
|
||||
self.returncode = self._final_returncode
|
||||
return self.returncode
|
||||
|
||||
def terminate(self):
|
||||
self.returncode = -15
|
||||
|
||||
def kill(self):
|
||||
self.returncode = -9
|
||||
|
||||
|
||||
def test_snapshot_sync_atomically_replaces_only_valid_complete_export(tmp_path):
|
||||
async def scenario():
|
||||
destination = tmp_path / "source-snapshot.jsonl"
|
||||
destination.write_bytes(_document("old") + b"\n")
|
||||
|
||||
invalid = FakeProcess([b"not-json\n"])
|
||||
|
||||
async def invalid_factory(*_args, **_kwargs):
|
||||
return invalid
|
||||
|
||||
synchronizer = SSHSnapshotSynchronizer(
|
||||
"host",
|
||||
destination,
|
||||
process_factory=invalid_factory,
|
||||
fingerprint=lambda source_id: f"key:{source_id}",
|
||||
)
|
||||
assert not await synchronizer.sync_once()
|
||||
assert destination.read_bytes() == _document("old") + b"\n"
|
||||
|
||||
rejected = FakeProcess([_document("rejected") + b"\n"], returncode=1)
|
||||
|
||||
async def rejected_factory(*_args, **_kwargs):
|
||||
return rejected
|
||||
|
||||
synchronizer.process_factory = rejected_factory
|
||||
assert not await synchronizer.sync_once()
|
||||
assert destination.read_bytes() == _document("old") + b"\n"
|
||||
|
||||
empty = FakeProcess([])
|
||||
|
||||
async def empty_factory(*_args, **_kwargs):
|
||||
return empty
|
||||
|
||||
synchronizer.process_factory = empty_factory
|
||||
assert not await synchronizer.sync_once()
|
||||
assert destination.read_bytes() == _document("old") + b"\n"
|
||||
|
||||
valid = FakeProcess([_document("new-a") + b"\n", _document("new-b") + b"\n"])
|
||||
|
||||
async def valid_factory(*_args, **_kwargs):
|
||||
return valid
|
||||
|
||||
synchronizer.process_factory = valid_factory
|
||||
assert await synchronizer.sync_once()
|
||||
assert synchronizer.snapshot_fingerprints == frozenset(
|
||||
{"key:new-a", "key:new-b"}
|
||||
)
|
||||
assert destination.read_bytes() == _document("new-a") + b"\n" + _document("new-b") + b"\n"
|
||||
assert destination.stat().st_mode & 0o777 == 0o600
|
||||
assert not list(tmp_path.glob(".source-snapshot.jsonl.*.tmp"))
|
||||
|
||||
synchronizer.process_factory = invalid_factory
|
||||
assert not await synchronizer.sync_once()
|
||||
assert synchronizer.snapshot_fingerprints == frozenset(
|
||||
{"key:new-a", "key:new-b"}
|
||||
)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_local_snapshot_accepts_empty_input_and_anchors_exporter_code(tmp_path):
|
||||
async def scenario():
|
||||
registration_root = tmp_path / "registration"
|
||||
destination = tmp_path / "auth" / "source-snapshot.jsonl"
|
||||
captured = []
|
||||
process = FakeProcess([])
|
||||
|
||||
async def factory(*args, **_kwargs):
|
||||
captured.append(args)
|
||||
return process
|
||||
|
||||
synchronizer = LocalSnapshotSynchronizer(
|
||||
registration_root,
|
||||
destination,
|
||||
process_factory=factory,
|
||||
)
|
||||
|
||||
assert await synchronizer.sync_once()
|
||||
assert destination.read_bytes() == b""
|
||||
assert destination.stat().st_mode & 0o777 == 0o600
|
||||
assert synchronizer.snapshot_fingerprints == frozenset()
|
||||
assert captured[0][0]
|
||||
assert captured[0][1] != str(registration_root / "scripts" / "export_registered_sessions.py")
|
||||
assert captured[0][-2:] == (
|
||||
str(registration_root / "keys" / "auth-sessions.jsonl"),
|
||||
str(registration_root / "keys" / "accounts.txt"),
|
||||
)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_local_snapshot_rejects_candidate_when_either_input_changes(tmp_path):
|
||||
async def scenario():
|
||||
registration_root = tmp_path / "registration"
|
||||
keys = registration_root / "keys"
|
||||
keys.mkdir(parents=True)
|
||||
sessions = keys / "auth-sessions.jsonl"
|
||||
accounts = keys / "accounts.txt"
|
||||
sessions.write_bytes(_document("old") + b"\n")
|
||||
accounts.write_text("", encoding="utf-8")
|
||||
destination = tmp_path / "auth" / "source-snapshot.jsonl"
|
||||
destination.parent.mkdir()
|
||||
destination.write_bytes(_document("stable") + b"\n")
|
||||
|
||||
process = FakeProcess([_document("candidate") + b"\n"])
|
||||
|
||||
async def factory(*_args, **_kwargs):
|
||||
accounts.write_text("new@example.test:p:sso\n", encoding="utf-8")
|
||||
return process
|
||||
|
||||
synchronizer = LocalSnapshotSynchronizer(
|
||||
registration_root,
|
||||
destination,
|
||||
process_factory=factory,
|
||||
)
|
||||
|
||||
assert not await synchronizer.sync_once()
|
||||
assert destination.read_bytes() == _document("stable") + b"\n"
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_snapshot_consumer_finishes_open_generation_before_replacement(tmp_path):
|
||||
async def scenario():
|
||||
destination = tmp_path / "source-snapshot.jsonl"
|
||||
destination.write_bytes(_document("old-a") + b"\n" + _document("old-b") + b"\n")
|
||||
source = DiskSnapshotSource(destination, poll_seconds=0.01)
|
||||
records = source.records()
|
||||
first = await anext(records)
|
||||
assert isinstance(first, SourceRecord)
|
||||
assert first.source_id == "old-a"
|
||||
|
||||
replacement = tmp_path / "replacement"
|
||||
replacement.write_bytes(_document("new-a") + b"\n")
|
||||
os.replace(replacement, destination)
|
||||
|
||||
assert (await anext(records)).source_id == "old-b"
|
||||
assert (await asyncio.wait_for(anext(records), 1)).source_id == "new-a"
|
||||
await source.close()
|
||||
await records.aclose()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_snapshot_source_reports_only_aggregate_new_records(tmp_path):
|
||||
class Synchronizer:
|
||||
def __init__(self):
|
||||
self.snapshot_fingerprints = None
|
||||
self.calls = 0
|
||||
self.source = None
|
||||
|
||||
async def sync_once(self):
|
||||
self.calls += 1
|
||||
if self.calls == 1:
|
||||
self.snapshot_fingerprints = frozenset({"a", "b"})
|
||||
else:
|
||||
self.snapshot_fingerprints = frozenset({"a", "b", "c"})
|
||||
self.source._closed = True
|
||||
return True
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
async def scenario():
|
||||
synchronizer = Synchronizer()
|
||||
events = []
|
||||
source = DiskSnapshotSource(
|
||||
tmp_path / "source-snapshot.jsonl",
|
||||
synchronizer=synchronizer,
|
||||
sync_seconds=0.001,
|
||||
event_callback=lambda kind, data: events.append((kind, data)),
|
||||
)
|
||||
synchronizer.source = source
|
||||
|
||||
await source._sync_loop()
|
||||
|
||||
assert events == [
|
||||
("source_connected", {}),
|
||||
("source_updated", {"new": 1, "total": 3}),
|
||||
]
|
||||
|
||||
asyncio.run(scenario())
|
||||
100
tools/GrokSession2CPAAndSub2API/README.md
Normal file
100
tools/GrokSession2CPAAndSub2API/README.md
Normal file
@@ -0,0 +1,100 @@
|
||||
# Grok Session → CPA / sub2api
|
||||
|
||||
纯前端单页工具：把 grok-free-register 产出的凭证，转成 CPA、sub2api 导入包、batch 建号、或 sso-to-oauth 请求体。
|
||||
|
||||
对照思路：GPT session → CPA / sub2api 纯前端转换器（同类工具，无外部依赖）。
|
||||
|
||||
## 本地使用
|
||||
|
||||
直接打开：
|
||||
|
||||
```text
|
||||
tools/GrokSession2CPAAndSub2API/index.html
|
||||
```
|
||||
|
||||
所有解析和转换都在浏览器本地完成，不上传 token。
|
||||
|
||||
## 支持输入
|
||||
|
||||
| 来源 | 说明 |
|
||||
|------|------|
|
||||
| `authenticated/xai-*.json` | 注册成功 OAuth（base_url 常是 api.x.ai） |
|
||||
| `cpa_ready/xai-*.json` | watchdog 终态（带 headers/email） |
|
||||
| `sub2api-data` / accounts 数组 | 已转换过的包，可再导成其它格式 |
|
||||
| `auth-sessions.jsonl` | `email` + `cookies.sso` |
|
||||
| `accounts.txt` | `email:password:sso` |
|
||||
| `grok.txt` | 每行一个 SSO JWT |
|
||||
|
||||
可粘贴 JSON / 数组 / jsonl，或拖入多个文件。
|
||||
|
||||
## 输出格式
|
||||
|
||||
### CPA
|
||||
|
||||
对齐注册机 cpa_ready：
|
||||
|
||||
- `type: "xai"`
|
||||
- `access_token` / `refresh_token` / `id_token`
|
||||
- `base_url` **强制** `https://cli-chat-proxy.grok.com/v1`（忽略 authenticated 里的 `api.x.ai`）
|
||||
- `headers`：保留源 headers，否则写 grok-cli 默认头
|
||||
- `email` / `sub` / `expired` / `expires_in`
|
||||
|
||||
单条输出对象，多条输出数组。
|
||||
|
||||
### sub2api 导入
|
||||
|
||||
`type: "sub2api-data"` 包，可走 sub2api 数据导入：
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "sub2api-data",
|
||||
"version": 1,
|
||||
"exported_at": "...",
|
||||
"proxies": [],
|
||||
"accounts": [ ... ]
|
||||
}
|
||||
```
|
||||
|
||||
账号字段：
|
||||
|
||||
- `platform: "grok"`, `type: "oauth"`
|
||||
- `credentials.access_token` / `refresh_token` / `id_token` / `expires_at` / `email` / `sub` / `base_url`
|
||||
- **有 `refresh_token`**：不写账号顶层 `expires_at` / `auto_pause_on_expired`（避免 access 6h 到期被 pause）
|
||||
- **无 `refresh_token`**：写 access JWT exp + `auto_pause_on_expired: true`
|
||||
|
||||
### batch 建号
|
||||
|
||||
`POST /api/v1/admin/accounts/batch` 体：
|
||||
|
||||
```json
|
||||
{ "accounts": [ /* 同上 sub2api 账号对象 */ ] }
|
||||
```
|
||||
|
||||
### sso-to-oauth
|
||||
|
||||
`POST /api/v1/admin/grok/sso-to-oauth` 体：
|
||||
|
||||
```json
|
||||
{
|
||||
"sso_tokens": ["eyJ...", "..."],
|
||||
"_meta": [ /* 仅本地对照，API 可忽略 */ ]
|
||||
}
|
||||
```
|
||||
|
||||
纯 SSO 输入（grok.txt / accounts.txt / auth-sessions）**只能**走这个格式；CPA / sub2api / batch 需要 OAuth。
|
||||
|
||||
## 关键规则（别踩坑）
|
||||
|
||||
1. free 号 API 走 `cli-chat-proxy.grok.com`，不是 `api.x.ai`。
|
||||
2. free OAuth access ≈ 6h（`expires_in=21600`）。有 refresh 时让 sub2api 自己刷；别把账号级 `expires_at` 钉成 access exp。
|
||||
3. 当前线上常见：refresh 被整批 revoke，access 在 exp 前仍可用。这种号导入后到期无法续命，只能重铸。
|
||||
4. 注册机 `grok.txt` / `accounts.txt` / `auth-sessions.jsonl` 只有 SSO，没有 OAuth；要进 sub2api 要么走 sso-to-oauth，要么先走 xai_enroller 出 `authenticated/` 再转。
|
||||
|
||||
## 与注册机衔接
|
||||
|
||||
```
|
||||
register → grok.txt / accounts.txt / auth-sessions.jsonl (SSO)
|
||||
→ xai_enroller → authenticated/xai-*.json (OAuth)
|
||||
→ acpa_watchdog → cpa_ready/xai-*.json + acc.md
|
||||
→ 本工具 → sub2api-data / batch / CPA / sso-to-oauth
|
||||
```
|
||||
902
tools/GrokSession2CPAAndSub2API/index.html
Normal file
902
tools/GrokSession2CPAAndSub2API/index.html
Normal file
@@ -0,0 +1,902 @@
|
||||
<!doctype html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<title>Grok Session → CPA / sub2api</title>
|
||||
<meta name="theme-color" content="#111827" />
|
||||
<style>
|
||||
:root {
|
||||
color-scheme: light dark;
|
||||
--bg: #0b1220;
|
||||
--surface: #121a2b;
|
||||
--surface-soft: #1a2438;
|
||||
--line: #2a3650;
|
||||
--text: #e8eefc;
|
||||
--muted: #93a0bc;
|
||||
--accent: #38bdf8;
|
||||
--accent-strong: #0ea5e9;
|
||||
--success: #34d399;
|
||||
--danger: #f87171;
|
||||
--warning: #fbbf24;
|
||||
--mono: "SFMono-Regular", Consolas, "Liberation Mono", Menlo, monospace;
|
||||
--sans: "Inter", "Avenir Next", "Helvetica Neue", "PingFang SC", "Noto Sans SC", sans-serif;
|
||||
--radius: 10px;
|
||||
--shadow: 0 18px 50px rgba(0, 0, 0, 0.35);
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
html, body { min-height: 100%; margin: 0; }
|
||||
body {
|
||||
font-family: var(--sans);
|
||||
color: var(--text);
|
||||
background:
|
||||
radial-gradient(circle at top left, rgba(56, 189, 248, 0.12), transparent 28%),
|
||||
linear-gradient(180deg, #0b1220, #0a101b 40%, #070b14);
|
||||
}
|
||||
button, textarea, input { font: inherit; }
|
||||
button { border: 0; cursor: pointer; }
|
||||
.app { width: min(1280px, calc(100vw - 28px)); margin: 0 auto; padding: 22px 0 36px; }
|
||||
.topbar { display: flex; justify-content: space-between; gap: 16px; align-items: flex-end; margin-bottom: 16px; }
|
||||
.eyebrow { margin: 0 0 8px; color: var(--accent); font-size: 0.78rem; font-weight: 800; letter-spacing: 0.12em; text-transform: uppercase; }
|
||||
h1 { margin: 0; font-size: clamp(1.5rem, 3vw, 2.3rem); line-height: 1.1; }
|
||||
.subtitle { margin: 10px 0 0; color: var(--muted); line-height: 1.5; max-width: 70ch; }
|
||||
.chips { display: flex; flex-wrap: wrap; gap: 8px; justify-content: flex-end; }
|
||||
.chip {
|
||||
display: inline-flex; align-items: center; min-height: 32px; padding: 6px 10px;
|
||||
border: 1px solid var(--line); border-radius: 999px; background: rgba(18, 26, 43, 0.9); color: var(--muted); font-size: 0.84rem;
|
||||
}
|
||||
.toolbar { display: flex; flex-wrap: wrap; gap: 10px; justify-content: space-between; align-items: center; margin-bottom: 12px; }
|
||||
.segmented {
|
||||
display: inline-grid; grid-template-columns: repeat(4, minmax(0, 1fr)); gap: 4px;
|
||||
width: min(760px, 100%); padding: 4px; border: 1px solid var(--line); border-radius: var(--radius); background: var(--surface);
|
||||
}
|
||||
.segmented button {
|
||||
min-height: 40px; border-radius: 8px; background: transparent; color: var(--muted); padding: 0 10px;
|
||||
}
|
||||
.segmented button[aria-pressed="true"] { background: var(--accent-strong); color: #041018; font-weight: 800; }
|
||||
.actions { display: flex; flex-wrap: wrap; gap: 8px; }
|
||||
.button {
|
||||
display: inline-flex; align-items: center; justify-content: center; min-height: 40px; padding: 0 14px;
|
||||
border-radius: var(--radius); transition: 0.15s ease;
|
||||
}
|
||||
.button-primary { background: var(--accent-strong); color: #041018; font-weight: 800; }
|
||||
.button-primary:hover { filter: brightness(1.08); }
|
||||
.button-secondary { border: 1px solid var(--line); background: var(--surface); color: var(--text); }
|
||||
.button-secondary:hover { background: var(--surface-soft); }
|
||||
.button:disabled { opacity: 0.5; cursor: not-allowed; }
|
||||
.workspace { display: grid; grid-template-columns: minmax(0, 1fr) minmax(0, 1fr); gap: 12px; }
|
||||
.panel {
|
||||
min-width: 0; border: 1px solid var(--line); border-radius: var(--radius);
|
||||
background: rgba(18, 26, 43, 0.92); box-shadow: var(--shadow);
|
||||
}
|
||||
.panel-head { display: flex; justify-content: space-between; gap: 12px; padding: 14px 16px; border-bottom: 1px solid var(--line); }
|
||||
.panel-head h2 { margin: 0; font-size: 1rem; }
|
||||
.panel-head p { margin: 4px 0 0; color: var(--muted); font-size: 0.88rem; line-height: 1.45; }
|
||||
.panel-body { padding: 14px 16px 16px; }
|
||||
.guide {
|
||||
margin: 0 0 12px; padding: 12px 14px; border: 1px solid rgba(56, 189, 248, 0.28);
|
||||
border-left: 4px solid var(--accent); border-radius: var(--radius); background: rgba(14, 165, 233, 0.08);
|
||||
color: var(--muted); line-height: 1.5; font-size: 0.9rem;
|
||||
}
|
||||
.guide strong { color: var(--text); }
|
||||
.guide code { font-family: var(--mono); color: var(--accent); }
|
||||
textarea {
|
||||
width: 100%; min-height: 460px; resize: vertical; border: 1px solid var(--line); border-radius: var(--radius);
|
||||
padding: 12px; background: #0a1220; color: var(--text); font-family: var(--mono); font-size: 0.82rem; line-height: 1.5;
|
||||
}
|
||||
textarea:focus { outline: 2px solid rgba(56, 189, 248, 0.35); border-color: var(--accent); }
|
||||
.status-row { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 10px; }
|
||||
.status {
|
||||
min-height: 30px; display: inline-flex; align-items: center; padding: 4px 10px; border-radius: 999px;
|
||||
border: 1px solid var(--line); color: var(--muted); font-size: 0.84rem; background: rgba(0,0,0,0.15);
|
||||
}
|
||||
.status.ok { color: var(--success); border-color: rgba(52, 211, 153, 0.35); }
|
||||
.status.err { color: var(--danger); border-color: rgba(248, 113, 113, 0.35); }
|
||||
.status.warn { color: var(--warning); border-color: rgba(251, 191, 36, 0.35); }
|
||||
.accounts, .issues {
|
||||
margin-top: 12px; display: grid; gap: 8px; max-height: 180px; overflow: auto;
|
||||
}
|
||||
.card {
|
||||
border: 1px solid var(--line); border-radius: 8px; padding: 10px 12px; background: rgba(0,0,0,0.14);
|
||||
font-size: 0.86rem; line-height: 1.4;
|
||||
}
|
||||
.card .title { font-weight: 700; }
|
||||
.card .meta { color: var(--muted); margin-top: 4px; font-family: var(--mono); font-size: 0.78rem; word-break: break-all; }
|
||||
.drop-active textarea { border-color: var(--accent); box-shadow: inset 0 0 0 1px rgba(56,189,248,0.35); }
|
||||
.footer-note { margin-top: 14px; color: var(--muted); font-size: 0.84rem; line-height: 1.5; }
|
||||
@media (max-width: 960px) {
|
||||
.workspace { grid-template-columns: 1fr; }
|
||||
.segmented { grid-template-columns: repeat(2, minmax(0, 1fr)); width: 100%; }
|
||||
.topbar { flex-direction: column; align-items: flex-start; }
|
||||
.chips { justify-content: flex-start; }
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="app">
|
||||
<div class="topbar">
|
||||
<div>
|
||||
<p class="eyebrow">Grok Free Register toolkit</p>
|
||||
<h1>Grok Session → CPA / sub2api</h1>
|
||||
<p class="subtitle">
|
||||
纯前端转换：把注册机 <code>xai-*.json</code> / SSO / accounts / sessions
|
||||
转成 CPA、sub2api 导入包、batch 建号、或 sso-to-oauth 请求体。本地解析，不上传。
|
||||
</p>
|
||||
</div>
|
||||
<div class="chips">
|
||||
<span class="chip">platform: grok</span>
|
||||
<span class="chip">base: cli-chat-proxy</span>
|
||||
<span class="chip">有 refresh 不写账号 6h 死期</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="toolbar">
|
||||
<div class="segmented" role="tablist" aria-label="输出格式">
|
||||
<button type="button" data-format="cpa" aria-pressed="true">CPA</button>
|
||||
<button type="button" data-format="sub2api" aria-pressed="false">sub2api 导入</button>
|
||||
<button type="button" data-format="batch" aria-pressed="false">batch 建号</button>
|
||||
<button type="button" data-format="sso" aria-pressed="false">sso-to-oauth</button>
|
||||
</div>
|
||||
<div class="actions">
|
||||
<button class="button button-secondary" id="btn-sample" type="button">填示例</button>
|
||||
<button class="button button-secondary" id="btn-clear" type="button">清空</button>
|
||||
<button class="button button-secondary" id="btn-copy" type="button" disabled>复制输出</button>
|
||||
<button class="button button-primary" id="btn-download" type="button" disabled>下载 JSON</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="workspace">
|
||||
<section class="panel" id="input-panel">
|
||||
<div class="panel-head">
|
||||
<div>
|
||||
<h2>输入</h2>
|
||||
<p>粘贴 / 拖入 JSON、JSON 数组、jsonl、accounts 行、纯 SSO 行</p>
|
||||
</div>
|
||||
</div>
|
||||
<div class="panel-body">
|
||||
<div class="guide">
|
||||
<strong>支持输入</strong><br />
|
||||
1. <code>authenticated/xai-*.json</code> 或 <code>cpa_ready/xai-*.json</code><br />
|
||||
2. sub2api account / accounts 数组 / sub2api-data 包<br />
|
||||
3. <code>auth-sessions.jsonl</code>（email + cookies.sso）<br />
|
||||
4. <code>accounts.txt</code>：<code>email:password:sso</code><br />
|
||||
5. <code>grok.txt</code>：每行一个 SSO JWT
|
||||
</div>
|
||||
<textarea id="session-input" spellcheck="false" placeholder="在此粘贴 xai-*.json / SSO / accounts 行…"></textarea>
|
||||
<div class="status-row">
|
||||
<span class="status" id="input-status">等待输入</span>
|
||||
<span class="status" id="convert-status">未转换</span>
|
||||
</div>
|
||||
<div class="accounts" id="accounts"></div>
|
||||
<div class="issues" id="issues"></div>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<section class="panel">
|
||||
<div class="panel-head">
|
||||
<div>
|
||||
<h2>输出</h2>
|
||||
<p id="output-help">CPA：type=xai，对齐注册机 cpa_ready / new-api 风格</p>
|
||||
</div>
|
||||
</div>
|
||||
<div class="panel-body">
|
||||
<textarea id="output" readonly spellcheck="false" placeholder="转换结果会显示在这里"></textarea>
|
||||
<div class="status-row">
|
||||
<span class="status" id="output-status">空</span>
|
||||
</div>
|
||||
<p class="footer-note">
|
||||
规则：<br />
|
||||
• free 号 <code>base_url</code> 固定为 <code>https://cli-chat-proxy.grok.com/v1</code>（忽略 authenticated 里的 api.x.ai）<br />
|
||||
• 有 <code>refresh_token</code> 时：sub2api 账号<strong>不写</strong>顶层 <code>expires_at/auto_pause</code>，避免 6h 被 pause<br />
|
||||
• 无 refresh 时：写 access 的 exp + auto_pause，到期停用<br />
|
||||
• 纯 SSO 输入只能出 <code>sso-to-oauth</code> 请求体；CPA/sub2api 需要 OAuth token
|
||||
</p>
|
||||
</div>
|
||||
</section>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
(() => {
|
||||
const CLI_BASE = "https://cli-chat-proxy.grok.com/v1";
|
||||
const TOKEN_ENDPOINT = "https://auth.x.ai/oauth2/token";
|
||||
const HELP = {
|
||||
cpa: "CPA：type=xai，对齐注册机 cpa_ready / cli-proxy 风格",
|
||||
sub2api: "sub2api 数据导入：type=sub2api-data / accounts[] platform=grok",
|
||||
batch: "batch 建号：POST /api/v1/admin/accounts/batch 的 {accounts:[]}",
|
||||
sso: "sso-to-oauth：POST /api/v1/admin/grok/sso-to-oauth 的 {sso_tokens:[]}",
|
||||
};
|
||||
|
||||
const els = {
|
||||
input: document.querySelector("#session-input"),
|
||||
output: document.querySelector("#output"),
|
||||
inputStatus: document.querySelector("#input-status"),
|
||||
convertStatus: document.querySelector("#convert-status"),
|
||||
outputStatus: document.querySelector("#output-status"),
|
||||
accounts: document.querySelector("#accounts"),
|
||||
issues: document.querySelector("#issues"),
|
||||
outputHelp: document.querySelector("#output-help"),
|
||||
inputPanel: document.querySelector("#input-panel"),
|
||||
btnCopy: document.querySelector("#btn-copy"),
|
||||
btnDownload: document.querySelector("#btn-download"),
|
||||
btnClear: document.querySelector("#btn-clear"),
|
||||
btnSample: document.querySelector("#btn-sample"),
|
||||
formatButtons: [...document.querySelectorAll("[data-format]")],
|
||||
};
|
||||
|
||||
const state = {
|
||||
format: "cpa",
|
||||
converted: [],
|
||||
ssoTokens: [],
|
||||
skipped: [],
|
||||
timer: null,
|
||||
outputText: "",
|
||||
};
|
||||
|
||||
function isPlainObject(value) {
|
||||
return Boolean(value) && typeof value === "object" && !Array.isArray(value);
|
||||
}
|
||||
|
||||
function firstNonEmpty(...values) {
|
||||
for (const value of values) {
|
||||
if (value === undefined || value === null) continue;
|
||||
if (typeof value === "string" && value.trim() === "") continue;
|
||||
return typeof value === "string" ? value.trim() : value;
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function setStatus(el, text, tone = "") {
|
||||
el.textContent = text;
|
||||
el.className = "status" + (tone ? ` ${tone}` : "");
|
||||
}
|
||||
|
||||
function decodeBase64Url(value) {
|
||||
const normalized = String(value).replace(/-/g, "+").replace(/_/g, "/");
|
||||
const padded = normalized.padEnd(Math.ceil(normalized.length / 4) * 4, "=");
|
||||
try {
|
||||
return decodeURIComponent(
|
||||
Array.from(atob(padded), (c) => `%${c.charCodeAt(0).toString(16).padStart(2, "0")}`).join("")
|
||||
);
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
function parseJwtPayload(token) {
|
||||
if (!token || typeof token !== "string" || token.split(".").length < 2) return null;
|
||||
try {
|
||||
const raw = decodeBase64Url(token.split(".")[1]);
|
||||
return raw ? JSON.parse(raw) : null;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
function looksLikeJwt(value) {
|
||||
return typeof value === "string" && /^[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+$/.test(value.trim());
|
||||
}
|
||||
|
||||
function normalizeTimestamp(value) {
|
||||
if (value === undefined || value === null || value === "") return undefined;
|
||||
if (value instanceof Date && !Number.isNaN(value.getTime())) return value.toISOString().replace(/\.\d{3}Z$/, "Z");
|
||||
if (typeof value === "number" && Number.isFinite(value)) {
|
||||
const ms = value > 1e12 ? value : value * 1000;
|
||||
return new Date(ms).toISOString().replace(/\.\d{3}Z$/, "Z");
|
||||
}
|
||||
if (typeof value === "string") {
|
||||
const n = Number(value);
|
||||
if (Number.isFinite(n) && /^\d+(\.\d+)?$/.test(value.trim())) {
|
||||
const ms = n > 1e12 ? n : n * 1000;
|
||||
return new Date(ms).toISOString().replace(/\.\d{3}Z$/, "Z");
|
||||
}
|
||||
const d = new Date(value);
|
||||
if (!Number.isNaN(d.getTime())) return d.toISOString().replace(/\.\d{3}Z$/, "Z");
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function unixSecondsFromValue(value) {
|
||||
if (value === undefined || value === null || value === "") return undefined;
|
||||
if (typeof value === "number" && Number.isFinite(value)) return Math.floor(value > 1e12 ? value / 1000 : value);
|
||||
if (typeof value === "string") {
|
||||
const n = Number(value);
|
||||
if (Number.isFinite(n) && /^\d+(\.\d+)?$/.test(value.trim())) return Math.floor(n > 1e12 ? n / 1000 : n);
|
||||
const d = new Date(value);
|
||||
if (!Number.isNaN(d.getTime())) return Math.floor(d.getTime() / 1000);
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function stripUndefined(obj) {
|
||||
if (Array.isArray(obj)) return obj.map(stripUndefined).filter((v) => v !== undefined);
|
||||
if (!isPlainObject(obj)) return obj;
|
||||
const out = {};
|
||||
for (const [k, v] of Object.entries(obj)) {
|
||||
if (v === undefined || v === null || v === "") continue;
|
||||
const next = stripUndefined(v);
|
||||
if (next === undefined) continue;
|
||||
if (isPlainObject(next) && !Object.keys(next).length) continue;
|
||||
out[k] = next;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
function extractSsoFromCookies(cookies) {
|
||||
if (!Array.isArray(cookies)) return undefined;
|
||||
let sso = "";
|
||||
let fallback = "";
|
||||
for (const cookie of cookies) {
|
||||
if (!isPlainObject(cookie)) continue;
|
||||
const name = String(cookie.name || "");
|
||||
const value = String(cookie.value || "").trim();
|
||||
if (!value) continue;
|
||||
if (name === "sso" && !sso) sso = value;
|
||||
if (name === "sso-rw" && !fallback) fallback = value;
|
||||
}
|
||||
return sso || fallback || undefined;
|
||||
}
|
||||
|
||||
function collectOAuthLikeObjects(value, sourceName = "pasted-json") {
|
||||
const found = [];
|
||||
const visited = new WeakSet();
|
||||
function visit(item, path) {
|
||||
if (!isPlainObject(item) && !Array.isArray(item)) return;
|
||||
if (Array.isArray(item)) {
|
||||
item.forEach((child, i) => visit(child, `${path}[${i}]`));
|
||||
return;
|
||||
}
|
||||
if (visited.has(item)) return;
|
||||
visited.add(item);
|
||||
|
||||
const access = firstNonEmpty(
|
||||
item.access_token, item.accessToken,
|
||||
item.tokens?.access_token, item.tokens?.accessToken,
|
||||
item.credentials?.access_token, item.credentials?.accessToken,
|
||||
);
|
||||
const refresh = firstNonEmpty(
|
||||
item.refresh_token, item.refreshToken,
|
||||
item.tokens?.refresh_token, item.tokens?.refreshToken,
|
||||
item.credentials?.refresh_token, item.credentials?.refreshToken,
|
||||
);
|
||||
const idToken = firstNonEmpty(
|
||||
item.id_token, item.idToken,
|
||||
item.tokens?.id_token, item.tokens?.idToken,
|
||||
item.credentials?.id_token, item.credentials?.idToken,
|
||||
);
|
||||
const email = firstNonEmpty(item.email, item.user?.email, item.credentials?.email, item.name);
|
||||
const sub = firstNonEmpty(item.sub, item.credentials?.sub, item.subject);
|
||||
const sso = firstNonEmpty(item.sso, item.sso_token, item.ssoToken, extractSsoFromCookies(item.cookies));
|
||||
const isXai = item.type === "xai" || item.auth_kind === "oauth" || item.platform === "grok";
|
||||
const isSub2 = item.platform === "grok" && item.credentials;
|
||||
|
||||
if (access || (refresh && (email || sub || isXai || isSub2)) || (sso && (email || item.cookies))) {
|
||||
found.push({ value: item, sourceName, path, kind: access || refresh ? "oauth" : "sso" });
|
||||
return;
|
||||
}
|
||||
|
||||
// sub2api-data envelope
|
||||
if (Array.isArray(item.accounts)) {
|
||||
visit(item.accounts, `${path}.accounts`);
|
||||
return;
|
||||
}
|
||||
|
||||
for (const [key, child] of Object.entries(item)) {
|
||||
if (["access_token", "accessToken", "refresh_token", "id_token", "cookies"].includes(key)) continue;
|
||||
visit(child, `${path}.${key}`);
|
||||
}
|
||||
}
|
||||
visit(value, "$");
|
||||
return found;
|
||||
}
|
||||
|
||||
function parseTextInput(text) {
|
||||
const raw = String(text || "").trim();
|
||||
if (!raw) return { oauthSources: [], ssoTokens: [], skipped: [] };
|
||||
|
||||
// Try full JSON first
|
||||
try {
|
||||
const parsed = JSON.parse(raw);
|
||||
const oauthSources = collectOAuthLikeObjects(parsed);
|
||||
const ssoTokens = [];
|
||||
// also collect pure sso fields if any
|
||||
for (const src of oauthSources) {
|
||||
if (src.kind === "sso") {
|
||||
const sso = firstNonEmpty(src.value.sso, src.value.sso_token, extractSsoFromCookies(src.value.cookies));
|
||||
if (sso) ssoTokens.push({ token: sso, email: firstNonEmpty(src.value.email, src.value.user?.email), path: src.path });
|
||||
}
|
||||
}
|
||||
if (oauthSources.length || ssoTokens.length) {
|
||||
return { oauthSources: oauthSources.filter((s) => s.kind === "oauth"), ssoTokens, skipped: [] };
|
||||
}
|
||||
} catch {
|
||||
// line mode below
|
||||
}
|
||||
|
||||
const oauthSources = [];
|
||||
const ssoTokens = [];
|
||||
const skipped = [];
|
||||
const lines = raw.split(/\r?\n/).map((l) => l.trim()).filter(Boolean);
|
||||
|
||||
lines.forEach((line, index) => {
|
||||
// jsonl
|
||||
if (line.startsWith("{") || line.startsWith("[")) {
|
||||
try {
|
||||
const parsed = JSON.parse(line);
|
||||
const found = collectOAuthLikeObjects(parsed, `line-${index + 1}`);
|
||||
for (const src of found) {
|
||||
if (src.kind === "oauth") oauthSources.push(src);
|
||||
else {
|
||||
const sso = firstNonEmpty(src.value.sso, src.value.sso_token, extractSsoFromCookies(src.value.cookies));
|
||||
if (sso) ssoTokens.push({ token: sso, email: firstNonEmpty(src.value.email), path: src.path });
|
||||
}
|
||||
}
|
||||
if (!found.length) skipped.push({ path: `line-${index + 1}`, reason: "JSON 行未识别为 Grok 凭证" });
|
||||
return;
|
||||
} catch (error) {
|
||||
skipped.push({ path: `line-${index + 1}`, reason: `JSON 行解析失败：${error.message}` });
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// accounts.txt: email:password:sso
|
||||
if (line.includes(":") && !looksLikeJwt(line)) {
|
||||
const parts = line.split(":");
|
||||
if (parts.length >= 3) {
|
||||
const email = parts[0];
|
||||
const sso = parts.slice(2).join(":");
|
||||
if (looksLikeJwt(sso)) {
|
||||
ssoTokens.push({ token: sso.trim(), email: email.trim(), path: `accounts:${index + 1}` });
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// pure SSO jwt line
|
||||
if (looksLikeJwt(line)) {
|
||||
ssoTokens.push({ token: line.trim(), path: `sso:${index + 1}` });
|
||||
return;
|
||||
}
|
||||
|
||||
skipped.push({ path: `line-${index + 1}`, reason: "无法识别该行" });
|
||||
});
|
||||
|
||||
return { oauthSources, ssoTokens, skipped };
|
||||
}
|
||||
|
||||
function convertOAuthRecord(record, options = {}) {
|
||||
if (!isPlainObject(record)) throw new Error("记录不是对象");
|
||||
|
||||
// unwrap sub2api account
|
||||
const credsIn = isPlainObject(record.credentials) ? record.credentials : {};
|
||||
const accessToken = firstNonEmpty(
|
||||
record.access_token, record.accessToken,
|
||||
credsIn.access_token, credsIn.accessToken,
|
||||
record.tokens?.access_token,
|
||||
);
|
||||
if (!accessToken) throw new Error("缺少 access_token");
|
||||
|
||||
const refreshToken = firstNonEmpty(
|
||||
record.refresh_token, record.refreshToken,
|
||||
credsIn.refresh_token, credsIn.refreshToken,
|
||||
record.tokens?.refresh_token,
|
||||
);
|
||||
const idToken = firstNonEmpty(
|
||||
record.id_token, record.idToken,
|
||||
credsIn.id_token, credsIn.idToken,
|
||||
record.tokens?.id_token,
|
||||
);
|
||||
const tokenType = firstNonEmpty(record.token_type, credsIn.token_type, "Bearer");
|
||||
const accessClaims = parseJwtPayload(accessToken) || {};
|
||||
const idClaims = parseJwtPayload(idToken) || {};
|
||||
|
||||
const email = firstNonEmpty(
|
||||
record.email, credsIn.email, record.user?.email, record.name,
|
||||
idClaims.email, accessClaims.email, accessClaims.preferred_username,
|
||||
);
|
||||
const sub = firstNonEmpty(
|
||||
record.sub, credsIn.sub, record.subject,
|
||||
idClaims.sub, accessClaims.sub, accessClaims.principal_id,
|
||||
);
|
||||
|
||||
const expiresIn = firstNonEmpty(
|
||||
record.expires_in, credsIn.expires_in,
|
||||
accessClaims.exp && accessClaims.iat ? Number(accessClaims.exp) - Number(accessClaims.iat) : undefined,
|
||||
21600,
|
||||
);
|
||||
|
||||
const expiredIso = firstNonEmpty(
|
||||
normalizeTimestamp(record.expired),
|
||||
normalizeTimestamp(record.expires_at),
|
||||
normalizeTimestamp(credsIn.expires_at),
|
||||
normalizeTimestamp(record.expiresAt),
|
||||
accessClaims.exp ? normalizeTimestamp(accessClaims.exp) : undefined,
|
||||
);
|
||||
const expiresUnix = firstNonEmpty(
|
||||
unixSecondsFromValue(expiredIso),
|
||||
unixSecondsFromValue(accessClaims.exp),
|
||||
);
|
||||
|
||||
const lastRefresh = firstNonEmpty(
|
||||
normalizeTimestamp(record.last_refresh),
|
||||
normalizeTimestamp(credsIn.last_refresh),
|
||||
normalizeTimestamp(options.now || new Date()),
|
||||
);
|
||||
|
||||
const baseUrl = CLI_BASE; // free CLI 强制
|
||||
const hasRefresh = Boolean(refreshToken);
|
||||
const name = firstNonEmpty(email, sub && `grok-${String(sub).slice(0, 8)}`, options.sourceName, "Grok Account");
|
||||
|
||||
const headers = isPlainObject(record.headers) ? record.headers : {
|
||||
"x-grok-client-version": "0.2.93",
|
||||
"x-xai-token-auth": "xai-grok-cli",
|
||||
"X-XAI-Token-Auth": "xai-grok-cli",
|
||||
"x-authenticateresponse": "authenticate-response",
|
||||
"x-grok-client-identifier": "grok-shell",
|
||||
"x-compaction-at": "400000",
|
||||
"User-Agent": "grok-shell/0.2.93 (linux; x86_64)",
|
||||
};
|
||||
|
||||
const cpa = stripUndefined({
|
||||
type: "xai",
|
||||
access_token: accessToken,
|
||||
refresh_token: refreshToken,
|
||||
id_token: idToken,
|
||||
token_type: tokenType,
|
||||
expires_in: Number(expiresIn) || undefined,
|
||||
expired: expiredIso,
|
||||
last_refresh: lastRefresh,
|
||||
sub,
|
||||
base_url: baseUrl,
|
||||
token_endpoint: firstNonEmpty(record.token_endpoint, TOKEN_ENDPOINT),
|
||||
auth_kind: "oauth",
|
||||
headers,
|
||||
email,
|
||||
});
|
||||
|
||||
// sub2api credentials always keep token-level expires_at
|
||||
const credentials = stripUndefined({
|
||||
access_token: accessToken,
|
||||
refresh_token: refreshToken,
|
||||
id_token: idToken,
|
||||
token_type: tokenType,
|
||||
expires_at: expiredIso,
|
||||
email,
|
||||
sub,
|
||||
base_url: baseUrl,
|
||||
});
|
||||
|
||||
// 有 refresh：不写账号级 expires_at，避免 6h auto-pause
|
||||
// 无 refresh：写 access exp + auto_pause
|
||||
const sub2apiAccount = stripUndefined({
|
||||
name,
|
||||
platform: "grok",
|
||||
type: "oauth",
|
||||
credentials,
|
||||
concurrency: Number.isFinite(Number(record.concurrency)) ? Number(record.concurrency) : 1,
|
||||
priority: Number.isFinite(Number(record.priority)) ? Number(record.priority) : 0,
|
||||
expires_at: hasRefresh ? undefined : expiresUnix,
|
||||
auto_pause_on_expired: hasRefresh ? undefined : (expiresUnix ? true : undefined),
|
||||
notes: firstNonEmpty(record.notes, record.note, record.remark),
|
||||
});
|
||||
|
||||
return {
|
||||
sourceName: options.sourceName,
|
||||
sourcePath: options.sourcePath,
|
||||
email,
|
||||
name,
|
||||
sub,
|
||||
hasRefresh,
|
||||
expiresAt: expiredIso,
|
||||
expiresUnix,
|
||||
cpa,
|
||||
sub2apiAccount,
|
||||
};
|
||||
}
|
||||
|
||||
function convertFromText(text) {
|
||||
const { oauthSources, ssoTokens, skipped } = parseTextInput(text);
|
||||
const converted = [];
|
||||
const allSkipped = [...skipped];
|
||||
const now = new Date();
|
||||
|
||||
oauthSources.forEach((item, index) => {
|
||||
try {
|
||||
converted.push(convertOAuthRecord(item.value, {
|
||||
now,
|
||||
sourceName: item.sourceName,
|
||||
sourcePath: item.path || `$[${index}]`,
|
||||
}));
|
||||
} catch (error) {
|
||||
allSkipped.push({
|
||||
path: item.path || `$[${index}]`,
|
||||
reason: error instanceof Error ? error.message : "无法转换",
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
// de-dupe sso
|
||||
const seenSso = new Set();
|
||||
const uniqueSso = [];
|
||||
for (const item of ssoTokens) {
|
||||
const token = String(item.token || "").trim();
|
||||
if (!token || seenSso.has(token)) continue;
|
||||
seenSso.add(token);
|
||||
uniqueSso.push(item);
|
||||
}
|
||||
|
||||
if (!converted.length && !uniqueSso.length) {
|
||||
allSkipped.push({ path: "$", reason: "未找到可识别的 Grok OAuth / SSO 记录" });
|
||||
}
|
||||
|
||||
return { converted, ssoTokens: uniqueSso, skipped: allSkipped };
|
||||
}
|
||||
|
||||
function buildSub2apiDocument(converted, now = new Date()) {
|
||||
return {
|
||||
type: "sub2api-data",
|
||||
version: 1,
|
||||
exported_at: normalizeTimestamp(now),
|
||||
proxies: [],
|
||||
accounts: converted.map((item) => item.sub2apiAccount),
|
||||
};
|
||||
}
|
||||
|
||||
function buildBatchDocument(converted) {
|
||||
return { accounts: converted.map((item) => item.sub2apiAccount) };
|
||||
}
|
||||
|
||||
function buildSsoDocument(ssoTokens) {
|
||||
return {
|
||||
sso_tokens: ssoTokens.map((item) => item.token),
|
||||
// 便于对照；API 实际只吃 sso_tokens
|
||||
_meta: ssoTokens.map((item, i) => stripUndefined({
|
||||
index: i + 1,
|
||||
email: item.email,
|
||||
path: item.path,
|
||||
token_head: String(item.token).slice(0, 24) + "...",
|
||||
})),
|
||||
};
|
||||
}
|
||||
|
||||
function buildCpaDocument(converted) {
|
||||
return converted.length === 1
|
||||
? converted[0].cpa
|
||||
: converted.map((item) => item.cpa);
|
||||
}
|
||||
|
||||
function buildOutputDocument() {
|
||||
const now = new Date();
|
||||
if (state.format === "cpa") {
|
||||
if (!state.converted.length) throw new Error("CPA 需要 OAuth 凭证（access_token）。纯 SSO 请切换到 sso-to-oauth");
|
||||
return buildCpaDocument(state.converted);
|
||||
}
|
||||
if (state.format === "sub2api") {
|
||||
if (!state.converted.length) throw new Error("sub2api 导入需要 OAuth 凭证。纯 SSO 请切换到 sso-to-oauth");
|
||||
return buildSub2apiDocument(state.converted, now);
|
||||
}
|
||||
if (state.format === "batch") {
|
||||
if (!state.converted.length) throw new Error("batch 建号需要 OAuth 凭证。纯 SSO 请切换到 sso-to-oauth");
|
||||
return buildBatchDocument(state.converted);
|
||||
}
|
||||
if (state.format === "sso") {
|
||||
if (!state.ssoTokens.length) {
|
||||
// also allow extracting nothing useful message
|
||||
throw new Error("未找到 SSO token。可粘贴 grok.txt / accounts.txt / auth-sessions.jsonl");
|
||||
}
|
||||
return buildSsoDocument(state.ssoTokens);
|
||||
}
|
||||
return buildSub2apiDocument(state.converted, now);
|
||||
}
|
||||
|
||||
function renderAccounts() {
|
||||
if (!state.converted.length && !state.ssoTokens.length) {
|
||||
els.accounts.innerHTML = "";
|
||||
return;
|
||||
}
|
||||
const oauthCards = state.converted.map((item) => `
|
||||
<div class="card">
|
||||
<div class="title">${escapeHtml(item.name || item.email || "Grok Account")}</div>
|
||||
<div class="meta">
|
||||
email=${escapeHtml(item.email || "-")}
|
||||
· sub=${escapeHtml(item.sub || "-")}
|
||||
· refresh=${item.hasRefresh ? "yes" : "no"}
|
||||
· exp=${escapeHtml(item.expiresAt || "-")}
|
||||
</div>
|
||||
</div>
|
||||
`).join("");
|
||||
const ssoCards = state.ssoTokens.slice(0, 20).map((item, i) => `
|
||||
<div class="card">
|
||||
<div class="title">SSO #${i + 1}${item.email ? " · " + escapeHtml(item.email) : ""}</div>
|
||||
<div class="meta">${escapeHtml(String(item.token).slice(0, 48))}...</div>
|
||||
</div>
|
||||
`).join("");
|
||||
els.accounts.innerHTML = oauthCards + ssoCards;
|
||||
}
|
||||
|
||||
function renderIssues() {
|
||||
if (!state.skipped.length) {
|
||||
els.issues.innerHTML = "";
|
||||
return;
|
||||
}
|
||||
els.issues.innerHTML = state.skipped.map((item) => `
|
||||
<div class="card">
|
||||
<div class="title">跳过 ${escapeHtml(item.path || "?")}</div>
|
||||
<div class="meta">${escapeHtml(item.reason || "")}</div>
|
||||
</div>
|
||||
`).join("");
|
||||
}
|
||||
|
||||
function escapeHtml(value) {
|
||||
return String(value ?? "")
|
||||
.replaceAll("&", "&")
|
||||
.replaceAll("<", "<")
|
||||
.replaceAll(">", ">")
|
||||
.replaceAll('"', """);
|
||||
}
|
||||
|
||||
function updateOutput() {
|
||||
els.outputHelp.textContent = HELP[state.format] || "";
|
||||
try {
|
||||
if (!els.input.value.trim()) {
|
||||
state.outputText = "";
|
||||
els.output.value = "";
|
||||
els.btnCopy.disabled = true;
|
||||
els.btnDownload.disabled = true;
|
||||
setStatus(els.outputStatus, "空");
|
||||
setStatus(els.convertStatus, "未转换");
|
||||
return;
|
||||
}
|
||||
const doc = buildOutputDocument();
|
||||
state.outputText = JSON.stringify(doc, null, 2);
|
||||
els.output.value = state.outputText;
|
||||
els.btnCopy.disabled = false;
|
||||
els.btnDownload.disabled = false;
|
||||
const count = state.format === "sso" ? state.ssoTokens.length : state.converted.length;
|
||||
setStatus(els.outputStatus, `已生成 · ${count} 条`, "ok");
|
||||
setStatus(els.convertStatus, `OAuth ${state.converted.length} / SSO ${state.ssoTokens.length}`, "ok");
|
||||
} catch (error) {
|
||||
state.outputText = "";
|
||||
els.output.value = "";
|
||||
els.btnCopy.disabled = true;
|
||||
els.btnDownload.disabled = true;
|
||||
setStatus(els.outputStatus, error instanceof Error ? error.message : "生成失败", "err");
|
||||
setStatus(els.convertStatus, `OAuth ${state.converted.length} / SSO ${state.ssoTokens.length}`, state.converted.length || state.ssoTokens.length ? "warn" : "err");
|
||||
}
|
||||
}
|
||||
|
||||
function scheduleConvert() {
|
||||
clearTimeout(state.timer);
|
||||
state.timer = setTimeout(() => {
|
||||
try {
|
||||
const result = convertFromText(els.input.value);
|
||||
state.converted = result.converted;
|
||||
state.ssoTokens = result.ssoTokens;
|
||||
state.skipped = result.skipped;
|
||||
const total = state.converted.length + state.ssoTokens.length;
|
||||
setStatus(
|
||||
els.inputStatus,
|
||||
total ? `识别 OAuth ${state.converted.length} · SSO ${state.ssoTokens.length}` : "未识别到凭证",
|
||||
total ? "ok" : "warn",
|
||||
);
|
||||
renderAccounts();
|
||||
renderIssues();
|
||||
updateOutput();
|
||||
} catch (error) {
|
||||
state.converted = [];
|
||||
state.ssoTokens = [];
|
||||
state.skipped = [{ path: "$", reason: error instanceof Error ? error.message : "转换失败" }];
|
||||
setStatus(els.inputStatus, error instanceof Error ? error.message : "转换失败", "err");
|
||||
renderAccounts();
|
||||
renderIssues();
|
||||
updateOutput();
|
||||
}
|
||||
}, 120);
|
||||
}
|
||||
|
||||
function downloadOutput() {
|
||||
if (!state.outputText) return;
|
||||
const stamp = new Date().toISOString().replace(/[:.]/g, "-").slice(0, 19);
|
||||
const name = `grok-${state.format}-${stamp}.json`;
|
||||
const blob = new Blob([state.outputText], { type: "application/json;charset=utf-8" });
|
||||
const url = URL.createObjectURL(blob);
|
||||
const a = document.createElement("a");
|
||||
a.href = url;
|
||||
a.download = name;
|
||||
a.click();
|
||||
URL.revokeObjectURL(url);
|
||||
}
|
||||
|
||||
async function copyOutput() {
|
||||
if (!state.outputText) return;
|
||||
try {
|
||||
await navigator.clipboard.writeText(state.outputText);
|
||||
setStatus(els.outputStatus, "已复制", "ok");
|
||||
} catch {
|
||||
els.output.select();
|
||||
document.execCommand("copy");
|
||||
setStatus(els.outputStatus, "已复制", "ok");
|
||||
}
|
||||
}
|
||||
|
||||
function setFormat(format) {
|
||||
state.format = format;
|
||||
els.formatButtons.forEach((btn) => {
|
||||
btn.setAttribute("aria-pressed", btn.dataset.format === format ? "true" : "false");
|
||||
});
|
||||
updateOutput();
|
||||
}
|
||||
|
||||
function fillSample() {
|
||||
const now = Math.floor(Date.now() / 1000);
|
||||
const exp = now + 21600;
|
||||
// fake jwt-like payload (not signed; only for local demo decode)
|
||||
const header = btoa(JSON.stringify({ alg: "none", typ: "JWT" })).replace(/=+$/, "").replace(/\+/g, "-").replace(/\//g, "_");
|
||||
const payload = btoa(JSON.stringify({
|
||||
sub: "55c526cb-9f70-4486-a232-3710e6f3df38",
|
||||
email: "demo@example.com",
|
||||
iat: now,
|
||||
exp,
|
||||
})).replace(/=+$/, "").replace(/\+/g, "-").replace(/\//g, "_");
|
||||
const access = `${header}.${payload}.demo`;
|
||||
const sample = {
|
||||
type: "xai",
|
||||
access_token: access,
|
||||
refresh_token: "demo-refresh-token-replace-me",
|
||||
id_token: access,
|
||||
token_type: "Bearer",
|
||||
expires_in: 21600,
|
||||
expired: new Date(exp * 1000).toISOString().replace(/\.\d{3}Z$/, "Z"),
|
||||
last_refresh: new Date().toISOString().replace(/\.\d{3}Z$/, "Z"),
|
||||
sub: "55c526cb-9f70-4486-a232-3710e6f3df38",
|
||||
base_url: "https://api.x.ai/v1",
|
||||
token_endpoint: TOKEN_ENDPOINT,
|
||||
auth_kind: "oauth",
|
||||
email: "demo@example.com",
|
||||
};
|
||||
els.input.value = JSON.stringify(sample, null, 2);
|
||||
scheduleConvert();
|
||||
}
|
||||
|
||||
// events
|
||||
els.input.addEventListener("input", scheduleConvert);
|
||||
els.btnClear.addEventListener("click", () => {
|
||||
els.input.value = "";
|
||||
scheduleConvert();
|
||||
});
|
||||
els.btnSample.addEventListener("click", fillSample);
|
||||
els.btnCopy.addEventListener("click", copyOutput);
|
||||
els.btnDownload.addEventListener("click", downloadOutput);
|
||||
els.formatButtons.forEach((btn) => {
|
||||
btn.addEventListener("click", () => setFormat(btn.dataset.format));
|
||||
});
|
||||
|
||||
// drag-drop
|
||||
["dragenter", "dragover"].forEach((type) => {
|
||||
els.inputPanel.addEventListener(type, (event) => {
|
||||
event.preventDefault();
|
||||
els.inputPanel.classList.add("drop-active");
|
||||
});
|
||||
});
|
||||
["dragleave", "drop"].forEach((type) => {
|
||||
els.inputPanel.addEventListener(type, (event) => {
|
||||
event.preventDefault();
|
||||
if (type === "drop") {
|
||||
const files = [...(event.dataTransfer?.files || [])];
|
||||
if (!files.length) return;
|
||||
Promise.all(files.map((file) => file.text())).then((parts) => {
|
||||
const existing = els.input.value.trim();
|
||||
const incoming = parts.join("\n");
|
||||
els.input.value = existing ? `${existing}\n${incoming}` : incoming;
|
||||
scheduleConvert();
|
||||
});
|
||||
}
|
||||
els.inputPanel.classList.remove("drop-active");
|
||||
});
|
||||
});
|
||||
|
||||
setFormat("cpa");
|
||||
})();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
1
tools/__init__.py
Normal file
1
tools/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Developer-only analysis and stress tools."""
|
||||
419
tools/runtime_log_analyzer.py
Normal file
419
tools/runtime_log_analyzer.py
Normal file
@@ -0,0 +1,419 @@
|
||||
"""
|
||||
Parse real runtime monitor logs from both CSP and legacy state-machine runs.
|
||||
|
||||
The analyzer is intentionally read-only. It does not tune parameters or modify
|
||||
runtime state; it only turns production logs into comparable stage rates.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
import re
|
||||
from typing import Iterable
|
||||
|
||||
|
||||
_CSP_RE = re.compile(
|
||||
r"^\[\*\] T:(?P<t>\d+) Q:(?P<q>\d+) phys:(?P<phys>\d+) "
|
||||
r"(?:p_send:(?P<p_send>\d+) )?t_slot:(?P<t_slot>\d+) q_slot:(?P<q_slot>\d+) q_pend:(?P<q_pend>\d+)"
|
||||
r"(?P<rest>.*)rate:(?P<rate>[0-9.]+)/min #(?P<ok>\d+)"
|
||||
)
|
||||
|
||||
_STATE_RE = re.compile(
|
||||
r"^\[\*\] slots:(?P<slots>\d+)/(?P<max_slots>\d+) act:(?P<active>\d+) "
|
||||
r"cpu:(?P<cpu>[0-9.]+)% avg:(?P<cpu_avg>[0-9.]+)%/(?P<cpu_target>[0-9.]+) "
|
||||
r"mem:(?P<mem>\d+)M T:(?P<t>\d+) Q:(?P<q>\d+) "
|
||||
r"sent:(?P<q_sent>\d+) got:(?P<q_ret>\d+)\((?P<q_hit>[0-9.]+)%\) "
|
||||
r"rate:(?P<rate>[0-9.]+)/min #(?P<ok>\d+)"
|
||||
)
|
||||
|
||||
_SOLVER_TIMELINE_PREFIX = "[solver_timeline] "
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MonitorRow:
|
||||
kind: str
|
||||
t: int
|
||||
q: int
|
||||
ok: int
|
||||
rate: float
|
||||
phys: int | None = None
|
||||
p_send: int | None = None
|
||||
t_slot: int | None = None
|
||||
q_slot: int | None = None
|
||||
q_pend: int | None = None
|
||||
p_batch: float | None = None
|
||||
t_prog: int | None = None
|
||||
q_inflight: int | None = None
|
||||
t_prod: int | None = None
|
||||
q_sent: int | None = None
|
||||
q_ret: int | None = None
|
||||
q_adm: int | None = None
|
||||
pair: int | None = None
|
||||
fail: int | None = None
|
||||
s_phys_wait: float | None = None
|
||||
s_phys_hold: float | None = None
|
||||
p_phys_wait: float | None = None
|
||||
p_phys_hold: float | None = None
|
||||
c_phys_wait: float | None = None
|
||||
c_phys_hold: float | None = None
|
||||
p_email_create: float | None = None
|
||||
p_page_prepare: float | None = None
|
||||
p_send_stage: float | None = None
|
||||
c_page_acquire: float | None = None
|
||||
c_verify: float | None = None
|
||||
c_register: float | None = None
|
||||
c_hot_hits: int | None = None
|
||||
c_hot_misses: int | None = None
|
||||
solver_goto: float | None = None
|
||||
solver_inject: float | None = None
|
||||
solver_initial: float | None = None
|
||||
solver_click: float | None = None
|
||||
solver_wait: float | None = None
|
||||
solver_reuse: float | None = None
|
||||
solver_visible: float | None = None
|
||||
slots: int | None = None
|
||||
max_slots: int | None = None
|
||||
active: int | None = None
|
||||
|
||||
@property
|
||||
def elapsed_min(self) -> float | None:
|
||||
if self.ok <= 0 or self.rate <= 0:
|
||||
return None
|
||||
return self.ok / self.rate
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SolverTimeline:
|
||||
events: list[dict]
|
||||
|
||||
|
||||
def _int_field(rest: str, name: str) -> int | None:
|
||||
match = re.search(rf"\b{name}:(\d+)", rest)
|
||||
return int(match.group(1)) if match else None
|
||||
|
||||
|
||||
def _float_field(rest: str, name: str) -> float | None:
|
||||
match = re.search(rf"\b{name}:([0-9.]+)", rest)
|
||||
return float(match.group(1)) if match else None
|
||||
|
||||
|
||||
def _float_pair_field(rest: str, name: str) -> tuple[float | None, float | None]:
|
||||
match = re.search(rf"\b{name}:([0-9.]+)/([0-9.]+)", rest)
|
||||
if not match:
|
||||
return None, None
|
||||
return float(match.group(1)), float(match.group(2))
|
||||
|
||||
|
||||
def _float_triple_field(rest: str, name: str) -> tuple[float | None, float | None, float | None]:
|
||||
match = re.search(rf"\b{name}:([0-9.]+)/([0-9.]+)/([0-9.]+)", rest)
|
||||
if not match:
|
||||
return None, None, None
|
||||
return float(match.group(1)), float(match.group(2)), float(match.group(3))
|
||||
|
||||
|
||||
def _int_pair_field(rest: str, name: str) -> tuple[int | None, int | None]:
|
||||
match = re.search(rf"\b{name}:(\d+)/(\d+)", rest)
|
||||
if not match:
|
||||
return None, None
|
||||
return int(match.group(1)), int(match.group(2))
|
||||
|
||||
|
||||
def parse_monitor_lines(text_or_lines: str | Iterable[str]) -> list[MonitorRow]:
|
||||
if isinstance(text_or_lines, str):
|
||||
lines = text_or_lines.splitlines()
|
||||
else:
|
||||
lines = list(text_or_lines)
|
||||
|
||||
rows: list[MonitorRow] = []
|
||||
for line in lines:
|
||||
csp = _CSP_RE.match(line)
|
||||
if csp:
|
||||
rest = csp.group("rest")
|
||||
s_phys_wait, s_phys_hold = _float_pair_field(rest, "s_phys")
|
||||
p_phys_wait, p_phys_hold = _float_pair_field(rest, "p_phys")
|
||||
c_phys_wait, c_phys_hold = _float_pair_field(rest, "c_phys")
|
||||
p_email_create, p_page_prepare, p_send_stage = _float_triple_field(rest, "p_stage")
|
||||
c_page_acquire, c_verify, c_register = _float_triple_field(rest, "c_stage")
|
||||
c_hot_hits, c_hot_misses = _int_pair_field(rest, "c_hot")
|
||||
rows.append(
|
||||
MonitorRow(
|
||||
kind="csp",
|
||||
t=int(csp.group("t")),
|
||||
q=int(csp.group("q")),
|
||||
ok=int(csp.group("ok")),
|
||||
rate=float(csp.group("rate")),
|
||||
phys=int(csp.group("phys")),
|
||||
p_send=int(csp.group("p_send")) if csp.group("p_send") is not None else None,
|
||||
t_slot=int(csp.group("t_slot")),
|
||||
q_slot=int(csp.group("q_slot")),
|
||||
q_pend=int(csp.group("q_pend")),
|
||||
p_batch=_float_field(rest, "p_batch"),
|
||||
t_prog=_int_field(rest, "t_prog"),
|
||||
q_inflight=_int_field(rest, "q_inflight"),
|
||||
t_prod=_int_field(rest, "t_prod"),
|
||||
q_sent=_int_field(rest, "q_sent"),
|
||||
q_ret=_int_field(rest, "q_ret"),
|
||||
q_adm=_int_field(rest, "q_adm"),
|
||||
pair=_int_field(rest, "pair"),
|
||||
fail=_int_field(rest, "fail"),
|
||||
s_phys_wait=s_phys_wait,
|
||||
s_phys_hold=s_phys_hold,
|
||||
p_phys_wait=p_phys_wait,
|
||||
p_phys_hold=p_phys_hold,
|
||||
c_phys_wait=c_phys_wait,
|
||||
c_phys_hold=c_phys_hold,
|
||||
p_email_create=p_email_create,
|
||||
p_page_prepare=p_page_prepare,
|
||||
p_send_stage=p_send_stage,
|
||||
c_page_acquire=c_page_acquire,
|
||||
c_verify=c_verify,
|
||||
c_register=c_register,
|
||||
c_hot_hits=c_hot_hits,
|
||||
c_hot_misses=c_hot_misses,
|
||||
solver_goto=_float_field(rest, "solver_goto"),
|
||||
solver_inject=_float_field(rest, "solver_inject"),
|
||||
solver_initial=_float_field(rest, "solver_initial"),
|
||||
solver_click=_float_field(rest, "solver_click"),
|
||||
solver_wait=_float_field(rest, "solver_wait"),
|
||||
solver_reuse=_float_field(rest, "solver_reuse"),
|
||||
solver_visible=_float_field(rest, "solver_visible"),
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
state = _STATE_RE.match(line)
|
||||
if state:
|
||||
rows.append(
|
||||
MonitorRow(
|
||||
kind="state_machine",
|
||||
t=int(state.group("t")),
|
||||
q=int(state.group("q")),
|
||||
ok=int(state.group("ok")),
|
||||
rate=float(state.group("rate")),
|
||||
q_sent=int(state.group("q_sent")),
|
||||
q_ret=int(state.group("q_ret")),
|
||||
slots=int(state.group("slots")),
|
||||
max_slots=int(state.group("max_slots")),
|
||||
active=int(state.group("active")),
|
||||
)
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def _rate_delta(rows: list[MonitorRow], attr: str) -> float | None:
|
||||
candidates = [r for r in rows if r.elapsed_min is not None and getattr(r, attr) is not None]
|
||||
if len(candidates) < 2:
|
||||
return None
|
||||
first = candidates[0]
|
||||
last = candidates[-1]
|
||||
dt = (last.elapsed_min or 0) - (first.elapsed_min or 0)
|
||||
if dt <= 0:
|
||||
return None
|
||||
return (getattr(last, attr) - getattr(first, attr)) / dt
|
||||
|
||||
|
||||
def _leader(values: dict[str, float | None]) -> str | None:
|
||||
if any(value is None for value in values.values()):
|
||||
return None
|
||||
return max(values, key=lambda key: values[key] or 0)
|
||||
|
||||
|
||||
def _avg(values: list[float]) -> float | None:
|
||||
if not values:
|
||||
return None
|
||||
return round(sum(values) / len(values), 2)
|
||||
|
||||
|
||||
def _ratio(numerator: int, denominator: int) -> float | None:
|
||||
if denominator <= 0:
|
||||
return None
|
||||
return round(numerator / denominator, 3)
|
||||
|
||||
|
||||
def summarize_monitor_rows(rows: list[MonitorRow], recent_count: int = 6) -> dict[str, float | int | str | None]:
|
||||
if not rows:
|
||||
return {"rows": 0}
|
||||
|
||||
last = rows[-1]
|
||||
recent = rows[-recent_count:]
|
||||
summary: dict[str, float | int | str | None] = {
|
||||
"rows": len(rows),
|
||||
"kind": last.kind,
|
||||
"last_ok": last.ok,
|
||||
"last_cumulative_rate": last.rate,
|
||||
"last_q_minus_t": last.q - last.t,
|
||||
"last_phys": last.phys,
|
||||
"last_p_send_sem": last.p_send,
|
||||
"last_t_slot": last.t_slot,
|
||||
"last_q_slot": last.q_slot,
|
||||
"last_q_pend": last.q_pend,
|
||||
"last_p_batch": last.p_batch,
|
||||
"last_t_prog": last.t_prog,
|
||||
"last_q_inflight": last.q_inflight,
|
||||
"last_q_return_minus_t_prod": (
|
||||
last.q_ret - last.t_prod
|
||||
if last.q_ret is not None and last.t_prod is not None
|
||||
else None
|
||||
),
|
||||
"last_slots": last.slots,
|
||||
"last_t_prod": last.t_prod,
|
||||
"last_q_ret": last.q_ret,
|
||||
"last_q_adm": last.q_adm,
|
||||
"last_pair": last.pair,
|
||||
"last_fail": last.fail,
|
||||
"last_s_phys_wait": last.s_phys_wait,
|
||||
"last_s_phys_hold": last.s_phys_hold,
|
||||
"last_p_phys_wait": last.p_phys_wait,
|
||||
"last_p_phys_hold": last.p_phys_hold,
|
||||
"last_c_phys_wait": last.c_phys_wait,
|
||||
"last_c_phys_hold": last.c_phys_hold,
|
||||
"last_p_email_create": last.p_email_create,
|
||||
"last_p_page_prepare": last.p_page_prepare,
|
||||
"last_p_send": last.p_send_stage,
|
||||
"last_c_page_acquire": last.c_page_acquire,
|
||||
"last_c_verify": last.c_verify,
|
||||
"last_c_register": last.c_register,
|
||||
"last_c_hot_hits": last.c_hot_hits,
|
||||
"last_c_hot_misses": last.c_hot_misses,
|
||||
"last_physical_hold_leader": _leader({
|
||||
"s": last.s_phys_hold,
|
||||
"p": last.p_phys_hold,
|
||||
"c": last.c_phys_hold,
|
||||
}),
|
||||
"last_physical_wait_leader": _leader({
|
||||
"s": last.s_phys_wait,
|
||||
"p": last.p_phys_wait,
|
||||
"c": last.c_phys_wait,
|
||||
}),
|
||||
"last_solver_goto": last.solver_goto,
|
||||
"last_solver_inject": last.solver_inject,
|
||||
"last_solver_initial": last.solver_initial,
|
||||
"last_solver_click": last.solver_click,
|
||||
"last_solver_wait": last.solver_wait,
|
||||
"last_solver_reuse": last.solver_reuse,
|
||||
"last_solver_visible": last.solver_visible,
|
||||
"recent_ok_per_min": _rate_delta(recent, "ok"),
|
||||
"recent_t_prod_per_min": _rate_delta(recent, "t_prod"),
|
||||
"recent_q_ret_per_min": _rate_delta(recent, "q_ret"),
|
||||
"recent_pair_per_min": _rate_delta(recent, "pair"),
|
||||
}
|
||||
return summary
|
||||
|
||||
|
||||
def parse_solver_timelines(text_or_lines: str | Iterable[str]) -> list[SolverTimeline]:
|
||||
if isinstance(text_or_lines, str):
|
||||
lines = text_or_lines.splitlines()
|
||||
else:
|
||||
lines = list(text_or_lines)
|
||||
|
||||
timelines: list[SolverTimeline] = []
|
||||
for line in lines:
|
||||
if not line.startswith(_SOLVER_TIMELINE_PREFIX):
|
||||
continue
|
||||
payload = line[len(_SOLVER_TIMELINE_PREFIX):]
|
||||
try:
|
||||
events = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(events, list):
|
||||
timelines.append(SolverTimeline(events=[e for e in events if isinstance(e, dict)]))
|
||||
return timelines
|
||||
|
||||
|
||||
def _last_page_trace(events: list[dict]) -> dict | None:
|
||||
for event in reversed(events):
|
||||
page_trace = event.get("page_trace")
|
||||
if isinstance(page_trace, dict):
|
||||
return page_trace
|
||||
return None
|
||||
|
||||
|
||||
def summarize_solver_timelines(timelines: list[SolverTimeline]) -> dict[str, float | int | None]:
|
||||
if not timelines:
|
||||
return {"solver_timeline_count": 0}
|
||||
|
||||
ok_count = 0
|
||||
click_calls: list[float] = []
|
||||
click_move_ms: list[float] = []
|
||||
click_down_up_ms: list[float] = []
|
||||
render_to_token_ms: list[float] = []
|
||||
token_write_to_poll_done_ms: list[float] = []
|
||||
poll_attempts: list[float] = []
|
||||
poll_read_avg_ms: list[float] = []
|
||||
click_before_count = 0
|
||||
center_iframe_hits = 0
|
||||
turnstile_iframe_seen = 0
|
||||
widget_seen = 0
|
||||
|
||||
for timeline in timelines:
|
||||
events = timeline.events
|
||||
page_trace = _last_page_trace(events)
|
||||
if page_trace:
|
||||
render_called = page_trace.get("render_called_at")
|
||||
token_written = page_trace.get("token_written_at")
|
||||
if isinstance(render_called, (int, float)) and isinstance(token_written, (int, float)):
|
||||
render_to_token_ms.append(float(token_written) - float(render_called))
|
||||
|
||||
for event in events:
|
||||
if event.get("event") == "click_before":
|
||||
click_before_count += 1
|
||||
dom = event.get("dom") if isinstance(event.get("dom"), dict) else {}
|
||||
widget = dom.get("widget") if isinstance(dom.get("widget"), dict) else {}
|
||||
center = dom.get("element_at_center") if isinstance(dom.get("element_at_center"), dict) else {}
|
||||
if widget.get("present") and widget.get("visible"):
|
||||
widget_seen += 1
|
||||
if center.get("is_iframe"):
|
||||
center_iframe_hits += 1
|
||||
if (dom.get("turnstile_iframe_count") or 0) > 0:
|
||||
turnstile_iframe_seen += 1
|
||||
|
||||
if event.get("event") == "click_after":
|
||||
call_ms = event.get("click_call_ms")
|
||||
if isinstance(call_ms, (int, float)):
|
||||
click_calls.append(float(call_ms))
|
||||
trace = event.get("click_trace") if isinstance(event.get("click_trace"), dict) else {}
|
||||
move1 = trace.get("mouse_move1_ms")
|
||||
move2 = trace.get("mouse_move2_ms")
|
||||
down = trace.get("mouse_down_ms")
|
||||
up = trace.get("mouse_up_ms")
|
||||
move_total = sum(float(v) for v in (move1, move2) if isinstance(v, (int, float)))
|
||||
down_up_total = sum(float(v) for v in (down, up) if isinstance(v, (int, float)))
|
||||
if move_total:
|
||||
click_move_ms.append(move_total)
|
||||
if down_up_total:
|
||||
click_down_up_ms.append(down_up_total)
|
||||
|
||||
if event.get("event") == "poll_done":
|
||||
if event.get("ok"):
|
||||
ok_count += 1
|
||||
attempts = event.get("poll_attempts")
|
||||
if isinstance(attempts, (int, float)):
|
||||
poll_attempts.append(float(attempts))
|
||||
read_avg = event.get("poll_read_ms_avg")
|
||||
if isinstance(read_avg, (int, float)):
|
||||
poll_read_avg_ms.append(float(read_avg))
|
||||
page_trace = event.get("page_trace") if isinstance(event.get("page_trace"), dict) else {}
|
||||
token_written = page_trace.get("token_written_at")
|
||||
event_t = event.get("t")
|
||||
created_at = page_trace.get("created_at")
|
||||
if all(isinstance(v, (int, float)) for v in (token_written, event_t, created_at)):
|
||||
token_write_to_poll_done_ms.append((float(event_t) * 1000.0) - (float(token_written) - float(created_at)))
|
||||
|
||||
return {
|
||||
"solver_timeline_count": len(timelines),
|
||||
"ok_count": ok_count,
|
||||
"avg_click_call_ms": _avg(click_calls),
|
||||
"avg_click_mouse_move_ms": _avg(click_move_ms),
|
||||
"avg_click_down_up_ms": _avg(click_down_up_ms),
|
||||
"avg_render_to_token_ms": _avg(render_to_token_ms),
|
||||
"avg_token_write_to_poll_done_ms": _avg(token_write_to_poll_done_ms),
|
||||
"avg_poll_attempts": _avg(poll_attempts),
|
||||
"avg_poll_read_ms": _avg(poll_read_avg_ms),
|
||||
"widget_seen_ratio": _ratio(widget_seen, click_before_count),
|
||||
"center_iframe_hit_ratio": _ratio(center_iframe_hits, click_before_count),
|
||||
"turnstile_iframe_seen_ratio": _ratio(turnstile_iframe_seen, click_before_count),
|
||||
}
|
||||
|
||||
|
||||
def analyze_text(text: str) -> dict[str, float | int | str | None]:
|
||||
return summarize_monitor_rows(parse_monitor_lines(text))
|
||||
3
xai_enroller/__init__.py
Normal file
3
xai_enroller/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""Standalone xAI OAuth Device Flow enroller."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
111
xai_enroller/__main__.py
Normal file
111
xai_enroller/__main__.py
Normal file
@@ -0,0 +1,111 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
|
||||
import httpx
|
||||
|
||||
from .config import Settings
|
||||
from .coordinator import EnrollmentCoordinator
|
||||
from .executors import HTTPProbeExecutor, PlaywrightExecutor
|
||||
from .protocol import XAIProfile, XAIProtocol
|
||||
from .sinks import CPAAuthFileSink
|
||||
from .sources import FileSourceAdapter, SQLiteSourceAdapter
|
||||
|
||||
|
||||
def build_parser():
|
||||
parser = argparse.ArgumentParser(description="Bounded xAI OAuth Device Flow enroller")
|
||||
parser.add_argument("--source", choices=("file", "sqlite"))
|
||||
parser.add_argument("--target", type=int)
|
||||
parser.add_argument("--concurrency", type=int)
|
||||
parser.add_argument("--retry-attempts", type=int)
|
||||
parser.add_argument("--executor", choices=("http", "playwright"))
|
||||
parser.add_argument("--sink", choices=("cpa",))
|
||||
return parser
|
||||
|
||||
|
||||
async def main_async(args=None):
|
||||
parsed = build_parser().parse_args(args)
|
||||
env = dict()
|
||||
if parsed.source:
|
||||
env["XAI_ENROLLER_SOURCE_KIND"] = parsed.source
|
||||
if parsed.target is not None:
|
||||
env["XAI_ENROLLER_TARGET"] = str(parsed.target)
|
||||
if parsed.concurrency is not None:
|
||||
env["XAI_ENROLLER_CONCURRENCY"] = str(parsed.concurrency)
|
||||
if parsed.retry_attempts is not None:
|
||||
env["XAI_ENROLLER_RETRY_ATTEMPTS"] = str(parsed.retry_attempts)
|
||||
if parsed.executor:
|
||||
env["XAI_ENROLLER_AUTH_EXECUTOR"] = parsed.executor
|
||||
if parsed.sink:
|
||||
env["XAI_ENROLLER_SINK"] = parsed.sink
|
||||
import os
|
||||
|
||||
merged = dict(os.environ)
|
||||
merged.update(env)
|
||||
settings = Settings.from_environ(merged)
|
||||
if settings.source_kind == "remote":
|
||||
raise ValueError("remote source requires the local auth service entrypoint")
|
||||
source = (
|
||||
FileSourceAdapter(settings.source_file)
|
||||
if settings.source_kind == "file"
|
||||
else SQLiteSourceAdapter(settings.source_db, settings.source_salt)
|
||||
)
|
||||
|
||||
# CF prewarm (external FlareSolverr) + register-side proxy for httpx/Playwright
|
||||
try:
|
||||
from grok_register.clearance import (
|
||||
format_prewarm_log,
|
||||
httpx_proxy_mounts,
|
||||
prewarm_clearance,
|
||||
)
|
||||
|
||||
print(format_prewarm_log(prewarm_clearance()))
|
||||
proxy_url = httpx_proxy_mounts()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f"[clearance] skip: {exc}")
|
||||
proxy_url = None
|
||||
|
||||
client = httpx.AsyncClient(proxy=proxy_url) if proxy_url else httpx.AsyncClient()
|
||||
protocol = XAIProtocol(
|
||||
client,
|
||||
XAIProfile.default(),
|
||||
default_poll_interval=settings.poll_interval,
|
||||
)
|
||||
executor = (
|
||||
HTTPProbeExecutor(client)
|
||||
if settings.executor == "http"
|
||||
else PlaywrightExecutor(settings.concurrency)
|
||||
)
|
||||
sink = None
|
||||
sink_client = None
|
||||
if settings.sink == "cpa":
|
||||
sink_client = (
|
||||
httpx.AsyncClient(proxy=proxy_url) if proxy_url else httpx.AsyncClient()
|
||||
)
|
||||
sink = CPAAuthFileSink(settings.cpa_base_url, settings.cpa_management_secret, sink_client)
|
||||
try:
|
||||
coordinator = EnrollmentCoordinator(
|
||||
source=source,
|
||||
protocol=protocol,
|
||||
executor=executor,
|
||||
sink=sink,
|
||||
ledger_path=settings.ledger_path,
|
||||
ledger_salt=settings.source_salt,
|
||||
concurrency=settings.concurrency,
|
||||
timeout=settings.timeout_sec,
|
||||
retry_attempts=settings.retry_attempts,
|
||||
)
|
||||
results = await coordinator.run(settings.target)
|
||||
for index, result in enumerate(results, 1):
|
||||
print(f"job {index}: {result.status.value} ({result.reason_code})")
|
||||
finally:
|
||||
await client.aclose()
|
||||
if sink_client:
|
||||
await sink_client.aclose()
|
||||
|
||||
|
||||
def main():
|
||||
asyncio.run(main_async())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
1035
xai_enroller/auth_pipeline.py
Normal file
1035
xai_enroller/auth_pipeline.py
Normal file
File diff suppressed because it is too large
Load Diff
139
xai_enroller/config.py
Normal file
139
xai_enroller/config.py
Normal file
@@ -0,0 +1,139 @@
|
||||
import os
|
||||
import stat
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
||||
def _int(env, key, default):
|
||||
try:
|
||||
return int(env.get(key, str(default)))
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"{key} must be an integer") from exc
|
||||
|
||||
|
||||
def _float(env, key, default):
|
||||
try:
|
||||
return float(env.get(key, str(default)))
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"{key} must be numeric") from exc
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Settings:
|
||||
source_kind: str
|
||||
source_file: Path | None
|
||||
source_db: Path | None
|
||||
source_salt: bytes
|
||||
ledger_path: Path
|
||||
target: int = 1
|
||||
concurrency: int = 1
|
||||
retry_attempts: int = 0
|
||||
timeout_sec: float = 1800.0
|
||||
poll_interval: float = 5.0
|
||||
executor: str = "http"
|
||||
sink: str | None = None
|
||||
cpa_base_url: str | None = None
|
||||
cpa_management_secret: str | None = None
|
||||
local_auth_dir: Path | None = None
|
||||
|
||||
@classmethod
|
||||
def from_environ(cls, env=None):
|
||||
env = dict(os.environ if env is None else env)
|
||||
source_kind = env.get("XAI_ENROLLER_SOURCE_KIND", "file")
|
||||
if source_kind not in {"file", "sqlite", "remote"}:
|
||||
raise ValueError("source kind must be file, sqlite, or remote")
|
||||
source_file = Path(env["XAI_ENROLLER_SOURCE_FILE"]) if env.get("XAI_ENROLLER_SOURCE_FILE") else None
|
||||
source_db = Path(env["XAI_ENROLLER_SOURCE_DB"]) if env.get("XAI_ENROLLER_SOURCE_DB") else None
|
||||
if source_kind == "file" and source_file is None:
|
||||
raise ValueError("source file is required")
|
||||
if source_kind == "sqlite" and source_db is None:
|
||||
raise ValueError("source db is required")
|
||||
salt = env.get("XAI_ENROLLER_SOURCE_SALT")
|
||||
if not salt:
|
||||
raise ValueError("source salt is required")
|
||||
ledger_path = Path(env.get("XAI_ENROLLER_LEDGER_PATH", "xai-enroller-ledger.db"))
|
||||
target = _int(env, "XAI_ENROLLER_TARGET", 1)
|
||||
concurrency = _int(env, "XAI_ENROLLER_CONCURRENCY", 1)
|
||||
retry_attempts = _int(env, "XAI_ENROLLER_RETRY_ATTEMPTS", 0)
|
||||
timeout_sec = _float(env, "XAI_ENROLLER_TIMEOUT_SEC", 1800)
|
||||
poll_interval = _float(env, "XAI_ENROLLER_POLL_SEC", 5)
|
||||
executor = env.get("XAI_ENROLLER_AUTH_EXECUTOR", "http")
|
||||
sink = env.get("XAI_ENROLLER_SINK") or None
|
||||
if not 1 <= target <= 100:
|
||||
raise ValueError("target must be between 1 and 100")
|
||||
if not 1 <= concurrency <= 4:
|
||||
raise ValueError("concurrency must be between 1 and 4")
|
||||
if not 0 <= retry_attempts <= 3:
|
||||
raise ValueError("retry attempts must be between 0 and 3")
|
||||
if timeout_sec <= 0:
|
||||
raise ValueError("timeout must be positive")
|
||||
if poll_interval <= 0:
|
||||
raise ValueError("poll interval must be positive")
|
||||
if executor not in {"http", "playwright"}:
|
||||
raise ValueError("executor must be http or playwright")
|
||||
if sink not in {None, "cpa", "local"}:
|
||||
raise ValueError("sink must be cpa or local")
|
||||
cpa_url = env.get("XAI_ENROLLER_CPA_BASE_URL") or None
|
||||
cpa_secret = env.get("XAI_ENROLLER_CPA_MANAGEMENT_SECRET") or None
|
||||
local_auth_dir = (
|
||||
Path(env["XAI_ENROLLER_LOCAL_AUTH_DIR"]).expanduser()
|
||||
if env.get("XAI_ENROLLER_LOCAL_AUTH_DIR")
|
||||
else None
|
||||
)
|
||||
if sink == "cpa":
|
||||
if not cpa_url or not cpa_secret:
|
||||
raise ValueError("CPA base URL and management secret are required")
|
||||
parsed = urlparse(cpa_url)
|
||||
is_private = parsed.hostname in {"localhost", "127.0.0.1", "::1"}
|
||||
if parsed.scheme != "https" and not is_private:
|
||||
raise ValueError("CPA base URL must use HTTPS")
|
||||
if sink == "local" and local_auth_dir is None:
|
||||
raise ValueError("local auth directory is required")
|
||||
return cls(
|
||||
source_kind=source_kind,
|
||||
source_file=source_file,
|
||||
source_db=source_db,
|
||||
source_salt=salt.encode(),
|
||||
ledger_path=ledger_path,
|
||||
target=target,
|
||||
concurrency=concurrency,
|
||||
retry_attempts=retry_attempts,
|
||||
timeout_sec=timeout_sec,
|
||||
poll_interval=poll_interval,
|
||||
executor=executor,
|
||||
sink=sink,
|
||||
cpa_base_url=cpa_url,
|
||||
cpa_management_secret=cpa_secret,
|
||||
local_auth_dir=local_auth_dir,
|
||||
)
|
||||
|
||||
def redacted_dict(self):
|
||||
return {
|
||||
"source_kind": self.source_kind,
|
||||
"source_file": str(self.source_file) if self.source_file else None,
|
||||
"source_db": str(self.source_db) if self.source_db else None,
|
||||
"ledger_path": str(self.ledger_path),
|
||||
"target": self.target,
|
||||
"concurrency": self.concurrency,
|
||||
"retry_attempts": self.retry_attempts,
|
||||
"timeout_sec": self.timeout_sec,
|
||||
"poll_interval": self.poll_interval,
|
||||
"executor": self.executor,
|
||||
"sink": self.sink,
|
||||
"cpa_base_url": self.cpa_base_url,
|
||||
"local_auth_dir": str(self.local_auth_dir) if self.local_auth_dir else None,
|
||||
}
|
||||
|
||||
|
||||
def require_private_regular_file(path: Path, *, require_0600=False):
|
||||
try:
|
||||
mode = path.stat()
|
||||
except OSError as exc:
|
||||
raise ValueError(f"source path is unavailable: {path}") from exc
|
||||
if not stat.S_ISREG(mode.st_mode):
|
||||
raise ValueError("source path must be a regular file")
|
||||
if mode.st_mode & (stat.S_IWGRP | stat.S_IWOTH):
|
||||
raise ValueError("source path must not be group/world writable")
|
||||
if require_0600 and mode.st_mode & (stat.S_IRGRP | stat.S_IROTH):
|
||||
raise ValueError("file source requires mode 0600 or stricter")
|
||||
177
xai_enroller/coordinator.py
Normal file
177
xai_enroller/coordinator.py
Normal file
@@ -0,0 +1,177 @@
|
||||
import asyncio
|
||||
from contextlib import suppress
|
||||
|
||||
import httpx
|
||||
|
||||
from .ledger import Ledger
|
||||
from .models import (
|
||||
AuthorizationStatus,
|
||||
JobResult,
|
||||
JobStatus,
|
||||
)
|
||||
|
||||
|
||||
class EnrollmentCoordinator:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
source,
|
||||
protocol,
|
||||
executor,
|
||||
sink,
|
||||
ledger_path,
|
||||
ledger_salt,
|
||||
concurrency=1,
|
||||
timeout=1800,
|
||||
retry_attempts=0,
|
||||
event_callback=None,
|
||||
):
|
||||
self.source = source
|
||||
self.protocol = protocol
|
||||
self.executor = executor
|
||||
self.sink = sink
|
||||
self.ledger = Ledger(ledger_path, ledger_salt)
|
||||
self.concurrency = concurrency
|
||||
self.timeout = timeout
|
||||
self.retry_attempts = max(0, retry_attempts)
|
||||
self.event_callback = event_callback
|
||||
self._semaphore = asyncio.Semaphore(concurrency)
|
||||
|
||||
@staticmethod
|
||||
def _is_pre_device_transport_failure(error):
|
||||
return isinstance(error, (httpx.TransportError, OSError))
|
||||
|
||||
def _emit(self, kind, data):
|
||||
if self.event_callback is None:
|
||||
return
|
||||
try:
|
||||
self.event_callback(kind, data)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def run(self, target=1):
|
||||
return await self.run_records(self.source.records(), target=target)
|
||||
|
||||
async def run_records(self, records, target=100):
|
||||
if not 1 <= target <= 100:
|
||||
raise ValueError("target must be between 1 and 100")
|
||||
self.ledger.recover_pending()
|
||||
results = []
|
||||
tasks = []
|
||||
index = 0
|
||||
if hasattr(records, "__aiter__"):
|
||||
async for source in records:
|
||||
if index >= target:
|
||||
break
|
||||
index += 1
|
||||
tasks.append(asyncio.create_task(self._run_one(source)))
|
||||
else:
|
||||
for source in records:
|
||||
if index >= target:
|
||||
break
|
||||
index += 1
|
||||
tasks.append(asyncio.create_task(self._run_one(source)))
|
||||
if hasattr(self.executor, "start"):
|
||||
await self.executor.start()
|
||||
try:
|
||||
if tasks:
|
||||
results = await asyncio.gather(*tasks)
|
||||
return results
|
||||
finally:
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
if hasattr(self.executor, "close"):
|
||||
with suppress(Exception):
|
||||
await self.executor.close()
|
||||
|
||||
async def _run_one(self, source):
|
||||
async with self._semaphore:
|
||||
for attempt in range(1, self.retry_attempts + 2):
|
||||
job_id = self.ledger.start(source.source_id, attempt=attempt)
|
||||
try:
|
||||
async with asyncio.timeout(self.timeout):
|
||||
try:
|
||||
flow = await self.protocol.start_device_flow()
|
||||
except Exception as error:
|
||||
result = JobResult(
|
||||
source.source_id,
|
||||
JobStatus.TRANSPORT_FAILED,
|
||||
"device_flow_failed",
|
||||
attempt,
|
||||
)
|
||||
self.ledger.finish(job_id, result.status, result.reason_code)
|
||||
if (
|
||||
self._is_pre_device_transport_failure(error)
|
||||
and attempt <= self.retry_attempts
|
||||
):
|
||||
continue
|
||||
return result
|
||||
self._emit(
|
||||
"device_flow",
|
||||
{
|
||||
"source_id": source.source_id,
|
||||
"user_code": flow.user_code,
|
||||
"verification_url": flow.verification_url,
|
||||
},
|
||||
)
|
||||
return await self._attempt_after_flow(source, job_id, flow, attempt)
|
||||
except asyncio.TimeoutError:
|
||||
result = JobResult(source.source_id, JobStatus.TIMEOUT, "timeout", attempt)
|
||||
self.ledger.finish(job_id, result.status, result.reason_code)
|
||||
return result
|
||||
except asyncio.CancelledError:
|
||||
result = JobResult(source.source_id, JobStatus.CANCELLED, "cancelled", attempt)
|
||||
self.ledger.finish(job_id, result.status, result.reason_code)
|
||||
raise
|
||||
except Exception:
|
||||
result = JobResult(
|
||||
source.source_id, JobStatus.TRANSPORT_FAILED, "transport_failed", attempt
|
||||
)
|
||||
self.ledger.finish(job_id, result.status, result.reason_code)
|
||||
return result
|
||||
|
||||
async def _attempt_after_flow(self, source, job_id, flow, attempt):
|
||||
authorization = await self.executor.confirm(source, flow)
|
||||
if authorization.status is not AuthorizationStatus.AUTHORIZED:
|
||||
status = JobStatus(authorization.status.value)
|
||||
result = JobResult(source.source_id, status, authorization.reason_code, attempt)
|
||||
self.ledger.finish(job_id, result.status, result.reason_code)
|
||||
return result
|
||||
if self.sink is None:
|
||||
result = JobResult(source.source_id, JobStatus.SINK_FAILED, "sink_unconfigured", attempt)
|
||||
self.ledger.finish(job_id, result.status, result.reason_code)
|
||||
return result
|
||||
try:
|
||||
credential = await self.protocol.poll_token(
|
||||
endpoint=flow.token_endpoint,
|
||||
flow=flow,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
except RuntimeError as error:
|
||||
reason = str(error)
|
||||
mapping = {
|
||||
"oauth_denied": JobStatus.OAUTH_DENIED,
|
||||
"oauth_expired": JobStatus.OAUTH_EXPIRED,
|
||||
"oauth_rejected": JobStatus.OAUTH_REJECTED,
|
||||
}
|
||||
result = JobResult(
|
||||
source.source_id, mapping.get(reason, JobStatus.TRANSPORT_FAILED), reason, attempt
|
||||
)
|
||||
self.ledger.finish(job_id, result.status, result.reason_code)
|
||||
return result
|
||||
try:
|
||||
receipt = await self.sink.store(credential)
|
||||
except Exception:
|
||||
result = JobResult(source.source_id, JobStatus.SINK_FAILED, "sink_failed", attempt)
|
||||
self.ledger.finish(job_id, result.status, result.reason_code)
|
||||
return result
|
||||
result = JobResult(
|
||||
source.source_id,
|
||||
JobStatus.IMPORTED,
|
||||
"imported",
|
||||
attempt,
|
||||
sink_receipt_fingerprint=receipt.fingerprint,
|
||||
)
|
||||
self.ledger.finish(job_id, result.status, result.reason_code, receipt.fingerprint)
|
||||
return result
|
||||
595
xai_enroller/executors.py
Normal file
595
xai_enroller/executors.py
Normal file
@@ -0,0 +1,595 @@
|
||||
import asyncio
|
||||
import glob
|
||||
import importlib
|
||||
import os
|
||||
from contextlib import suppress
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
from grok_register.clearance import (
|
||||
apply_clearance_to_context,
|
||||
cached_user_agent,
|
||||
playwright_proxy_settings,
|
||||
)
|
||||
|
||||
from .models import AuthorizationResult, AuthorizationStatus, DeviceFlow, SourceRecord
|
||||
|
||||
|
||||
class HTTPProbeExecutor:
|
||||
def __init__(self, client: httpx.AsyncClient, max_body=64 * 1024):
|
||||
self.client = client
|
||||
self.max_body = max_body
|
||||
|
||||
async def confirm(self, source: SourceRecord, flow: DeviceFlow):
|
||||
parsed = urlparse(flow.verification_url)
|
||||
if parsed.scheme != "https" or parsed.hostname != "accounts.x.ai":
|
||||
return AuthorizationResult(AuthorizationStatus.NEEDS_INTERACTION, "invalid_url")
|
||||
response = await self.client.get(
|
||||
flow.verification_url,
|
||||
headers={"Cookie": f"sso={source.sso_token}"},
|
||||
follow_redirects=False,
|
||||
)
|
||||
body = await response.aread()
|
||||
if len(body) > self.max_body:
|
||||
return AuthorizationResult(AuthorizationStatus.NEEDS_INTERACTION, "response_too_large")
|
||||
text = body.decode("utf-8", errors="replace").lower()
|
||||
if response.status_code == 403 and any(
|
||||
marker in text for marker in ("challenge", "cf-chl", "captcha")
|
||||
):
|
||||
return AuthorizationResult(AuthorizationStatus.NEEDS_BROWSER, "challenge")
|
||||
if response.status_code in {301, 302, 303, 307, 308}:
|
||||
return AuthorizationResult(AuthorizationStatus.NEEDS_INTERACTION, "redirect")
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION, "http_probe_inconclusive"
|
||||
)
|
||||
|
||||
|
||||
class PlaywrightExecutor:
|
||||
ALLOWED_CONTROLS = frozenset({"Authorize", "Allow", "Continue", "Confirm"})
|
||||
ATTEMPT_TIMEOUT_SECONDS = 55.0
|
||||
CLOSE_TIMEOUT_SECONDS = 5.0
|
||||
# Keep confirm loops tight: most successful device verifies finish under ~15s.
|
||||
LOOP_DEADLINE_SECONDS = 28.0
|
||||
COOKIE_WAIT_MS = 250
|
||||
CONSENT_WAIT_MS = 350
|
||||
CODE_WAIT_MS = 300
|
||||
CHALLENGE_WAIT_MS = 700
|
||||
POLL_WAIT_MS = 250
|
||||
# Drop heavy assets that never participate in device verification.
|
||||
BLOCKED_RESOURCE_TYPES = frozenset({"image", "media", "font", "stylesheet"})
|
||||
BLOCKED_URL_MARKERS = (
|
||||
"google-analytics",
|
||||
"googletagmanager",
|
||||
"doubleclick",
|
||||
"facebook.net",
|
||||
"hotjar",
|
||||
"segment.io",
|
||||
)
|
||||
|
||||
def __init__(self, concurrency=1, playwright_factory=None, executable_path=None):
|
||||
self.concurrency = concurrency
|
||||
self.playwright_factory = playwright_factory
|
||||
self.executable_path = executable_path or self._find_executable_path()
|
||||
self._playwright = None
|
||||
self._browser = None
|
||||
self._semaphore = asyncio.Semaphore(concurrency)
|
||||
self._lifecycle_lock = asyncio.Lock()
|
||||
|
||||
@staticmethod
|
||||
def _find_executable_path():
|
||||
configured = os.environ.get("XAI_ENROLLER_BROWSER_EXECUTABLE")
|
||||
if configured:
|
||||
return configured
|
||||
candidates = glob.glob(os.path.expanduser("~/.cloakbrowser/chromium-*/chrome"))
|
||||
candidates.extend(
|
||||
glob.glob(
|
||||
os.path.expanduser(
|
||||
"~/.cloakbrowser/chromium-*/Chromium.app/Contents/MacOS/Chromium"
|
||||
)
|
||||
)
|
||||
)
|
||||
return sorted(candidates)[-1] if candidates else None
|
||||
|
||||
async def start(self):
|
||||
async with self._lifecycle_lock:
|
||||
if self._browser is not None:
|
||||
is_connected = getattr(self._browser, "is_connected", None)
|
||||
if not callable(is_connected) or is_connected():
|
||||
return
|
||||
await self._close_unlocked()
|
||||
factory = self.playwright_factory
|
||||
if factory is None:
|
||||
module = importlib.import_module("playwright.async_api")
|
||||
factory = module.async_playwright
|
||||
playwright = await factory().start()
|
||||
options = {"headless": True}
|
||||
if self.executable_path:
|
||||
options["executable_path"] = self.executable_path
|
||||
# REGISTER_PROXY / HTTP_PROXY / HTTPS_PROXY → this process only
|
||||
proxy = playwright_proxy_settings()
|
||||
if proxy:
|
||||
options["proxy"] = proxy
|
||||
try:
|
||||
browser = await playwright.chromium.launch(**options)
|
||||
except BaseException:
|
||||
await self._close_transport_resource(playwright.stop())
|
||||
raise
|
||||
self._playwright = playwright
|
||||
self._browser = browser
|
||||
|
||||
async def close(self):
|
||||
async with self._lifecycle_lock:
|
||||
await self._close_unlocked()
|
||||
|
||||
async def _close_unlocked(self):
|
||||
browser = self._browser
|
||||
playwright = self._playwright
|
||||
self._browser = None
|
||||
self._playwright = None
|
||||
try:
|
||||
if browser is not None:
|
||||
await self._close_transport_resource(browser.close())
|
||||
finally:
|
||||
if playwright is not None:
|
||||
await self._close_transport_resource(playwright.stop())
|
||||
|
||||
@staticmethod
|
||||
def _consume_task_result(task):
|
||||
with suppress(BaseException):
|
||||
task.result()
|
||||
|
||||
@classmethod
|
||||
async def _close_transport_resource(cls, awaitable):
|
||||
task = asyncio.create_task(awaitable)
|
||||
try:
|
||||
done, _pending = await asyncio.wait(
|
||||
{task}, timeout=cls.CLOSE_TIMEOUT_SECONDS
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
task.cancel()
|
||||
task.add_done_callback(cls._consume_task_result)
|
||||
raise
|
||||
if task in done:
|
||||
cls._consume_task_result(task)
|
||||
return True
|
||||
task.cancel()
|
||||
task.add_done_callback(cls._consume_task_result)
|
||||
return False
|
||||
|
||||
async def _recycle_browser(self, browser):
|
||||
async with self._lifecycle_lock:
|
||||
if self._browser is browser:
|
||||
await self._close_unlocked()
|
||||
|
||||
async def __aenter__(self):
|
||||
await self.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
await self.close()
|
||||
|
||||
@classmethod
|
||||
async def _close_attempt_resources(cls, page, context):
|
||||
async def close_session():
|
||||
closers = []
|
||||
if page is not None:
|
||||
closers.append(page.close())
|
||||
if context is not None:
|
||||
closers.append(context.close())
|
||||
if closers:
|
||||
await asyncio.gather(*closers, return_exceptions=True)
|
||||
|
||||
cleanup = asyncio.create_task(close_session())
|
||||
cancelled = False
|
||||
deadline = asyncio.get_running_loop().time() + cls.CLOSE_TIMEOUT_SECONDS
|
||||
while not cleanup.done():
|
||||
remaining = deadline - asyncio.get_running_loop().time()
|
||||
if remaining <= 0:
|
||||
break
|
||||
try:
|
||||
await asyncio.wait({cleanup}, timeout=remaining)
|
||||
except asyncio.CancelledError:
|
||||
cancelled = True
|
||||
if cleanup.done():
|
||||
cls._consume_task_result(cleanup)
|
||||
else:
|
||||
# asyncio.wait_for() is not a hard deadline: it cancels the child and
|
||||
# then waits for cancellation to finish, which a wedged Playwright
|
||||
# transport may never do. Abandon the cleanup task after one real
|
||||
# wall-clock deadline and consume its eventual result asynchronously.
|
||||
cleanup.cancel()
|
||||
cleanup.add_done_callback(cls._consume_task_result)
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
return cleanup.done()
|
||||
|
||||
@staticmethod
|
||||
async def _click_visible_exact(page, names):
|
||||
for name in names:
|
||||
try:
|
||||
button = page.get_by_role("button", name=name, exact=True)
|
||||
except TypeError:
|
||||
button = page.get_by_role("button", name=name)
|
||||
for index in range(await button.count()):
|
||||
candidate = button.nth(index) if hasattr(button, "nth") else button.first
|
||||
is_visible = getattr(candidate, "is_visible", None)
|
||||
if is_visible is None or await is_visible():
|
||||
await candidate.click()
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def _click_visible_turnstile(page):
|
||||
"""Click a visible Turnstile widget using the registration solver motion."""
|
||||
for selector in (
|
||||
".cf-turnstile",
|
||||
"iframe[src*='challenges.cloudflare.com']",
|
||||
"iframe[src*='turnstile']",
|
||||
):
|
||||
locator = page.locator(selector).first
|
||||
if not await locator.count() or not await locator.is_visible():
|
||||
continue
|
||||
box = await locator.bounding_box()
|
||||
if not box:
|
||||
continue
|
||||
x = box["x"] + box["width"] / 2
|
||||
y = box["y"] + box["height"] / 2
|
||||
await page.mouse.move(max(0, x - 25), max(0, y - 8))
|
||||
await page.mouse.move(x, y, steps=8)
|
||||
await page.mouse.down()
|
||||
await asyncio.sleep(0.05)
|
||||
await page.mouse.up()
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _expanded_cookies(source):
|
||||
allowed = {
|
||||
"name",
|
||||
"value",
|
||||
"url",
|
||||
"domain",
|
||||
"path",
|
||||
"expires",
|
||||
"httpOnly",
|
||||
"secure",
|
||||
"sameSite",
|
||||
}
|
||||
cookies = []
|
||||
for source_cookie in source.cookies:
|
||||
cookie = {
|
||||
key: value for key, value in source_cookie.items() if key in allowed
|
||||
}
|
||||
if cookie.get("domain"):
|
||||
cookie.pop("url", None)
|
||||
elif cookie.get("url"):
|
||||
cookie.pop("domain", None)
|
||||
cookie.pop("path", None)
|
||||
cookies.append(cookie)
|
||||
if not cookies:
|
||||
cookies = [
|
||||
{
|
||||
"name": "sso",
|
||||
"value": source.sso_token,
|
||||
"domain": "accounts.x.ai",
|
||||
"path": "/",
|
||||
"secure": True,
|
||||
"httpOnly": True,
|
||||
}
|
||||
]
|
||||
seen = {
|
||||
(cookie.get("name"), cookie.get("domain"), cookie.get("path", "/"))
|
||||
for cookie in cookies
|
||||
}
|
||||
for cookie in list(cookies) if source.cookies else ():
|
||||
name = cookie.get("name", "")
|
||||
if not name.startswith("sso"):
|
||||
continue
|
||||
key = (name, ".x.ai", cookie.get("path", "/"))
|
||||
if key in seen:
|
||||
continue
|
||||
clone = dict(cookie)
|
||||
clone.pop("url", None)
|
||||
clone["domain"] = ".x.ai"
|
||||
clone.setdefault("path", "/")
|
||||
cookies.append(clone)
|
||||
seen.add(key)
|
||||
return cookies
|
||||
|
||||
async def confirm(self, source: SourceRecord, flow: DeviceFlow):
|
||||
await self.start()
|
||||
async with self._semaphore:
|
||||
browser = self._browser
|
||||
attempt = asyncio.create_task(
|
||||
self._confirm_once(browser, source, flow)
|
||||
)
|
||||
try:
|
||||
done, _pending = await asyncio.wait(
|
||||
{attempt}, timeout=self.ATTEMPT_TIMEOUT_SECONDS
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
attempt.cancel()
|
||||
attempt.add_done_callback(self._consume_task_result)
|
||||
raise
|
||||
if attempt in done:
|
||||
return attempt.result()
|
||||
|
||||
# A Playwright transport can swallow cancellation while one of its
|
||||
# RPCs is wedged. Do not let asyncio.wait_for() extend the nominal
|
||||
# timeout indefinitely: detach the attempt and replace the complete
|
||||
# browser transport before admitting the next account.
|
||||
attempt.cancel()
|
||||
attempt.add_done_callback(self._consume_task_result)
|
||||
await self._recycle_browser(browser)
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION, "confirmation_timeout"
|
||||
)
|
||||
|
||||
async def _install_request_filters(self, page):
|
||||
"""Skip non-essential resources so device verification spends less wall time."""
|
||||
|
||||
async def route_handler(route):
|
||||
request = route.request
|
||||
resource_type = getattr(request, "resource_type", "") or ""
|
||||
url = getattr(request, "url", "") or ""
|
||||
if resource_type in self.BLOCKED_RESOURCE_TYPES or any(
|
||||
marker in url for marker in self.BLOCKED_URL_MARKERS
|
||||
):
|
||||
await route.abort()
|
||||
return
|
||||
await route.continue_()
|
||||
|
||||
with suppress(Exception):
|
||||
await page.route("**/*", route_handler)
|
||||
|
||||
async def _confirm_once(self, browser, source: SourceRecord, flow: DeviceFlow):
|
||||
context = None
|
||||
page = None
|
||||
code_submitted = False
|
||||
consent_submitted = False
|
||||
challenge_clicks = 0
|
||||
try:
|
||||
context_kwargs = {}
|
||||
ua = cached_user_agent()
|
||||
if ua:
|
||||
context_kwargs["user_agent"] = ua
|
||||
context = await browser.new_context(**context_kwargs)
|
||||
# CF cookies from external FlareSolverr prewarm (host-level)
|
||||
await apply_clearance_to_context(context)
|
||||
page = await context.new_page()
|
||||
await self._install_request_filters(page)
|
||||
await context.add_cookies(self._expanded_cookies(source))
|
||||
await page.goto(flow.verification_url, wait_until="domcontentloaded")
|
||||
deadline = (
|
||||
asyncio.get_running_loop().time() + self.LOOP_DEADLINE_SECONDS
|
||||
)
|
||||
unknown_since = None
|
||||
while asyncio.get_running_loop().time() < deadline:
|
||||
parsed = urlparse(page.url)
|
||||
if parsed.scheme != "https" or not parsed.hostname or not (
|
||||
parsed.hostname == "x.ai"
|
||||
or parsed.hostname.endswith(".x.ai")
|
||||
):
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION, "unsafe_page"
|
||||
)
|
||||
query_error = parse_qs(parsed.query).get("error", [None])[0]
|
||||
text = await page.locator("body").inner_text()
|
||||
text_lower = text.lower()
|
||||
page_title = getattr(page, "title", None)
|
||||
title_lower = (
|
||||
(await page_title()).lower() if page_title is not None else ""
|
||||
)
|
||||
|
||||
if query_error == "rate_limited":
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION, "rate_limited"
|
||||
)
|
||||
if query_error:
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION, "device_verify_failed"
|
||||
)
|
||||
if any(
|
||||
marker in text_lower
|
||||
for marker in ("rate limit", "too many requests", "请求过于频繁")
|
||||
):
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION, "rate_limited"
|
||||
)
|
||||
if (
|
||||
"attention required" in title_lower
|
||||
and "cloudflare" in title_lower
|
||||
) or any(
|
||||
marker in text_lower
|
||||
for marker in (
|
||||
"sorry, you have been blocked",
|
||||
"unable to access x.ai",
|
||||
"cloudflare ray id",
|
||||
)
|
||||
):
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION,
|
||||
"challenge_required",
|
||||
)
|
||||
if "/oauth2/device/done" in parsed.path or any(
|
||||
marker in text_lower
|
||||
for marker in ("device authorized", "设备已授权")
|
||||
):
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.AUTHORIZED, "confirmed"
|
||||
)
|
||||
|
||||
cookie_clicked = await self._click_visible_exact(
|
||||
page,
|
||||
(
|
||||
"全部拒绝",
|
||||
"拒绝全部",
|
||||
"Reject all",
|
||||
"Reject All",
|
||||
),
|
||||
)
|
||||
if cookie_clicked:
|
||||
unknown_since = None
|
||||
await page.wait_for_timeout(self.COOKIE_WAIT_MS)
|
||||
continue
|
||||
|
||||
if "/oauth2/device/consent" in parsed.path or any(
|
||||
marker in text_lower
|
||||
for marker in ("authorize grok build", "授权 grok build")
|
||||
):
|
||||
if not consent_submitted:
|
||||
allowed = await self._click_visible_exact(
|
||||
page, ("允许", "Allow", "Authorize", "Approve")
|
||||
)
|
||||
if allowed:
|
||||
consent_submitted = True
|
||||
unknown_since = None
|
||||
await page.wait_for_timeout(self.CONSENT_WAIT_MS)
|
||||
continue
|
||||
else:
|
||||
await page.wait_for_timeout(self.POLL_WAIT_MS)
|
||||
continue
|
||||
|
||||
code_input = page.locator('input[name="user_code"]')
|
||||
try:
|
||||
code_input_count = await code_input.count()
|
||||
except AttributeError:
|
||||
if await self._click_visible_exact(
|
||||
page, self.ALLOWED_CONTROLS
|
||||
):
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.AUTHORIZED,
|
||||
"confirmed",
|
||||
)
|
||||
code_input_count = 0
|
||||
if code_input_count and await code_input.first.is_visible():
|
||||
if code_submitted:
|
||||
await page.wait_for_timeout(self.POLL_WAIT_MS)
|
||||
continue
|
||||
unknown_since = None
|
||||
current = await code_input.first.input_value()
|
||||
# fill() is far cheaper than press_sequentially and the
|
||||
# device-code form does not require keystroke events.
|
||||
if flow.user_code.replace("-", "") not in current.replace(
|
||||
"-", ""
|
||||
):
|
||||
await code_input.first.fill(flow.user_code)
|
||||
submit = page.locator(
|
||||
'button[type="submit"], input[type="submit"]'
|
||||
)
|
||||
if not await submit.count():
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION,
|
||||
"submit_missing",
|
||||
)
|
||||
predicate = lambda response: (
|
||||
urlparse(response.url).hostname == "auth.x.ai"
|
||||
and urlparse(response.url).path == "/oauth2/device/verify"
|
||||
)
|
||||
async with page.expect_response(
|
||||
predicate, timeout=12000
|
||||
) as response_info:
|
||||
code_submitted = True
|
||||
await submit.first.click()
|
||||
response = await response_info.value
|
||||
location = urlparse(response.headers.get("location", ""))
|
||||
error = parse_qs(location.query).get("error", [None])[0]
|
||||
if error == "rate_limited":
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION,
|
||||
"rate_limited",
|
||||
)
|
||||
if error:
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION,
|
||||
"device_verify_failed",
|
||||
)
|
||||
await page.wait_for_timeout(self.CODE_WAIT_MS)
|
||||
continue
|
||||
|
||||
if any(
|
||||
marker in text_lower
|
||||
for marker in (
|
||||
"continue with email",
|
||||
"sign in with email",
|
||||
"使用邮箱登录",
|
||||
)
|
||||
) or await page.locator('input[type="password"]').count():
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION,
|
||||
"login_required",
|
||||
)
|
||||
|
||||
if any(
|
||||
marker in text_lower
|
||||
for marker in (
|
||||
"captcha",
|
||||
"verify you are human",
|
||||
"security challenge",
|
||||
"turnstile",
|
||||
"人机验证",
|
||||
"安全验证",
|
||||
)
|
||||
):
|
||||
if (
|
||||
challenge_clicks < 3
|
||||
and await self._click_visible_turnstile(page)
|
||||
):
|
||||
challenge_clicks += 1
|
||||
unknown_since = None
|
||||
await page.wait_for_timeout(self.CHALLENGE_WAIT_MS)
|
||||
continue
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION,
|
||||
"challenge_required",
|
||||
)
|
||||
|
||||
if any(
|
||||
marker in text_lower
|
||||
for marker in (
|
||||
"multi-factor",
|
||||
"two-factor",
|
||||
"authentication code",
|
||||
"verification code",
|
||||
"双重验证",
|
||||
"验证码",
|
||||
)
|
||||
):
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION,
|
||||
"mfa_required",
|
||||
)
|
||||
|
||||
if await self._click_visible_exact(page, self.ALLOWED_CONTROLS):
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.AUTHORIZED,
|
||||
"confirmed",
|
||||
)
|
||||
|
||||
now = asyncio.get_running_loop().time()
|
||||
if unknown_since is None:
|
||||
unknown_since = now
|
||||
elif now - unknown_since >= 2.5:
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION, "unknown_page"
|
||||
)
|
||||
|
||||
await page.wait_for_timeout(self.POLL_WAIT_MS)
|
||||
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION, "confirmation_timeout"
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as error:
|
||||
is_timeout = (
|
||||
error.__class__.__name__ == "TimeoutError"
|
||||
and error.__class__.__module__.startswith("playwright")
|
||||
)
|
||||
return AuthorizationResult(
|
||||
AuthorizationStatus.NEEDS_INTERACTION,
|
||||
"confirmation_timeout" if is_timeout else "browser_error",
|
||||
)
|
||||
finally:
|
||||
await self._close_attempt_resources(page, context)
|
||||
105
xai_enroller/inventory.py
Normal file
105
xai_enroller/inventory.py
Normal file
@@ -0,0 +1,105 @@
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class InventoryError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClaimBatch:
|
||||
batch_id: str
|
||||
directory: Path
|
||||
moved: int
|
||||
note: str = ""
|
||||
|
||||
|
||||
class CredentialInventory:
|
||||
def __init__(self, ledger, available_directory, claimed_directory):
|
||||
self.ledger = ledger
|
||||
self.available_directory = Path(available_directory)
|
||||
self.claimed_directory = Path(claimed_directory)
|
||||
|
||||
def take(self, limit, note=""):
|
||||
try:
|
||||
limit = int(limit)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError("claim count must be an integer") from exc
|
||||
if limit <= 0:
|
||||
raise ValueError("claim count must be positive")
|
||||
note = str(note or "").strip()
|
||||
batch_id = self._new_batch_id()
|
||||
receipts = self.ledger.claim_available(limit, batch_id)
|
||||
return self._complete_batch(batch_id, receipts, note)
|
||||
|
||||
def recover(self):
|
||||
rows = self.ledger.pending_claims()
|
||||
batches = {}
|
||||
for row in rows:
|
||||
batches.setdefault(row["batch_id"], []).append(
|
||||
row["sink_receipt_fingerprint"]
|
||||
)
|
||||
recovered = 0
|
||||
for batch_id, receipts in batches.items():
|
||||
recovered += self._complete_batch(batch_id, receipts, "").moved
|
||||
return recovered
|
||||
|
||||
@staticmethod
|
||||
def _new_batch_id():
|
||||
timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
|
||||
return f"{timestamp}-{secrets.token_hex(3)}"
|
||||
|
||||
@staticmethod
|
||||
def _validate_receipt(receipt):
|
||||
if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}", receipt):
|
||||
raise InventoryError("invalid credential receipt")
|
||||
|
||||
@staticmethod
|
||||
def _fsync_directory(directory):
|
||||
descriptor = os.open(directory, os.O_RDONLY)
|
||||
try:
|
||||
os.fsync(descriptor)
|
||||
finally:
|
||||
os.close(descriptor)
|
||||
|
||||
def _complete_batch(self, batch_id, receipts, note):
|
||||
destination_directory = self.claimed_directory / batch_id
|
||||
if not receipts:
|
||||
return ClaimBatch(batch_id, destination_directory, 0, note)
|
||||
self.available_directory.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
self.claimed_directory.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
destination_directory.mkdir(mode=0o700, parents=False, exist_ok=True)
|
||||
for directory in (
|
||||
self.available_directory,
|
||||
self.claimed_directory,
|
||||
destination_directory,
|
||||
):
|
||||
os.chmod(directory, 0o700)
|
||||
|
||||
moved = 0
|
||||
for receipt in receipts:
|
||||
self._validate_receipt(receipt)
|
||||
source = self.available_directory / f"{receipt}.json"
|
||||
destination = destination_directory / source.name
|
||||
source_exists = source.is_file()
|
||||
destination_exists = destination.is_file()
|
||||
if source_exists and destination_exists:
|
||||
raise InventoryError("credential exists in both inventory locations")
|
||||
if not source_exists and not destination_exists:
|
||||
raise InventoryError("credential file is missing")
|
||||
if source_exists:
|
||||
os.replace(source, destination)
|
||||
os.chmod(destination, 0o600)
|
||||
moved += 1
|
||||
|
||||
self._fsync_directory(self.available_directory)
|
||||
self._fsync_directory(destination_directory)
|
||||
self._fsync_directory(self.claimed_directory)
|
||||
marked = self.ledger.mark_claimed(batch_id, note=note)
|
||||
if marked != len(receipts):
|
||||
raise InventoryError("inventory batch did not commit completely")
|
||||
return ClaimBatch(batch_id, destination_directory, moved, note)
|
||||
275
xai_enroller/ledger.py
Normal file
275
xai_enroller/ledger.py
Normal file
@@ -0,0 +1,275 @@
|
||||
import hashlib
|
||||
import hmac
|
||||
import os
|
||||
import sqlite3
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
from .models import JobStatus
|
||||
|
||||
|
||||
class Ledger:
|
||||
def __init__(self, path: Path, salt: bytes):
|
||||
self.path = Path(path)
|
||||
self.salt = bytes(salt)
|
||||
self._init()
|
||||
|
||||
def _connect(self):
|
||||
connection = sqlite3.connect(self.path)
|
||||
connection.row_factory = sqlite3.Row
|
||||
return connection
|
||||
|
||||
def _init(self):
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS jobs (
|
||||
job_id INTEGER PRIMARY KEY,
|
||||
source_fingerprint TEXT NOT NULL,
|
||||
attempt_number INTEGER NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
started_at TEXT NOT NULL,
|
||||
finished_at TEXT,
|
||||
reason_code TEXT,
|
||||
sink_receipt_fingerprint TEXT
|
||||
)
|
||||
"""
|
||||
)
|
||||
columns = {
|
||||
row["name"] for row in connection.execute("PRAGMA table_info(jobs)")
|
||||
}
|
||||
if "authorization_started" not in columns:
|
||||
connection.execute(
|
||||
"ALTER TABLE jobs ADD COLUMN authorization_started INTEGER NOT NULL DEFAULT 0"
|
||||
)
|
||||
connection.execute(
|
||||
"CREATE INDEX IF NOT EXISTS jobs_source_status_idx "
|
||||
"ON jobs(source_fingerprint, status)"
|
||||
)
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS credential_inventory (
|
||||
sink_receipt_fingerprint TEXT PRIMARY KEY,
|
||||
state TEXT NOT NULL CHECK(state IN ('available', 'claiming', 'claimed')),
|
||||
claimed_at TEXT NOT NULL DEFAULT '',
|
||||
batch_id TEXT NOT NULL DEFAULT '',
|
||||
note TEXT NOT NULL DEFAULT ''
|
||||
)
|
||||
"""
|
||||
)
|
||||
connection.execute(
|
||||
"CREATE INDEX IF NOT EXISTS credential_inventory_state_idx "
|
||||
"ON credential_inventory(state)"
|
||||
)
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT OR IGNORE INTO credential_inventory(
|
||||
sink_receipt_fingerprint, state
|
||||
)
|
||||
SELECT sink_receipt_fingerprint, 'available'
|
||||
FROM jobs
|
||||
WHERE status=?
|
||||
AND COALESCE(TRIM(sink_receipt_fingerprint), '') <> ''
|
||||
""",
|
||||
(JobStatus.IMPORTED.value,),
|
||||
)
|
||||
os.chmod(self.path, 0o600)
|
||||
|
||||
def fingerprint(self, source_id):
|
||||
return hmac.new(self.salt, source_id.encode(), hashlib.sha256).hexdigest()
|
||||
|
||||
def _fingerprint(self, source_id):
|
||||
return self.fingerprint(source_id)
|
||||
|
||||
def start(self, source_id, attempt=1):
|
||||
return self.start_fingerprint(self.fingerprint(source_id), attempt=attempt)
|
||||
|
||||
def start_fingerprint(self, source_fingerprint, attempt=None):
|
||||
if attempt is None:
|
||||
attempt = self.next_attempt(source_fingerprint)
|
||||
now = datetime.now(timezone.utc).isoformat()
|
||||
with self._connect() as connection:
|
||||
cursor = connection.execute(
|
||||
"INSERT INTO jobs(source_fingerprint, attempt_number, status, started_at) "
|
||||
"VALUES (?, ?, ?, ?)",
|
||||
(source_fingerprint, attempt, "pending", now),
|
||||
)
|
||||
return cursor.lastrowid
|
||||
|
||||
def next_attempt(self, source_fingerprint):
|
||||
with self._connect() as connection:
|
||||
row = connection.execute(
|
||||
"SELECT COALESCE(MAX(attempt_number), 0) + 1 AS next_attempt "
|
||||
"FROM jobs WHERE source_fingerprint=?",
|
||||
(source_fingerprint,),
|
||||
).fetchone()
|
||||
return int(row["next_attempt"])
|
||||
|
||||
def mark_authorization_started(self, job_id):
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"UPDATE jobs SET authorization_started=1 "
|
||||
"WHERE job_id=? AND status='pending'",
|
||||
(job_id,),
|
||||
)
|
||||
|
||||
def finish(self, job_id, status, reason_code, sink_receipt_fingerprint=None):
|
||||
status_value = status.value if isinstance(status, JobStatus) else str(status)
|
||||
now = datetime.now(timezone.utc).isoformat()
|
||||
with self._connect() as connection:
|
||||
cursor = connection.execute(
|
||||
"UPDATE jobs SET status=?, finished_at=?, reason_code=?, "
|
||||
"sink_receipt_fingerprint=? "
|
||||
"WHERE job_id=? AND status='pending'",
|
||||
(status_value, now, reason_code, sink_receipt_fingerprint, job_id),
|
||||
)
|
||||
if (
|
||||
cursor.rowcount == 1
|
||||
and status_value == JobStatus.IMPORTED.value
|
||||
and sink_receipt_fingerprint
|
||||
):
|
||||
connection.execute(
|
||||
"INSERT OR IGNORE INTO credential_inventory("
|
||||
"sink_receipt_fingerprint, state) VALUES (?, 'available')",
|
||||
(sink_receipt_fingerprint,),
|
||||
)
|
||||
|
||||
def claim_available(self, limit, batch_id):
|
||||
limit = int(limit)
|
||||
if limit <= 0:
|
||||
return []
|
||||
batch_id = str(batch_id).strip()
|
||||
if not batch_id:
|
||||
raise ValueError("batch_id must not be empty")
|
||||
with self._connect() as connection:
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
rows = connection.execute(
|
||||
"""
|
||||
WITH latest_import AS (
|
||||
SELECT sink_receipt_fingerprint, MAX(finished_at) AS finished_at
|
||||
FROM jobs
|
||||
WHERE status=?
|
||||
AND COALESCE(TRIM(sink_receipt_fingerprint), '') <> ''
|
||||
GROUP BY sink_receipt_fingerprint
|
||||
)
|
||||
SELECT inventory.sink_receipt_fingerprint
|
||||
FROM credential_inventory AS inventory
|
||||
JOIN latest_import USING (sink_receipt_fingerprint)
|
||||
WHERE inventory.state='available'
|
||||
ORDER BY latest_import.finished_at DESC,
|
||||
inventory.sink_receipt_fingerprint ASC
|
||||
LIMIT ?
|
||||
""",
|
||||
(JobStatus.IMPORTED.value, limit),
|
||||
).fetchall()
|
||||
fingerprints = [row["sink_receipt_fingerprint"] for row in rows]
|
||||
connection.executemany(
|
||||
"UPDATE credential_inventory SET state='claiming', batch_id=?, "
|
||||
"claimed_at='', note='' "
|
||||
"WHERE sink_receipt_fingerprint=? AND state='available'",
|
||||
((batch_id, fingerprint) for fingerprint in fingerprints),
|
||||
)
|
||||
return fingerprints
|
||||
|
||||
def mark_claimed(self, batch_id, note=""):
|
||||
batch_id = str(batch_id).strip()
|
||||
if not batch_id:
|
||||
raise ValueError("batch_id must not be empty")
|
||||
with self._connect() as connection:
|
||||
cursor = connection.execute(
|
||||
"UPDATE credential_inventory SET state='claimed', claimed_at=?, note=? "
|
||||
"WHERE state='claiming' AND batch_id=?",
|
||||
(datetime.now(timezone.utc).isoformat(), str(note), batch_id),
|
||||
)
|
||||
return cursor.rowcount
|
||||
|
||||
def inventory_counts(self):
|
||||
counts = {"available": 0, "claiming": 0, "claimed": 0}
|
||||
with self._connect() as connection:
|
||||
rows = connection.execute(
|
||||
"SELECT state, COUNT(*) AS count FROM credential_inventory GROUP BY state"
|
||||
)
|
||||
for row in rows:
|
||||
counts[row["state"]] = int(row["count"])
|
||||
return counts
|
||||
|
||||
def pending_claims(self, batch_id=None):
|
||||
query = (
|
||||
"SELECT sink_receipt_fingerprint, batch_id, note "
|
||||
"FROM credential_inventory WHERE state='claiming'"
|
||||
)
|
||||
params = ()
|
||||
if batch_id is not None:
|
||||
query += " AND batch_id=?"
|
||||
params = (str(batch_id),)
|
||||
query += " ORDER BY batch_id, sink_receipt_fingerprint"
|
||||
with self._connect() as connection:
|
||||
rows = connection.execute(query, params)
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
def recover_claiming(self, batch_id=None, *, note=""):
|
||||
query = (
|
||||
"UPDATE credential_inventory SET state='available', claimed_at='', "
|
||||
"batch_id='', note=? WHERE state='claiming'"
|
||||
)
|
||||
params = [str(note)]
|
||||
if batch_id is not None:
|
||||
query += " AND batch_id=?"
|
||||
params.append(str(batch_id))
|
||||
with self._connect() as connection:
|
||||
cursor = connection.execute(query, params)
|
||||
return cursor.rowcount
|
||||
|
||||
def recover_pending(self, *, reason="recovered_pending"):
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"UPDATE jobs SET status=?, finished_at=?, reason_code=? WHERE status='pending'",
|
||||
(JobStatus.CANCELLED.value, datetime.now(timezone.utc).isoformat(), reason),
|
||||
)
|
||||
|
||||
def has_imported(self, source_id):
|
||||
with self._connect() as connection:
|
||||
row = connection.execute(
|
||||
"SELECT 1 FROM jobs WHERE source_fingerprint=? AND status=? LIMIT 1",
|
||||
(self.fingerprint(source_id), JobStatus.IMPORTED.value),
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
||||
def imported_fingerprints(self):
|
||||
with self._connect() as connection:
|
||||
rows = connection.execute(
|
||||
"SELECT DISTINCT source_fingerprint FROM jobs WHERE status=?",
|
||||
(JobStatus.IMPORTED.value,),
|
||||
)
|
||||
return {row["source_fingerprint"] for row in rows}
|
||||
|
||||
def aggregate_counts(self):
|
||||
with self._connect() as connection:
|
||||
row = connection.execute(
|
||||
"""
|
||||
SELECT
|
||||
COUNT(DISTINCT CASE WHEN authorization_started=1 OR status=?
|
||||
THEN source_fingerprint END) AS attempted_unique,
|
||||
COUNT(DISTINCT CASE WHEN status=?
|
||||
THEN source_fingerprint END) AS imported_unique,
|
||||
COUNT(CASE WHEN finished_at IS NOT NULL THEN 1 END) AS finalized_attempts,
|
||||
COUNT(CASE WHEN status=? THEN 1 END) AS imported_attempts,
|
||||
COUNT(CASE WHEN reason_code='rate_limited' THEN 1 END) AS rate_limited
|
||||
FROM jobs
|
||||
""",
|
||||
(
|
||||
JobStatus.IMPORTED.value,
|
||||
JobStatus.IMPORTED.value,
|
||||
JobStatus.IMPORTED.value,
|
||||
),
|
||||
).fetchone()
|
||||
return {key: int(row[key] or 0) for key in row.keys()}
|
||||
|
||||
def get(self, job_id):
|
||||
with self._connect() as connection:
|
||||
row = connection.execute(
|
||||
"SELECT source_fingerprint, attempt_number, status, started_at, finished_at, "
|
||||
"reason_code, sink_receipt_fingerprint FROM jobs WHERE job_id=?",
|
||||
(job_id,),
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
128
xai_enroller/models.py
Normal file
128
xai_enroller/models.py
Normal file
@@ -0,0 +1,128 @@
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SourceRecord:
|
||||
source_id: str
|
||||
sso_token: str
|
||||
cookies: tuple[dict, ...] = ()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "SourceRecord(<redacted>)"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DeviceFlow:
|
||||
device_code: str
|
||||
user_code: str
|
||||
verification_url: str
|
||||
expires_in: int
|
||||
interval: float
|
||||
token_endpoint: str = "https://auth.x.ai/oauth2/token"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "DeviceFlow(<redacted>)"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OAuthCredential:
|
||||
access_token: str
|
||||
refresh_token: str
|
||||
id_token: Optional[str]
|
||||
token_type: Optional[str]
|
||||
expires_in: int
|
||||
expires_at: str
|
||||
last_refresh: str
|
||||
subject: Optional[str]
|
||||
token_endpoint: str
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "OAuthCredential(<redacted>)"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SinkReceipt:
|
||||
fingerprint: str
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "SinkReceipt(<redacted>)"
|
||||
|
||||
|
||||
class AuthorizationStatus(str, Enum):
|
||||
AUTHORIZED = "authorized"
|
||||
NEEDS_BROWSER = "needs_browser"
|
||||
NEEDS_INTERACTION = "needs_interaction"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AuthorizationResult:
|
||||
status: AuthorizationStatus
|
||||
reason_code: str
|
||||
|
||||
|
||||
class JobStatus(str, Enum):
|
||||
IMPORTED = "imported"
|
||||
NEEDS_BROWSER = "needs_browser"
|
||||
NEEDS_INTERACTION = "needs_interaction"
|
||||
SOURCE_INVALID = "source_invalid"
|
||||
OAUTH_DENIED = "oauth_denied"
|
||||
OAUTH_EXPIRED = "oauth_expired"
|
||||
OAUTH_REJECTED = "oauth_rejected"
|
||||
SINK_FAILED = "sink_failed"
|
||||
TRANSPORT_FAILED = "transport_failed"
|
||||
TIMEOUT = "timeout"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class JobResult:
|
||||
source_id: str
|
||||
status: JobStatus
|
||||
reason_code: str
|
||||
attempt_number: int = 1
|
||||
sink_receipt_fingerprint: Optional[str] = None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"JobResult(source_id=<redacted>, status={self.status.value!r}, "
|
||||
f"reason_code={self.reason_code!r}, attempt_number={self.attempt_number})"
|
||||
)
|
||||
|
||||
|
||||
class PipelineState(str, Enum):
|
||||
QUEUED = "queued"
|
||||
PREPARED = "prepared"
|
||||
ACTIVE = "active"
|
||||
RETRY_WAITING = "retry_waiting"
|
||||
IMPORTED = "imported"
|
||||
TERMINAL = "terminal"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PreparedJob:
|
||||
source: SourceRecord
|
||||
source_fingerprint: str
|
||||
flow: DeviceFlow
|
||||
flow_created_monotonic: float
|
||||
job_id: int
|
||||
attempt_number: int
|
||||
task_number: Optional[int] = None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
"PreparedJob(source=<redacted>, flow=<redacted>, "
|
||||
f"attempt_number={self.attempt_number})"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CompletionJob:
|
||||
prepared: PreparedJob
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
"CompletionJob(source=<redacted>, flow=<redacted>, "
|
||||
f"attempt_number={self.prepared.attempt_number})"
|
||||
)
|
||||
181
xai_enroller/protocol.py
Normal file
181
xai_enroller/protocol.py
Normal file
@@ -0,0 +1,181 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import binascii
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Awaitable, Callable
|
||||
from urllib.parse import urlencode
|
||||
|
||||
import httpx
|
||||
|
||||
from .models import DeviceFlow, OAuthCredential
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class XAIProfile:
|
||||
client_id: str
|
||||
scope: str
|
||||
version: str = "xai-device-v1"
|
||||
|
||||
@classmethod
|
||||
def default(cls):
|
||||
return cls(
|
||||
client_id="b1a00492-073a-47ea-816f-4c329264a828",
|
||||
scope="openid profile email offline_access grok-cli:access api:access",
|
||||
version="grok-cli-device-v1",
|
||||
)
|
||||
|
||||
|
||||
class XAIProtocol:
|
||||
DISCOVERY_URL = "https://auth.x.ai/.well-known/openid-configuration"
|
||||
MAX_BODY = 64 * 1024
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: httpx.AsyncClient,
|
||||
profile: XAIProfile,
|
||||
sleep=None,
|
||||
default_poll_interval=5.0,
|
||||
):
|
||||
self.client = client
|
||||
self.profile = profile
|
||||
self.sleep = sleep or asyncio.sleep
|
||||
self.default_poll_interval = float(default_poll_interval)
|
||||
self._discovered = None
|
||||
|
||||
@staticmethod
|
||||
def _allowed_url(value):
|
||||
from urllib.parse import urlparse
|
||||
|
||||
parsed = urlparse(value)
|
||||
hostname = parsed.hostname
|
||||
return bool(hostname) and parsed.scheme == "https" and (
|
||||
hostname == "x.ai" or hostname.endswith(".x.ai")
|
||||
)
|
||||
|
||||
async def _json(self, response):
|
||||
body = await response.aread()
|
||||
if len(body) > self.MAX_BODY:
|
||||
raise RuntimeError("response body exceeds limit")
|
||||
try:
|
||||
return json.loads(body)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise RuntimeError("invalid JSON response") from exc
|
||||
|
||||
async def discover(self):
|
||||
response = await self.client.get(self.DISCOVERY_URL, follow_redirects=False)
|
||||
if response.status_code // 100 != 2:
|
||||
raise RuntimeError("discovery rejected")
|
||||
document = await self._json(response)
|
||||
device_endpoint = document.get("device_authorization_endpoint")
|
||||
token_endpoint = document.get("token_endpoint")
|
||||
if not device_endpoint or not token_endpoint or not all(
|
||||
self._allowed_url(url) for url in (device_endpoint, token_endpoint)
|
||||
):
|
||||
raise ValueError("discovery endpoint failed x.ai HTTPS allowlist")
|
||||
self._discovered = (device_endpoint, token_endpoint)
|
||||
return self._discovered
|
||||
|
||||
async def start_device_flow(self):
|
||||
device_endpoint, token_endpoint = await self.discover()
|
||||
response = await self.client.post(
|
||||
device_endpoint,
|
||||
data={"client_id": self.profile.client_id, "scope": self.profile.scope},
|
||||
follow_redirects=False,
|
||||
)
|
||||
if response.status_code // 100 != 2:
|
||||
raise RuntimeError("device authorization rejected")
|
||||
document = await self._json(response)
|
||||
try:
|
||||
device_code = document["device_code"]
|
||||
user_code = document["user_code"]
|
||||
base_url = document.get("verification_uri") or document["verification_url"]
|
||||
expires_in = int(document["expires_in"])
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise RuntimeError("invalid device authorization response") from exc
|
||||
if not self._allowed_url(base_url):
|
||||
raise ValueError("verification URL failed x.ai HTTPS allowlist")
|
||||
interval = float(document.get("interval", self.default_poll_interval))
|
||||
verification_url = document.get("verification_uri_complete")
|
||||
if verification_url is None:
|
||||
verification_url = (
|
||||
f"{base_url}{'&' if '?' in base_url else '?'}"
|
||||
f"{urlencode({'user_code': user_code})}"
|
||||
)
|
||||
if not self._allowed_url(verification_url):
|
||||
raise ValueError("verification URL failed x.ai HTTPS allowlist")
|
||||
return DeviceFlow(device_code, user_code, verification_url, expires_in, interval, token_endpoint)
|
||||
|
||||
async def poll_token(self, *, endpoint, flow, timeout):
|
||||
deadline = time.monotonic() + timeout
|
||||
interval = max(0.0, float(flow.interval))
|
||||
while time.monotonic() < deadline:
|
||||
response = await self.client.post(
|
||||
endpoint,
|
||||
data={
|
||||
"client_id": self.profile.client_id,
|
||||
"device_code": flow.device_code,
|
||||
"grant_type": "urn:ietf:params:oauth:grant-type:device_code",
|
||||
},
|
||||
follow_redirects=False,
|
||||
)
|
||||
document = await self._json(response)
|
||||
if response.status_code // 100 == 2:
|
||||
return self._credential(document, endpoint)
|
||||
error = document.get("error")
|
||||
if error in {"authorization_pending", "slow_down"}:
|
||||
interval += 1 if error == "slow_down" else 0
|
||||
await self.sleep(interval)
|
||||
continue
|
||||
if error == "access_denied":
|
||||
raise RuntimeError("oauth_denied")
|
||||
if error == "expired_token":
|
||||
raise RuntimeError("oauth_expired")
|
||||
raise RuntimeError("oauth_rejected")
|
||||
raise RuntimeError("oauth_expired")
|
||||
|
||||
@staticmethod
|
||||
def _jwt_subject(token):
|
||||
if not isinstance(token, str):
|
||||
return None
|
||||
parts = token.split(".")
|
||||
if len(parts) != 3 or not parts[1]:
|
||||
return None
|
||||
payload = parts[1] + "=" * (-len(parts[1]) % 4)
|
||||
try:
|
||||
claims = json.loads(base64.urlsafe_b64decode(payload).decode("utf-8"))
|
||||
except (binascii.Error, UnicodeDecodeError, ValueError):
|
||||
return None
|
||||
if not isinstance(claims, dict):
|
||||
return None
|
||||
for claim in ("sub", "principal_id"):
|
||||
value = claims.get(claim)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _credential(document, endpoint):
|
||||
if not document.get("access_token") or not document.get("refresh_token"):
|
||||
raise RuntimeError("oauth_rejected")
|
||||
now = datetime.now(timezone.utc)
|
||||
expires_in = int(document.get("expires_in", 3600))
|
||||
expires_at = datetime.fromtimestamp(now.timestamp() + expires_in, timezone.utc).isoformat().replace("+00:00", "Z")
|
||||
subject = (
|
||||
XAIProtocol._jwt_subject(document.get("id_token"))
|
||||
or XAIProtocol._jwt_subject(document["access_token"])
|
||||
or document.get("sub")
|
||||
)
|
||||
return OAuthCredential(
|
||||
access_token=document["access_token"],
|
||||
refresh_token=document["refresh_token"],
|
||||
id_token=document.get("id_token"),
|
||||
token_type=document.get("token_type"),
|
||||
expires_in=expires_in,
|
||||
expires_at=expires_at,
|
||||
last_refresh=now.isoformat().replace("+00:00", "Z"),
|
||||
subject=subject,
|
||||
token_endpoint=endpoint,
|
||||
)
|
||||
630
xai_enroller/remote_stream.py
Normal file
630
xai_enroller/remote_stream.py
Normal file
@@ -0,0 +1,630 @@
|
||||
"""Atomic local snapshots of password-free registration sessions."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import sys
|
||||
import tempfile
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
|
||||
from .models import SourceRecord
|
||||
|
||||
MAX_SESSION_RECORD_BYTES = 256 * 1024
|
||||
|
||||
|
||||
class RemoteStreamError(RuntimeError):
|
||||
"""A classified stream failure whose text never includes remote output."""
|
||||
|
||||
|
||||
class SSHSnapshotSynchronizer:
|
||||
"""Atomically refresh one validated, password-free local JSONL snapshot."""
|
||||
|
||||
MAX_STDERR_BYTES = 16 * 1024
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host,
|
||||
destination,
|
||||
*,
|
||||
remote_root="/opt/grok-free-register",
|
||||
identity_file=None,
|
||||
process_factory=asyncio.create_subprocess_exec,
|
||||
fingerprint=None,
|
||||
):
|
||||
self.host = host
|
||||
self.destination = Path(destination)
|
||||
self.remote_root = remote_root
|
||||
self.identity_file = identity_file
|
||||
self.process_factory = process_factory
|
||||
self.fingerprint = fingerprint or (lambda source_id: source_id)
|
||||
self.snapshot_fingerprints = None
|
||||
self._process = None
|
||||
|
||||
def _input_generation(self):
|
||||
return None
|
||||
|
||||
def _input_generation_unchanged(self, generation):
|
||||
return True
|
||||
|
||||
def _empty_snapshot_allowed(self):
|
||||
return False
|
||||
|
||||
def _command(self):
|
||||
return (
|
||||
f"cd {shlex.quote(self.remote_root)} && "
|
||||
"python3 scripts/export_registered_sessions.py "
|
||||
"keys/auth-sessions.jsonl keys/accounts.txt"
|
||||
)
|
||||
|
||||
def _args(self):
|
||||
args = [
|
||||
"ssh",
|
||||
"-C",
|
||||
"-T",
|
||||
"-o",
|
||||
"BatchMode=yes",
|
||||
"-o",
|
||||
"ConnectTimeout=15",
|
||||
"-o",
|
||||
"ServerAliveInterval=15",
|
||||
"-o",
|
||||
"ServerAliveCountMax=3",
|
||||
]
|
||||
if self.identity_file:
|
||||
args.extend(["-i", self.identity_file])
|
||||
args.extend(["--", self.host, self._command()])
|
||||
return args
|
||||
|
||||
async def _read_stderr(self, stream):
|
||||
retained = 0
|
||||
while True:
|
||||
chunk = await stream.read(4096)
|
||||
if not chunk:
|
||||
return
|
||||
retained = min(self.MAX_STDERR_BYTES, retained + len(chunk))
|
||||
|
||||
async def _terminate(self, process):
|
||||
if process is None or process.returncode is not None:
|
||||
return
|
||||
with suppress(ProcessLookupError):
|
||||
process.terminate()
|
||||
try:
|
||||
await asyncio.wait_for(process.wait(), timeout=3)
|
||||
except TimeoutError:
|
||||
with suppress(ProcessLookupError):
|
||||
process.kill()
|
||||
await process.wait()
|
||||
|
||||
async def close(self):
|
||||
await self._terminate(self._process)
|
||||
|
||||
async def sync_once(self):
|
||||
self.destination.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
os.chmod(self.destination.parent, 0o700)
|
||||
fd, temporary_name = tempfile.mkstemp(
|
||||
prefix=f".{self.destination.name}.",
|
||||
suffix=".tmp",
|
||||
dir=self.destination.parent,
|
||||
)
|
||||
process = None
|
||||
stderr_task = None
|
||||
try:
|
||||
os.fchmod(fd, 0o600)
|
||||
input_generation = self._input_generation()
|
||||
process = await self.process_factory(
|
||||
*self._args(),
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
limit=MAX_SESSION_RECORD_BYTES + 1,
|
||||
)
|
||||
self._process = process
|
||||
stderr_task = asyncio.create_task(self._read_stderr(process.stderr))
|
||||
with os.fdopen(fd, "wb") as stream:
|
||||
fd = -1
|
||||
record_count = 0
|
||||
snapshot_fingerprints = set()
|
||||
while True:
|
||||
try:
|
||||
raw = await process.stdout.readline()
|
||||
except (ValueError, asyncio.LimitOverrunError) as exc:
|
||||
raise ValueError("invalid remote session snapshot") from exc
|
||||
if not raw:
|
||||
break
|
||||
if (
|
||||
not raw.endswith(b"\n")
|
||||
or len(raw) - 1 > MAX_SESSION_RECORD_BYTES
|
||||
):
|
||||
raise ValueError("invalid remote session snapshot")
|
||||
record = parse_session_document(raw[:-1])
|
||||
snapshot_fingerprints.add(self.fingerprint(record.source_id))
|
||||
stream.write(raw)
|
||||
record_count += 1
|
||||
returncode = await process.wait()
|
||||
if stderr_task is not None:
|
||||
await stderr_task
|
||||
stderr_task = None
|
||||
if returncode != 0:
|
||||
raise RemoteStreamError("remote snapshot export failed")
|
||||
if record_count == 0 and not self._empty_snapshot_allowed():
|
||||
raise RemoteStreamError("remote snapshot export was empty")
|
||||
if not self._input_generation_unchanged(input_generation):
|
||||
raise RemoteStreamError("snapshot inputs changed during export")
|
||||
stream.flush()
|
||||
os.fsync(stream.fileno())
|
||||
os.replace(temporary_name, self.destination)
|
||||
os.chmod(self.destination, 0o600)
|
||||
directory_fd = os.open(self.destination.parent, os.O_RDONLY)
|
||||
try:
|
||||
os.fsync(directory_fd)
|
||||
finally:
|
||||
os.close(directory_fd)
|
||||
self.snapshot_fingerprints = frozenset(snapshot_fingerprints)
|
||||
return True
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
return False
|
||||
finally:
|
||||
if fd >= 0:
|
||||
os.close(fd)
|
||||
if stderr_task is not None:
|
||||
if not stderr_task.done():
|
||||
stderr_task.cancel()
|
||||
await asyncio.gather(stderr_task, return_exceptions=True)
|
||||
with suppress(Exception):
|
||||
await self._terminate(process)
|
||||
if self._process is process:
|
||||
self._process = None
|
||||
with suppress(FileNotFoundError):
|
||||
os.unlink(temporary_name)
|
||||
|
||||
|
||||
class LocalSnapshotSynchronizer(SSHSnapshotSynchronizer):
|
||||
"""Export registration data from one local project into an atomic snapshot."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
register_root,
|
||||
destination,
|
||||
*,
|
||||
process_factory=asyncio.create_subprocess_exec,
|
||||
fingerprint=None,
|
||||
python_executable=sys.executable,
|
||||
exporter_path=None,
|
||||
):
|
||||
super().__init__(
|
||||
"",
|
||||
destination,
|
||||
process_factory=process_factory,
|
||||
fingerprint=fingerprint,
|
||||
)
|
||||
self.register_root = Path(register_root)
|
||||
self.python_executable = str(python_executable)
|
||||
self.exporter_path = Path(exporter_path or (
|
||||
Path(__file__).resolve().parents[1]
|
||||
/ "scripts"
|
||||
/ "export_registered_sessions.py"
|
||||
))
|
||||
self.sessions_path = self.register_root / "keys" / "auth-sessions.jsonl"
|
||||
self.accounts_path = self.register_root / "keys" / "accounts.txt"
|
||||
|
||||
@staticmethod
|
||||
def _path_generation(path):
|
||||
try:
|
||||
stat_result = path.stat()
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
return (
|
||||
stat_result.st_dev,
|
||||
stat_result.st_ino,
|
||||
stat_result.st_size,
|
||||
stat_result.st_mtime_ns,
|
||||
)
|
||||
|
||||
def _input_generation(self):
|
||||
return (
|
||||
self._path_generation(self.sessions_path),
|
||||
self._path_generation(self.accounts_path),
|
||||
)
|
||||
|
||||
def _input_generation_unchanged(self, generation):
|
||||
return self._input_generation() == generation
|
||||
|
||||
def _empty_snapshot_allowed(self):
|
||||
return True
|
||||
|
||||
def _args(self):
|
||||
return [
|
||||
self.python_executable,
|
||||
str(self.exporter_path),
|
||||
str(self.sessions_path),
|
||||
str(self.accounts_path),
|
||||
]
|
||||
|
||||
|
||||
class DiskSnapshotSource:
|
||||
"""Consume immutable local snapshot generations under queue backpressure."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
path,
|
||||
*,
|
||||
synchronizer=None,
|
||||
sync_seconds=30.0,
|
||||
poll_seconds=0.25,
|
||||
sleep=asyncio.sleep,
|
||||
event_callback=None,
|
||||
fingerprint=None,
|
||||
):
|
||||
self.path = Path(path)
|
||||
self.synchronizer = synchronizer
|
||||
self.sync_seconds = float(sync_seconds)
|
||||
self.poll_seconds = float(poll_seconds)
|
||||
self.sleep = sleep
|
||||
self.event_callback = event_callback
|
||||
self.fingerprint = fingerprint or (lambda source_id: source_id)
|
||||
self._snapshot_fingerprints = None
|
||||
self._last_reported_snapshot_fingerprints = None
|
||||
self._closed = False
|
||||
self._sync_task = None
|
||||
self._last_sync_ok = None
|
||||
|
||||
def _emit(self, kind, data):
|
||||
if self.event_callback is not None:
|
||||
with suppress(Exception):
|
||||
self.event_callback(kind, data)
|
||||
|
||||
@property
|
||||
def snapshot_fingerprints(self):
|
||||
if self.synchronizer is not None:
|
||||
return self.synchronizer.snapshot_fingerprints
|
||||
return self._snapshot_fingerprints
|
||||
|
||||
async def _sync_loop(self):
|
||||
while not self._closed:
|
||||
refreshed = await self.synchronizer.sync_once()
|
||||
current = self.synchronizer.snapshot_fingerprints
|
||||
if refreshed != self._last_sync_ok:
|
||||
self._emit(
|
||||
"source_connected" if refreshed else "source_disconnected",
|
||||
{} if refreshed else {"reason": "snapshot_sync_failed"},
|
||||
)
|
||||
self._last_sync_ok = refreshed
|
||||
if refreshed and current is not None:
|
||||
previous = self._last_reported_snapshot_fingerprints
|
||||
if previous is not None:
|
||||
added = current - previous
|
||||
if added:
|
||||
self._emit(
|
||||
"source_updated",
|
||||
{"new": len(added), "total": len(current)},
|
||||
)
|
||||
self._last_reported_snapshot_fingerprints = current
|
||||
try:
|
||||
await asyncio.wait_for(self._wait_closed(), timeout=self.sync_seconds)
|
||||
except TimeoutError:
|
||||
pass
|
||||
|
||||
async def _wait_closed(self):
|
||||
while not self._closed:
|
||||
await self.sleep(min(0.25, self.sync_seconds))
|
||||
|
||||
async def close(self):
|
||||
self._closed = True
|
||||
if self._sync_task is not None:
|
||||
self._sync_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await self._sync_task
|
||||
self._sync_task = None
|
||||
if self.synchronizer is not None:
|
||||
await self.synchronizer.close()
|
||||
|
||||
@staticmethod
|
||||
def _generation(stat_result):
|
||||
return (
|
||||
stat_result.st_dev,
|
||||
stat_result.st_ino,
|
||||
stat_result.st_mtime_ns,
|
||||
stat_result.st_size,
|
||||
)
|
||||
|
||||
async def records(self):
|
||||
if self.synchronizer is not None and self._sync_task is None:
|
||||
self._sync_task = asyncio.create_task(self._sync_loop())
|
||||
consumed_generation = None
|
||||
while not self._closed:
|
||||
try:
|
||||
current = self.path.stat()
|
||||
except FileNotFoundError:
|
||||
await self.sleep(self.poll_seconds)
|
||||
continue
|
||||
generation = self._generation(current)
|
||||
if generation == consumed_generation:
|
||||
await self.sleep(self.poll_seconds)
|
||||
continue
|
||||
try:
|
||||
stream = self.path.open("rb")
|
||||
except FileNotFoundError:
|
||||
continue
|
||||
with stream:
|
||||
opened_generation = self._generation(os.fstat(stream.fileno()))
|
||||
generation_fingerprints = set()
|
||||
valid_generation = True
|
||||
while not self._closed:
|
||||
raw = await asyncio.to_thread(
|
||||
stream.readline, MAX_SESSION_RECORD_BYTES + 2
|
||||
)
|
||||
if not raw:
|
||||
break
|
||||
if (
|
||||
not raw.endswith(b"\n")
|
||||
or len(raw) - 1 > MAX_SESSION_RECORD_BYTES
|
||||
):
|
||||
self._emit(
|
||||
"source_record_rejected", {"reason": "invalid_record"}
|
||||
)
|
||||
valid_generation = False
|
||||
break
|
||||
try:
|
||||
record = parse_session_document(raw[:-1])
|
||||
except ValueError:
|
||||
self._emit(
|
||||
"source_record_rejected", {"reason": "invalid_record"}
|
||||
)
|
||||
valid_generation = False
|
||||
continue
|
||||
generation_fingerprints.add(self.fingerprint(record.source_id))
|
||||
yield record
|
||||
if valid_generation:
|
||||
self._snapshot_fingerprints = frozenset(generation_fingerprints)
|
||||
consumed_generation = opened_generation
|
||||
|
||||
|
||||
def parse_session_document(raw: bytes | str) -> SourceRecord:
|
||||
if len(raw) > MAX_SESSION_RECORD_BYTES:
|
||||
raise ValueError("invalid remote session record")
|
||||
try:
|
||||
document = json.loads(raw)
|
||||
source_id = document["email"]
|
||||
raw_cookies = document["cookies"]
|
||||
except (UnicodeDecodeError, TypeError, ValueError, KeyError) as exc:
|
||||
raise ValueError("invalid remote session record") from exc
|
||||
if not isinstance(source_id, str) or not source_id:
|
||||
raise ValueError("invalid remote session record")
|
||||
try:
|
||||
source_id.encode("utf-8")
|
||||
except UnicodeEncodeError as exc:
|
||||
raise ValueError("invalid remote session record") from exc
|
||||
if not isinstance(raw_cookies, list) or not raw_cookies:
|
||||
raise ValueError("invalid remote session record")
|
||||
|
||||
cookies = []
|
||||
sso_token = ""
|
||||
fallback_sso_token = ""
|
||||
allowed = {
|
||||
"name",
|
||||
"value",
|
||||
"url",
|
||||
"domain",
|
||||
"path",
|
||||
"expires",
|
||||
"httpOnly",
|
||||
"secure",
|
||||
"sameSite",
|
||||
}
|
||||
for raw_cookie in raw_cookies:
|
||||
if not isinstance(raw_cookie, dict):
|
||||
raise ValueError("invalid remote session record")
|
||||
cookie = {key: raw_cookie[key] for key in allowed if key in raw_cookie}
|
||||
name = cookie.get("name")
|
||||
value = cookie.get("value")
|
||||
if not isinstance(name, str) or not name or not isinstance(value, str) or not value:
|
||||
raise ValueError("invalid remote session record")
|
||||
scope = cookie.get("url") or cookie.get("domain")
|
||||
if not isinstance(scope, str) or not scope:
|
||||
raise ValueError("invalid remote session record")
|
||||
try:
|
||||
name.encode("utf-8")
|
||||
value.encode("utf-8")
|
||||
scope.encode("utf-8")
|
||||
except UnicodeEncodeError as exc:
|
||||
raise ValueError("invalid remote session record") from exc
|
||||
if name == "sso" and not sso_token:
|
||||
sso_token = value
|
||||
elif name == "sso-rw" and not fallback_sso_token:
|
||||
fallback_sso_token = value
|
||||
cookies.append(cookie)
|
||||
sso_token = sso_token or fallback_sso_token
|
||||
if not sso_token:
|
||||
raise ValueError("invalid remote session record")
|
||||
return SourceRecord(source_id, sso_token, tuple(cookies))
|
||||
|
||||
|
||||
class RemoteSessionStream:
|
||||
"""Yield full snapshots and appends from one reconnecting SSH child."""
|
||||
|
||||
MAX_STDERR_BYTES = 16 * 1024
|
||||
MAX_RECORD_BYTES = MAX_SESSION_RECORD_BYTES
|
||||
RECONNECT_DELAYS = (1.0, 2.0, 5.0, 10.0, 30.0)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str,
|
||||
*,
|
||||
remote_root: str = "/opt/grok-free-register",
|
||||
identity_file: str | None = None,
|
||||
process_factory=asyncio.create_subprocess_exec,
|
||||
sleep=asyncio.sleep,
|
||||
event_callback=None,
|
||||
):
|
||||
self.host = host
|
||||
self.remote_root = remote_root
|
||||
self.identity_file = identity_file
|
||||
self.process_factory = process_factory
|
||||
self.sleep = sleep
|
||||
self.event_callback = event_callback
|
||||
self._process = None
|
||||
self._closed = False
|
||||
self._last_disconnect_reason = None
|
||||
|
||||
def _emit(self, kind, data):
|
||||
if self.event_callback is None:
|
||||
return
|
||||
try:
|
||||
self.event_callback(kind, data)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _command(self):
|
||||
return (
|
||||
f"cd {shlex.quote(self.remote_root)} && "
|
||||
"python3 -u scripts/export_registered_sessions.py --follow "
|
||||
"keys/auth-sessions.jsonl keys/accounts.txt"
|
||||
)
|
||||
|
||||
def _args(self):
|
||||
args = [
|
||||
"ssh",
|
||||
"-C",
|
||||
"-T",
|
||||
"-o",
|
||||
"BatchMode=yes",
|
||||
"-o",
|
||||
"ConnectTimeout=15",
|
||||
"-o",
|
||||
"ServerAliveInterval=15",
|
||||
"-o",
|
||||
"ServerAliveCountMax=3",
|
||||
]
|
||||
if self.identity_file:
|
||||
args.extend(["-i", self.identity_file])
|
||||
args.extend(["--", self.host, self._command()])
|
||||
return args
|
||||
|
||||
async def _read_stderr(self, stream):
|
||||
retained = bytearray()
|
||||
while True:
|
||||
chunk = await stream.read(4096)
|
||||
if not chunk:
|
||||
break
|
||||
remaining = self.MAX_STDERR_BYTES - len(retained)
|
||||
if remaining > 0:
|
||||
retained.extend(chunk[:remaining])
|
||||
return bytes(retained)
|
||||
|
||||
@staticmethod
|
||||
def _classify_disconnect(returncode, stderr):
|
||||
normalized = stderr.lower()
|
||||
if b"permission denied" in normalized or b"host key verification failed" in normalized:
|
||||
return "ssh_auth_failed"
|
||||
if b"could not resolve" in normalized or b"name or service not known" in normalized:
|
||||
return "ssh_resolution_failed"
|
||||
if any(
|
||||
marker in normalized
|
||||
for marker in (b"connection refused", b"connection timed out", b"no route to host")
|
||||
):
|
||||
return "ssh_connection_failed"
|
||||
if returncode == 3:
|
||||
return "remote_snapshot_changed"
|
||||
if returncode == 4:
|
||||
return "remote_data_invalid"
|
||||
return "remote_stream_closed"
|
||||
|
||||
async def _terminate_process(self, process):
|
||||
if process is None or process.returncode is not None:
|
||||
return
|
||||
with suppress(ProcessLookupError):
|
||||
process.terminate()
|
||||
try:
|
||||
await asyncio.wait_for(process.wait(), timeout=3)
|
||||
except TimeoutError:
|
||||
with suppress(ProcessLookupError):
|
||||
process.kill()
|
||||
await process.wait()
|
||||
|
||||
async def close(self):
|
||||
self._closed = True
|
||||
await self._terminate_process(self._process)
|
||||
|
||||
async def records(self):
|
||||
reconnect_index = 0
|
||||
while not self._closed:
|
||||
process = None
|
||||
stderr_task = None
|
||||
yielded = False
|
||||
reason = "remote_stream_closed"
|
||||
try:
|
||||
process = await self.process_factory(
|
||||
*self._args(),
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
limit=self.MAX_RECORD_BYTES + 1,
|
||||
)
|
||||
self._process = process
|
||||
stderr_task = asyncio.create_task(self._read_stderr(process.stderr))
|
||||
self._emit("source_connected", {})
|
||||
while not self._closed:
|
||||
try:
|
||||
raw = await process.stdout.readline()
|
||||
except (ValueError, asyncio.LimitOverrunError):
|
||||
reason = "remote_record_too_large"
|
||||
break
|
||||
if not raw:
|
||||
break
|
||||
if (
|
||||
not raw.endswith(b"\n")
|
||||
or len(raw) - 1 > self.MAX_RECORD_BYTES
|
||||
):
|
||||
reason = "remote_record_too_large"
|
||||
break
|
||||
try:
|
||||
record = parse_session_document(raw[:-1])
|
||||
except ValueError:
|
||||
self._emit("source_record_rejected", {"reason": "invalid_record"})
|
||||
continue
|
||||
if not yielded:
|
||||
self._last_disconnect_reason = None
|
||||
yielded = True
|
||||
reconnect_index = 0
|
||||
yield record
|
||||
|
||||
# A follow exporter is intentionally long-lived. Once framing is
|
||||
# invalid, waiting for it to exit on its own would hang this source
|
||||
# forever, so terminate it before collecting the exit status.
|
||||
if reason != "remote_stream_closed":
|
||||
await self._terminate_process(process)
|
||||
if process.returncode is None:
|
||||
await process.wait()
|
||||
stderr = await stderr_task
|
||||
if reason == "remote_stream_closed":
|
||||
reason = self._classify_disconnect(process.returncode, stderr)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except (OSError, RuntimeError):
|
||||
reason = "ssh_start_failed"
|
||||
finally:
|
||||
if stderr_task is not None and not stderr_task.done():
|
||||
stderr_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await stderr_task
|
||||
await self._terminate_process(process)
|
||||
if self._process is process:
|
||||
self._process = None
|
||||
|
||||
if self._closed:
|
||||
break
|
||||
if reason != self._last_disconnect_reason:
|
||||
self._emit("source_disconnected", {"reason": reason})
|
||||
self._last_disconnect_reason = reason
|
||||
if reason == "remote_snapshot_changed":
|
||||
delay = 0.1
|
||||
else:
|
||||
delay = self.RECONNECT_DELAYS[min(reconnect_index, len(self.RECONNECT_DELAYS) - 1)]
|
||||
if not yielded:
|
||||
reconnect_index += 1
|
||||
await self.sleep(delay)
|
||||
|
||||
|
||||
PersistentSSHSource = RemoteSessionStream
|
||||
883
xai_enroller/service.py
Normal file
883
xai_enroller/service.py
Normal file
@@ -0,0 +1,883 @@
|
||||
"""Interactive composition for the asynchronous local authentication service."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import secrets
|
||||
import sys
|
||||
import tempfile
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from .inventory import InventoryError
|
||||
|
||||
|
||||
DEFAULT_LOCAL_AUTH_DIR = Path.home() / "Downloads" / "grok-free-register-auth"
|
||||
AUTHENTICATED_DIRNAME = "authenticated"
|
||||
CLAIMED_DIRNAME = "claimed"
|
||||
|
||||
|
||||
class AuthServiceConfigurationError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
class InteractiveCommandPrompt:
|
||||
"""Stdlib-only line editor that keeps commands visible during events."""
|
||||
|
||||
CLEAR_LINE = "\r\x1b[2K"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
input_stream=None,
|
||||
output=None,
|
||||
interactive=None,
|
||||
prompt="认证> ",
|
||||
):
|
||||
self.input_stream = input_stream or sys.stdin
|
||||
self.output = output or sys.stdout
|
||||
self.prompt = prompt
|
||||
if interactive is None:
|
||||
interactive = bool(
|
||||
getattr(self.input_stream, "isatty", lambda: False)()
|
||||
and getattr(self.output, "isatty", lambda: False)()
|
||||
)
|
||||
self.interactive = interactive
|
||||
self.buffer = ""
|
||||
self.active = False
|
||||
self._callback = None
|
||||
self._fd = None
|
||||
self._original_terminal = None
|
||||
|
||||
def _write(self, text):
|
||||
try:
|
||||
self.output.write(text)
|
||||
self.output.flush()
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
|
||||
def _redraw(self):
|
||||
if self.active and self.interactive:
|
||||
self._write(f"{self.CLEAR_LINE}{self.prompt}{self.buffer}")
|
||||
|
||||
def start(self, callback):
|
||||
self._callback = callback
|
||||
self.active = True
|
||||
self._redraw()
|
||||
|
||||
def feed(self, text):
|
||||
for character in text:
|
||||
if character in {"\r", "\n"}:
|
||||
command = self.buffer
|
||||
self.buffer = ""
|
||||
self._write("\n")
|
||||
if command.strip() and self._callback is not None:
|
||||
self._callback(command)
|
||||
if self.active and self.interactive:
|
||||
self._write(self.prompt)
|
||||
elif character in {"\b", "\x7f"}:
|
||||
self.buffer = self.buffer[:-1]
|
||||
self._redraw()
|
||||
elif character == "\x04":
|
||||
if self._callback is not None:
|
||||
self._callback(None)
|
||||
elif character.isprintable():
|
||||
self.buffer += character
|
||||
self._redraw()
|
||||
|
||||
def write_event(self, message):
|
||||
if self.active and self.interactive:
|
||||
self._write(f"{self.CLEAR_LINE}{message}\n{self.prompt}{self.buffer}")
|
||||
else:
|
||||
self._write(f"{message}\n")
|
||||
|
||||
def attach(self, loop, callback):
|
||||
self.start(callback)
|
||||
self._fd = self.input_stream.fileno()
|
||||
if self.interactive:
|
||||
import termios
|
||||
import tty
|
||||
|
||||
self._original_terminal = termios.tcgetattr(self._fd)
|
||||
tty.setcbreak(self._fd)
|
||||
|
||||
def stdin_ready():
|
||||
data = os.read(self._fd, 256)
|
||||
if data:
|
||||
self.feed(data.decode(errors="ignore"))
|
||||
elif self._callback is not None:
|
||||
loop.remove_reader(self._fd)
|
||||
self._callback(None)
|
||||
|
||||
else:
|
||||
|
||||
def stdin_ready():
|
||||
line = self.input_stream.readline()
|
||||
if line == "":
|
||||
if self._callback is not None:
|
||||
loop.remove_reader(self._fd)
|
||||
self._callback(None)
|
||||
return
|
||||
if self._callback is not None:
|
||||
self._callback(line)
|
||||
|
||||
loop.add_reader(self._fd, stdin_ready)
|
||||
|
||||
def close(self, loop):
|
||||
if self._fd is not None:
|
||||
with suppress(Exception):
|
||||
loop.remove_reader(self._fd)
|
||||
if self._original_terminal is not None and self._fd is not None:
|
||||
import termios
|
||||
|
||||
with suppress(Exception):
|
||||
termios.tcsetattr(
|
||||
self._fd, termios.TCSADRAIN, self._original_terminal
|
||||
)
|
||||
if self.active and self.interactive:
|
||||
self._write(f"{self.CLEAR_LINE}\n")
|
||||
self.active = False
|
||||
|
||||
|
||||
def resolve_auth_log_mode(argv=None, env=None):
|
||||
argv = list(sys.argv[1:] if argv is None else argv)
|
||||
env = dict(os.environ if env is None else env)
|
||||
unknown = [argument for argument in argv if argument != "--debug"]
|
||||
if unknown:
|
||||
raise AuthServiceConfigurationError(
|
||||
f"unsupported auth service argument: {unknown[0]}"
|
||||
)
|
||||
mode = (env.get("XAI_AUTH_SERVICE_LOG_MODE") or "user").strip().lower()
|
||||
if "--debug" in argv:
|
||||
mode = "debug"
|
||||
if mode not in {"user", "debug"}:
|
||||
raise AuthServiceConfigurationError(
|
||||
"XAI_AUTH_SERVICE_LOG_MODE must be user or debug"
|
||||
)
|
||||
return mode
|
||||
|
||||
|
||||
def prepare_local_service_environment(env=None):
|
||||
"""Create private local persistence defaults without storing secrets in the repo."""
|
||||
merged = dict(os.environ if env is None else env)
|
||||
destination = Path(
|
||||
merged.get("XAI_ENROLLER_LOCAL_AUTH_DIR", DEFAULT_LOCAL_AUTH_DIR)
|
||||
).expanduser()
|
||||
destination.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
os.chmod(destination, 0o700)
|
||||
merged["XAI_ENROLLER_SINK"] = "local"
|
||||
merged["XAI_ENROLLER_LOCAL_AUTH_DIR"] = str(destination)
|
||||
merged.setdefault(
|
||||
"XAI_ENROLLER_LEDGER_PATH", str(destination / "enrollment-ledger.db")
|
||||
)
|
||||
if not merged.get("XAI_ENROLLER_SOURCE_SALT"):
|
||||
salt_file = destination / ".ledger-salt"
|
||||
try:
|
||||
salt = salt_file.read_text(encoding="utf-8").strip()
|
||||
except FileNotFoundError:
|
||||
salt = secrets.token_hex(32)
|
||||
fd, temporary_name = tempfile.mkstemp(
|
||||
prefix=".ledger-salt.", suffix=".tmp", dir=destination, text=True
|
||||
)
|
||||
try:
|
||||
os.fchmod(fd, 0o600)
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
||||
stream.write(salt + "\n")
|
||||
stream.flush()
|
||||
os.fsync(stream.fileno())
|
||||
os.replace(temporary_name, salt_file)
|
||||
finally:
|
||||
if os.path.exists(temporary_name):
|
||||
os.unlink(temporary_name)
|
||||
os.chmod(salt_file, 0o600)
|
||||
merged["XAI_ENROLLER_SOURCE_SALT"] = salt
|
||||
return merged
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AuthServiceSettings:
|
||||
source_kind: str
|
||||
ssh_host: str | None
|
||||
register_root: str
|
||||
remote_root: str
|
||||
identity_file: str | None
|
||||
sync_seconds: int = 30
|
||||
retry_seconds: int = 60
|
||||
min_authorization_interval_seconds: float = 10.0
|
||||
|
||||
@classmethod
|
||||
def from_environ(cls, env=None):
|
||||
env = dict(os.environ if env is None else env)
|
||||
ssh_host = (env.get("XAI_AUTH_SERVICE_SSH_HOST") or "").strip()
|
||||
requested_source = (env.get("XAI_AUTH_SERVICE_SOURCE") or "auto").strip().lower()
|
||||
if requested_source not in {"auto", "local", "ssh"}:
|
||||
raise AuthServiceConfigurationError(
|
||||
"XAI_AUTH_SERVICE_SOURCE must be auto, local, or ssh"
|
||||
)
|
||||
source_kind = (
|
||||
"ssh" if requested_source == "auto" and ssh_host else
|
||||
"local" if requested_source == "auto" else
|
||||
requested_source
|
||||
)
|
||||
if source_kind == "ssh" and not ssh_host:
|
||||
raise AuthServiceConfigurationError("XAI_AUTH_SERVICE_SSH_HOST is required")
|
||||
register_root = (
|
||||
env.get("XAI_AUTH_SERVICE_REGISTER_ROOT")
|
||||
or str(Path(__file__).resolve().parents[1])
|
||||
).strip()
|
||||
remote_root = (
|
||||
env.get("XAI_AUTH_SERVICE_REMOTE_ROOT") or "/opt/grok-free-register"
|
||||
).strip()
|
||||
identity_file = (env.get("XAI_AUTH_SERVICE_SSH_IDENTITY") or "").strip() or None
|
||||
try:
|
||||
sync_seconds = int(env.get("XAI_AUTH_SERVICE_SYNC_SEC", "30"))
|
||||
retry_seconds = int(env.get("XAI_AUTH_SERVICE_RETRY_SEC", "60"))
|
||||
min_authorization_interval_seconds = float(
|
||||
env.get("XAI_AUTH_SERVICE_MIN_INTERVAL_SEC", "10")
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise AuthServiceConfigurationError(
|
||||
"auth service intervals must be numeric"
|
||||
) from exc
|
||||
if not 5 <= sync_seconds <= 3600:
|
||||
raise AuthServiceConfigurationError(
|
||||
"XAI_AUTH_SERVICE_SYNC_SEC must be between 5 and 3600"
|
||||
)
|
||||
if not 30 <= retry_seconds <= 86400:
|
||||
raise AuthServiceConfigurationError(
|
||||
"XAI_AUTH_SERVICE_RETRY_SEC must be between 30 and 86400"
|
||||
)
|
||||
if not 0 <= min_authorization_interval_seconds <= 3600:
|
||||
raise AuthServiceConfigurationError(
|
||||
"XAI_AUTH_SERVICE_MIN_INTERVAL_SEC must be between 0 and 3600"
|
||||
)
|
||||
return cls(
|
||||
source_kind=source_kind,
|
||||
ssh_host=ssh_host or None,
|
||||
register_root=register_root,
|
||||
remote_root=remote_root,
|
||||
identity_file=identity_file,
|
||||
sync_seconds=sync_seconds,
|
||||
retry_seconds=retry_seconds,
|
||||
min_authorization_interval_seconds=min_authorization_interval_seconds,
|
||||
)
|
||||
|
||||
|
||||
class EventTerminal:
|
||||
"""Mode-aware aggregate renderer; identifiers and credentials stay hidden."""
|
||||
|
||||
_SAFE_REASONS = {
|
||||
"already_imported",
|
||||
"authorization_transport_failed",
|
||||
"browser_error",
|
||||
"challenge_required",
|
||||
"confirmation_timeout",
|
||||
"device_flow_failed",
|
||||
"device_flow_invalid",
|
||||
"device_flow_refresh_failed",
|
||||
"imported",
|
||||
"internal_error",
|
||||
"oauth_denied",
|
||||
"oauth_expired",
|
||||
"oauth_rejected",
|
||||
"operator_cancelled",
|
||||
"rate_limited",
|
||||
"sink_failed",
|
||||
"sink_unconfigured",
|
||||
"snapshot_sync_failed",
|
||||
"token_failed",
|
||||
"token_transport_failed",
|
||||
"unknown_page",
|
||||
}
|
||||
_TRANSIENT_REASONS = {
|
||||
"authorization_transport_failed",
|
||||
"browser_error",
|
||||
"challenge_required",
|
||||
"confirmation_timeout",
|
||||
"device_flow_failed",
|
||||
"device_flow_refresh_failed",
|
||||
"sink_failed",
|
||||
"token_failed",
|
||||
"token_transport_failed",
|
||||
"unknown_page",
|
||||
}
|
||||
|
||||
def __init__(self, mode="user", output=None):
|
||||
if mode not in {"user", "debug"}:
|
||||
raise ValueError("terminal mode must be user or debug")
|
||||
self.mode = mode
|
||||
self.output = output or self._print
|
||||
|
||||
@staticmethod
|
||||
def _print(message):
|
||||
print(message, flush=True)
|
||||
|
||||
@staticmethod
|
||||
def _percentage(value):
|
||||
return f"{100.0 * float(value):.1f}%"
|
||||
|
||||
@staticmethod
|
||||
def _rate(value):
|
||||
return "—" if value is None else f"{float(value):.2f}/分"
|
||||
|
||||
@staticmethod
|
||||
def _pending(value):
|
||||
return "—" if value is None else str(value)
|
||||
|
||||
@staticmethod
|
||||
def _task_suffix(value):
|
||||
return "" if value is None else f" #{value}"
|
||||
|
||||
def _format_user(self, kind, data):
|
||||
if kind == "startup":
|
||||
source = {"local": "本机", "ssh": "远端"}.get(
|
||||
data.get("source_kind"), "等待同步"
|
||||
)
|
||||
return (
|
||||
f"[✓] 本地认证服务已启动 | 来源 {source} | "
|
||||
f"输出 {data['destination']} | 待处理 — | 可用 {data['available']}"
|
||||
)
|
||||
if kind == "service_started":
|
||||
return "[▶] 认证流水线运行中"
|
||||
if kind == "service_stopped":
|
||||
return "[■] 认证服务已停止"
|
||||
if kind == "source_connected":
|
||||
return "[✓] 账号来源已连接 | 本地快照可用"
|
||||
if kind == "source_updated":
|
||||
return f"[↻] 发现新账号 {data['new']} | 快照共 {data['total']}"
|
||||
if kind == "source_disconnected":
|
||||
return "[!] 账号来源暂时断开 | 继续使用上一份有效快照"
|
||||
if kind == "source_record_rejected":
|
||||
return "[!] 已忽略一条无效来源记录"
|
||||
if kind == "rate_limited":
|
||||
pace = data.get("min_authorization_interval_seconds")
|
||||
pace_text = (
|
||||
f" | 恢复后间隔 {pace:.0f}秒" if pace is not None else ""
|
||||
)
|
||||
return (
|
||||
f"[⏸] 触发限流 | {data['wait_seconds']}秒后单次探测{pace_text}"
|
||||
)
|
||||
if kind == "rate_limit_cleared":
|
||||
return f"[▶] 限流解除 | 实际等待 {data['elapsed_seconds']}秒"
|
||||
if kind == "pacing_adjusted":
|
||||
return (
|
||||
f"[•] 认证间隔调整为 {data['min_authorization_interval_seconds']:.0f}秒"
|
||||
)
|
||||
if kind == "authorization_started":
|
||||
return (
|
||||
f"[→] 开始认证 #{data['task_number']} | "
|
||||
f"待处理 {self._pending(data.get('pending_total'))}"
|
||||
)
|
||||
if kind == "result":
|
||||
if data.get("status") == "imported":
|
||||
return (
|
||||
f"[✓] 认证成功{self._task_suffix(data.get('task_number'))} | "
|
||||
f"运行平均 {self._rate(data.get('lifetime_imports_per_minute'))} | "
|
||||
f"累计 {data['imported_unique']} | 可用 {data.get('available', '—')}"
|
||||
)
|
||||
reason = data.get("reason")
|
||||
if reason == "rate_limited":
|
||||
return None
|
||||
action = "暂时失败，将自动重试" if reason in self._TRANSIENT_REASONS else "已跳过"
|
||||
return f"[✗] 认证未完成{self._task_suffix(data.get('task_number'))} | {action}"
|
||||
if kind == "control":
|
||||
states = {
|
||||
"paused": "[⏸] 认证服务已暂停",
|
||||
"running": "[▶] 认证服务已恢复",
|
||||
"cancelling": "[■] 正在取消当前任务",
|
||||
"idle": "[•] 当前没有可取消的任务",
|
||||
"stopping": "[■] 正在安全退出",
|
||||
"usage: take <positive-count>": "[!] 用法：take N，N 必须是正整数",
|
||||
"inventory unavailable": "[!] 当前凭据库存不可用",
|
||||
"help": (
|
||||
"[•] 命令 s 状态 | take N 取用 | p 暂停 | "
|
||||
"r 恢复 | c 取消 | q 退出"
|
||||
),
|
||||
"unknown": "[!] 未知命令 | 输入 help 查看命令",
|
||||
}
|
||||
return states.get(data.get("state"), "[•] 控制命令已处理")
|
||||
if kind == "status":
|
||||
state = {"running": "运行中", "paused": "已暂停", "stopping": "停止中"}.get(
|
||||
data.get("state"), "未知"
|
||||
)
|
||||
cooldown = (
|
||||
f"限流冷却 {data.get('cooldown_remaining_seconds', 0):.0f}秒"
|
||||
if data.get("cooldown")
|
||||
else "不限流"
|
||||
)
|
||||
return (
|
||||
f"[•] 状态 {state} | 待处理 {self._pending(data.get('pending_total'))} | "
|
||||
f"阶段 {data.get('active_stage', 'idle')} | "
|
||||
f"运行平均 {self._rate(data.get('lifetime_imports_per_minute'))} | "
|
||||
f"累计 {data.get('imported_unique', 0)} | "
|
||||
f"可用 {data.get('available', 0)} | 已取用 {data.get('claimed', 0)} | {cooldown}"
|
||||
)
|
||||
if kind == "inventory_taken":
|
||||
return (
|
||||
f"[✓] 已取用 {data['moved']} 个凭据 | "
|
||||
f"剩余可用 {data['available']} | "
|
||||
f"位置 claimed/{data['batch_id']}/"
|
||||
)
|
||||
if kind == "inventory_error":
|
||||
return (
|
||||
f"[!] 凭据取用失败 | 可用 {data['available']} | "
|
||||
f"处理中 {data['claiming']} | 已取用 {data['claimed']}"
|
||||
)
|
||||
if kind == "pipeline_error":
|
||||
return f"[!] 认证流水线异常 | 阶段 {data.get('stage', 'unknown')}"
|
||||
return None
|
||||
|
||||
def _format_debug(self, kind, data):
|
||||
if kind == "service_started":
|
||||
return (
|
||||
"• service: running; "
|
||||
f"min_interval={data['min_authorization_interval_seconds']:.1f}s"
|
||||
)
|
||||
if kind == "service_stopped":
|
||||
return "• service: stopped"
|
||||
if kind == "source_connected":
|
||||
return "• source: local snapshot updated"
|
||||
if kind == "source_updated":
|
||||
return f"• source: new={data['new']}; total={data['total']}"
|
||||
if kind == "source_disconnected":
|
||||
return "⚠ source: snapshot_sync_failed; keeping previous snapshot"
|
||||
if kind == "source_record_rejected":
|
||||
return "⚠ source: invalid record rejected"
|
||||
if kind == "rate_limited":
|
||||
pace = data.get("min_authorization_interval_seconds")
|
||||
pace_text = f"; resume_interval={pace:.1f}s" if pace is not None else ""
|
||||
return (
|
||||
"⏸ authentication rate limited; "
|
||||
f"next single probe in {data['wait_seconds']}s{pace_text}"
|
||||
)
|
||||
if kind == "rate_limit_cleared":
|
||||
return (
|
||||
"▶ authentication rate limit cleared after "
|
||||
f"{data['elapsed_seconds']}s"
|
||||
)
|
||||
if kind == "pacing_adjusted":
|
||||
return (
|
||||
"• pacing: min_interval="
|
||||
f"{data['min_authorization_interval_seconds']:.1f}s"
|
||||
f" ({data.get('reason', 'adjusted')})"
|
||||
)
|
||||
if kind == "authorization_started":
|
||||
return (
|
||||
"→ authentication: next task started; "
|
||||
f"task={data['task_number']}; attempt={data['attempt_number']}; "
|
||||
f"queued={data['source_queue']}; "
|
||||
f"pending={self._pending(data.get('pending_total'))}"
|
||||
)
|
||||
if kind == "result":
|
||||
if data.get("status") == "imported":
|
||||
return (
|
||||
"✓ authentication: imported; "
|
||||
f"task={data.get('task_number')}; total={data['imported_unique']}; "
|
||||
f"5m_rate={self._rate(data.get('five_minute_imports_per_minute'))}; "
|
||||
f"attempt_success={self._percentage(data['attempt_success'])}; "
|
||||
f"eventual_success={self._percentage(data['eventual_success'])}"
|
||||
)
|
||||
reason = data.get("reason")
|
||||
safe_reason = reason if reason in self._SAFE_REASONS else "unknown"
|
||||
return (
|
||||
f"⚠ authentication: {data.get('status', 'unknown')} ({safe_reason}); "
|
||||
f"task={data.get('task_number')}; attempt={data.get('attempt_number')}"
|
||||
)
|
||||
if kind == "control":
|
||||
state = data.get("state", "unknown")
|
||||
safe_state = state if str(state).replace(" ", "").isalnum() else "unknown"
|
||||
return f"• service: {safe_state}"
|
||||
if kind == "status":
|
||||
rate = data.get("five_minute_imports_per_minute")
|
||||
rate_text = "unknown" if rate is None else f"{rate:.2f}/min"
|
||||
lifetime_rate = data.get("lifetime_imports_per_minute")
|
||||
lifetime_rate_text = (
|
||||
"unknown"
|
||||
if lifetime_rate is None
|
||||
else f"{lifetime_rate:.2f}/min"
|
||||
)
|
||||
return (
|
||||
f"• status: {data['state']}; "
|
||||
f"queues={data['source_queue']}/{data['prepared_queue']}/"
|
||||
f"{data['completion_queue']}; active={data['active_stage']}; "
|
||||
f"retry_waiting={data['retry_waiting']}; "
|
||||
f"next_retry={data['next_retry_seconds']:.1f}s; "
|
||||
f"started={data['authorization_starts']}; "
|
||||
f"cooldown={str(data['cooldown']).lower()}; "
|
||||
f"cooldown_remaining={data['cooldown_remaining_seconds']:.1f}s; "
|
||||
f"probe={str(data['probe_in_flight']).lower()}; "
|
||||
f"min_interval={data['min_authorization_interval_seconds']:.1f}s; "
|
||||
f"pace_remaining={data['pacing_remaining_seconds']:.1f}s; "
|
||||
f"imported={data['imported_unique']}; attempted={data['attempted_unique']}; "
|
||||
f"rate_limited={data['rate_limited']}; 5m_rate={rate_text}; "
|
||||
f"lifetime_rate={lifetime_rate_text}; "
|
||||
f"available={data['available']}; claiming={data['claiming']}; "
|
||||
f"claimed={data['claimed']}"
|
||||
)
|
||||
if kind == "inventory_taken":
|
||||
return (
|
||||
f"✓ inventory: claimed={data['moved']}; "
|
||||
f"available={data['available']}; batch={data['batch_id']}"
|
||||
)
|
||||
if kind == "inventory_error":
|
||||
return (
|
||||
f"⚠ inventory: available={data['available']}; "
|
||||
f"claiming={data['claiming']}; claimed={data['claimed']}"
|
||||
)
|
||||
if kind == "pipeline_error":
|
||||
reason = data.get("reason")
|
||||
safe_reason = reason if reason in self._SAFE_REASONS else "unknown"
|
||||
stage = data.get("stage", "unknown")
|
||||
safe_stage = stage if str(stage).replace("_", "").isalnum() else "unknown"
|
||||
return f"⚠ service: {safe_stage} stage stopped ({safe_reason})"
|
||||
known = self._format_user(kind, data)
|
||||
if known is not None:
|
||||
return known
|
||||
reason = data.get("reason")
|
||||
suffix = f" reason={reason}" if reason in self._SAFE_REASONS else ""
|
||||
safe_kind = kind if kind.replace("_", "").isalnum() else "unknown"
|
||||
return f"• debug event={safe_kind}{suffix}"
|
||||
|
||||
def emit(self, event):
|
||||
kind, data = event
|
||||
try:
|
||||
message = (
|
||||
self._format_debug(kind, data)
|
||||
if self.mode == "debug"
|
||||
else self._format_user(kind, data)
|
||||
)
|
||||
if message is not None:
|
||||
self.output(message)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
class AuthPipelineRunner:
|
||||
"""Map interactive controls onto the persistent pipeline lifecycle."""
|
||||
|
||||
def __init__(self, pipeline, emit, *, interval_seconds=None, inventory=None):
|
||||
self.pipeline = pipeline
|
||||
self.emit = emit
|
||||
self.inventory = inventory
|
||||
self.paused = False
|
||||
self.current_cycle = None
|
||||
|
||||
async def handle_command(self, command):
|
||||
command = command.strip().lower()
|
||||
if command == "p":
|
||||
self.paused = True
|
||||
self.pipeline.pause()
|
||||
self.emit(("control", {"state": "paused"}))
|
||||
elif command == "r":
|
||||
self.paused = False
|
||||
self.pipeline.resume()
|
||||
self.emit(("control", {"state": "running"}))
|
||||
elif command == "s":
|
||||
status = self.pipeline.status()
|
||||
status.update(self.pipeline.ledger.inventory_counts())
|
||||
self.emit(("status", status))
|
||||
elif command.startswith("take "):
|
||||
parts = command.split()
|
||||
if len(parts) != 2 or not parts[1].isdigit() or int(parts[1]) <= 0:
|
||||
self.emit(("control", {"state": "usage: take <positive-count>"}))
|
||||
return True
|
||||
if self.inventory is None:
|
||||
self.emit(("control", {"state": "inventory unavailable"}))
|
||||
return True
|
||||
try:
|
||||
batch = await asyncio.to_thread(self.inventory.take, int(parts[1]))
|
||||
except InventoryError as exc:
|
||||
self.emit(
|
||||
(
|
||||
"inventory_error",
|
||||
{
|
||||
"reason": str(exc),
|
||||
**self.pipeline.ledger.inventory_counts(),
|
||||
},
|
||||
)
|
||||
)
|
||||
return True
|
||||
counts = self.pipeline.ledger.inventory_counts()
|
||||
self.emit(
|
||||
(
|
||||
"inventory_taken",
|
||||
{
|
||||
"batch_id": batch.batch_id,
|
||||
"directory": str(batch.directory),
|
||||
"moved": batch.moved,
|
||||
**counts,
|
||||
},
|
||||
)
|
||||
)
|
||||
elif command == "c":
|
||||
cancelled = await self.pipeline.cancel_active()
|
||||
self.emit(
|
||||
("control", {"state": "cancelling" if cancelled else "idle"})
|
||||
)
|
||||
elif command in {"help", "h", "?"}:
|
||||
self.emit(("control", {"state": "help"}))
|
||||
elif command in {"q", "quit", "exit"}:
|
||||
self.pipeline.request_stop()
|
||||
self.emit(("control", {"state": "stopping"}))
|
||||
return False
|
||||
elif command:
|
||||
self.emit(("control", {"state": "unknown"}))
|
||||
return True
|
||||
|
||||
async def run(self):
|
||||
self.current_cycle = asyncio.create_task(self.pipeline.run())
|
||||
try:
|
||||
await self.current_cycle
|
||||
finally:
|
||||
self.current_cycle = None
|
||||
|
||||
|
||||
async def _run_interactive(runner, terminal):
|
||||
loop = asyncio.get_running_loop()
|
||||
commands = asyncio.Queue()
|
||||
prompt = InteractiveCommandPrompt()
|
||||
previous_output = terminal.output
|
||||
terminal.output = prompt.write_event
|
||||
prompt.write_event(
|
||||
"命令：s 状态 | take N 取用凭据 | p 暂停 | r 恢复 | c 取消当前任务 | q 退出"
|
||||
)
|
||||
try:
|
||||
prompt.attach(loop, commands.put_nowait)
|
||||
except Exception:
|
||||
prompt.close(loop)
|
||||
terminal.output = previous_output
|
||||
raise
|
||||
worker = asyncio.create_task(runner.run())
|
||||
try:
|
||||
while not worker.done():
|
||||
command_task = asyncio.create_task(commands.get())
|
||||
done, _pending = await asyncio.wait(
|
||||
{worker, command_task}, return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
if worker in done:
|
||||
command_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await command_task
|
||||
await worker
|
||||
break
|
||||
command = command_task.result()
|
||||
if command is None:
|
||||
command = "q"
|
||||
if not await runner.handle_command(command):
|
||||
break
|
||||
except asyncio.CancelledError:
|
||||
runner.pipeline.request_stop()
|
||||
raise
|
||||
finally:
|
||||
prompt.close(loop)
|
||||
terminal.output = previous_output
|
||||
runner.pipeline.request_stop()
|
||||
await asyncio.gather(worker, return_exceptions=True)
|
||||
|
||||
|
||||
async def main_async(*, log_mode="user"):
|
||||
import httpx
|
||||
|
||||
from .auth_pipeline import AuthPipeline
|
||||
from .config import Settings
|
||||
from .executors import PlaywrightExecutor
|
||||
from .inventory import CredentialInventory
|
||||
from .ledger import Ledger
|
||||
from .protocol import XAIProfile, XAIProtocol
|
||||
from .remote_stream import (
|
||||
DiskSnapshotSource,
|
||||
LocalSnapshotSynchronizer,
|
||||
SSHSnapshotSynchronizer,
|
||||
)
|
||||
from .sinks import LocalAuthFileSink
|
||||
|
||||
service_settings = AuthServiceSettings.from_environ()
|
||||
merged = prepare_local_service_environment()
|
||||
merged["XAI_ENROLLER_SOURCE_KIND"] = "remote"
|
||||
merged["XAI_ENROLLER_AUTH_EXECUTOR"] = "playwright"
|
||||
merged["XAI_ENROLLER_CONCURRENCY"] = "1"
|
||||
settings = Settings.from_environ(merged)
|
||||
terminal = EventTerminal(mode=log_mode)
|
||||
|
||||
# CF prewarm via external FlareSolverr (:8191). Separate from REGISTER_PROXY.
|
||||
try:
|
||||
from grok_register.clearance import (
|
||||
format_prewarm_log,
|
||||
httpx_proxy_mounts,
|
||||
prewarm_clearance,
|
||||
register_proxy_url,
|
||||
)
|
||||
|
||||
prewarm = await asyncio.to_thread(prewarm_clearance)
|
||||
terminal.emit(("control", {"state": format_prewarm_log(prewarm)}))
|
||||
reg_proxy = register_proxy_url()
|
||||
if reg_proxy:
|
||||
terminal.emit(("control", {"state": f"register_proxy={reg_proxy}"}))
|
||||
proxy_url = httpx_proxy_mounts()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
terminal.emit(("control", {"state": f"clearance_prewarm_skip: {exc}"}))
|
||||
proxy_url = None
|
||||
|
||||
# httpx honors REGISTER_PROXY/HTTP_PROXY only when passed explicitly
|
||||
client = httpx.AsyncClient(proxy=proxy_url) if proxy_url else httpx.AsyncClient()
|
||||
pipeline = None
|
||||
try:
|
||||
ledger = Ledger(settings.ledger_path, settings.source_salt)
|
||||
snapshot_path = Path(settings.local_auth_dir) / "source-snapshot.jsonl"
|
||||
if service_settings.source_kind == "ssh":
|
||||
synchronizer = SSHSnapshotSynchronizer(
|
||||
service_settings.ssh_host,
|
||||
snapshot_path,
|
||||
remote_root=service_settings.remote_root,
|
||||
identity_file=service_settings.identity_file,
|
||||
fingerprint=ledger.fingerprint,
|
||||
)
|
||||
else:
|
||||
synchronizer = LocalSnapshotSynchronizer(
|
||||
service_settings.register_root,
|
||||
snapshot_path,
|
||||
fingerprint=ledger.fingerprint,
|
||||
)
|
||||
source = DiskSnapshotSource(
|
||||
snapshot_path,
|
||||
synchronizer=synchronizer,
|
||||
sync_seconds=service_settings.sync_seconds,
|
||||
event_callback=lambda kind, data: terminal.emit((kind, data)),
|
||||
fingerprint=ledger.fingerprint,
|
||||
)
|
||||
protocol = XAIProtocol(
|
||||
client,
|
||||
XAIProfile.default(),
|
||||
default_poll_interval=settings.poll_interval,
|
||||
)
|
||||
executor = PlaywrightExecutor(concurrency=1)
|
||||
sink = LocalAuthFileSink(
|
||||
Path(settings.local_auth_dir) / AUTHENTICATED_DIRNAME,
|
||||
name_secret=settings.source_salt,
|
||||
)
|
||||
inventory = CredentialInventory(
|
||||
ledger,
|
||||
Path(settings.local_auth_dir) / AUTHENTICATED_DIRNAME,
|
||||
Path(settings.local_auth_dir) / CLAIMED_DIRNAME,
|
||||
)
|
||||
recovered = await asyncio.to_thread(inventory.recover)
|
||||
if recovered:
|
||||
terminal.emit(("control", {"state": f"recovered {recovered} claims"}))
|
||||
terminal.emit(
|
||||
(
|
||||
"startup",
|
||||
{
|
||||
"destination": f"{AUTHENTICATED_DIRNAME}/",
|
||||
"source_kind": service_settings.source_kind,
|
||||
**ledger.inventory_counts(),
|
||||
},
|
||||
)
|
||||
)
|
||||
pipeline = AuthPipeline(
|
||||
source=source,
|
||||
protocol=protocol,
|
||||
executor=executor,
|
||||
sink=sink,
|
||||
ledger=ledger,
|
||||
timeout=settings.timeout_sec,
|
||||
min_authorization_interval=(
|
||||
service_settings.min_authorization_interval_seconds
|
||||
),
|
||||
event_callback=lambda kind, data: terminal.emit((kind, data)),
|
||||
)
|
||||
# Keep base cooldown configurable; the gate grows it after consecutive trips.
|
||||
pipeline.rate_gate.base_cooldown = float(service_settings.retry_seconds)
|
||||
pipeline.rate_gate.COOLDOWN_SECONDS = float(service_settings.retry_seconds)
|
||||
pipeline.rate_gate._current_cooldown = float(service_settings.retry_seconds)
|
||||
runner = AuthPipelineRunner(pipeline, terminal.emit, inventory=inventory)
|
||||
await _run_interactive(runner, terminal)
|
||||
finally:
|
||||
if pipeline is not None:
|
||||
pipeline.request_stop()
|
||||
await client.aclose()
|
||||
|
||||
|
||||
def _known_configuration_key(error):
|
||||
message = str(error)
|
||||
keys = (
|
||||
"XAI_AUTH_SERVICE_LOG_MODE",
|
||||
"XAI_AUTH_SERVICE_SOURCE",
|
||||
"XAI_AUTH_SERVICE_SSH_HOST",
|
||||
"XAI_AUTH_SERVICE_REGISTER_ROOT",
|
||||
"XAI_AUTH_SERVICE_SYNC_SEC",
|
||||
"XAI_AUTH_SERVICE_RETRY_SEC",
|
||||
"XAI_AUTH_SERVICE_MIN_INTERVAL_SEC",
|
||||
"XAI_ENROLLER_TARGET",
|
||||
"XAI_ENROLLER_RETRY_ATTEMPTS",
|
||||
"XAI_ENROLLER_TIMEOUT_SEC",
|
||||
"XAI_ENROLLER_POLL_SEC",
|
||||
"XAI_ENROLLER_AUTH_EXECUTOR",
|
||||
"XAI_ENROLLER_SINK",
|
||||
"XAI_ENROLLER_CPA_BASE_URL",
|
||||
"XAI_ENROLLER_CPA_MANAGEMENT_SECRET",
|
||||
"XAI_ENROLLER_LOCAL_AUTH_DIR",
|
||||
"XAI_ENROLLER_SOURCE_SALT",
|
||||
)
|
||||
explicit = next((key for key in keys if key in message), None)
|
||||
if explicit is not None:
|
||||
return explicit
|
||||
message_keys = {
|
||||
"target must be between 1 and 100": "XAI_ENROLLER_TARGET",
|
||||
"retry attempts must be between 0 and 3": "XAI_ENROLLER_RETRY_ATTEMPTS",
|
||||
"timeout must be positive": "XAI_ENROLLER_TIMEOUT_SEC",
|
||||
"poll interval must be positive": "XAI_ENROLLER_POLL_SEC",
|
||||
"executor must be http or playwright": "XAI_ENROLLER_AUTH_EXECUTOR",
|
||||
"sink must be cpa or local": "XAI_ENROLLER_SINK",
|
||||
"CPA base URL must use HTTPS": "XAI_ENROLLER_CPA_BASE_URL",
|
||||
"CPA base URL and management secret are required": (
|
||||
"XAI_ENROLLER_CPA_BASE_URL/XAI_ENROLLER_CPA_MANAGEMENT_SECRET"
|
||||
),
|
||||
"local auth directory is required": "XAI_ENROLLER_LOCAL_AUTH_DIR",
|
||||
"source salt is required": "XAI_ENROLLER_SOURCE_SALT",
|
||||
}
|
||||
return message_keys.get(message)
|
||||
|
||||
|
||||
def _print_unexpected_service_failure(error, *, log_mode):
|
||||
if log_mode == "debug":
|
||||
detail = f"异常类别 {type(error).__name__}"
|
||||
else:
|
||||
detail = "使用 bash auth-service.sh --debug 查看异常类别"
|
||||
print(f"[✗] 认证服务异常终止 | {detail}", file=sys.stderr)
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
load_dotenv(dotenv_path=Path(__file__).resolve().parents[1] / ".env")
|
||||
mode = "user"
|
||||
try:
|
||||
mode = resolve_auth_log_mode(argv)
|
||||
asyncio.run(main_async(log_mode=mode))
|
||||
return 0
|
||||
except AuthServiceConfigurationError as error:
|
||||
key = _known_configuration_key(error) or "认证服务参数"
|
||||
print(
|
||||
f"[✗] 配置错误：{key} | 教程 docs/guides/auth-service.md#配置远端同步",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 2
|
||||
except ValueError as error:
|
||||
key = _known_configuration_key(error)
|
||||
if key is not None:
|
||||
print(
|
||||
f"[✗] 配置错误：{key} | 教程 docs/guides/auth-service.md#配置远端同步",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 2
|
||||
_print_unexpected_service_failure(error, log_mode=mode)
|
||||
return 1
|
||||
except KeyboardInterrupt:
|
||||
return 130
|
||||
except Exception as error:
|
||||
_print_unexpected_service_failure(error, log_mode=mode)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
106
xai_enroller/sinks.py
Normal file
106
xai_enroller/sinks.py
Normal file
@@ -0,0 +1,106 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Protocol
|
||||
from urllib.parse import urlencode
|
||||
|
||||
import httpx
|
||||
|
||||
from .models import OAuthCredential, SinkReceipt
|
||||
|
||||
|
||||
class SinkError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class CredentialSink(Protocol):
|
||||
async def store(self, credential: OAuthCredential) -> SinkReceipt: ...
|
||||
|
||||
|
||||
def cpa_document(credential: OAuthCredential):
|
||||
return {
|
||||
"type": "xai",
|
||||
"access_token": credential.access_token,
|
||||
"refresh_token": credential.refresh_token,
|
||||
"id_token": credential.id_token,
|
||||
"token_type": credential.token_type,
|
||||
"expires_in": credential.expires_in,
|
||||
"expired": credential.expires_at,
|
||||
"last_refresh": credential.last_refresh,
|
||||
"sub": credential.subject,
|
||||
"base_url": "https://api.x.ai/v1",
|
||||
"token_endpoint": credential.token_endpoint,
|
||||
"auth_kind": "oauth",
|
||||
}
|
||||
|
||||
|
||||
def credential_filename(credential: OAuthCredential, name_secret: bytes):
|
||||
subject = credential.subject or credential.refresh_token
|
||||
digest = hmac.new(name_secret, subject.encode(), hashlib.sha256).hexdigest()[:16]
|
||||
return f"xai-{digest}.json"
|
||||
|
||||
|
||||
class CPAAuthFileSink:
|
||||
def __init__(self, base_url, management_secret, client: httpx.AsyncClient, name_secret=None):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.management_secret = management_secret
|
||||
self.client = client
|
||||
self.name_secret = name_secret or management_secret.encode()
|
||||
|
||||
async def store(self, credential: OAuthCredential):
|
||||
filename = credential_filename(credential, self.name_secret)
|
||||
document = cpa_document(credential)
|
||||
response = await self.client.post(
|
||||
f"{self.base_url}/v0/management/auth-files?{urlencode({'name': filename})}",
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.management_secret}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json=document,
|
||||
follow_redirects=False,
|
||||
)
|
||||
if response.status_code // 100 != 2:
|
||||
raise SinkError("CPA upload rejected")
|
||||
return SinkReceipt(filename.removesuffix(".json"))
|
||||
|
||||
|
||||
class LocalAuthFileSink:
|
||||
"""Atomically persist canonical CPA-compatible documents locally."""
|
||||
|
||||
def __init__(self, directory, *, name_secret: bytes):
|
||||
self.directory = Path(directory).expanduser()
|
||||
self.name_secret = name_secret
|
||||
|
||||
async def store(self, credential: OAuthCredential):
|
||||
return await asyncio.to_thread(self._store_sync, credential)
|
||||
|
||||
def _store_sync(self, credential: OAuthCredential):
|
||||
self.directory.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
os.chmod(self.directory, 0o700)
|
||||
filename = credential_filename(credential, self.name_secret)
|
||||
destination = self.directory / filename
|
||||
payload = json.dumps(cpa_document(credential), ensure_ascii=False, indent=2) + "\n"
|
||||
fd, temporary_name = tempfile.mkstemp(
|
||||
prefix=f".{filename}.", suffix=".tmp", dir=self.directory, text=True
|
||||
)
|
||||
try:
|
||||
os.fchmod(fd, 0o600)
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
||||
stream.write(payload)
|
||||
stream.flush()
|
||||
os.fsync(stream.fileno())
|
||||
os.replace(temporary_name, destination)
|
||||
os.chmod(destination, 0o600)
|
||||
directory_fd = os.open(self.directory, os.O_RDONLY)
|
||||
try:
|
||||
os.fsync(directory_fd)
|
||||
finally:
|
||||
os.close(directory_fd)
|
||||
finally:
|
||||
if os.path.exists(temporary_name):
|
||||
os.unlink(temporary_name)
|
||||
return SinkReceipt(filename.removesuffix(".json"))
|
||||
61
xai_enroller/sources.py
Normal file
61
xai_enroller/sources.py
Normal file
@@ -0,0 +1,61 @@
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from typing import Iterator, Protocol
|
||||
|
||||
from .config import require_private_regular_file
|
||||
from .models import SourceRecord
|
||||
|
||||
|
||||
class SourceAdapter(Protocol):
|
||||
def records(self) -> Iterator[SourceRecord]: ...
|
||||
|
||||
|
||||
class FileSourceAdapter:
|
||||
def __init__(self, path: Path):
|
||||
self.path = Path(path)
|
||||
|
||||
def records(self) -> Iterator[SourceRecord]:
|
||||
require_private_regular_file(self.path, require_0600=True)
|
||||
seen = set()
|
||||
with self.path.open("r", encoding="utf-8") as handle:
|
||||
for line_number, line in enumerate(handle, 1):
|
||||
raw = line.rstrip("\r\n")
|
||||
if not raw:
|
||||
continue
|
||||
source_id, separator, token = raw.partition("\t")
|
||||
if not separator or not source_id or not token:
|
||||
raise ValueError(f"invalid source line {line_number}")
|
||||
if source_id in seen:
|
||||
continue
|
||||
seen.add(source_id)
|
||||
yield SourceRecord(source_id, token)
|
||||
|
||||
|
||||
class SQLiteSourceAdapter:
|
||||
def __init__(self, path: Path, salt: bytes):
|
||||
self.path = Path(path)
|
||||
self.salt = bytes(salt)
|
||||
|
||||
def records(self) -> Iterator[SourceRecord]:
|
||||
require_private_regular_file(self.path)
|
||||
uri = f"file:{self.path.absolute()}?mode=ro"
|
||||
connection = sqlite3.connect(uri, uri=True)
|
||||
try:
|
||||
rows = connection.execute(
|
||||
"SELECT token FROM accounts "
|
||||
"WHERE status = 'active' AND deleted_at IS NULL "
|
||||
"ORDER BY updated_at DESC"
|
||||
)
|
||||
seen = set()
|
||||
for (token,) in rows:
|
||||
fingerprint = hmac.new(self.salt, token.encode(), hashlib.sha256).digest()
|
||||
source_id = base64.urlsafe_b64encode(fingerprint).rstrip(b"=").decode()
|
||||
if source_id in seen:
|
||||
continue
|
||||
seen.add(source_id)
|
||||
yield SourceRecord(source_id, token)
|
||||
finally:
|
||||
connection.close()
|
||||
Reference in New Issue
Block a user