refactor: integrate reverse request handlers for NSFW

app/services/grok/services/nsfw.py

@@ -6,23 +6,15 @@ NSFW (Unhinged) 模式服务
 
 from dataclasses import dataclass
 from typing import Optional
-import datetime
-import random
 
 from curl_cffi.requests import AsyncSession
 
-from app.core.config import get_config
 from app.core.logger import logger
 from app.services.grok.protocols.grpc_web import (
-    encode_grpc_web_payload,
-    parse_grpc_web_response,
-    get_grpc_status,
-)
-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"
-
+from app.core.config import get_config
+from app.core.exceptions import UpstreamException
+from app.services.reverse import NsfwMgmtReverse, SetBirthReverse
+from app.services.reverse.utils.grpc import GrpcStatus
 
 @dataclass
 class NSFWResult:
@@ -38,139 +30,33 @@ class NSFWResult:
 class NSFWService:
     """NSFW 模式服务"""
 
-    def __init__(self, proxy: str = None):
-        self.proxy = proxy or get_config("network.base_proxy_url")
-        self.timeout = float(get_config("network.timeout"))
-
-    def _build_proxies(self) -> Optional[dict]:
-        """构建代理配置"""
-        return {"http": self.proxy, "https": self.proxy} if self.proxy else None
-
-    @staticmethod
-    def _random_birth_date() -> str:
-        """生成随机出生日期（20-40岁）"""
-        today = datetime.date.today()
-        birth_year = today.year - random.randint(20, 40)
-        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)
-        user_agent = get_config("security.user_agent")
-        return {
-            "accept": "*/*",
-            "content-type": "application/grpc-web+proto",
-            "origin": "https://grok.com",
-            "referer": "https://grok.com/",
-            "user-agent": user_agent,
-            "x-grpc-web": "1",
-            "x-user-agent": "connect-es/2.1.1",
-            "cookie": cookie,
-        }
-
-    def _build_birth_headers(self, token: str) -> dict:
-        """构造设置出生日期请求头"""
-        cookie = build_sso_cookie(token, include_rw=True)
-        user_agent = get_config("security.user_agent")
-        return {
-            "accept": "*/*",
-            "content-type": "application/json",
-            "origin": "https://grok.com",
-            "referer": "https://grok.com/?_s=account",
-            "user-agent": user_agent,
-            "cookie": cookie,
-        }
-
-    @staticmethod
-    def _build_payload() -> bytes:
-        """构造请求 payload"""
-        # protobuf (match captured HAR):
-        # 0a 02 10 01                   -> field 1 (len=2) with inner bool=true
-        # 12 1a                         -> field 2, length 26
-        #   0a 18 <name>                -> nested message with name string
-        name = b"always_show_nsfw_content"
-        inner = b"\x0a" + bytes([len(name)]) + name
-        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()}
-
-        try:
-            response = await session.post(
-                BIRTH_DATE_API,
-                json=payload,
-                headers=headers,
-                timeout=self.timeout,
-                proxies=self._build_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)
-        payload = self._build_payload()
-        logger.debug(f"NSFW payload: len={len(payload)} hex={payload.hex()}")
-
         try:
             browser = get_config("security.browser")
             async with AsyncSession(impersonate=browser) as session:
                 # 先设置出生日期
-                ok, birth_status, birth_err = await self._set_birth_date(session, token)
-                if not ok:
+                try:
+                    await SetBirthReverse.request(session, token)
+                except UpstreamException as e:
+                    status = None
+                    if e.details and "status" in e.details:
+                        status = e.details["status"]
+                    else:
+                        status = getattr(e, "status_code", None)
                     return NSFWResult(
                         success=False,
-                        http_status=birth_status,
-                        error=f"Set birth date failed: {birth_err}",
+                        http_status=status or 0,
+                        error=f"Set birth date failed: {str(e)}",
                     )
 
                 # 开启 NSFW
