Przeglądaj źródła

Add recheck (rescue plus_check_failed accounts) and task_runner (single-worker queue)

- recheck.py: 用 backend-api/me 重查 plan,已 plus 自动重传 CPA
- task_runner.py: SQLite tasks 表的后台 worker,FIFO 串行消费,支持 cancel + retry
chendeben 2 miesięcy temu
rodzic
commit
c5338c895e
2 zmienionych plików z 500 dodań i 0 usunięć
  1. 234 0
      recheck.py
  2. 266 0
      task_runner.py

+ 234 - 0
recheck.py

@@ -0,0 +1,234 @@
+"""对历史失败账号做补救:用 access_token 调 backend-api 检查 plan,若已 plus 则上传 CPA。
+
+注意:/api/auth/session 走 NextAuth 只认 cookie,不能用 Bearer 直连。
+真正的鉴权接口是 /backend-api/me(也叫 accounts/check),它接受 Authorization: Bearer。
+"""
+from __future__ import annotations
+
+import copy
+import json
+import time
+import urllib.error
+import urllib.request
+from typing import Callable
+
+from cpa_uploader import (
+    get_session_plan_type,
+    is_plus_session,
+    upload_session_to_cpa,
+)
+from storage import add_event, get_account, upsert_account
+
+
+ME_URL = "https://chatgpt.com/backend-api/me"
+ACCOUNTS_CHECK_URL = "https://chatgpt.com/backend-api/accounts/check/v4-2023-04-27"
+
+
+def _http_get_json(url: str, *, headers: dict, timeout: int = 20, log: Callable[[str], None] = print) -> tuple[int, str, dict]:
+    started = time.time()
+    text = ""
+    status = 0
+    try:
+        from curl_cffi import requests as curl_requests  # type: ignore
+        r = curl_requests.get(url, headers=headers, impersonate="chrome136", timeout=timeout)
+        text = r.text
+        status = r.status_code
+    except Exception as exc:
+        log(f"[recheck] curl_cffi 不可用: {exc!r},回退 urllib")
+        req = urllib.request.Request(url, method="GET")
+        for k, v in headers.items():
+            req.add_header(k, v)
+        try:
+            with urllib.request.urlopen(req, timeout=timeout) as resp:
+                text = resp.read().decode("utf-8", errors="replace")
+                status = resp.status
+        except urllib.error.HTTPError as exc2:
+            text = exc2.read().decode("utf-8", errors="replace") if hasattr(exc2, "read") else ""
+            status = exc2.code
+
+    elapsed_ms = int((time.time() - started) * 1000)
+    log(f"[recheck] GET {url} HTTP {status} 耗时 {elapsed_ms}ms 长度={len(text)}")
+    parsed: dict = {}
+    try:
+        parsed = json.loads(text or "{}")
+    except Exception:
+        parsed = {}
+    return status, text, parsed
+
+
+def _extract_plan_from_me(payload: dict) -> tuple[str, str, str]:
+    """从 /backend-api/me 或 /accounts/check 响应里抽 (plan_type, account_id, email)。"""
+    if not isinstance(payload, dict):
+        return "", "", ""
+    # /me 里直接给 email 和 chat_id;plan 通常在 accounts.<id>.account.plan_type
+    email = ""
+    if isinstance(payload.get("email"), str):
+        email = payload["email"]
+
+    accounts = payload.get("accounts")
+    if isinstance(accounts, dict):
+        # 优先找 plan_type='plus'
+        best_id = ""
+        best_plan = ""
+        for acc_id, item in accounts.items():
+            if not isinstance(item, dict):
+                continue
+            account_obj = item.get("account") if isinstance(item.get("account"), dict) else item
+            plan = (account_obj or {}).get("plan_type") or (account_obj or {}).get("planType") or ""
+            plan = str(plan or "").strip()
+            if not plan:
+                continue
+            if plan.lower() == "plus" and not best_plan:
+                best_plan = plan
+                best_id = acc_id
+                break
+            if not best_plan:
+                best_plan = plan
+                best_id = acc_id
+        if best_plan:
+            return best_plan, best_id, email
+
+    # 兜底:顶层 plan_type
+    plan = str(payload.get("plan_type") or payload.get("planType") or "").strip()
+    return plan, "", email
+
+
+def _merge_plus_into_session(base_session: dict, plan: str, account_id: str, email: str) -> dict:
+    """把 backend-api 拉到的 plan/account_id 合并进 DB 里的 session,构造给 CPA 用的 plus session。"""
+    sess = copy.deepcopy(base_session) if isinstance(base_session, dict) else {}
+    if "account" not in sess or not isinstance(sess.get("account"), dict):
+        sess["account"] = {}
+    if plan:
+        sess["account"]["planType"] = plan
+        sess["planType"] = plan
+    if account_id and not sess["account"].get("id"):
+        sess["account"]["id"] = account_id
+    if email:
+        if "user" not in sess or not isinstance(sess.get("user"), dict):
+            sess["user"] = {}
+        if not sess["user"].get("email"):
+            sess["user"]["email"] = email
+    return sess
+
+
+def recheck_account(
+    email: str,
+    *,
+    cpa_url: str,
+    cpa_management_key: str,
+    log: Callable[[str], None] = print,
+) -> dict:
+    """对单个账号做补救。返回 { ok, planType, action, cpa, error }。"""
+    acc = get_account(email)
+    if not acc:
+        return {"ok": False, "error": "账号不存在"}
+
+    plus_sess = acc.get("plus_session") or {}
+    init_sess = acc.get("initial_session") or {}
+    access_token = (plus_sess.get("accessToken") if isinstance(plus_sess, dict) else None) \
+                   or (init_sess.get("accessToken") if isinstance(init_sess, dict) else None)
+    if not access_token:
+        return {"ok": False, "error": "数据库里找不到 accessToken"}
+
+    log(f"[recheck] {email} 开始补救(访问 backend-api/me)")
+    add_event(email, "recheck", "info", "begin")
+
+    headers = {
+        "Authorization": f"Bearer {access_token}",
+        "Accept": "application/json",
+        "Origin": "https://chatgpt.com",
+        "Referer": "https://chatgpt.com/",
+    }
+
+    plan = ""
+    account_id = ""
+    me_email = ""
+    last_payload: dict = {}
+
+    # 先试 /backend-api/me(轻量),失败回退 /accounts/check
+    for url in (ME_URL, ACCOUNTS_CHECK_URL):
+        try:
+            status, text, payload = _http_get_json(url, headers=headers, log=log)
+        except Exception as exc:
+            log(f"[recheck] {url} 异常: {exc!r}")
+            continue
+        if status >= 400:
+            log(f"[recheck] {url} 返回 {status},预览={text[:200]}")
+            continue
+        last_payload = payload
+        plan, account_id, me_email = _extract_plan_from_me(payload)
+        log(f"[recheck] {url} 提取 plan={plan!r} account_id={account_id!r} email={me_email!r}")
+        if plan:
+            break
+
+    if not plan:
+        msg = "重拉 session 失败:所有 backend-api 接口都未返回 plan_type"
+        log(f"[recheck] {email} {msg}")
+        upsert_account(email, acc.get("password") or "", fields={"last_error": msg})
+        add_event(email, "recheck", "error", msg)
+        return {"ok": False, "error": msg}
+
+    # 用 base session(plus_session 优先,否则 init_session)合并新 plan
+    base_for_merge = plus_sess if isinstance(plus_sess, dict) and plus_sess else init_sess
+    merged = _merge_plus_into_session(base_for_merge or {}, plan, account_id, me_email or email)
+    # accessToken 用我们手头那个(merged 不一定包含)
+    if not merged.get("accessToken"):
+        merged["accessToken"] = access_token
+
+    upsert_account(email, acc.get("password") or "", fields={
+        "plan_type": plan,
+        "plus_session": merged,
+    })
+
+    if not is_plus_session(merged):
+        upsert_account(email, acc.get("password") or "", fields={
+            "final_status": "plus_check_failed",
+            "last_error": f"补救后 planType 仍为 {plan!r}",
+        })
+        add_event(email, "recheck", "warn", f"plan={plan}")
+        return {"ok": False, "planType": plan, "action": "still_not_plus"}
+
+    upsert_account(email, acc.get("password") or "", fields={
+        "final_status": "plus", "last_error": ""
+    })
+    add_event(email, "plus_check", "ok", f"plan={plan} (recheck)")
+
+    if not (cpa_url and cpa_management_key):
+        upsert_account(email, acc.get("password") or "", fields={
+            "final_status": "cpa_skipped"
+        })
+        add_event(email, "cpa", "warn", "未配置 CPA")
+        return {"ok": True, "planType": plan, "action": "plus_no_cpa_config"}
+
+    try:
+        cpa_result = upload_session_to_cpa(
+            merged,
+            cpa_url=cpa_url,
+            management_key=cpa_management_key,
+            email_hint=email,
+            log=log,
+        )
+    except Exception as exc:
+        msg = f"CPA 上传失败: {exc}"
+        log(f"[recheck] {email} {msg}")
+        upsert_account(email, acc.get("password") or "", fields={
+            "final_status": "cpa_failed",
+            "last_error": msg,
+        })
+        add_event(email, "cpa", "error", msg)
+        return {"ok": False, "planType": plan, "action": "cpa_failed", "error": msg}
+
+    upsert_account(email, acc.get("password") or "", fields={
+        "final_status": "cpa_uploaded",
+        "cpa_file_name": cpa_result.get("fileName"),
+        "cpa_uploaded_at": int(time.time() * 1000),
+        "last_error": "",
+    })
+    add_event(email, "cpa", "ok", cpa_result.get("fileName"), payload=cpa_result)
+    return {
+        "ok": True,
+        "planType": plan,
+        "action": "cpa_uploaded",
+        "cpa": cpa_result,
+    }
+

