Prefer SSO OAuth PKCE for CPA minting and sync latest accounts.

Switch CPA export to SSO→Authorization Code with referrer=grok-build, keep device-code as optional fallback, and capture new register/auth artifacts.

oidc_mint/oauth_code.py

@@ -0,0 +1,434 @@
+"""xAI OAuth Authorization Code + PKCE (SSO cookie → CPA token).
+
+对齐最新可用流程：authorize / consent 必须带 referrer=grok-build，
+否则 access_token JWT 缺少 referrer 字段，cli-chat-proxy / grok-build 不可用。
+
+参考实现：sso -> oauth2/authorize(referrer=grok-build) -> consent allow
+        -> oauth2/token (authorization_code + PKCE)
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import json
+import re
+import secrets
+import time
+from dataclasses import dataclass
+from typing import Any, Callable
+from urllib.parse import parse_qs, urlencode, urljoin, urlparse
+
+from .proxyutil import resolve_proxy
+
+CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828"
+ISSUER = "https://auth.x.ai"
+TOKEN_URL = f"{ISSUER}/oauth2/token"
+AUTHORIZE_URL = f"{ISSUER}/oauth2/authorize"
+REDIRECT_URI = "http://127.0.0.1:56121/callback"
+# 比旧 device-code scope 多 conversations:*，对齐 grok-build
+SCOPE = (
+    "openid profile email offline_access "
+    "grok-cli:access api:access conversations:read conversations:write"
+)
+GROK_REFERRER = "grok-build"
+GROK_VERSION = "0.2.93"
+GROK_TOKEN_UA = (
+    f"grok-pager/{GROK_VERSION} grok-shell/{GROK_VERSION} (linux; x86_64)"
+)
+# Next.js Server Action id（consent 页 POST 需要）
+NEXT_ACTION_ID = "4005315a1d7e426de592990bb54bb37471f39dd6d2"
+BROWSER_UA = (
+    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
+    "(KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36"
+)
+
+LogFn = Callable[[str], None]
+
+
+def _noop_log(_: str) -> None:
+    return None
+
+
+class OAuthCodeError(RuntimeError):
+    pass
+
+
+@dataclass
+class AuthCodeFlow:
+    state: str
+    nonce: str
+    code_verifier: str
+    code_challenge: str
+
+
+@dataclass
+class TokenResult:
+    access_token: str
+    refresh_token: str
+    id_token: str | None
+    token_type: str
+    expires_in: int
+    raw: dict[str, Any]
+    referrer: str = ""
+
+
+def _b64url(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
+
+
+def new_auth_code_flow() -> AuthCodeFlow:
+    verifier = _b64url(secrets.token_bytes(32))
+    state = _b64url(secrets.token_bytes(16))
+    nonce = _b64url(secrets.token_bytes(16))
+    challenge = _b64url(hashlib.sha256(verifier.encode("ascii")).digest())
+    return AuthCodeFlow(
+        state=state,
+        nonce=nonce,
+        code_verifier=verifier,
+        code_challenge=challenge,
+    )
+
+
+def normalize_sso_cookie(raw: str) -> str:
+    token = (raw or "").strip()
+    if token.lower().startswith("sso="):
+        token = token[4:].strip()
+    return token
+
+
+def jwt_payload(token: str) -> dict[str, Any]:
+    parts = (token or "").split(".")
+    if len(parts) < 2:
+        raise ValueError("invalid JWT")
+    seg = parts[1]
+    seg += "=" * (-len(seg) % 4)
+    return json.loads(base64.urlsafe_b64decode(seg.encode("ascii")))
+
+
+def _short(value: str, limit: int = 240) -> str:
+    value = value or ""
+    return value if len(value) <= limit else value[:limit]
+
+
+def _make_session(proxy: str | None = None):
+    """优先 curl_cffi（Chrome TLS），否则回退标准 requests。"""
+    resolved = resolve_proxy(proxy)
+    proxies = {"http": resolved, "https": resolved} if resolved else None
+
+    try:
+        from curl_cffi import requests as crequests
+
+        for impersonate in ("chrome131", "chrome124", "chrome120"):
+            try:
+                return crequests.Session(impersonate=impersonate, proxies=proxies)
+            except Exception:
+                continue
+        return crequests.Session(proxies=proxies)
+    except ImportError:
+        pass
+
+    try:
+        import requests as std_requests
+    except ImportError as e:  # pragma: no cover
+        raise OAuthCodeError(
+            "需要 curl_cffi 或 requests 才能执行 SSO→OAuth 转换"
+        ) from e
+    sess = std_requests.Session()
+    if proxies:
+        sess.proxies.update(proxies)
+    return sess
+
+
+def _set_sso_cookies(session: Any, sso: str) -> None:
+    sso = normalize_sso_cookie(sso)
+    if not sso:
+        raise OAuthCodeError("sso cookie 为空")
+    # curl_cffi / requests cookie jar
+    for domain in ("accounts.x.ai", "auth.x.ai", ".x.ai"):
+        for name in ("sso", "sso-rw"):
+            try:
+                session.cookies.set(name, sso, domain=domain, path="/")
+            except Exception:
+                try:
+                    session.cookies.set(name, sso)
+                except Exception:
+                    pass
+
+
+def _browser_headers(method: str, url: str, next_action: str = "") -> dict[str, str]:
+    headers = {
+        "User-Agent": BROWSER_UA,
+        "Sec-CH-UA": '"Not(A:Brand";v="99", "Google Chrome";v="133", "Chromium";v="133"',
+        "Sec-CH-UA-Mobile": "?0",
+        "Sec-CH-UA-Platform": '"Linux"',
+        "Accept-Language": "en-US,en;q=0.9",
+    }
+    if method.upper() == "POST":
+        headers.update(
+            {
+                "Accept": "text/x-component",
+                "Content-Type": "text/plain;charset=UTF-8",
+                "Origin": "https://accounts.x.ai",
+                "Referer": url,
+                "Sec-Fetch-Site": "same-origin",
+                "Sec-Fetch-Mode": "cors",
+                "Sec-Fetch-Dest": "empty",
+            }
+        )
+        if next_action:
+            headers["Next-Action"] = next_action
+    else:
+        headers.update(
+            {
+                "Accept": (
+                    "text/html,application/xhtml+xml,application/xml;q=0.9,"
+                    "application/json;q=0.8,*/*;q=0.7"
+                ),
+                "Upgrade-Insecure-Requests": "1",
+                "Sec-Fetch-Site": "none",
+                "Sec-Fetch-Mode": "navigate",
+                "Sec-Fetch-Dest": "document",
+            }
+        )
+    return headers
+
+
+def _token_headers() -> dict[str, str]:
+    return {
+        "User-Agent": GROK_TOKEN_UA,
+        "Accept": "*/*",
+        "X-Grok-Client-Version": GROK_VERSION,
+        "Content-Type": "application/x-www-form-urlencoded",
+    }
+
+
+def _final_url(resp: Any) -> str:
+    try:
+        return str(getattr(resp, "url", "") or "")
+    except Exception:
+        return ""
+
+
+def open_authorize_page(session: Any, flow: AuthCodeFlow) -> str:
+    params = {
+        "response_type": "code",
+        "client_id": CLIENT_ID,
+        "redirect_uri": REDIRECT_URI,
+        "scope": SCOPE,
+        "code_challenge": flow.code_challenge,
+        "code_challenge_method": "S256",
+        "state": flow.state,
+        "nonce": flow.nonce,
+        "referrer": GROK_REFERRER,
+    }
+    url = f"{AUTHORIZE_URL}?{urlencode(params)}"
+    resp = session.get(
+        url,
+        headers=_browser_headers("GET", url),
+        allow_redirects=True,
+        timeout=30,
+    )
+    body = resp.text or ""
+    final = _final_url(resp)
+    if resp.status_code < 200 or resp.status_code >= 300:
+        raise OAuthCodeError(
+            f"authorize HTTP {resp.status_code}: {_short(body)}"
+        )
+    if "sign-in" in final or "sign-up" in final:
+        raise OAuthCodeError("sso 无效（authorize 跳转登录页）")
+    if "/oauth2/consent" not in final:
+        # 少数情况 consent 在 body 的 redirect 里
+        m = re.search(r'https?://[^"\']+/oauth2/consent[^"\']*', body)
+        if m:
+            final = m.group(0)
+        else:
+            raise OAuthCodeError(f"authorize 未进入 consent: {final or _short(body)}")
+    return final
+
+
+def parse_consent_code(body: str) -> str:
+    """从 Next.js RSC / text-x-component 响应中解析 code。"""
+    text = body or ""
+    # 1) 逐行 JSON（Go 实现路径）
+    for line in text.splitlines():
+        idx = line.find("{")
+        if idx < 0:
+            continue
+        try:
+            obj = json.loads(line[idx:])
+        except Exception:
+            continue
+        if isinstance(obj, dict) and obj.get("code"):
+            if obj.get("success") is False:
+                raise OAuthCodeError(
+                    f"consent 失败: {obj.get('error') or obj.get('action')}"
+                )
+            return str(obj["code"]).strip()
+        if isinstance(obj, list):
+            for item in obj:
+                if isinstance(item, dict) and item.get("code"):
+                    return str(item["code"]).strip()
+
+    # 2) 宽松正则
+    m = re.search(r'"code"\s*:\s*"([A-Za-z0-9._~\-]+)"', text)
+    if m:
+        return m.group(1)
+
+    # 3) redirect 里带 code=
+    m = re.search(r"[?&]code=([A-Za-z0-9._~\-]+)", text)
+    if m:
+        return m.group(1)
+
+    raise OAuthCodeError(f"consent 响应缺少 code: {_short(text, 300)}")
+
+
+def approve_authorization(session: Any, consent_url: str, flow: AuthCodeFlow) -> str:
+    payload = [
+        {
+            "action": "allow",
+            "clientId": CLIENT_ID,
+            "redirectUri": REDIRECT_URI,
+            "scope": SCOPE,
+            "state": flow.state,
+            "codeChallenge": flow.code_challenge,
+            "codeChallengeMethod": "S256",
+            "nonce": flow.nonce,
+            "principalType": "User",
+            "principalId": "",
+            "referrer": GROK_REFERRER,
+        }
+    ]
+    body = json.dumps(payload, separators=(",", ":"))
+    resp = session.post(
+        consent_url,
+        data=body.encode("utf-8"),
+        headers=_browser_headers("POST", consent_url, NEXT_ACTION_ID),
+        allow_redirects=True,
+        timeout=30,
+    )
+    text = resp.text or ""
+    if resp.status_code < 200 or resp.status_code >= 300:
+        raise OAuthCodeError(f"consent HTTP {resp.status_code}: {_short(text, 300)}")
+
+    # 有时 302 到 redirect_uri?code=
+    final = _final_url(resp)
+    if "code=" in final:
+        qs = parse_qs(urlparse(final).query)
+        code = (qs.get("code") or [""])[0]
+        if code:
+            return code
+
+    return parse_consent_code(text)
+
+
+def exchange_auth_code(session: Any, code: str, flow: AuthCodeFlow) -> TokenResult:
+    form = {
+        "grant_type": "authorization_code",
+        "code": code,
+        "redirect_uri": REDIRECT_URI,
+        "client_id": CLIENT_ID,
+        "code_verifier": flow.code_verifier,
+    }
+    resp = session.post(
+        TOKEN_URL,
+        data=urlencode(form),
+        headers=_token_headers(),
+        timeout=30,
+    )
+    text = resp.text or ""
+    if resp.status_code < 200 or resp.status_code >= 300:
+        raise OAuthCodeError(f"token HTTP {resp.status_code}: {_short(text, 300)}")
+    try:
+        data = resp.json()
+    except Exception as e:
+        raise OAuthCodeError(f"token 响应非 JSON: {_short(text)}") from e
+    if not isinstance(data, dict) or not data.get("access_token"):
+        raise OAuthCodeError(f"token 响应缺少 access_token: {data!r}")
+
+    access = str(data["access_token"]).strip()
+    refresh = str(data.get("refresh_token") or "").strip()
+    if not refresh:
+        raise OAuthCodeError("token 响应缺少 refresh_token")
+
+    referrer = ""
+    try:
+        pl = jwt_payload(access)
+        referrer = str(pl.get("referrer") or "")
+    except Exception:
+        pl = {}
+
+    expires_in = int(data.get("expires_in") or 21600)
+    return TokenResult(
+        access_token=access,
+        refresh_token=refresh,
+        id_token=(str(data["id_token"]).strip() if data.get("id_token") else None),
+        token_type=str(data.get("token_type") or "Bearer"),
+        expires_in=expires_in,
+        raw=data,
+        referrer=referrer,
+    )
+
+
+def sso_to_token(
+    sso_cookie: str,
+    *,
+    proxy: str | None = None,
+    log: LogFn | None = None,
+    require_referrer: bool = True,
+) -> TokenResult:
+    """SSO cookie → 带 referrer=grok-build 的 OAuth token。"""
+    log = log or _noop_log
+    sso = normalize_sso_cookie(sso_cookie)
+    if not sso:
+        raise OAuthCodeError("sso cookie 为空")
+
+    flow = new_auth_code_flow()
+    session = _make_session(proxy)
+    try:
+        _set_sso_cookies(session, sso)
+        log(f"Authorization Code Flow referrer={GROK_REFERRER}")
+        consent_url = open_authorize_page(session, flow)
+        log(f"authorize -> consent: {_short(consent_url, 120)}")
+        code = approve_authorization(session, consent_url, flow)
+        log("consent allow ok")
+        token = exchange_auth_code(session, code, flow)
+        if token.referrer != GROK_REFERRER:
+            msg = f"access_token 未包含预期 referrer（got={token.referrer!r}）"
+            if require_referrer:
+                raise OAuthCodeError(msg)
+            log(f"WARN {msg}")
+        else:
+            log("access_token referrer=grok-build ok")
+        log(f"token ok expires_in={token.expires_in} refresh=yes")
+        return token
+    finally:
+        try:
+            session.close()
+        except Exception:
+            pass
+
+
+def mint_from_sso(
+    sso_cookie: str,
+    *,
+    proxy: str | None = None,
+    log: LogFn | None = None,
+    require_referrer: bool = True,
+) -> dict[str, Any]:
+    """上层统一返回 dict，兼容 cpa_export。"""
+    tr = sso_to_token(
+        sso_cookie,
+        proxy=proxy,
+        log=log,
+        require_referrer=require_referrer,
+    )
+    return {
+        "access_token": tr.access_token,
+        "refresh_token": tr.refresh_token,
+        "id_token": tr.id_token,
+        "token_type": tr.token_type,
+        "expires_in": tr.expires_in,
+        "referrer": tr.referrer,
+        "sso": normalize_sso_cookie(sso_cookie),
+    }