-                response = await session.post(
-                    NSFW_API,
-                    data=payload,
-                    headers=headers,
-                    timeout=self.timeout,
-                    proxies=self._build_proxies(),
-                )
-
-                if response.status_code != 200:
-                    return NSFWResult(
-                        success=False,
-                        http_status=response.status_code,
-                        error=f"HTTP {response.status_code}",
-                    )
-
-                # 解析 gRPC-Web 响应
-                _, trailers = parse_grpc_web_response(
-                    response.content, content_type=response.headers.get("content-type")
-                )
-
-                grpc_status = get_grpc_status(trailers)
-                logger.debug(
-                    f"NSFW response: http={response.status_code} grpc={grpc_status.code} "
-                    f"msg={grpc_status.message} trailers={trailers}"
-                )
-
-                # HTTP 200 且无 grpc-status（空响应）或 grpc-status=0 都算成功
-                success = grpc_status.code == -1 or grpc_status.ok
+                grpc_status: GrpcStatus = await NsfwMgmtReverse.request(session, token)
+                success = grpc_status.code in (-1, 0)
 
                 return NSFWResult(
                     success=success,
-                    http_status=response.status_code,
+                    http_status=200,
                     grpc_status=grpc_status.code,
                     grpc_message=grpc_status.message or None,
                 )

app/services/reverse/__init__.py

@@ -6,7 +6,11 @@ from .assets_download import AssetsDownloadReverse
 from .assets_list import AssetsListReverse
 from .assets_upload import AssetsUploadReverse
 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 .utils.headers import build_headers
 from .utils.statsig import StatsigGenerator
 
@@ -17,7 +21,11 @@ __all__ = [
     "AssetsListReverse",
     "AssetsUploadReverse",
     "MediaPostReverse",
+    "NsfwMgmtReverse",
     "RateLimitsReverse",
+    "SetBirthReverse",
+    "LivekitTokenReverse",
+    "LivekitWebSocketReverse",
     "StatsigGenerator",
     "build_headers",
 ]

app/services/reverse/nsfw_mgmt.py

@@ -0,0 +1,132 @@
+"""
+Reverse interface: NSFW feature controls (gRPC-Web).
+"""
+
+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.grpc import GrpcClient, GrpcStatus
+
+NSFW_MGMT_API = "https://grok.com/auth_mgmt.AuthManagement/UpdateUserFeatureControls"
+
+
+class NsfwMgmtReverse:
+    """/auth_mgmt.AuthManagement/UpdateUserFeatureControls reverse interface."""
+
+    @staticmethod
+    async def request(session: AsyncSession, token: str) -> GrpcStatus:
+        """Enable NSFW feature control via gRPC-Web.
+
+        Args:
+            session: AsyncSession, the session to use for the request.
+            token: str, the SSO token.
+
+        Returns:
+            GrpcStatus: Parsed gRPC status.
+        """
+        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,
+                origin="https://grok.com",
+                referer="https://grok.com/?_s=data",
+            )
+            headers["Content-Type"] = "application/grpc-web+proto"
+            headers["Accept"] = "*/*"
+            headers["Sec-Fetch-Dest"] = "empty"
+            headers["x-grpc-web"] = "1"
+            headers["x-user-agent"] = "connect-es/2.1.1"
+            headers["Cache-Control"] = "no-cache"
+            headers["Pragma"] = "no-cache"
+
+            # Build payload
+            name = "always_show_nsfw_content".encode("utf-8")
+            inner = b"\x0a" + bytes([len(name)]) + name
+            protobuf = b"\x0a\x02\x10\x01\x12" + bytes([len(inner)]) + inner
+            payload = GrpcClient.encode_payload(protobuf)
+
+            # Curl Config
+            timeout = get_config("network.timeout")
+            browser = get_config("security.browser")
+
+            async def _do_request():
+                response = await session.post(
+                    NSFW_MGMT_API,
+                    headers=headers,
+                    data=payload,
+                    timeout=timeout,
+                    proxies=proxies,
+                    impersonate=browser,
+                )
+
+                if response.status_code != 200:
+                    logger.error(
+                        f"NsfwMgmtReverse: Request failed, {response.status_code}",
+                        extra={"error_type": "UpstreamException"},
+                    )
+                    raise UpstreamException(
+                        message=f"NsfwMgmtReverse: Request failed, {response.status_code}",
+                        details={"status": response.status_code},
+                    )
+
+                return response
+
+            response = await retry_on_status(_do_request)
+
+            _, trailers = GrpcClient.parse_response(
+                response.content,
+                content_type=response.headers.get("content-type"),
+                headers=response.headers,
+            )
+            grpc_status = GrpcClient.get_status(trailers)
+
+            if grpc_status.code not in (-1, 0):
+                raise UpstreamException(
+                    message=f"NsfwMgmtReverse: gRPC failed, {grpc_status.code}",
+                    details={
+                        "status": grpc_status.http_equiv,
+                        "grpc_status": grpc_status.code,
+                        "grpc_message": grpc_status.message,
+                    },
+                )
+
+            return grpc_status
+
+        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, "nsfw_mgmt_auth_failed"
+                        )
+                    except Exception:
+                        pass
+                raise
+
+            # Handle other non-upstream exceptions
+            logger.error(
+                f"NsfwMgmtReverse: Request failed, {str(e)}",
+                extra={"error_type": type(e).__name__},
+            )
+            raise UpstreamException(
+                message=f"NsfwMgmtReverse: Request failed, {str(e)}",
+                details={"status": 502, "error": str(e)},
+            )
+
+
+__all__ = ["NsfwMgmtReverse"]

