app/upstream/auth.py

"""anuma.ai AuthProvider：Privy JWT（authorization: Bearer）+ 平台头。

凭据（account/<name>.json）：
- ``privy_access_token``：Privy **identity token**（上游 portal.anuma.ai 的
  authorization: Bearer；~1h 过期，exp - iat = 3600s）
- ``privy_pat``：Privy **persisted access token**（``privy:pat``，仅刷新端点用；
  刷新时 Bearer 必须用 pat，identity token 会被拒）
- ``privy_app_id``：Privy 应用 ID（JWT aud，如 cmjrfihuc03h8l10ca0bi9o2y）
- ``privy_refresh_token``：Privy refresh token（一次性，刷新后轮换，自动写回）
- ``privy_cf_cookie``：浏览器 auth.privy.io 的 Cloudflare clearance cookie 串
  （httpOnly，从 DevTools Network 请求头复制；绑定本机 IP 与 Chrome UA，
  过期需重新抓取——见 README「自动刷新」节）
- ``conversation_id``：会话 UUID（x-conversation-id 头；首次生图自动生成并持久化，
  后续复用同一会话——对齐网页「同一会话连续生图」行为）

过期处理：JWT 到期前 ``token_refresh_margin`` 秒主动刷新
（``POST auth.privy.io/api/v1/sessions``，实测于 2026-08-01：
需带 origin/referer/user-agent/cf cookie + **Bearer pat**（identity token 会报
missing_or_invalid_token）；刷新返回新 pat + 新 id_token + 新 rt，
**新 rt 必须写回账号文件**，Privy 的 refresh token 一次性轮换）。
刷新失败时抛明确错误（不判死账号），提示重新登录更新凭据。
"""
from __future__ import annotations

import base64
import json
import time
from pathlib import Path
from typing import Any

import httpx

from app.account import FailReason
from app.upstream.base import AuthProvider

APP_VERSION = "41cf240"  # x-anuma-app-version（2026-08-06 从浏览器真实请求抓取；旧 7e5848e 已失效）

# 刷新请求必须复刻浏览器特征（cf_clearance 绑定 UA + IP）：
# UA 需与抓取 cf_clearance 时浏览器的 UA 一致（实测 Chrome 150）。
_CHROME_UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/150.0.0.0 Safari/537.36"
)
_PRIVY_REFRESH_URL = "https://auth.privy.io/api/v1/sessions"  # 实测端点（非 /token、非 /refresh）


def _decode_jwt_payload(token: str) -> dict[str, Any]:
    """解析 JWT payload（不验签，仅取过期时间等元数据）。"""
    try:
        payload = token.split(".")[1]
        payload += "=" * (-len(payload) % 4)
        return json.loads(base64.urlsafe_b64decode(payload))
    except Exception:  # noqa: BLE001
        return {}


