From 413493d5b524de1eeea029f40eba8cbe1087e766 Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Tue, 10 Feb 2026 10:39:55 +0800 Subject: [PATCH] refactor: integrate reverse request handlers for image and voice services --- app/services/grok/services/image.py | 14 +- app/services/grok/services/voice.py | 104 ++-------- app/services/reverse/__init__.py | 3 +- app/services/reverse/assets_upload.py | 2 +- app/services/reverse/utils/grpc.py | 5 +- app/services/reverse/utils/headers.py | 4 +- app/services/reverse/utils/statsig.py | 7 +- app/services/reverse/utils/websocket.py | 144 ++++++++++++++ app/services/reverse/ws_livekit.py | 253 ++++++++++++++++++++++++ 9 files changed, 417 insertions(+), 119 deletions(-) create mode 100644 app/services/reverse/utils/websocket.py create mode 100644 app/services/reverse/ws_livekit.py diff --git a/app/services/grok/services/image.py b/app/services/grok/services/image.py index 573e1a2..218c334 100644 --- a/app/services/grok/services/image.py +++ b/app/services/grok/services/image.py @@ -10,14 +10,13 @@ import ssl import time import uuid from typing import AsyncGenerator, Dict, Optional -from urllib.parse import urlparse import aiohttp -from aiohttp_socks import ProxyConnector from app.core.config import get_config from app.core.logger import logger from app.services.grok.utils.headers import build_sso_cookie +from app.services.reverse.utils.websocket import resolve_proxy WS_URL = "wss://grok.com/ws/imagine/listen" @@ -36,16 +35,7 @@ class ImageService: def _resolve_proxy(self) -> tuple[aiohttp.BaseConnector, Optional[str]]: proxy_url = get_config("network.base_proxy_url") - if not proxy_url: - return aiohttp.TCPConnector(ssl=self._ssl_context), None - - scheme = urlparse(proxy_url).scheme.lower() - if scheme.startswith("socks"): - logger.info(f"Using SOCKS proxy: {proxy_url}") - return ProxyConnector.from_url(proxy_url, ssl=self._ssl_context), None - - logger.info(f"Using HTTP proxy: {proxy_url}") - return aiohttp.TCPConnector(ssl=self._ssl_context), proxy_url + return resolve_proxy(proxy_url, self._ssl_context) def _get_ws_headers(self, token: str) -> Dict[str, str]: cookie = build_sso_cookie(token, include_rw=True) diff --git a/app/services/grok/services/voice.py b/app/services/grok/services/voice.py index 006fc54..208bf95 100644 --- a/app/services/grok/services/voice.py +++ b/app/services/grok/services/voice.py @@ -2,25 +2,17 @@ Grok Voice Mode Service """ -import orjson -from typing import Dict, Any +from typing import Any, Dict from curl_cffi.requests import AsyncSession -from app.core.logger import logger from app.core.config import get_config -from app.core.exceptions import UpstreamException -from app.services.grok.utils.headers import apply_statsig, build_sso_cookie - -LIVEKIT_TOKEN_API = "https://grok.com/rest/livekit/tokens" +from app.services.reverse import LivekitTokenReverse class VoiceService: """Voice Mode Service (LiveKit)""" - def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("network.base_proxy_url") - async def get_token( self, token: str, @@ -28,86 +20,12 @@ class VoiceService: personality: str = "assistant", speed: float = 1.0, ) -> Dict[str, Any]: - """ - Get LiveKit token - - Args: - token: Auth token - Returns: - Dict containing token and livekitUrl - """ - logger.debug( - f"Voice token request: voice={voice}, personality={personality}, speed={speed}" - ) - headers = self._build_headers(token) - payload = self._build_payload(voice, personality, speed) - - proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None - - try: - browser = get_config("security.browser") - timeout = get_config("network.timeout") - async with AsyncSession(impersonate=browser) as session: - response = await session.post( - LIVEKIT_TOKEN_API, - headers=headers, - data=orjson.dumps(payload), - timeout=timeout, - proxies=proxies, - ) - - if response.status_code != 200: - body = response.text[:200] - logger.error( - f"Voice token failed: status={response.status_code}, body={body}" - ) - raise UpstreamException( - message=f"Failed to get voice token: {response.status_code}", - details={"status": response.status_code, "body": response.text}, - ) - - result = response.json() - logger.info(f"Voice token obtained: voice={voice}") - return result - - except Exception as e: - logger.error(f"Voice service error: {e}") - if isinstance(e, UpstreamException): - raise - raise UpstreamException(f"Voice service error: {str(e)}") - - def _build_headers(self, token: str) -> Dict[str, str]: - headers = { - "Accept": "*/*", - "Content-Type": "application/json", - "Origin": "https://grok.com", - "Referer": "https://grok.com/", - # Statsig ID is crucial - } - - apply_statsig(headers) - headers["Cookie"] = build_sso_cookie(token) - - return headers - - def _build_payload( - self, - voice: str = "ara", - personality: str = "assistant", - speed: float = 1.0, - ) -> Dict[str, Any]: - """Construct payload with voice settings""" - return { - "sessionPayload": orjson.dumps( - { - "voice": voice, - "personality": personality, - "playback_speed": speed, - "enable_vision": False, - "turn_detection": {"type": "server_vad"}, - } - ).decode(), - "requestAgentDispatch": False, - "livekitUrl": "wss://livekit.grok.com", - "params": {"enable_markdown_transcript": "true"}, - } + browser = get_config("security.browser") + async with AsyncSession(impersonate=browser) as session: + return await LivekitTokenReverse.request( + session, + token=token, + voice=voice, + personality=personality, + speed=speed, + ) diff --git a/app/services/reverse/__init__.py b/app/services/reverse/__init__.py index af7a4fd..69594b5 100644 --- a/app/services/reverse/__init__.py +++ b/app/services/reverse/__init__.py @@ -9,8 +9,7 @@ from .media_post import MediaPostReverse from .nsfw_mgmt import NsfwMgmtReverse from .rate_limits import RateLimitsReverse from .set_birth import SetBirthReverse -from .livekit_tokens import LivekitTokenReverse -from .ws_livekit import LivekitWebSocketReverse +from .ws_livekit import LivekitTokenReverse, LivekitWebSocketReverse from .utils.headers import build_headers from .utils.statsig import StatsigGenerator diff --git a/app/services/reverse/assets_upload.py b/app/services/reverse/assets_upload.py index e4f65ac..0466be6 100644 --- a/app/services/reverse/assets_upload.py +++ b/app/services/reverse/assets_upload.py @@ -6,9 +6,9 @@ from typing import Any from curl_cffi.requests import AsyncSession from app.core.logger import logger +from app.core.config import get_config from app.core.exceptions import UpstreamException from app.services.token.service import TokenService -from app.core.config import get_config from app.services.reverse.utils.headers import build_headers from app.services.reverse.utils.retry import retry_on_status diff --git a/app/services/reverse/utils/grpc.py b/app/services/reverse/utils/grpc.py index 74b745e..2cb2688 100644 --- a/app/services/reverse/utils/grpc.py +++ b/app/services/reverse/utils/grpc.py @@ -10,9 +10,6 @@ from typing import Dict, List, Mapping, Tuple from urllib.parse import unquote -_B64_RE = re.compile(rb"^[A-Za-z0-9+/=\r\n]+$") - - @dataclass(frozen=True) class GrpcStatus: code: int @@ -51,7 +48,7 @@ class GrpcClient: return base64.b64decode(compact, validate=False) head = body[: min(len(body), 2048)] - if head and _B64_RE.fullmatch(head): + if head and re.compile(rb"^[A-Za-z0-9+/=\r\n]+$").fullmatch(head): compact = b"".join(body.split()) try: return base64.b64decode(compact, validate=True) diff --git a/app/services/reverse/utils/headers.py b/app/services/reverse/utils/headers.py index fd7e74c..e26de01 100644 --- a/app/services/reverse/utils/headers.py +++ b/app/services/reverse/utils/headers.py @@ -5,8 +5,8 @@ import orjson from urllib.parse import urlparse from typing import Dict, Optional -from app.core.config import get_config from app.core.logger import logger +from app.core.config import get_config from app.services.reverse.utils.statsig import StatsigGenerator @@ -14,7 +14,7 @@ def _build_sso_cookie(sso_token: str) -> str: """ Build SSO Cookie string. """ - # Remove "sso=" prefix if present + # Format sso_token = sso_token[4:] if sso_token.startswith("sso=") else sso_token # SSO Cookie diff --git a/app/services/reverse/utils/statsig.py b/app/services/reverse/utils/statsig.py index 28280a0..69e8196 100644 --- a/app/services/reverse/utils/statsig.py +++ b/app/services/reverse/utils/statsig.py @@ -6,11 +6,8 @@ import base64 import random import string -from app.core.config import get_config from app.core.logger import logger - - -STATIC_STATSIG_ID = "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=" +from app.core.config import get_config class StatsigGenerator: @@ -53,7 +50,7 @@ class StatsigGenerator: # Static Statsig ID logger.debug("Generating static Statsig ID") - return STATIC_STATSIG_ID + return "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=" __all__ = ["StatsigGenerator"] diff --git a/app/services/reverse/utils/websocket.py b/app/services/reverse/utils/websocket.py new file mode 100644 index 0000000..a4a8ffb --- /dev/null +++ b/app/services/reverse/utils/websocket.py @@ -0,0 +1,144 @@ +""" +WebSocket helpers for reverse interfaces. +""" + +import ssl +import certifi +import aiohttp +from aiohttp_socks import ProxyConnector +from typing import Mapping, Optional +from urllib.parse import urlparse + +from app.core.logger import logger +from app.core.config import get_config + + +def _default_ssl_context() -> ssl.SSLContext: + context = ssl.create_default_context() + context.load_verify_locations(certifi.where()) + return context + + +def _normalize_socks_proxy(proxy_url: str) -> tuple[str, Optional[bool]]: + scheme = urlparse(proxy_url).scheme.lower() + rdns: Optional[bool] = None + base_scheme = scheme + + if scheme == "socks5h": + base_scheme = "socks5" + rdns = True + elif scheme == "socks4a": + base_scheme = "socks4" + rdns = True + + if base_scheme != scheme: + proxy_url = proxy_url.replace(f"{scheme}://", f"{base_scheme}://", 1) + + return proxy_url, rdns + + +def resolve_proxy( + proxy_url: str | None, ssl_context: ssl.SSLContext +) -> tuple[aiohttp.BaseConnector, Optional[str]]: + """Resolve proxy connector. + + Args: + proxy_url: str, the proxy URL. + ssl_context: ssl.SSLContext, the SSL context. + + Returns: + tuple[aiohttp.BaseConnector, Optional[str]]: The proxy connector and the proxy URL. + """ + if not proxy_url: + return aiohttp.TCPConnector(ssl=ssl_context), None + + scheme = urlparse(proxy_url).scheme.lower() + if scheme.startswith("socks"): + normalized, rdns = _normalize_socks_proxy(proxy_url) + logger.info(f"Using SOCKS proxy: {proxy_url}") + try: + if rdns is not None: + return ( + ProxyConnector.from_url(normalized, rdns=rdns, ssl=ssl_context), + None, + ) + except TypeError: + return ProxyConnector.from_url(normalized, ssl=ssl_context), None + return ProxyConnector.from_url(normalized, ssl=ssl_context), None + + logger.info(f"Using HTTP proxy: {proxy_url}") + return aiohttp.TCPConnector(ssl=ssl_context), proxy_url + + +class WebSocketConnection: + """WebSocket connection wrapper.""" + + def __init__( + self, + session: aiohttp.ClientSession, + ws: aiohttp.ClientWebSocketResponse, + ) -> None: + self.session = session + self.ws = ws + + async def close(self) -> None: + if not self.ws.closed: + await self.ws.close() + await self.session.close() + + async def __aenter__(self) -> aiohttp.ClientWebSocketResponse: + return self.ws + + async def __aexit__(self, exc_type, exc, tb) -> None: + await self.close() + + +class WebSocketClient: + """WebSocket client with proxy support.""" + + def __init__(self, proxy: str | None = None) -> None: + self.proxy = proxy or get_config("network.base_proxy_url") + self._ssl_context = _default_ssl_context() + + async def connect( + self, + url: str, + headers: Mapping[str, str] | None = None, + timeout: float | aiohttp.ClientTimeout | None = None, + ) -> WebSocketConnection: + """Connect to the WebSocket. + + Args: + url: str, the URL to connect to. + headers: Mapping[str, str], the headers to send. + timeout: float | aiohttp.ClientTimeout | None, the timeout. + + Returns: + WebSocketConnection: The WebSocket connection. + """ + # Resolve proxy + connector, proxy = resolve_proxy(self.proxy, self._ssl_context) + + # Build client timeout + client_timeout = ( + timeout + if isinstance(timeout, aiohttp.ClientTimeout) + else aiohttp.ClientTimeout(total=timeout) + ) + + # Create session + session = aiohttp.ClientSession(connector=connector, timeout=client_timeout) + try: + ws = await session.ws_connect( + url, + headers=headers, + proxy=proxy, + ssl=self._ssl_context, + ) + return WebSocketConnection(session, ws) + except Exception: + await session.close() + raise + + +__all__ = ["WebSocketClient", "WebSocketConnection", "resolve_proxy"] diff --git a/app/services/reverse/ws_livekit.py b/app/services/reverse/ws_livekit.py new file mode 100644 index 0000000..52eb84e --- /dev/null +++ b/app/services/reverse/ws_livekit.py @@ -0,0 +1,253 @@ +""" +Reverse interface: LiveKit token + WebSocket. +""" + +import orjson +from typing import Any, Dict +from urllib.parse import urlencode +from curl_cffi.requests import AsyncSession + +from app.core.logger import logger +from app.core.config import get_config +from app.core.exceptions import UpstreamException +from app.services.token.service import TokenService +from app.services.reverse.utils.headers import build_headers +from app.services.reverse.utils.retry import retry_on_status +from app.services.reverse.utils.websocket import WebSocketClient, WebSocketConnection + +LIVEKIT_TOKEN_API = "https://grok.com/rest/livekit/tokens" +LIVEKIT_WS_URL = "wss://livekit.grok.com" + + +class LivekitTokenReverse: + """/rest/livekit/tokens reverse interface.""" + + @staticmethod + async def request( + session: AsyncSession, + token: str, + voice: str = "ara", + personality: str = "assistant", + speed: float = 1.0, + livekit_url: str = LIVEKIT_WS_URL, + ) -> Dict[str, Any]: + """Fetch LiveKit token. + + Args: + session: AsyncSession, the session to use for the request. + token: str, the SSO token. + voice: str, the voice to use for the request. + personality: str, the personality to use for the request. + speed: float, the speed to use for the request. + livekit_url: str, the LiveKit URL to use for the request. + + Returns: + Dict[str, Any]: The LiveKit token. + """ + try: + # Get proxies + base_proxy = get_config("network.base_proxy_url") + proxies = {"http": base_proxy, "https": base_proxy} if base_proxy else None + + # Build headers + headers = build_headers( + cookie_token=token, + content_type="application/json", + origin="https://grok.com", + referer="https://grok.com/", + ) + + # Build payload + payload = { + "sessionPayload": orjson.dumps( + { + "voice": voice, + "personality": personality, + "playback_speed": speed, + "enable_vision": False, + "turn_detection": {"type": "server_vad"}, + } + ).decode(), + "requestAgentDispatch": False, + "livekitUrl": livekit_url, + "params": {"enable_markdown_transcript": "true"}, + } + + # Curl Config + timeout = get_config("network.timeout") + browser = get_config("security.browser") + + async def _do_request(): + response = await session.post( + LIVEKIT_TOKEN_API, + headers=headers, + data=orjson.dumps(payload), + timeout=timeout, + proxies=proxies, + impersonate=browser, + ) + + if response.status_code != 200: + body = response.text[:200] + logger.error( + f"LivekitTokenReverse: Request failed, {response.status_code}, body={body}" + ) + raise UpstreamException( + message=f"LivekitTokenReverse: Request failed, {response.status_code}", + details={"status": response.status_code, "body": response.text}, + ) + + return response + + return await retry_on_status(_do_request) + + + except Exception as e: + # Handle upstream exception + if isinstance(e, UpstreamException): + status = None + if e.details and "status" in e.details: + status = e.details["status"] + else: + status = getattr(e, "status_code", None) + if status == 401: + try: + await TokenService.record_fail( + token, status, "livekit_token_auth_failed" + ) + except Exception: + pass + raise + + # Handle other non-upstream exceptions + logger.error( + f"LivekitTokenReverse: Request failed, {str(e)}", + extra={"error_type": type(e).__name__}, + ) + raise UpstreamException( + message=f"LivekitTokenReverse: Request failed, {str(e)}", + details={"status": 502, "error": str(e)}, + ) + + +class LivekitWebSocketReverse: + """LiveKit WebSocket reverse interface.""" + + def __init__(self) -> None: + self._client = WebSocketClient() + + def build_url( + self, + access_token: str, + *, + livekit_url: str = LIVEKIT_WS_URL, + auto_subscribe: bool = True, + sdk: str = "js", + version: str = "2.11.4", + protocol: int = 15, + ) -> str: + """Build LiveKit WebSocket URL. + + Args: + access_token: str, the LiveKit access token. + livekit_url: str, the LiveKit URL to use for the request. + auto_subscribe: bool, whether to auto subscribe to the WebSocket. + sdk: str, the SDK to use for the request. + version: str, the version to use for the request. + protocol: int, the protocol to use for the request. + + Returns: + str: The LiveKit WebSocket URL. + """ + # Build base URL + base = livekit_url.rstrip("/") + if not base.endswith("/rtc"): + base = f"{base}/rtc" + + # Build parameters + params = { + "access_token": access_token, + "auto_subscribe": str(int(auto_subscribe)), + "sdk": sdk, + "version": version, + "protocol": str(protocol), + } + + return f"{base}?{urlencode(params)}" + + def _build_headers(self, extra: Dict[str, str] | None = None) -> Dict[str, str]: + """Build LiveKit WebSocket headers.""" + # Build headers + headers = { + "Origin": "https://grok.com", + "User-Agent": get_config("security.user_agent"), + "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8", + "Cache-Control": "no-cache", + "Pragma": "no-cache", + } + + # Update headers + if extra: + headers.update(extra) + return headers + + async def connect( + self, + access_token: str, + *, + livekit_url: str = LIVEKIT_WS_URL, + auto_subscribe: bool = True, + sdk: str = "js", + version: str = "2.11.4", + protocol: int = 15, + headers: Dict[str, str] | None = None, + timeout: float | None = None, + ) -> WebSocketConnection: + """Connect to the LiveKit WebSocket. + + Args: + access_token: str, the LiveKit access token. + livekit_url: str, the LiveKit URL to use for the request. + auto_subscribe: bool, whether to auto subscribe to the WebSocket. + sdk: str, the SDK to use for the request. + version: str, the version to use for the request. + protocol: int, the protocol to use for the request. + headers: Dict[str, str], the headers to send. + timeout: float, the timeout to use for the request. + + Returns: + WebSocketConnection: The LiveKit WebSocket connection. + """ + # Build URL + url = self.build_url( + access_token, + livekit_url=livekit_url, + auto_subscribe=auto_subscribe, + sdk=sdk, + version=version, + protocol=protocol, + ) + + # Build WebSocket headers + ws_headers = self._build_headers(headers) + + # Build timeout + if timeout is None: + timeout = get_config("network.timeout") + + # Connect to the LiveKit WebSocket + try: + return await self._client.connect(url, headers=ws_headers, timeout=timeout) + except Exception as e: + logger.error(f"LivekitWebSocketReverse: Connect failed, {e}") + raise UpstreamException( + f"LivekitWebSocketReverse: Connect failed, {str(e)}" + ) + + +__all__ = [ + "LivekitTokenReverse", + "LivekitWebSocketReverse", + "LIVEKIT_TOKEN_API", + "LIVEKIT_WS_URL", +]