app/services/reverse/set_birth.py

@@ -0,0 +1,117 @@
+"""
+Reverse interface: set birth date.
+"""
+
+import datetime
+import random
+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.services.reverse.utils.headers import build_headers
+from app.services.reverse.utils.retry import retry_on_status
+
+SET_BIRTH_API = "https://grok.com/rest/auth/set-birth-date"
+
+
+class SetBirthReverse:
+    """/rest/auth/set-birth-date reverse interface."""
+
+    @staticmethod
+    async def request(session: AsyncSession, token: str) -> Any:
+        """Set birth date in Grok.
+
+        Args:
+            session: AsyncSession, the session to use for the request.
+            token: str, the SSO token.
+
+        Returns:
+            Any: The response from the request.
+        """
+        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/?_s=home",
+            )
+
+            # Build payload
+            today = datetime.date.today()
+            birth_year = today.year - random.randint(20, 40)
+            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)
+            payload = {
+                "birthDate": f"{birth_year:04d}-{birth_month:02d}-{birth_day:02d}"
+                f"T{hour:02d}:{minute:02d}:{second:02d}.{microsecond:03d}Z"
+            }
+
+            # Curl Config
+            timeout = get_config("network.timeout")
+            browser = get_config("security.browser")
+
+            async def _do_request():
+                response = await session.post(
+                    SET_BIRTH_API,
+                    headers=headers,
+                    json=payload,
+                    timeout=timeout,
+                    proxies=proxies,
+                    impersonate=browser,
+                )
+
+                if response.status_code not in (200, 204):
+                    logger.error(
+                        f"SetBirthReverse: Request failed, {response.status_code}",
+                        extra={"error_type": "UpstreamException"},
+                    )
+                    raise UpstreamException(
+                        message=f"SetBirthReverse: Request failed, {response.status_code}",
+                        details={"status": response.status_code},
+                    )
+
+                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, "set_birth_auth_failed"
+                        )
+                    except Exception:
+                        pass
+                raise
+
+            # Handle other non-upstream exceptions
+            logger.error(
+                f"SetBirthReverse: Request failed, {str(e)}",
+                extra={"error_type": type(e).__name__},
+            )
+            raise UpstreamException(
+                message=f"SetBirthReverse: Request failed, {str(e)}",
+                details={"status": 502, "error": str(e)},
+            )
+
+
+__all__ = ["SetBirthReverse"]

app/services/reverse/utils/grpc.py