class DefaultAuthProvider(AuthProvider):
    def __init__(self, account, settings, http_client, account_file: Path | None = None) -> None:
        self._account = account
        self._settings = settings
        self._http: httpx.AsyncClient = http_client
        self._account_file = account_file

    async def get_auth(self) -> dict[str, str]:
        # 2026-08-01 实测：portal.anuma.ai 用 Privy **identity token**（privy:id_token，
        # payload 含 linked_accounts 的 eyJjciI6... 开头 JWT）做 Bearer，不是 privy:pat。
        # 账号文件里 privy_access_token 存 identity token；旧 pat 放 privy_pat 字段备用。
        token = str(
            getattr(self._account, "privy_identity_token", "")
            or self._account.privy_access_token
            or ""
        )
        if not token:
            raise RuntimeError("account missing privy_access_token (identity token)")
        exp = _decode_jwt_payload(token).get("exp") or 0
        now = time.time()
        margin = float(getattr(self._settings, "token_refresh_margin", 300))
        if exp and exp - now <= margin:
            refreshed = await self._try_refresh()
            if refreshed:
                token = refreshed
            elif exp <= now:
                # token 已过期且无 refresh token 可换：抛明确错误而非打上游吃 401
                # （401 会被 classify_failure 判 AUTH_FAILED → 账号永久 disabled）。
                # RuntimeError 不命中账号级失效分类 → 原样抛出，账号保留可用状态，
                # 用户重新登录更新 privy_access_token 后即可恢复。
                raise RuntimeError(
                    f"privy access token expired (exp={int(exp)}); re-login and update "
                    "privy_access_token in account/<name>.json"
                )
        # 2026-08-08：生图时 client.py 已把 conversation_id 写回 account，带上同值头：
        # 只有显式存过/生成过才带（避免普通对话开垃圾会话）。
        conv_id = str(getattr(self._account, "conversation_id", "") or "")
        headers: dict[str, str] = {
            "authorization": f"Bearer {token}",
            # 2026-08-07 实测：portal.anuma.ai 走 Cloudflare，缺浏览器 UA（httpx 默认
            # python-httpx/x.x）一律 403 -> 判 auth_failed -> 账号阵亡。必须带 Chrome UA。
            "user-agent": _CHROME_UA,
            # 2026-08-07 实测：origin/referer 必须指向 chat.anuma.ai（同源校验），
            # 缺失时 portal 403 authorization_error。
            "origin": "https://chat.anuma.ai",
            "referer": "https://chat.anuma.ai/",
            "x-anuma-feature": "chat",
            "x-anuma-platform": "web",
            "x-anuma-app-version": APP_VERSION,
            "x-privacy-mode": "standard",
            "x-anuma-preprocessors-enabled": "true",
            # 生图时 client.py 已把 conversation_id 写回 account，带上同值头：
            # 只有显式存过/生成过才带（避免普通对话开垃圾会话）。
            **({"x-conversation-id": conv_id} if conv_id else {}),
        }
        # 2026-08-07 实测：portal.anuma.ai 走 Cloudflare，无有效 cf_clearance 一律 403
        # （11 个注册机号全 auth_failed 根因）。cf_cookie 绑定 UA+IP，同一浏览器
        # 同一 IP 导出的一份可多号共用；缺失时请求仍发（刷新端点会带，打 portal 不带会 403）。
        cf_cookie = str(getattr(self._account, "privy_cf_cookie", "") or "")
        if cf_cookie:
            headers["cookie"] = cf_cookie
        return headers

    async def _try_refresh(self) -> str | None:
        """用 Privy refresh token 换新 access token（实测端点 auth.privy.io/api/v1/sessions）。

        请求需复刻浏览器特征：origin/referer 指向 anuma、Chrome UA、cf_clearance cookie
        （绑定 UA+IP）、Bearer 旧 pat（过期也带上，服务器只校验会话归属）。
        响应返回新 pat + **新 rt**（一次性轮换）→ 两者都写回账号文件。
        """
        refresh = str(getattr(self._account, "privy_refresh_token", "") or "")
        app_id = str(getattr(self._account, "privy_app_id", "") or "")
        # 2026-08-01 实测：刷新端点 Bearer 必须用 privy:pat（persisted access token），
        # 传 identity token 会 400 missing_or_invalid_token → 自动刷新永远失败。
        pat = str(getattr(self._account, "privy_pat", "") or "")
        if not refresh or not app_id or not pat:
            return None
        headers: dict[str, str] = {
            "content-type": "application/json",
            "privy-app-id": app_id,
            "authorization": f"Bearer {pat}",
            "origin": "https://chat.anuma.ai",
            "referer": "https://chat.anuma.ai/",
            "user-agent": _CHROME_UA,
            "privy-client": "react-auth:3.14.1",
        }
        cf_cookie = str(getattr(self._account, "privy_cf_cookie", "") or "")
        if cf_cookie:
            headers["cookie"] = cf_cookie
        try:
            resp = await self._http.post(
                _PRIVY_REFRESH_URL, json={"refresh_token": refresh}, headers=headers,
            )
            if resp.status_code == 200:
                data = resp.json()
                # 2026-08-01 实测响应字段：privy_access_token（新 pat）、
                # identity_token（新 identity）、refresh_token（新 rt，一次性轮换）
                new_pat = data.get("privy_access_token") or ""
                new_id = data.get("identity_token") or ""
                if new_pat:
                    self._account.privy_pat = new_pat
                if new_id:
                    self._account.privy_identity_token = new_id
                new_refresh = data.get("refresh_token") or ""
                if new_refresh:
                    # Privy refresh token 一次性：必须持久化新 rt，否则下次刷新失败
                    self._account.privy_refresh_token = new_refresh
                if new_pat or new_id:
                    self._persist()
                    return new_id or new_pat
        except Exception:  # noqa: BLE001
            pass
        return None

    def _persist(self) -> None:
        """把内存 account 状态（含轮换后的 token）原子写回账号文件。失败不阻断请求。"""
        if self._account_file is None:
            return
        try:
            self._account_file.parent.mkdir(parents=True, exist_ok=True)
            tmp = self._account_file.with_suffix(".json.tmp")
            tmp.write_text(
                json.dumps(self._account.model_dump(mode="json"), ensure_ascii=False, indent=2),
                encoding="utf-8",
            )
            tmp.replace(self._account_file)
        except Exception:  # noqa: BLE001
            pass

    def is_auth_failure(self, exc: BaseException) -> bool:
        if isinstance(exc, httpx.HTTPStatusError):
            status = exc.response.status_code
            if status in (401, 403):
                body = exc.response.text.lower()
                # 2026-08-07 实测：model_tier_required（订阅限制）不是认证失败，
                # 不应罚账号冷却/停用。仅 authorization_error / invalid token 判认证失败。
                if "model_tier_required" in body:
                    return False
                # 403 也可能是配额/内容限制；仅当明确 authorization_error 时判认证失败
                if status == 401 or "authorization_error" in body or "invalid token" in body:
                    return True
        return False

    async def check_balance(self) -> int | None:
        """查账号可用额度；返回 None = 查不到（无 token/401/网络错/未配置 http）。

        2026-08-08：GET /api/v1/credits/balance 实测返回
        {available_credits, lifetime_credits, subscription_tier, ...}；
        token 失效 → 401 invalid or expired token。仅返回**起效**额度，供
        失败账号判定是否<=0 删除。超时 15s，失败不抛（None 即可删）。

        token 复用 ``get_auth()``（含自动刷新）——静态存储的 identity token 1h 即过期，
        直接读字段查余额几乎必失败；走 get_auth 才能拿到刷新后的有效 token。
        """
        try:
            if self._http is None:
                return -1  # 未配置 http client（测试环境）：不可判定，调用方留池
            # 优先用 get_auth（含刷新）拿有效 token；失败再回退静态字段。
            token = ""
            try:
                headers = await self.get_auth()
                auth_val = headers.get("authorization", "")
                if auth_val.lower().startswith("bearer "):
                    token = auth_val[7:]
            except Exception:  # noqa: BLE001
                token = ""
            if not token:
                token = str(
                    getattr(self._account, "privy_identity_token", "")
                    or self._account.privy_access_token or ""
                )
            if not token:
                return -1  # 无 token 不可判定，留池
            r = await self._http.get(
                "https://portal.anuma.ai/api/v1/credits/balance",
                headers={
                    "authorization": f"Bearer {token}",
                    "user-agent": _CHROME_UA,
                    "origin": "https://chat.anuma.ai",
                    "referer": "https://chat.anuma.ai/",
                    "x-anuma-feature": "chat",
                    "x-anuma-platform": "web",
                    "x-anuma-app-version": APP_VERSION,
                    "x-privacy-mode": "standard",
                    "x-anuma-preprocessors-enabled": "true",
                },
                timeout=15.0,
            )
            if r.status_code != 200:
                return None
            return int((r.json().get("available_credits") or 0))
        except Exception:  # noqa: BLE001
            return None

    def classify_failure(self, exc: BaseException) -> FailReason | None:
        if self.is_auth_failure(exc):
            return FailReason.AUTH_FAILED
        if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code == 429:
            return FailReason.QUOTA_EXHAUSTED
        return None