From f937d57d1e681ba3ae218ba24ce3dc1cc640b5fe Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Fri, 6 Feb 2026 17:12:50 +0800 Subject: [PATCH] feat: enhance token management with normalization and improved error handling for token updates --- app/api/v1/admin.py | 56 ++++++++++++++++++++++++++++- app/services/grok/services/nsfw.py | 57 ++++++++++++++++++++++++++++++ app/static/token/token.js | 24 ++++++++++--- 3 files changed, 132 insertions(+), 5 deletions(-) diff --git a/app/api/v1/admin.py b/app/api/v1/admin.py index 438b399..09346ee 100644 --- a/app/api/v1/admin.py +++ b/app/api/v1/admin.py @@ -242,9 +242,63 @@ async def update_tokens_api(data: dict): storage = get_storage() try: from app.services.token.manager import get_token_manager + from app.services.token.models import TokenInfo async with storage.acquire_lock("tokens_save", timeout=10): - await storage.save_tokens(data) + existing = await storage.load_tokens() or {} + normalized = {} + allowed_fields = set(TokenInfo.model_fields.keys()) + existing_map = {} + for pool_name, tokens in existing.items(): + if not isinstance(tokens, list): + continue + pool_map = {} + for item in tokens: + if isinstance(item, str): + token_data = {"token": item} + elif isinstance(item, dict): + token_data = dict(item) + else: + continue + raw_token = token_data.get("token") + if isinstance(raw_token, str) and raw_token.startswith("sso="): + token_data["token"] = raw_token[4:] + token_key = token_data.get("token") + if isinstance(token_key, str): + pool_map[token_key] = token_data + existing_map[pool_name] = pool_map + for pool_name, tokens in (data or {}).items(): + if not isinstance(tokens, list): + continue + pool_list = [] + for item in tokens: + if isinstance(item, str): + token_data = {"token": item} + elif isinstance(item, dict): + token_data = dict(item) + else: + continue + + raw_token = token_data.get("token") + if isinstance(raw_token, str) and raw_token.startswith("sso="): + token_data["token"] = raw_token[4:] + + base = existing_map.get(pool_name, {}).get(token_data.get("token"), {}) + merged = dict(base) + merged.update(token_data) + if merged.get("tags") is None: + merged["tags"] = [] + + filtered = {k: v for k, v in merged.items() if k in allowed_fields} + try: + info = TokenInfo(**filtered) + pool_list.append(info.model_dump()) + except Exception as e: + logger.warning(f"Skip invalid token in pool '{pool_name}': {e}") + continue + normalized[pool_name] = pool_list + + await storage.save_tokens(normalized) mgr = await get_token_manager() await mgr.reload() return {"status": "success", "message": "Token 已更新"} diff --git a/app/services/grok/services/nsfw.py b/app/services/grok/services/nsfw.py index 476793e..2041ec1 100644 --- a/app/services/grok/services/nsfw.py +++ b/app/services/grok/services/nsfw.py @@ -6,6 +6,8 @@ NSFW (Unhinged) 模式服务 from dataclasses import dataclass from typing import Optional +import datetime +import random from curl_cffi.requests import AsyncSession @@ -20,6 +22,7 @@ from app.services.grok.utils.headers import build_sso_cookie NSFW_API = "https://grok.com/auth_mgmt.AuthManagement/UpdateUserFeatureControls" +BIRTH_DATE_API = "https://grok.com/rest/auth/set-birth-date" BROWSER = "chrome136" TIMEOUT = 30 @@ -41,6 +44,20 @@ class NSFWService: def __init__(self, proxy: str = None): self.proxy = proxy or get_config("grok.base_proxy_url", "") + @staticmethod + def _random_birth_date() -> str: + """生成随机出生日期""" + today = datetime.date.today() + age = random.randint(20, 40) + birth_year = today.year - age + birth_month = random.randint(1, 12) + birth_day = random.randint(1, 28) + hour = random.randint(0, 23) + minute = random.randint(0, 59) + second = random.randint(0, 59) + microsecond = random.randint(0, 999) + return f"{birth_year:04d}-{birth_month:02d}-{birth_day:02d}T{hour:02d}:{minute:02d}:{second:02d}.{microsecond:03d}Z" + def _build_headers(self, token: str) -> dict: """构造 gRPC-Web 请求头""" cookie = build_sso_cookie(token, include_rw=True) @@ -55,6 +72,18 @@ class NSFWService: "cookie": cookie, } + def _build_birth_headers(self, token: str) -> dict: + """构造设置出生日期请求头""" + cookie = build_sso_cookie(token, include_rw=True) + return { + "accept": "*/*", + "content-type": "application/json", + "origin": "https://grok.com", + "referer": "https://grok.com/?_s=account", + "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36", + "cookie": cookie, + } + @staticmethod def _build_payload() -> bytes: """构造请求 payload""" @@ -67,6 +96,25 @@ class NSFWService: protobuf = b"\x0a\x02\x10\x01\x12" + bytes([len(inner)]) + inner return encode_grpc_web_payload(protobuf) + async def _set_birth_date(self, session: AsyncSession, token: str) -> tuple[bool, int, Optional[str]]: + """设置出生日期""" + headers = self._build_birth_headers(token) + payload = {"birthDate": self._random_birth_date()} + proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None + try: + response = await session.post( + BIRTH_DATE_API, + json=payload, + headers=headers, + timeout=TIMEOUT, + proxies=proxies, + ) + if response.status_code in (200, 204): + return True, response.status_code, None + return False, response.status_code, f"HTTP {response.status_code}" + except Exception as e: + return False, 0, str(e)[:100] + async def enable(self, token: str) -> NSFWResult: """为单个 token 开启 NSFW 模式""" headers = self._build_headers(token) @@ -80,6 +128,15 @@ class NSFWService: try: async with AsyncSession(impersonate=BROWSER) as session: + ok, birth_status, birth_err = await self._set_birth_date( + session, token + ) + if not ok: + return NSFWResult( + success=False, + http_status=birth_status, + error=f"Set birth date failed: {birth_err}", + ) response = await session.post( NSFW_API, data=payload, diff --git a/app/static/token/token.js b/app/static/token/token.js index d7ff84b..1ed1d0a 100644 --- a/app/static/token/token.js +++ b/app/static/token/token.js @@ -152,7 +152,13 @@ function processTokens(data) { note: t.note || '', fail_count: t.fail_count || 0, use_count: t.use_count || 0, - tags: t.tags || [] + tags: t.tags || [], + created_at: t.created_at, + last_used_at: t.last_used_at, + last_fail_at: t.last_fail_at, + last_fail_reason: t.last_fail_reason, + last_sync_at: t.last_sync_at, + last_asset_clear_at: t.last_asset_clear_at }; flatTokens.push({ ...tObj, pool: pool, _selected: false }); }); @@ -511,14 +517,22 @@ async function syncToServer() { const newTokens = {}; flatTokens.forEach(t => { if (!newTokens[t.pool]) newTokens[t.pool] = []; - newTokens[t.pool].push({ + const payload = { token: t.token, status: t.status, quota: t.quota, note: t.note, fail_count: t.fail_count, - use_count: t.use_count || 0 - }); + use_count: t.use_count || 0, + tags: Array.isArray(t.tags) ? t.tags : [] + }; + if (typeof t.created_at === 'number') payload.created_at = t.created_at; + if (typeof t.last_used_at === 'number') payload.last_used_at = t.last_used_at; + if (typeof t.last_fail_at === 'number') payload.last_fail_at = t.last_fail_at; + if (typeof t.last_sync_at === 'number') payload.last_sync_at = t.last_sync_at; + if (typeof t.last_asset_clear_at === 'number') payload.last_asset_clear_at = t.last_asset_clear_at; + if (typeof t.last_fail_reason === 'string' && t.last_fail_reason) payload.last_fail_reason = t.last_fail_reason; + newTokens[t.pool].push(payload); }); try { @@ -563,6 +577,8 @@ async function submitImport() { status: 'active', quota: defaultQuota, note: '', + tags: [], + fail_count: 0, use_count: 0, _selected: false });