@@ -0,0 +1,156 @@
+"""
+gRPC-Web helpers for reverse interfaces.
+"""
+
+import base64
+import re
+import struct
+from dataclasses import dataclass
+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
+    message: str = ""
+
+    @property
+    def ok(self) -> bool:
+        return self.code == 0
+
+    @property
+    def http_equiv(self) -> int:
+        mapping = {
+            0: 200,
+            16: 401,
+            7: 403,
+            8: 429,
+            4: 504,
+            14: 503,
+        }
+        return mapping.get(self.code, 502)
+
+
+class GrpcClient:
+    """gRPC-Web helpers wrapper."""
+
+    @staticmethod
+    def encode_payload(data: bytes) -> bytes:
+        """Encode gRPC-Web data frame."""
+        return b"\x00" + struct.pack(">I", len(data)) + data
+
+    @staticmethod
+    def _maybe_decode_grpc_web_text(body: bytes, content_type: str | None) -> bytes:
+        ct = (content_type or "").lower()
+        if "grpc-web-text" in ct:
+            compact = b"".join(body.split())
+            return base64.b64decode(compact, validate=False)
+
+        head = body[: min(len(body), 2048)]
+        if head and _B64_RE.fullmatch(head):
+            compact = b"".join(body.split())
+            try:
+                return base64.b64decode(compact, validate=True)
+            except Exception:
+                return body
+        return body
+
+    @staticmethod
+    def _parse_trailer_block(payload: bytes) -> Dict[str, str]:
+        text = payload.decode("utf-8", errors="replace")
+        lines = [ln for ln in re.split(r"\r\n|\n", text) if ln]
+
+        trailers: Dict[str, str] = {}
+        for ln in lines:
+            if ":" not in ln:
+                continue
+            k, v = ln.split(":", 1)
+            trailers[k.strip().lower()] = v.strip()
+
+        if "grpc-message" in trailers:
+            trailers["grpc-message"] = unquote(trailers["grpc-message"])
+
+        return trailers
+
+    @classmethod
+    def parse_response(
+        cls,
+        body: bytes,
+        content_type: str | None = None,
+        headers: Mapping[str, str] | None = None,
+    ) -> Tuple[List[bytes], Dict[str, str]]:
+        decoded = cls._maybe_decode_grpc_web_text(body, content_type)
+
+        messages: List[bytes] = []
+        trailers: Dict[str, str] = {}
+
+        i = 0
+        n = len(decoded)
+        while i < n:
+            if n - i < 5:
+                break
+
+            flag = decoded[i]
+            length = int.from_bytes(decoded[i + 1 : i + 5], "big")
+            i += 5
+
+            if n - i < length:
+                break
+
+            payload = decoded[i : i + length]
+            i += length
+
+            if flag & 0x80:
+                trailers.update(cls._parse_trailer_block(payload))
+            elif flag & 0x01:
+                raise ValueError("grpc-web compressed flag not supported")
+            else:
+                messages.append(payload)
+
+        if headers:
+            lower = {k.lower(): v for k, v in headers.items()}
+            if "grpc-status" in lower and "grpc-status" not in trailers:
+                trailers["grpc-status"] = str(lower["grpc-status"]).strip()
+            if "grpc-message" in lower and "grpc-message" not in trailers:
+                trailers["grpc-message"] = unquote(str(lower["grpc-message"]).strip())
+
+        return messages, trailers
+
+    @staticmethod
+    def get_status(trailers: Mapping[str, str]) -> GrpcStatus:
+        raw = str(trailers.get("grpc-status", "")).strip()
+        msg = str(trailers.get("grpc-message", "")).strip()
+        try:
+            code = int(raw)
+        except Exception:
+            code = -1
+        return GrpcStatus(code=code, message=msg)
+
+
+def encode_grpc_web_payload(data: bytes) -> bytes:
+    return GrpcClient.encode_payload(data)
+
+
+def parse_grpc_web_response(
+    body: bytes,
+    content_type: str | None = None,
+    headers: Mapping[str, str] | None = None,
+) -> Tuple[List[bytes], Dict[str, str]]:
+    return GrpcClient.parse_response(body, content_type=content_type, headers=headers)
+
+
+def get_grpc_status(trailers: Mapping[str, str]) -> GrpcStatus:
+    return GrpcClient.get_status(trailers)
+
+
+__all__ = [
+    "encode_grpc_web_payload",
+    "parse_grpc_web_response",
+    "get_grpc_status",
+    "GrpcStatus",
+    "GrpcClient",
+]