kyson/grok-register
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: finish shared registration flow hardening

Browse Source
This commit is contained in:
github-actions[bot]
2026-07-14 16:49:13 +00:00
parent 5c9a58df8e
commit 3ab094ea1c
8 changed files with 344 additions and 644 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/temp-final-flow-fixes.yml
-34
View File
@@ -1,34 +0,0 @@
name: Temporary final flow fixes
on:
push:
paths:
- tools/apply_final_flow_fixes.py
- tools/fix_final_flow_test_generation.py
- .github/workflows/temp-final-flow-fixes.yml
workflow_dispatch:
permissions:
contents: write
jobs:
apply:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Repair focused test generation
run: python tools/fix_final_flow_test_generation.py
- name: Apply final flow fixes
run: python tools/apply_final_flow_fixes.py
- name: Verify syntax
run: python -m compileall -q grok_register_ttk.py registration_flow.py cpa_xai tests
- name: Run focused unit tests
run: python -m unittest tests.test_grok2api_remote_pool tests.test_registration_flow
- name: Commit changes and cleanup
run: |
git config user.name "github-actions[bot]"
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
git rm -f tools/apply_final_flow_fixes.py tools/fix_final_flow_test_generation.py .github/workflows/temp-final-flow-fixes.yml
git add -A
git commit -m "fix: finish shared registration flow hardening"
git push
cpa_xai/browser_confirm.py
+1 -1
View File
@@ -211,7 +211,7 @@ def create_standalone_page(proxy: Optional[str] = None, headless: bool = False,
except Exception:
if browser is not None:
close_standalone(browser)
elif proxy_bridge is not None:
if proxy_bridge is not None:
try:
proxy_bridge.stop()
except Exception:
grok_register_ttk.py
+95
View File
@@ -216,6 +216,26 @@ def validate_config(raw):
if parsed.scheme not in ("http", "https") or not parsed.netloc:
raise ConfigError(f"配置项 {key} 必须是有效的 http/https URL")
provider = cfg["email_provider"]
if provider == "cloudflare" and not cfg["cloudflare_api_base"]:
raise ConfigError("Cloudflare 模式需要配置 cloudflare_api_base")
if provider == "cloudmail":
missing = [
key for key in ("cloudmail_api_base", "cloudmail_public_token", "cloudmail_domains")
if not cfg[key]
]
if missing:
raise ConfigError("Cloud Mail 模式缺少必需配置: " + ", ".join(missing))
if cfg["grok2api_auto_add_remote"]:
missing = [
key for key in ("grok2api_remote_base", "grok2api_remote_app_key")
if not cfg[key]
]
if missing:
raise ConfigError("远端 token 入池缺少必需配置: " + ", ".join(missing))
if cfg["cpa_copy_to_hotload"] and not cfg["cpa_hotload_dir"]:
raise ConfigError("启用 CPA 热加载复制时必须配置 cpa_hotload_dir")
for key in path_keys:
value = cfg[key]
if value.startswith("~"):
@@ -3278,6 +3298,65 @@ def _queue_unsaved_account(path, payload, error, log_callback=None):
return False
def retry_pending_file(pending_path, output_path=None, log_callback=None):
logger = log_callback or (lambda message: None)
pending_path = os.path.abspath(os.path.expanduser(str(pending_path)))
if not os.path.isfile(pending_path):
raise FileNotFoundError(f"pending 文件不存在: {pending_path}")
suffix = ".pending.jsonl"
if output_path:
target_path = os.path.abspath(os.path.expanduser(str(output_path)))
elif pending_path.endswith(suffix):
target_path = pending_path[:-len(suffix)]
else:
target_path = pending_path + ".recovered.txt"
unresolved = []
restored = 0
with open(pending_path, "r", encoding="utf-8") as handle:
lines = handle.readlines()
for line_number, raw_line in enumerate(lines, 1):
stripped = raw_line.strip()
if not stripped:
continue
try:
record = json.loads(stripped)
if not isinstance(record, dict):
raise ValueError("record must be a JSON object")
email = str(record.get("email") or "").strip()
password = str(record.get("password") or "")
sso = str(record.get("sso") or "").strip()
if not email or not sso:
raise ValueError("record missing email or sso")
_append_account_line(target_path, email, password, sso)
restored += 1
logger(f"[+] 已恢复 pending 账号: {email}")
except Exception as exc:
unresolved.append(raw_line if raw_line.endswith("\n") else raw_line + "\n")
logger(f"[!] pending 第 {line_number} 行恢复失败: {exc}")
directory = os.path.dirname(pending_path) or "."
fd, temp_path = tempfile.mkstemp(prefix=".pending-retry-", suffix=".jsonl.tmp", dir=directory)
try:
with os.fdopen(fd, "w", encoding="utf-8") as handle:
handle.writelines(unresolved)
handle.flush()
os.fsync(handle.fileno())
if unresolved:
os.replace(temp_path, pending_path)
temp_path = None
try:
os.chmod(pending_path, 0o600)
except Exception:
pass
else:
os.unlink(temp_path)
temp_path = None
os.unlink(pending_path)
finally:
if temp_path and os.path.exists(temp_path):
os.unlink(temp_path)
return {"restored": restored, "remaining": len(unresolved), "output_path": target_path}
def run_registration_common(count, log_callback, cancel_callback, accounts_output_file, observer):
from registration_flow import RegistrationCallbacks, RegistrationOperations, run_batch
callbacks = RegistrationCallbacks(log=log_callback, cancelled=cancel_callback)
@@ -3775,6 +3854,22 @@ def main_cli():
def main():
if len(sys.argv) > 1 and sys.argv[1].strip().lower() == "retry-pending":
if len(sys.argv) < 3:
print("用法: python grok_register_ttk.py retry-pending <pending文件> [输出文件]", file=sys.stderr)
return
try:
summary = retry_pending_file(
sys.argv[2],
output_path=sys.argv[3] if len(sys.argv) > 3 else None,
log_callback=cli_log,
)
cli_log(
f"[*] pending 恢复完成: 已恢复 {summary['restored']} | 剩余 {summary['remaining']} | 输出 {summary['output_path']}"
)
except Exception as exc:
log_exception("pending 恢复失败", exc, cli_log)
return
if len(sys.argv) > 1 and sys.argv[1].strip().lower() in ("start", "cli", "--cli"):
main_cli()
return
registration_flow.py
+62 -16
View File
@@ -52,11 +52,22 @@ class OutputResult:
cpa: Dict[str, Any] = field(default_factory=dict)
@dataclass
class RegistrationSettings:
count: int
enable_nsfw: bool = True
max_mail_retry: int = 3
max_slot_retry: int = 3
cleanup_interval: int = 5
@dataclass
class BatchResult:
success_count: int = 0
fail_count: int = 0
processed_count: int = 0
registered_unsaved_count: int = 0
postprocess_warning_count: int = 0
cancelled: bool = False
results: list = field(default_factory=list)
@@ -152,24 +163,41 @@ def persist_account_result(result, callbacks, ops):
)
def _notify_observer(observer, result, account, output, callbacks):
try:
observer(result, account, output)
except Exception as exc:
callbacks.log(f"[Debug] observer 执行失败: {exc}")
def run_batch(count, callbacks, observer, ops, enable_nsfw=True, cleanup_interval=5,
max_slot_retry=3, max_mail_retry=3):
settings = RegistrationSettings(
count=int(count),
enable_nsfw=bool(enable_nsfw),
cleanup_interval=int(cleanup_interval),
max_slot_retry=int(max_slot_retry),
max_mail_retry=int(max_mail_retry),
)
result = BatchResult()
retry_count_for_slot = 0
ops.start_browser()
callbacks.log("[*] 浏览器已启动")
last_cleanup_success_count = 0
try:
while result.processed_count < count:
ops.start_browser()
callbacks.log("[*] 浏览器已启动")
while result.processed_count < settings.count:
if callbacks.cancelled():
result.cancelled = True
break
callbacks.log(f"--- 开始第 {result.processed_count + 1}/{count} 个账号 ---")
callbacks.log(f"--- 开始第 {result.processed_count + 1}/{settings.count} 个账号 ---")
account = None
output = None
try:
account = register_one_account(
callbacks, ops, enable_nsfw=enable_nsfw,
max_mail_retry=max_mail_retry,
callbacks,
ops,
enable_nsfw=settings.enable_nsfw,
max_mail_retry=settings.max_mail_retry,
)
output = persist_account_result(account, callbacks, ops)
result.results.append({"registration": account, "output": output})
@@ -178,19 +206,35 @@ def run_batch(count, callbacks, observer, ops, enable_nsfw=True, cleanup_interva
if output.saved:
result.success_count += 1
callbacks.log(f"[+] 注册并保存成功: {account.email}")
if (
settings.cleanup_interval > 0
and result.success_count % settings.cleanup_interval == 0
and result.success_count != last_cleanup_success_count
and result.processed_count < settings.count
):
ops.cleanup(f"已成功 {result.success_count} 个账号，执行定期清理")
last_cleanup_success_count = result.success_count
else:
result.fail_count += 1
result.registered_unsaved_count += 1
callbacks.log(f"[-] 注册成功但持久化未完成: {account.email}")
if result.success_count > 0 and result.success_count % cleanup_interval == 0 and result.processed_count < count:
ops.cleanup(f"已成功 {result.success_count} 个账号，执行定期清理")
pool_warning = any(
state.get("enabled") and not state.get("ok")
for state in output.pools.values()
)
cpa_warning = bool(output.cpa and not output.cpa.get("ok") and not output.cpa.get("skipped"))
if pool_warning or cpa_warning:
result.postprocess_warning_count += 1
except ops.cancelled_exception:
result.cancelled = True
callbacks.log("[!] 注册被停止")
break
except ops.retry_exception as exc:
retry_count_for_slot += 1
if retry_count_for_slot <= max_slot_retry:
callbacks.log(f"[!] 当前账号流程卡住，重试第 {retry_count_for_slot}/{max_slot_retry} 次: {exc}")
if retry_count_for_slot <= settings.max_slot_retry:
callbacks.log(
f"[!] 当前账号流程卡住，重试第 {retry_count_for_slot}/{settings.max_slot_retry} 次: {exc}"
)
else:
result.fail_count += 1
result.processed_count += 1
@@ -202,15 +246,17 @@ def run_batch(count, callbacks, observer, ops, enable_nsfw=True, cleanup_interva
retry_count_for_slot = 0
callbacks.log(f"[-] 注册失败: {exc}")
finally:
observer(result, account, output)
_notify_observer(observer, result, account, output, callbacks)
if callbacks.cancelled():
result.cancelled = True
break
if ops.browser_missing():
ops.start_browser()
else:
ops.restart_browser()
ops.sleep(1)
if result.processed_count < settings.count:
if ops.browser_missing():
ops.start_browser()
else:
ops.restart_browser()
ops.sleep(1)
finally:
ops.cleanup("任务结束")
return result
tests/test_grok2api_remote_pool.py
+68 -15
View File
@@ -1,15 +1,30 @@
import sys
import types
import unittest
from unittest.mock import patch
# Keep this unit test independent from optional browser/network dependencies.
drission = types.ModuleType("DrissionPage")
drission.Chromium = type("Chromium", (), {})
drission.ChromiumOptions = type("ChromiumOptions", (), {})
drission_errors = types.ModuleType("DrissionPage.errors")
drission_errors.PageDisconnectedError = type("PageDisconnectedError", (Exception,), {})
curl_cffi = types.ModuleType("curl_cffi")
curl_cffi.requests = types.SimpleNamespace()
sys.modules.setdefault("DrissionPage", drission)
sys.modules.setdefault("DrissionPage.errors", drission_errors)
sys.modules.setdefault("curl_cffi", curl_cffi)
import grok_register_ttk as app
class DummyResponse:
def __init__(self, payload=None, status_code=200, reason=""):
def __init__(self, payload=None, status_code=200, reason="", headers=None, text=""):
self._payload = payload or {}
self.status_code = status_code
self.reason = reason
self.text = ""
self.headers = headers or {}
self.text = text
def raise_for_status(self):
if self.status_code >= 400:
@@ -26,12 +41,17 @@ class Grok2ApiRemotePoolTests(unittest.TestCase):
def tearDown(self):
app.config = self.original_config
def test_remote_pool_falls_back_to_admin_api_prefix_when_root_tokens_add_is_404(self):
def _configure(self, **overrides):
app.config.update({
"grok2api_remote_base": "https://grok.example.com",
"grok2api_remote_app_key": "app-secret",
"grok2api_pool_name": "ssoBasic",
"grok2api_allow_legacy_full_save": False,
**overrides,
})
def test_remote_pool_falls_back_to_admin_api_prefix_when_root_tokens_add_is_404(self):
self._configure()
calls = []
def fake_post(url, **kwargs):
@@ -56,11 +76,10 @@ class Grok2ApiRemotePoolTests(unittest.TestCase):
})
def test_remote_pool_does_not_duplicate_admin_api_prefix_when_base_already_points_to_admin_api(self):
app.config.update({
"grok2api_remote_base": "https://grok.example.com/admin/api",
"grok2api_remote_app_key": "app-secret",
"grok2api_pool_name": "ssoSuper",
})
self._configure(
grok2api_remote_base="https://grok.example.com/admin/api",
grok2api_pool_name="ssoSuper",
)
calls = []
def fake_post(url, **kwargs):
@@ -76,12 +95,8 @@ class Grok2ApiRemotePoolTests(unittest.TestCase):
])
self.assertEqual(calls[0][1]["json"]["pool"], "super")
def test_remote_pool_full_save_fallback_tries_admin_api_tokens_path(self):
app.config.update({
"grok2api_remote_base": "https://grok.example.com",
"grok2api_remote_app_key": "app-secret",
"grok2api_pool_name": "ssoBasic",
})
def test_remote_pool_full_save_fallback_requires_opt_in_and_uses_etag(self):
self._configure(grok2api_allow_legacy_full_save=True)
get_calls = []
post_calls = []
@@ -96,7 +111,10 @@ class Grok2ApiRemotePoolTests(unittest.TestCase):
def fake_get(url, **kwargs):
get_calls.append((url, kwargs))
if url == "https://grok.example.com/admin/api/tokens":
return DummyResponse({"tokens": {"ssoBasic": []}})
return DummyResponse(
{"tokens": {"ssoBasic": []}},
headers={"ETag": '"version-7"'},
)
return DummyResponse(status_code=404)
with patch.object(app, "http_post", side_effect=fake_post), \
@@ -109,10 +127,45 @@ class Grok2ApiRemotePoolTests(unittest.TestCase):
"https://grok.example.com/admin/api/tokens",
])
self.assertEqual(post_calls[-1][0], "https://grok.example.com/admin/api/tokens")
self.assertEqual(post_calls[-1][1]["headers"]["If-Match"], '"version-7"')
self.assertEqual(post_calls[-1][1]["json"], {
"ssoBasic": [{"token": "fallback123", "tags": ["auto-register"], "note": "a@example.com"}],
})
def test_remote_pool_legacy_fallback_is_disabled_by_default(self):
self._configure()
with patch.object(app, "http_post", return_value=DummyResponse(status_code=404)), \
patch.object(app, "http_get") as get_mock:
with self.assertRaises(app.RemoteTokenCompatibilityError):
app.add_token_to_grok2api_remote_pool("abc")
get_mock.assert_not_called()
def test_remote_pool_500_does_not_fallback(self):
self._configure(grok2api_allow_legacy_full_save=True)
with patch.object(app, "http_post", return_value=DummyResponse(status_code=500, text="boom")), \
patch.object(app, "http_get") as get_mock:
with self.assertRaises(app.RemoteTokenRequestError):
app.add_token_to_grok2api_remote_pool("abc")
get_mock.assert_not_called()
def test_remote_pool_legacy_fallback_rejects_missing_etag(self):
self._configure(grok2api_allow_legacy_full_save=True)
def fake_post(url, **kwargs):
if url.endswith("/tokens/add"):
return DummyResponse(status_code=404)
return DummyResponse({"status": "success"})
def fake_get(url, **kwargs):
if url.endswith("/tokens"):
return DummyResponse({"tokens": {"ssoBasic": []}})
return DummyResponse(status_code=404)
with patch.object(app, "http_post", side_effect=fake_post), \
patch.object(app, "http_get", side_effect=fake_get):
with self.assertRaises(app.RemoteTokenCompatibilityError):
app.add_token_to_grok2api_remote_pool("abc")
if __name__ == "__main__":
unittest.main()
tests/test_registration_flow.py
+118
View File
@@ -0,0 +1,118 @@
import unittest
from registration_flow import (
RegistrationCallbacks,
RegistrationOperations,
run_batch,
)
class Cancelled(Exception):
pass
class Retryable(Exception):
pass
class FakeOps:
def __init__(self, save_ok=True, observer_events=None):
self.events = []
self.save_ok = save_ok
self.observer_events = observer_events if observer_events is not None else []
self.account_no = 0
def operations(self):
return RegistrationOperations(
start_browser=lambda: self.events.append("start"),
restart_browser=lambda: self.events.append("restart"),
browser_missing=lambda: False,
open_signup_page=lambda: self.events.append("open"),
fill_email_and_submit=self._email,
save_mail_credential=lambda email, token: True,
fill_code_and_submit=lambda email, token: "123456",
fill_profile_and_submit=lambda: {"given_name": "A", "family_name": "B", "password": "pw"},
wait_for_sso_cookie=lambda: "sso-token",
enable_nsfw=lambda sso: (True, "ok"),
persist_account_line=self._persist,
queue_unsaved_result=lambda payload, error: True,
add_tokens=lambda sso, email: {
"local": {"enabled": False, "ok": None, "error": None},
"remote": {"enabled": False, "ok": None, "error": None},
},
export_cpa=lambda email, password, sso: {"ok": False, "skipped": True},
cleanup=lambda reason: self.events.append(("cleanup", reason)),
sleep=lambda seconds: self.events.append(("sleep", seconds)),
cancelled_exception=Cancelled,
retry_exception=Retryable,
)
def _email(self):
self.account_no += 1
return f"user{self.account_no}@example.com", "mail-token"
def _persist(self, email, password, sso):
if not self.save_ok:
raise OSError("disk full")
self.events.append(("persist", email))
class RegistrationFlowTests(unittest.TestCase):
def callbacks(self, logs=None):
logs = logs if logs is not None else []
return RegistrationCallbacks(log=logs.append, cancelled=lambda: False)
def test_start_failure_still_runs_cleanup(self):
fake = FakeOps()
ops = fake.operations()
ops.start_browser = lambda: (_ for _ in ()).throw(RuntimeError("start failed"))
with self.assertRaises(RuntimeError):
run_batch(1, self.callbacks(), lambda *args: None, ops)
self.assertEqual(fake.events, [("cleanup", "任务结束")])
def test_last_account_does_not_restart_browser(self):
fake = FakeOps()
batch = run_batch(1, self.callbacks(), lambda *args: None, fake.operations())
self.assertEqual(batch.success_count, 1)
self.assertNotIn("restart", fake.events)
self.assertEqual(fake.events[-1], ("cleanup", "任务结束"))
def test_cleanup_interval_does_not_repeat_after_unsaved_result(self):
fake = FakeOps(save_ok=True)
ops = fake.operations()
original_persist = ops.persist_account_line
calls = {"count": 0}
def persist(email, password, sso):
calls["count"] += 1
if calls["count"] == 2:
raise OSError("disk full")
original_persist(email, password, sso)
ops.persist_account_line = persist
batch = run_batch(2, self.callbacks(), lambda *args: None, ops, cleanup_interval=1)
interval_cleanups = [
event for event in fake.events
if isinstance(event, tuple)
and len(event) > 1
and isinstance(event[1], str)
and "已成功" in event[1]
]
self.assertEqual(len(interval_cleanups), 1)
self.assertEqual(batch.success_count, 1)
self.assertEqual(batch.registered_unsaved_count, 1)
def test_observer_failure_is_logged_and_batch_continues(self):
fake = FakeOps()
logs = []
def broken_observer(*args):
raise RuntimeError("ui broke")
batch = run_batch(1, self.callbacks(logs), broken_observer, fake.operations())
self.assertEqual(batch.success_count, 1)
self.assertTrue(any("observer 执行失败" in line for line in logs))
if __name__ == "__main__":
unittest.main()
tools/apply_final_flow_fixes.py
-532
View File
@@ -1,532 +0,0 @@
#!/usr/bin/env python3
from pathlib import Path
import ast
ROOT = Path(__file__).resolve().parents[1]
def replace_once(text, old, new, label):
count = text.count(old)
if count != 1:
raise RuntimeError(f"{label}: expected one match, got {count}")
return text.replace(old, new, 1)
def replace_between(text, start, end, replacement, label):
i = text.find(start)
if i < 0:
raise RuntimeError(f"{label}: start not found")
j = text.find(end, i)
if j < 0:
raise RuntimeError(f"{label}: end not found")
return text[:i] + replacement + text[j:]
# 1) Shared registration flow boundaries and result states.
flow_path = ROOT / "registration_flow.py"
flow = flow_path.read_text(encoding="utf-8")
flow = replace_once(
flow,
"@dataclass\nclass BatchResult:\n success_count: int = 0\n fail_count: int = 0\n processed_count: int = 0\n cancelled: bool = False\n results: list = field(default_factory=list)\n",
"@dataclass\nclass RegistrationSettings:\n count: int\n enable_nsfw: bool = True\n max_mail_retry: int = 3\n max_slot_retry: int = 3\n cleanup_interval: int = 5\n\n\n@dataclass\nclass BatchResult:\n success_count: int = 0\n fail_count: int = 0\n processed_count: int = 0\n registered_unsaved_count: int = 0\n postprocess_warning_count: int = 0\n cancelled: bool = False\n results: list = field(default_factory=list)\n",
"settings and batch states",
)
new_run_batch = r'''def _notify_observer(observer, result, account, output, callbacks):
try:
observer(result, account, output)
except Exception as exc:
callbacks.log(f"[Debug] observer 执行失败: {exc}")
def run_batch(count, callbacks, observer, ops, enable_nsfw=True, cleanup_interval=5,
max_slot_retry=3, max_mail_retry=3):
settings = RegistrationSettings(
count=int(count),
enable_nsfw=bool(enable_nsfw),
cleanup_interval=int(cleanup_interval),
max_slot_retry=int(max_slot_retry),
max_mail_retry=int(max_mail_retry),
)
result = BatchResult()
retry_count_for_slot = 0
last_cleanup_success_count = 0
try:
ops.start_browser()
callbacks.log("[*] 浏览器已启动")
while result.processed_count < settings.count:
if callbacks.cancelled():
result.cancelled = True
break
callbacks.log(f"--- 开始第 {result.processed_count + 1}/{settings.count} 个账号 ---")
account = None
output = None
try:
account = register_one_account(
callbacks,
ops,
enable_nsfw=settings.enable_nsfw,
max_mail_retry=settings.max_mail_retry,
)
output = persist_account_result(account, callbacks, ops)
result.results.append({"registration": account, "output": output})
retry_count_for_slot = 0
result.processed_count += 1
if output.saved:
result.success_count += 1
callbacks.log(f"[+] 注册并保存成功: {account.email}")
if (
settings.cleanup_interval > 0
and result.success_count % settings.cleanup_interval == 0
and result.success_count != last_cleanup_success_count
and result.processed_count < settings.count
):
ops.cleanup(f"已成功 {result.success_count} 个账号，执行定期清理")
last_cleanup_success_count = result.success_count
else:
result.fail_count += 1
result.registered_unsaved_count += 1
callbacks.log(f"[-] 注册成功但持久化未完成: {account.email}")
pool_warning = any(
state.get("enabled") and not state.get("ok")
for state in output.pools.values()
)
cpa_warning = bool(output.cpa and not output.cpa.get("ok") and not output.cpa.get("skipped"))
if pool_warning or cpa_warning:
result.postprocess_warning_count += 1
except ops.cancelled_exception:
result.cancelled = True
callbacks.log("[!] 注册被停止")
break
except ops.retry_exception as exc:
retry_count_for_slot += 1
if retry_count_for_slot <= settings.max_slot_retry:
callbacks.log(
f"[!] 当前账号流程卡住，重试第 {retry_count_for_slot}/{settings.max_slot_retry} 次: {exc}"
)
else:
result.fail_count += 1
result.processed_count += 1
retry_count_for_slot = 0
callbacks.log(f"[-] 当前账号已达到最大重试次数，跳过: {exc}")
except Exception as exc:
result.fail_count += 1
result.processed_count += 1
retry_count_for_slot = 0
callbacks.log(f"[-] 注册失败: {exc}")
finally:
_notify_observer(observer, result, account, output, callbacks)
if callbacks.cancelled():
result.cancelled = True
break
if result.processed_count < settings.count:
if ops.browser_missing():
ops.start_browser()
else:
ops.restart_browser()
ops.sleep(1)
finally:
ops.cleanup("任务结束")
return result
'''
flow = replace_between(flow, "def run_batch(", "\n return result\n", new_run_batch, "run_batch")
# remove the old trailing return consumed only through marker
if flow.endswith("\n return result\n"):
flow = flow[:-len("\n return result\n")] + "\n"
ast.parse(flow)
flow_path.write_text(flow, encoding="utf-8")
# 2) Main config dependencies and retry-pending CLI.
main_path = ROOT / "grok_register_ttk.py"
main = main_path.read_text(encoding="utf-8")
cross_field = r''' provider = cfg["email_provider"]
if provider == "cloudflare" and not cfg["cloudflare_api_base"]:
raise ConfigError("Cloudflare 模式需要配置 cloudflare_api_base")
if provider == "cloudmail":
missing = [
key for key in ("cloudmail_api_base", "cloudmail_public_token", "cloudmail_domains")
if not cfg[key]
]
if missing:
raise ConfigError("Cloud Mail 模式缺少必需配置: " + ", ".join(missing))
if cfg["grok2api_auto_add_remote"]:
missing = [
key for key in ("grok2api_remote_base", "grok2api_remote_app_key")
if not cfg[key]
]
if missing:
raise ConfigError("远端 token 入池缺少必需配置: " + ", ".join(missing))
if cfg["cpa_copy_to_hotload"] and not cfg["cpa_hotload_dir"]:
raise ConfigError("启用 CPA 热加载复制时必须配置 cpa_hotload_dir")
'''
main = replace_once(
main,
" for key in path_keys:\n value = cfg[key]\n if value.startswith(\"~\"):\n cfg[key] = os.path.expanduser(value)\n return cfg\n",
cross_field + " for key in path_keys:\n value = cfg[key]\n if value.startswith(\"~\"):\n cfg[key] = os.path.expanduser(value)\n return cfg\n",
"cross-field config validation",
)
retry_pending = r'''def retry_pending_file(pending_path, output_path=None, log_callback=None):
logger = log_callback or (lambda message: None)
pending_path = os.path.abspath(os.path.expanduser(str(pending_path)))
if not os.path.isfile(pending_path):
raise FileNotFoundError(f"pending 文件不存在: {pending_path}")
suffix = ".pending.jsonl"
if output_path:
target_path = os.path.abspath(os.path.expanduser(str(output_path)))
elif pending_path.endswith(suffix):
target_path = pending_path[:-len(suffix)]
else:
target_path = pending_path + ".recovered.txt"
unresolved = []
restored = 0
with open(pending_path, "r", encoding="utf-8") as handle:
lines = handle.readlines()
for line_number, raw_line in enumerate(lines, 1):
stripped = raw_line.strip()
if not stripped:
continue
try:
record = json.loads(stripped)
if not isinstance(record, dict):
raise ValueError("record must be a JSON object")
email = str(record.get("email") or "").strip()
password = str(record.get("password") or "")
sso = str(record.get("sso") or "").strip()
if not email or not sso:
raise ValueError("record missing email or sso")
_append_account_line(target_path, email, password, sso)
restored += 1
logger(f"[+] 已恢复 pending 账号: {email}")
except Exception as exc:
unresolved.append(raw_line if raw_line.endswith("\n") else raw_line + "\n")
logger(f"[!] pending 第 {line_number} 行恢复失败: {exc}")
directory = os.path.dirname(pending_path) or "."
fd, temp_path = tempfile.mkstemp(prefix=".pending-retry-", suffix=".jsonl.tmp", dir=directory)
try:
with os.fdopen(fd, "w", encoding="utf-8") as handle:
handle.writelines(unresolved)
handle.flush()
os.fsync(handle.fileno())
if unresolved:
os.replace(temp_path, pending_path)
temp_path = None
try:
os.chmod(pending_path, 0o600)
except Exception:
pass
else:
os.unlink(temp_path)
temp_path = None
os.unlink(pending_path)
finally:
if temp_path and os.path.exists(temp_path):
os.unlink(temp_path)
return {"restored": restored, "remaining": len(unresolved), "output_path": target_path}
'''
main = replace_once(main, "def run_registration_common(count, log_callback, cancel_callback, accounts_output_file, observer):", retry_pending + "def run_registration_common(count, log_callback, cancel_callback, accounts_output_file, observer):", "retry pending helper")
main = replace_once(
main,
"def main():\n if len(sys.argv) > 1 and sys.argv[1].strip().lower() in (\"start\", \"cli\", \"--cli\"):\n",
"def main():\n if len(sys.argv) > 1 and sys.argv[1].strip().lower() == \"retry-pending\":\n if len(sys.argv) < 3:\n print(\"用法: python grok_register_ttk.py retry-pending <pending文件> [输出文件]\", file=sys.stderr)\n return\n try:\n summary = retry_pending_file(\n sys.argv[2],\n output_path=sys.argv[3] if len(sys.argv) > 3 else None,\n log_callback=cli_log,\n )\n cli_log(\n f\"[*] pending 恢复完成: 已恢复 {summary['restored']} | 剩余 {summary['remaining']} | 输出 {summary['output_path']}\"\n )\n except Exception as exc:\n log_exception(\"pending 恢复失败\", exc, cli_log)\n return\n if len(sys.argv) > 1 and sys.argv[1].strip().lower() in (\"start\", \"cli\", \"--cli\"):\n",
"retry pending dispatch",
)
ast.parse(main)
main_path.write_text(main, encoding="utf-8")
# 3) CPA browser/proxy bridge cleanup boundary.
browser_path = ROOT / "cpa_xai" / "browser_confirm.py"
browser = browser_path.read_text(encoding="utf-8")
browser = replace_once(
browser,
" except Exception:\n if browser is not None:\n close_standalone(browser)\n elif proxy_bridge is not None:\n try:\n proxy_bridge.stop()\n except Exception:\n pass\n raise\n",
" except Exception:\n if browser is not None:\n close_standalone(browser)\n if proxy_bridge is not None:\n try:\n proxy_bridge.stop()\n except Exception:\n pass\n raise\n",
"proxy bridge cleanup",
)
ast.parse(browser)
browser_path.write_text(browser, encoding="utf-8")
# 4) Remote token pool tests aligned with ETag/legacy policy.
remote_tests = r'''import unittest
from unittest.mock import patch
import grok_register_ttk as app
class DummyResponse:
def __init__(self, payload=None, status_code=200, reason="", headers=None, text=""):
self._payload = payload or {}
self.status_code = status_code
self.reason = reason
self.headers = headers or {}
self.text = text
def raise_for_status(self):
if self.status_code >= 400:
raise RuntimeError(f"HTTP Error {self.status_code}: {self.reason}")
def json(self):
return self._payload
class Grok2ApiRemotePoolTests(unittest.TestCase):
def setUp(self):
self.original_config = app.config.copy()
def tearDown(self):
app.config = self.original_config
def _configure(self, **overrides):
app.config.update({
"grok2api_remote_base": "https://grok.example.com",
"grok2api_remote_app_key": "app-secret",
"grok2api_pool_name": "ssoBasic",
"grok2api_allow_legacy_full_save": False,
**overrides,
})
def test_remote_pool_falls_back_to_admin_api_prefix_when_root_tokens_add_is_404(self):
self._configure()
calls = []
def fake_post(url, **kwargs):
calls.append((url, kwargs))
if url == "https://grok.example.com/tokens/add":
return DummyResponse(status_code=404)
return DummyResponse({"status": "success", "count": 1})
with patch.object(app, "http_post", side_effect=fake_post):
ok = app.add_token_to_grok2api_remote_pool("sso=abc123", email="a@example.com")
self.assertTrue(ok)
self.assertEqual([url for url, _ in calls], [
"https://grok.example.com/tokens/add",
"https://grok.example.com/admin/api/tokens/add",
])
self.assertEqual(calls[-1][1]["params"], {"app_key": "app-secret"})
self.assertEqual(calls[-1][1]["json"], {
"tokens": ["abc123"],
"pool": "basic",
"tags": ["auto-register"],
})
def test_remote_pool_does_not_duplicate_admin_api_prefix_when_base_already_points_to_admin_api(self):
self._configure(
grok2api_remote_base="https://grok.example.com/admin/api",
grok2api_pool_name="ssoSuper",
)
calls = []
def fake_post(url, **kwargs):
calls.append((url, kwargs))
return DummyResponse({"status": "success", "count": 1})
with patch.object(app, "http_post", side_effect=fake_post):
ok = app.add_token_to_grok2api_remote_pool("sso=super123", email="a@example.com")
self.assertTrue(ok)
self.assertEqual([url for url, _ in calls], [
"https://grok.example.com/admin/api/tokens/add",
])
self.assertEqual(calls[0][1]["json"]["pool"], "super")
def test_remote_pool_full_save_fallback_requires_opt_in_and_uses_etag(self):
self._configure(grok2api_allow_legacy_full_save=True)
get_calls = []
post_calls = []
def fake_post(url, **kwargs):
post_calls.append((url, kwargs))
if url.endswith("/tokens/add"):
return DummyResponse(status_code=404)
if url == "https://grok.example.com/admin/api/tokens":
return DummyResponse({"status": "success"})
return DummyResponse(status_code=404)
def fake_get(url, **kwargs):
get_calls.append((url, kwargs))
if url == "https://grok.example.com/admin/api/tokens":
return DummyResponse(
{"tokens": {"ssoBasic": []}},
headers={"ETag": '"version-7"'},
)
return DummyResponse(status_code=404)
with patch.object(app, "http_post", side_effect=fake_post), \
patch.object(app, "http_get", side_effect=fake_get):
ok = app.add_token_to_grok2api_remote_pool("sso=fallback123", email="a@example.com")
self.assertTrue(ok)
self.assertEqual([url for url, _ in get_calls], [
"https://grok.example.com/tokens",
"https://grok.example.com/admin/api/tokens",
])
self.assertEqual(post_calls[-1][0], "https://grok.example.com/admin/api/tokens")
self.assertEqual(post_calls[-1][1]["headers"]["If-Match"], '"version-7"')
self.assertEqual(post_calls[-1][1]["json"], {
"ssoBasic": [{"token": "fallback123", "tags": ["auto-register"], "note": "a@example.com"}],
})
def test_remote_pool_legacy_fallback_is_disabled_by_default(self):
self._configure()
with patch.object(app, "http_post", return_value=DummyResponse(status_code=404)), \
patch.object(app, "http_get") as get_mock:
with self.assertRaises(app.RemoteTokenCompatibilityError):
app.add_token_to_grok2api_remote_pool("abc")
get_mock.assert_not_called()
def test_remote_pool_500_does_not_fallback(self):
self._configure(grok2api_allow_legacy_full_save=True)
with patch.object(app, "http_post", return_value=DummyResponse(status_code=500, text="boom")), \
patch.object(app, "http_get") as get_mock:
with self.assertRaises(app.RemoteTokenRequestError):
app.add_token_to_grok2api_remote_pool("abc")
get_mock.assert_not_called()
def test_remote_pool_legacy_fallback_rejects_missing_etag(self):
self._configure(grok2api_allow_legacy_full_save=True)
def fake_post(url, **kwargs):
if url.endswith("/tokens/add"):
return DummyResponse(status_code=404)
return DummyResponse({"status": "success"})
def fake_get(url, **kwargs):
if url.endswith("/tokens"):
return DummyResponse({"tokens": {"ssoBasic": []}})
return DummyResponse(status_code=404)
with patch.object(app, "http_post", side_effect=fake_post), \
patch.object(app, "http_get", side_effect=fake_get):
with self.assertRaises(app.RemoteTokenCompatibilityError):
app.add_token_to_grok2api_remote_pool("abc")
if __name__ == "__main__":
unittest.main()
'''
remote_path = ROOT / "tests" / "test_grok2api_remote_pool.py"
ast.parse(remote_tests)
remote_path.write_text(remote_tests, encoding="utf-8")
# 5) Shared flow tests with fake operations only.
flow_tests = r'''import unittest
from registration_flow import (
RegistrationCallbacks,
RegistrationOperations,
run_batch,
)
class Cancelled(Exception):
pass
class Retryable(Exception):
pass
class FakeOps:
def __init__(self, save_ok=True, observer_events=None):
self.events = []
self.save_ok = save_ok
self.observer_events = observer_events if observer_events is not None else []
self.account_no = 0
def operations(self):
return RegistrationOperations(
start_browser=lambda: self.events.append("start"),
restart_browser=lambda: self.events.append("restart"),
browser_missing=lambda: False,
open_signup_page=lambda: self.events.append("open"),
fill_email_and_submit=self._email,
save_mail_credential=lambda email, token: True,
fill_code_and_submit=lambda email, token: "123456",
fill_profile_and_submit=lambda: {"given_name": "A", "family_name": "B", "password": "pw"},
wait_for_sso_cookie=lambda: "sso-token",
enable_nsfw=lambda sso: (True, "ok"),
persist_account_line=self._persist,
queue_unsaved_result=lambda payload, error: True,
add_tokens=lambda sso, email: {
"local": {"enabled": False, "ok": None, "error": None},
"remote": {"enabled": False, "ok": None, "error": None},
},
export_cpa=lambda email, password, sso: {"ok": False, "skipped": True},
cleanup=lambda reason: self.events.append(("cleanup", reason)),
sleep=lambda seconds: self.events.append(("sleep", seconds)),
cancelled_exception=Cancelled,
retry_exception=Retryable,
)
def _email(self):
self.account_no += 1
return f"user{self.account_no}@example.com", "mail-token"
def _persist(self, email, password, sso):
if not self.save_ok:
raise OSError("disk full")
self.events.append(("persist", email))
class RegistrationFlowTests(unittest.TestCase):
def callbacks(self, logs=None):
logs = logs if logs is not None else []
return RegistrationCallbacks(log=logs.append, cancelled=lambda: False)
def test_start_failure_still_runs_cleanup(self):
fake = FakeOps()
ops = fake.operations()
ops.start_browser = lambda: (_ for _ in ()).throw(RuntimeError("start failed"))
with self.assertRaises(RuntimeError):
run_batch(1, self.callbacks(), lambda *args: None, ops)
self.assertEqual(fake.events, [("cleanup", "任务结束")])
def test_last_account_does_not_restart_browser(self):
fake = FakeOps()
batch = run_batch(1, self.callbacks(), lambda *args: None, fake.operations())
self.assertEqual(batch.success_count, 1)
self.assertNotIn("restart", fake.events)
self.assertEqual(fake.events[-1], ("cleanup", "任务结束"))
def test_cleanup_interval_does_not_repeat_after_unsaved_result(self):
fake = FakeOps(save_ok=True)
ops = fake.operations()
original_persist = ops.persist_account_line
calls = {"count": 0}
def persist(email, password, sso):
calls["count"] += 1
if calls["count"] == 2:
raise OSError("disk full")
original_persist(email, password, sso)
ops.persist_account_line = persist
batch = run_batch(2, self.callbacks(), lambda *args: None, ops, cleanup_interval=1)
interval_cleanups = [event for event in fake.events if isinstance(event, tuple) and "已成功" in event[1]]
self.assertEqual(len(interval_cleanups), 1)
self.assertEqual(batch.success_count, 1)
self.assertEqual(batch.registered_unsaved_count, 1)
def test_observer_failure_is_logged_and_batch_continues(self):
fake = FakeOps()
logs = []
def broken_observer(*args):
raise RuntimeError("ui broke")
batch = run_batch(1, self.callbacks(logs), broken_observer, fake.operations())
self.assertEqual(batch.success_count, 1)
self.assertTrue(any("observer 执行失败" in line for line in logs))
if __name__ == "__main__":
unittest.main()
'''
flow_test_path = ROOT / "tests" / "test_registration_flow.py"
ast.parse(flow_tests)
flow_test_path.write_text(flow_tests, encoding="utf-8")
for path in (flow_path, main_path, browser_path, remote_path, flow_test_path):
ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
print("final flow fixes applied")
tools/fix_final_flow_test_generation.py
-46
View File
@@ -1,46 +0,0 @@
#!/usr/bin/env python3
from pathlib import Path
script = Path(__file__).resolve().with_name("apply_final_flow_fixes.py")
text = script.read_text(encoding="utf-8")
old_import = '''remote_tests = r\'''import unittest
from unittest.mock import patch
import grok_register_ttk as app
'''
new_import = '''remote_tests = r\'''import sys
import types
import unittest
from unittest.mock import patch
# Keep this unit test independent from optional browser/network dependencies.
drission = types.ModuleType("DrissionPage")
drission.Chromium = type("Chromium", (), {})
drission.ChromiumOptions = type("ChromiumOptions", (), {})
drission_errors = types.ModuleType("DrissionPage.errors")
drission_errors.PageDisconnectedError = type("PageDisconnectedError", (Exception,), {})
curl_cffi = types.ModuleType("curl_cffi")
curl_cffi.requests = types.SimpleNamespace()
sys.modules.setdefault("DrissionPage", drission)
sys.modules.setdefault("DrissionPage.errors", drission_errors)
sys.modules.setdefault("curl_cffi", curl_cffi)
import grok_register_ttk as app
'''
if old_import not in text:
if new_import not in text:
raise RuntimeError("remote test import block not found")
else:
text = text.replace(old_import, new_import, 1)
old_cleanup = ' interval_cleanups = [event for event in fake.events if isinstance(event, tuple) and "已成功" in event[1]]\n'
new_cleanup = ' interval_cleanups = [\n event for event in fake.events\n if isinstance(event, tuple)\n and len(event) > 1\n and isinstance(event[1], str)\n and "已成功" in event[1]\n ]\n'
if old_cleanup not in text:
if new_cleanup not in text:
raise RuntimeError("cleanup test expression not found")
else:
text = text.replace(old_cleanup, new_cleanup, 1)
script.write_text(text, encoding="utf-8")
print("final flow test generation repaired")