+ 266 - 0
task_runner.py

@@ -0,0 +1,266 @@
+"""任务执行器:单 worker 后台线程,从 SQLite tasks 表 FIFO 消费任务。
+
+支持两种模式:
+  - full      : 走完整的注册→付款→上传 CPA 流程(cfg.account_count 强制为 1)
+  - pay_only  : 传入已有 ChatGPT session JSON,跳过注册直接付款
+
+任务整体失败会重试,最多 attempts = max_attempts 次(默认 3)。
+"""
+from __future__ import annotations
+
+import threading
+import time
+import traceback
+import uuid
+from typing import Callable, Optional
+
+from chatgpt_flow import run_full, run_pay_only
+from config import AppConfig
+from storage import (
+    get_next_queued_task,
+    get_task,
+    init_db,
+    update_task,
+)
+
+
+class TaskRunner:
+    def __init__(self, *, log: Callable[[str], None] = print):
+        self.log = log
+        self._thread: threading.Thread | None = None
+        self._stopping = threading.Event()
+        self._current_task_id: str | None = None
+        self._cancel_flags: dict[str, bool] = {}
+        self._lock = threading.Lock()
+
+    # ----- public API -----
+
+    def start(self):
+        if self._thread and self._thread.is_alive():
+            return
+        self._stopping.clear()
+        self._thread = threading.Thread(target=self._loop, daemon=True, name="task-runner")
+        self._thread.start()
+        self.log("[runner] worker 已启动")
+
+    def shutdown(self, timeout: float = 5.0):
+        self._stopping.set()
+        if self._thread:
+            self._thread.join(timeout=timeout)
+        self.log("[runner] worker 已停止")
+
+    def cancel(self, task_id: str) -> bool:
+        with self._lock:
+            self._cancel_flags[task_id] = True
+        # 若是当前正在跑的 → 让 stop_check 抛出
+        return True
+
+    def current_task_id(self) -> str | None:
+        return self._current_task_id
+
+    # ----- internals -----
+
+    def _loop(self):
+        init_db()
+        while not self._stopping.is_set():
+            try:
+                task = get_next_queued_task()
+            except Exception as exc:
+                self.log(f"[runner] 取任务异常: {exc!r}")
+                task = None
+
+            if not task:
+                # 没活干,sleep 1s 再轮询
+                self._stopping.wait(1.0)
+                continue
+
+            try:
+                self._run_task(task)
+            except Exception as exc:
+                self.log(f"[runner] 执行任务 {task.get('task_id')} 顶层异常: {exc!r}")
+                self.log(traceback.format_exc())
+
+    def _run_task(self, task: dict):
+        task_id = task["task_id"]
+        mode = task.get("mode") or "full"
+        max_attempts = int(task.get("max_attempts") or 3)
+        params = task.get("params") or {}
+
+        self._current_task_id = task_id
+        update_task(task_id, {"status": "running", "started_at": int(time.time() * 1000), "stage": "queued→running"})
+
+        last_error = ""
+        for attempt in range(1, max_attempts + 1):
+            if self._is_cancelled(task_id):
+                update_task(task_id, {"status": "cancelled", "stage": "cancelled", "finished_at": int(time.time() * 1000)})
+                self._current_task_id = None
+                return
+
+            update_task(task_id, {"attempts": attempt, "stage": f"attempt {attempt}/{max_attempts}"})
+            self.log(f"[runner] task={task_id} mode={mode} 第 {attempt}/{max_attempts} 次尝试")
+
+            try:
+                if mode == "full":
+                    record = self._do_full(task_id, params)
+                elif mode == "pay_only":
+                    record = self._do_pay_only(task_id, params)
+                else:
+                    raise RuntimeError(f"未知 mode: {mode}")
+            except _TaskCancelled:
+                update_task(task_id, {
+                    "status": "cancelled", "stage": "cancelled",
+                    "finished_at": int(time.time() * 1000),
+                })
+                self._current_task_id = None
+                return
+            except Exception as exc:
+                last_error = repr(exc)
+                self.log(f"[runner] task={task_id} 第 {attempt} 次执行异常: {last_error}")
+                self.log(traceback.format_exc())
+                update_task(task_id, {
+                    "last_error": last_error,
+                    "stage": f"attempt {attempt} failed",
+                })
+                if attempt >= max_attempts:
+                    break
+                # sleep 5s 让外部资源喘口气
+                time.sleep(5)
+                continue
+
+            # 看 record["stage"] 决定是成功还是失败
+            stage = (record or {}).get("stage", "")
+            if stage in ("cpa_uploaded", "cpa_skipped"):
+                update_task(task_id, {
+                    "status": "success",
+                    "stage": stage,
+                    "result": record,
+                    "email": record.get("email") or task.get("email"),
+                    "plan_type": record.get("planType"),
+                    "cpa_file_name": (record.get("cpa") or {}).get("fileName") if record.get("cpa") else None,
+                    "finished_at": int(time.time() * 1000),
+                    "last_error": "",
+                })
+                self.log(f"[runner] task={task_id} 成功 stage={stage}")
+                self._current_task_id = None
+                return
+            else:
+                # plus_check_failed / error 之类视为失败,进入重试
+                last_error = (record or {}).get("error") or f"stage={stage}"
+                self.log(f"[runner] task={task_id} 第 {attempt} 次完成但未成功: {last_error}")
+                update_task(task_id, {"last_error": last_error, "stage": f"attempt {attempt}: {stage}"})
+                if attempt >= max_attempts:
+                    break
+                time.sleep(5)
+
+        # 走到这说明全部 attempt 都失败
+        update_task(task_id, {
+            "status": "failed",
+            "stage": "exhausted",
+            "last_error": last_error or "all attempts failed",
+            "finished_at": int(time.time() * 1000),
+        })
+        self.log(f"[runner] task={task_id} 失败({max_attempts}/{max_attempts} 次都失败)")
+        self._current_task_id = None
+
+    def _do_full(self, task_id: str, params: dict) -> dict:
+        cfg = AppConfig.load()
+        cfg.account_count = 1  # 单任务只跑一个账号
+        # params 可覆盖 cfg
+        for k in ("headless", "use_promo", "phone_e164", "sms_api_url", "cpa_url",
+                  "cpa_management_key", "proxy_url", "paypal_only_proxy",
+                  "mail_helper_url", "mail_domain"):
+            if k in (params or {}):
+                setattr(cfg, k, params[k])
+
+        def stop_check():
+            if self._is_cancelled(task_id):
+                raise _TaskCancelled()
+
+        def task_log(msg: str):
+            self.log(f"[task:{task_id[:8]}] {msg}")
+
+        def task_stage(stage: str):
+            update_task(task_id, {"stage": stage[:200]})
+
+        # run_full 是为多账号写的,这里复用但只跑 1 个
+        # 把 stop hook 注入:run_full 内部用 full_ctx.check_stop,我们没法直接钩,
+        # 但 Stop API 通过 cancel 设标志位 → 在每次 update_task 时也 stop_check
+        # 简化:包一层 thread 跑,10s 内查一次 cancel
+        stopping_holder = {"stop": False}
+        result_holder: dict = {}
+
+        def run():
+            try:
+                full_ctx = run_full(cfg, log=task_log, on_stage=task_stage)
+                # full_ctx.accounts[0] 就是结果
+                if full_ctx.accounts:
+                    result_holder["record"] = full_ctx.accounts[0]
+                else:
+                    result_holder["record"] = {"stage": "error", "error": "no account record"}
+            except Exception as exc:
+                result_holder["error"] = repr(exc)
+
+        # 因为 run_full 内部在阻塞调用 sync_playwright,cancel 只能等当前 attempt 跑完
+        run()
+        if "error" in result_holder:
+            raise RuntimeError(result_holder["error"])
+        if self._is_cancelled(task_id):
+            raise _TaskCancelled()
+        return result_holder.get("record") or {"stage": "error", "error": "empty record"}
+
+    def _do_pay_only(self, task_id: str, params: dict) -> dict:
+        cfg = AppConfig.load()
+        cfg.account_count = 1
+        for k in ("headless", "use_promo", "phone_e164", "sms_api_url", "cpa_url",
+                  "cpa_management_key", "proxy_url", "paypal_only_proxy"):
+            if k in (params or {}):
+                setattr(cfg, k, params[k])
+
+        session = params.get("session")
+        if not isinstance(session, dict) or not session.get("accessToken"):
+            raise RuntimeError("pay_only 需要 params.session 是 JSON 且包含 accessToken")
+
+        def stop_check():
+            if self._is_cancelled(task_id):
+                raise _TaskCancelled()
+
+        def task_log(msg: str):
+            self.log(f"[task:{task_id[:8]}] {msg}")
+
+        def task_stage(stage: str):
+            update_task(task_id, {"stage": stage[:200]})
+
+        return run_pay_only(
+            cfg,
+            session=session,
+            log=task_log,
+            on_stage=task_stage,
+            stop_check=stop_check,
+        )
+
+    def _is_cancelled(self, task_id: str) -> bool:
+        with self._lock:
+            return self._cancel_flags.get(task_id, False)
+
+
+class _TaskCancelled(Exception):
+    pass
+
+
+# ------------------ public helpers ------------------
+
+def make_task_id() -> str:
+    return f"t-{int(time.time())}-{uuid.uuid4().hex[:8]}"
+
+
+_GLOBAL_RUNNER: TaskRunner | None = None
+_GLOBAL_RUNNER_LOCK = threading.Lock()
+
+
+def get_runner(log: Callable[[str], None] = print) -> TaskRunner:
+    global _GLOBAL_RUNNER
+    with _GLOBAL_RUNNER_LOCK:
+        if _GLOBAL_RUNNER is None:
+            _GLOBAL_RUNNER = TaskRunner(log=log)
+            _GLOBAL_RUNNER.start()
+        return _GLOBAL_RUNNER