110 lines
3.8 KiB
Python
110 lines
3.8 KiB
Python
"""Mint the short-lived JWT that the API needs, from a session cookie.
|
||
|
||
The web session cookie (``ory_kratos_session``, ~30 days) authenticates
|
||
``GET /v2/studio/studio/whoami?tokenizeAs=default_jwt`` over plain HTTP and
|
||
returns a ``tokenized`` JWT valid for ~15 minutes. Verified live: a cookie alone
|
||
is enough — no browser, no device id.
|
||
|
||
The endpoint occasionally answers a transient ``401`` under rapid repeated calls,
|
||
so minting retries a few times with a short backoff before giving up.
|
||
"""
|
||
|
||
import base64
|
||
import json
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
|
||
from .errors import TripoAPIError
|
||
|
||
ORIGIN = "https://studio.tripo3d.ai"
|
||
COOKIE_NAME = "ory_kratos_session"
|
||
WHOAMI = "/v2/studio/studio/whoami?tokenizeAs=default_jwt"
|
||
|
||
RETRY_ATTEMPTS = 3
|
||
RETRY_BACKOFF = 1.5
|
||
|
||
|
||
def _request(api_base, cookie, device_id=None, timeout=30):
|
||
headers = {
|
||
"Accept": "application/json",
|
||
"Origin": ORIGIN,
|
||
"Referer": ORIGIN + "/",
|
||
"User-Agent": "tripocli/2.1.0",
|
||
"Cookie": cookie,
|
||
}
|
||
if device_id:
|
||
headers["x-tripo-device-id"] = device_id
|
||
|
||
request = urllib.request.Request(f"{api_base}{WHOAMI}", headers=headers, method="GET")
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||
return json.loads(response.read().decode("utf-8"))
|
||
except urllib.error.HTTPError as error:
|
||
body = error.read().decode("utf-8", "replace")
|
||
raise TripoAPIError(
|
||
f"会话换取 JWT 失败 (HTTP {error.code}): {body[:200]}", code=error.code
|
||
) from None
|
||
except urllib.error.URLError as error:
|
||
raise TripoAPIError(f"网络错误: {error.reason}") from None
|
||
|
||
|
||
def decode_jwt_payload(token):
|
||
"""Return the JWT payload, or ``{}`` when it cannot be decoded."""
|
||
try:
|
||
parts = (token or "").split(".")
|
||
if len(parts) >= 2:
|
||
payload = parts[1] + "=" * (-len(parts[1]) % 4)
|
||
return json.loads(base64.urlsafe_b64decode(payload).decode("utf-8"))
|
||
except Exception:
|
||
pass
|
||
return {}
|
||
|
||
|
||
def normalize_cookie(cookie):
|
||
"""Accept either a bare cookie value or a full ``name=value; name2=value2`` string."""
|
||
cookie = (cookie or "").strip().strip('"').strip("'")
|
||
if not cookie:
|
||
raise ValueError("cookie 为空")
|
||
if "=" not in cookie:
|
||
cookie = f"{COOKIE_NAME}={cookie}"
|
||
return cookie
|
||
|
||
|
||
def fetch_identity(api_base, cookie, device_id=None):
|
||
"""Read ``(user_id, email)`` for a session cookie.
|
||
|
||
Note ``/studio/whoami`` authenticates with the *cookie*, not the Bearer JWT
|
||
(a JWT-only call answers 401), so identity lookup goes through the same
|
||
cookie path as minting.
|
||
"""
|
||
cookie = normalize_cookie(cookie)
|
||
data = _request(api_base, cookie, device_id)
|
||
traits = (data.get("identity") or {}).get("traits") or {}
|
||
user_id = data.get("id") or decode_jwt_payload(data.get("tokenized") or "").get("sub")
|
||
return user_id, traits.get("email")
|
||
|
||
|
||
def mint_token(api_base, cookie, device_id=None):
|
||
"""Exchange a session cookie for ``(token, session_expires_at)``."""
|
||
cookie = normalize_cookie(cookie)
|
||
|
||
last_error = None
|
||
for attempt in range(1, RETRY_ATTEMPTS + 1):
|
||
try:
|
||
data = _request(api_base, cookie, device_id)
|
||
except TripoAPIError as error:
|
||
last_error = error
|
||
# 401 here is usually a transient throttle; 5xx/network also worth retrying.
|
||
if error.code in (401, 403, 429, 500, 502, 503, 504, None) and attempt < RETRY_ATTEMPTS:
|
||
time.sleep(RETRY_BACKOFF * attempt)
|
||
continue
|
||
raise
|
||
|
||
token = data.get("tokenized")
|
||
if not token:
|
||
raise TripoAPIError("会话有效但未返回 tokenized，cookie 可能已失效", raw=data)
|
||
return token, data.get("expires_at")
|
||
|
||
raise last_error
|