diff --git a/app/control/proxy/__init__.py b/app/control/proxy/__init__.py index 1dbe949c3..9d7581d60 100644 --- a/app/control/proxy/__init__.py +++ b/app/control/proxy/__init__.py @@ -6,6 +6,7 @@ """ import asyncio +from urllib.parse import urlparse from app.platform.logging.logger import logger from app.platform.config.snapshot import get_config @@ -26,6 +27,14 @@ from .providers.manual import ManualClearanceProvider from .providers.flaresolverr import FlareSolverrClearanceProvider +_DEFAULT_CLEARANCE_ORIGIN = "https://grok.com" +BundleKey = tuple[str, str] + + +def _clearance_host(clearance_origin: str | None) -> str: + host = urlparse(clearance_origin or _DEFAULT_CLEARANCE_ORIGIN).hostname + return (host or "grok.com").lower() + class ProxyDirectory: """Owns egress nodes and clearance bundles. @@ -36,11 +45,11 @@ class ProxyDirectory: def __init__(self) -> None: self._nodes: list[EgressNode] = [] self._resource_nodes: list[EgressNode] = [] # for media downloads - self._bundles: dict[str, ClearanceBundle] = {} + self._bundles: dict[BundleKey, ClearanceBundle] = {} self._lock = asyncio.Lock() - # Single-flight guard: at most one FlareSolverr call per affinity key. + # Single-flight guard: at most one FlareSolverr call per proxy+host key. # Other coroutines wait on the Event until the active refresh completes. - self._refresh_events: dict[str, asyncio.Event] = {} + self._refresh_events: dict[BundleKey, asyncio.Event] = {} self._manual = ManualClearanceProvider() self._flare = FlareSolverrClearanceProvider() self._egress_mode: EgressMode = EgressMode.DIRECT @@ -117,12 +126,12 @@ async def load(self) -> None: self._bundles = { key: bundle.model_copy(update={"state": ClearanceBundleState.INVALID}) for key, bundle in self._bundles.items() - if key in valid_affinities + if key[0] in valid_affinities } self._refresh_events = { key: event for key, event in self._refresh_events.items() - if key in valid_affinities + if key[0] in valid_affinities } self._config_sig = config_sig @@ -144,6 +153,7 @@ async def acquire( scope: ProxyScope = ProxyScope.APP, kind: RequestKind = RequestKind.HTTP, resource: bool = False, + clearance_origin: str | None = None, ) -> ProxyLease: """Return a ProxyLease for the next request. @@ -151,9 +161,12 @@ async def acquire( """ proxy_url = await self._pick_proxy_url(resource=resource) affinity = proxy_url or "direct" + clearance_host = _clearance_host(clearance_origin) bundle = await self._get_or_build_bundle( - affinity_key=affinity, proxy_url=proxy_url or "" + affinity_key=affinity, + proxy_url=proxy_url or "", + clearance_origin=clearance_origin or _DEFAULT_CLEARANCE_ORIGIN, ) return ProxyLease( @@ -161,6 +174,7 @@ async def acquire( proxy_url=proxy_url, cf_cookies=bundle.cf_cookies if bundle else "", user_agent=bundle.user_agent if bundle else "", + clearance_host=clearance_host, scope=scope, kind=kind, acquired_at=now_ms(), @@ -173,13 +187,13 @@ async def feedback(self, lease: ProxyLease, result: ProxyFeedback) -> None: ProxyFeedbackKind.UNAUTHORIZED, ): # Invalidate associated clearance bundle. - affinity = lease.proxy_url or "direct" + key = (lease.proxy_url or "direct", lease.clearance_host) async with self._lock: - bundle = self._bundles.get(affinity) - if bundle: - from .models import ClearanceBundleState + from .models import ClearanceBundleState - self._bundles[affinity] = bundle.model_copy( + bundle = self._bundles.get(key) + if bundle: + self._bundles[key] = bundle.model_copy( update={"state": ClearanceBundleState.INVALID} ) @@ -233,41 +247,48 @@ async def _get_or_build_bundle( *, affinity_key: str, proxy_url: str, + clearance_origin: str, ) -> ClearanceBundle | None: if self._clearance_mode == ClearanceMode.NONE: return None + clearance_host = _clearance_host(clearance_origin) + key: BundleKey = (affinity_key, clearance_host) - # Single-flight: only one coroutine fetches clearance per affinity key. + # Single-flight: only one coroutine fetches clearance per proxy+host key. # Concurrent callers wait on the Event and retry once it fires. while True: async with self._lock: - bundle = self._bundles.get(affinity_key) + bundle = self._bundles.get(key) if bundle and bundle.state.value == 0: # VALID return bundle - event = self._refresh_events.get(affinity_key) + event = self._refresh_events.get(key) if event is None: # This coroutine wins the right to refresh. event = asyncio.Event() - self._refresh_events[affinity_key] = event + self._refresh_events[key] = event break # Another coroutine is already refreshing — wait for it, then retry. await event.wait() try: if self._clearance_mode == ClearanceMode.MANUAL: - bundle = self._manual.build_bundle(affinity_key=affinity_key) + bundle = self._manual.build_bundle( + affinity_key=affinity_key, + clearance_host=clearance_host, + ) else: bundle = await self._flare.refresh_bundle( affinity_key=affinity_key, proxy_url=proxy_url, + target_url=clearance_origin, ) if bundle: async with self._lock: - self._bundles[affinity_key] = bundle + self._bundles[key] = bundle return bundle finally: async with self._lock: - self._refresh_events.pop(affinity_key, None) + self._refresh_events.pop(key, None) event.set() # Wake all waiters so they retry with the new bundle. # ------------------------------------------------------------------ @@ -308,6 +329,7 @@ async def warm_up(self) -> None: await self._get_or_build_bundle( affinity_key=affinity, proxy_url=proxy_url, + clearance_origin=_DEFAULT_CLEARANCE_ORIGIN, ) async def refresh_clearance_safe(self) -> None: @@ -322,34 +344,45 @@ async def refresh_clearance_safe(self) -> None: return async with self._lock: nodes = list(self._nodes) - existing = set(self._bundles.keys()) + existing = list(self._bundles.keys()) - affinity_items = ( + refresh_targets: dict[BundleKey, tuple[str, str]] = {} + default_items = ( [(n.proxy_url or "direct", n.proxy_url or "") for n in nodes] if nodes else [("direct", "")] ) - # Also refresh bundles for keys that no longer have a matching node - # (e.g. pool was reconfigured) so stale entries get cleaned up. - all_keys = {a for a, _ in affinity_items} | existing + for affinity, proxy_url in default_items: + key: BundleKey = (affinity, _clearance_host(_DEFAULT_CLEARANCE_ORIGIN)) + refresh_targets[key] = (proxy_url, _DEFAULT_CLEARANCE_ORIGIN) + for key in existing: + affinity, clearance_host = key + refresh_targets.setdefault( + key, + ("" if affinity == "direct" else affinity, f"https://{clearance_host}"), + ) - for affinity in all_keys: - proxy_url = "" if affinity == "direct" else affinity + for key, (proxy_url, clearance_origin) in refresh_targets.items(): + affinity, clearance_host = key if self._clearance_mode == ClearanceMode.MANUAL: - new_bundle = self._manual.build_bundle(affinity_key=affinity) + new_bundle = self._manual.build_bundle( + affinity_key=affinity, + clearance_host=clearance_host, + ) else: new_bundle = await self._flare.refresh_bundle( affinity_key=affinity, proxy_url=proxy_url, + target_url=clearance_origin, ) if new_bundle: async with self._lock: - self._bundles[affinity] = new_bundle - logger.debug("clearance bundle refreshed: affinity={}", affinity) + self._bundles[key] = new_bundle + logger.debug("clearance bundle refreshed: bundle={}", key) else: logger.warning( - "clearance refresh failed, keeping old bundle: affinity={}", - affinity, + "clearance refresh failed, keeping old bundle: bundle={}", + key, ) # ------------------------------------------------------------------ @@ -374,7 +407,7 @@ def nodes(self) -> list[EgressNode]: return list(self._nodes) @property - def bundles(self) -> dict[str, ClearanceBundle]: + def bundles(self) -> dict[BundleKey, ClearanceBundle]: """Read-only snapshot of the current clearance bundles.""" return dict(self._bundles) diff --git a/app/control/proxy/models.py b/app/control/proxy/models.py index 73fcdf563..8edde22cc 100644 --- a/app/control/proxy/models.py +++ b/app/control/proxy/models.py @@ -1,9 +1,9 @@ """Control-plane proxy domain models.""" from enum import IntEnum, StrEnum -from typing import Any, Self +from typing import Self -from pydantic import BaseModel, Field +from pydantic import BaseModel class ProxyScope(StrEnum): @@ -76,6 +76,7 @@ class ClearanceBundle(BaseModel): user_agent: str = "" state: ClearanceBundleState = ClearanceBundleState.VALID affinity_key: str = "" # associates bundle with an egress node + clearance_host: str = "grok.com" last_refresh_at: int | None = None # ms @@ -84,6 +85,7 @@ class ProxyLease(BaseModel): proxy_url: str | None = None cf_cookies: str = "" user_agent: str = "" + clearance_host: str = "grok.com" scope: ProxyScope = ProxyScope.APP kind: RequestKind = RequestKind.HTTP acquired_at: int = 0 # ms diff --git a/app/control/proxy/providers/flaresolverr.py b/app/control/proxy/providers/flaresolverr.py index 23eb4dc1a..8f07d037e 100644 --- a/app/control/proxy/providers/flaresolverr.py +++ b/app/control/proxy/providers/flaresolverr.py @@ -2,9 +2,9 @@ import asyncio import json -import re from urllib import request as urllib_request from urllib.error import HTTPError, URLError +from urllib.parse import urlparse from app.platform.logging.logger import logger from app.platform.config.snapshot import get_config @@ -15,18 +15,6 @@ def _extract_all_cookies(cookies: list[dict]) -> str: return "; ".join(f"{c.get('name')}={c.get('value')}" for c in cookies) -def _extract_cookie_value(cookies: list[dict], name: str) -> str: - for c in cookies: - if c.get("name") == name: - return c.get("value") or "" - return "" - - -def _browser_profile(user_agent: str) -> str: - m = re.search(r"Chrome/(\d+)", user_agent) - return f"chrome{m.group(1)}" if m else "chrome120" - - class FlareSolverrClearanceProvider: """Refresh CF clearance bundles via a FlareSolverr instance.""" @@ -35,6 +23,7 @@ async def refresh_bundle( *, affinity_key: str, proxy_url: str, + target_url: str = "https://grok.com", ) -> ClearanceBundle | None: cfg = get_config() mode = ClearanceMode.parse(cfg.get_str("proxy.clearance.mode", "none")) @@ -49,19 +38,22 @@ async def refresh_bundle( fs_url = fs_url, proxy_url = proxy_url, timeout_sec = timeout_sec, + target_url = target_url, ) if not result: logger.warning( - "flaresolverr clearance refresh failed: affinity={} proxy={}", - affinity_key, proxy_url or "", + "flaresolverr clearance refresh failed: affinity={} proxy={} target={}", + affinity_key, proxy_url or "", target_url, ) return None + host = result.get("clearance_host", "grok.com") return ClearanceBundle( - bundle_id = f"flaresolverr:{affinity_key}", + bundle_id = f"flaresolverr:{affinity_key}@{host}", cf_cookies = result.get("cookies", ""), user_agent = result.get("user_agent", ""), affinity_key = affinity_key, + clearance_host = host, ) async def _solve( @@ -70,10 +62,12 @@ async def _solve( fs_url: str, proxy_url: str, timeout_sec: int, + target_url: str, ) -> dict[str, str] | None: + target = target_url.strip() or "https://grok.com" payload: dict = { "cmd": "request.get", - "url": "https://grok.com", + "url": target, "maxTimeout": timeout_sec * 1000, } if proxy_url: @@ -107,10 +101,16 @@ def _post() -> dict: return None ua = solution.get("userAgent", "") or "" + host = (urlparse(target).hostname or "").lower() + filtered = [ + cookie for cookie in cookies + if not host or not cookie.get("domain") or host.endswith(str(cookie.get("domain", "")).lstrip(".").lower()) + ] + chosen = filtered or cookies return { - "cookies": _extract_all_cookies(cookies), + "cookies": _extract_all_cookies(chosen), "user_agent": ua, - "browser": _browser_profile(ua), + "clearance_host": host or "grok.com", } except HTTPError as exc: diff --git a/app/control/proxy/providers/manual.py b/app/control/proxy/providers/manual.py index fcb8df8b2..0844a3673 100644 --- a/app/control/proxy/providers/manual.py +++ b/app/control/proxy/providers/manual.py @@ -8,17 +8,23 @@ class ManualClearanceProvider: """Build a ClearanceBundle from static config values.""" - def build_bundle(self, *, affinity_key: str) -> ClearanceBundle | None: + def build_bundle( + self, + *, + affinity_key: str, + clearance_host: str = "grok.com", + ) -> ClearanceBundle | None: cfg = get_config() mode = ClearanceMode.parse(cfg.get_str("proxy.clearance.mode", "none")) if mode != ClearanceMode.MANUAL: return None clearance = resolve_clearance_config(cfg) return ClearanceBundle( - bundle_id=f"manual:{affinity_key}", + bundle_id=f"manual:{affinity_key}@{clearance_host}", cf_cookies=clearance.cf_cookies, user_agent=clearance.user_agent, affinity_key=affinity_key, + clearance_host=clearance_host, ) diff --git a/app/dataplane/proxy/__init__.py b/app/dataplane/proxy/__init__.py index f7dc43eab..33dfbb6e9 100644 --- a/app/dataplane/proxy/__init__.py +++ b/app/dataplane/proxy/__init__.py @@ -23,8 +23,14 @@ async def acquire( scope: ProxyScope = ProxyScope.APP, kind: RequestKind = RequestKind.HTTP, resource: bool = False, + clearance_origin: str | None = None, ) -> ProxyLease: - return await self._dir.acquire(scope=scope, kind=kind, resource=resource) + return await self._dir.acquire( + scope=scope, + kind=kind, + resource=resource, + clearance_origin=clearance_origin, + ) async def feedback(self, lease: ProxyLease, result: ProxyFeedback) -> None: await self._dir.feedback(lease, result) diff --git a/app/dataplane/proxy/table.py b/app/dataplane/proxy/table.py index 8696b6e3d..3a83cb40d 100644 --- a/app/dataplane/proxy/table.py +++ b/app/dataplane/proxy/table.py @@ -6,11 +6,17 @@ """ from dataclasses import dataclass, field +from typing import TYPE_CHECKING from app.control.proxy.models import ( EgressNode, ClearanceBundle, EgressMode, ClearanceMode, ) +if TYPE_CHECKING: + from app.control.proxy import ProxyDirectory + +BundleKey = tuple[str, str] + @dataclass class ProxyRuntimeTable: @@ -19,7 +25,7 @@ class ProxyRuntimeTable: egress_mode: EgressMode = EgressMode.DIRECT clearance_mode: ClearanceMode = ClearanceMode.NONE nodes: list[EgressNode] = field(default_factory=list) - bundles: dict[str, ClearanceBundle] = field(default_factory=dict) + bundles: dict[BundleKey, ClearanceBundle] = field(default_factory=dict) @property def node_count(self) -> int: diff --git a/app/dataplane/reverse/protocol/xai_auth.py b/app/dataplane/reverse/protocol/xai_auth.py index c814a300c..de1879354 100644 --- a/app/dataplane/reverse/protocol/xai_auth.py +++ b/app/dataplane/reverse/protocol/xai_auth.py @@ -15,6 +15,12 @@ from app.control.proxy.models import ProxyFeedback, ProxyFeedbackKind, ProxyScope, RequestKind from app.dataplane.proxy import get_proxy_runtime from app.dataplane.proxy.adapters.session import ResettableSession, build_session_kwargs +from app.dataplane.reverse.runtime.endpoint_table import ( + ACCEPT_TOS as ACCEPT_TOS_URL, + BASE as GROK_ORIGIN, + NSFW_MGMT as NSFW_MGMT_URL, + SET_BIRTH as SET_BIRTH_URL, +) from app.dataplane.reverse.transport.grpc_web import post_grpc_web from app.dataplane.reverse.transport.http import post_json @@ -25,9 +31,7 @@ # Endpoint URLs # ------------------------------------------------------------------ -ACCEPT_TOS_URL = "https://accounts.x.ai/auth_mgmt.AuthManagement/SetTosAcceptedVersion" -NSFW_MGMT_URL = "https://grok.com/auth_mgmt.AuthManagement/UpdateUserFeatureControls" -SET_BIRTH_URL = "https://grok.com/rest/auth/set-birth-date" +ACCOUNTS_ORIGIN = "https://accounts.x.ai" # ------------------------------------------------------------------ # Payload builders @@ -94,7 +98,11 @@ async def _grpc_call( if not shared: proxy = await get_proxy_runtime() - lease = await proxy.acquire(scope=ProxyScope.APP, kind=RequestKind.HTTP) + lease = await proxy.acquire( + scope=ProxyScope.APP, + kind=RequestKind.HTTP, + clearance_origin=origin, + ) try: _, trailers = await post_grpc_web( @@ -142,8 +150,8 @@ async def accept_tos(token: str) -> GrpcStatus: token, build_accept_tos_payload(), label = "accept_tos", - origin = "https://accounts.x.ai", - referer = "https://accounts.x.ai/accept-tos", + origin = ACCOUNTS_ORIGIN, + referer = f"{ACCOUNTS_ORIGIN}/accept-tos", ) @@ -154,8 +162,8 @@ async def set_nsfw(token: str, enabled: bool) -> GrpcStatus: token, build_nsfw_mgmt_payload(enabled), label = "enable_nsfw" if enabled else "disable_nsfw", - origin = "https://grok.com", - referer = "https://grok.com/?_s=data", + origin = GROK_ORIGIN, + referer = f"{GROK_ORIGIN}/?_s=data", ) @@ -186,14 +194,18 @@ async def set_birth_date( if not shared: proxy = await get_proxy_runtime() - lease = await proxy.acquire(scope=ProxyScope.APP, kind=RequestKind.HTTP) + lease = await proxy.acquire( + scope=ProxyScope.APP, + kind=RequestKind.HTTP, + clearance_origin=GROK_ORIGIN, + ) payload = orjson.dumps(build_set_birth_payload()) try: result = await post_json( SET_BIRTH_URL, token, payload, lease=lease, timeout_s=timeout_s, - origin="https://grok.com", referer="https://grok.com/?_s=data", + origin=GROK_ORIGIN, referer=f"{GROK_ORIGIN}/?_s=data", session=session, ) except UpstreamError as exc: @@ -216,16 +228,20 @@ async def set_birth_date( async def nsfw_sequence(token: str) -> None: - """Run set_birth_date → enable_nsfw sharing one session + lease. + """Run accept_tos → set_birth_date → enable_nsfw. - Compared to calling the two functions individually this saves one TCP+TLS - handshake per token, which matters in high-concurrency batch scenarios. + accept_tos runs against ``accounts.x.ai`` with its own host-specific + clearance. The grok.com birth-date and NSFW update steps still share one + session + lease to avoid an extra handshake per token. """ - cfg = get_config() - timeout_s = cfg.get_float("nsfw.timeout", 30.0) + await accept_tos(token) proxy = await get_proxy_runtime() - lease = await proxy.acquire(scope=ProxyScope.APP, kind=RequestKind.HTTP) + lease = await proxy.acquire( + scope=ProxyScope.APP, + kind=RequestKind.HTTP, + clearance_origin=GROK_ORIGIN, + ) kwargs = build_session_kwargs(lease=lease) try: @@ -233,7 +249,7 @@ async def nsfw_sequence(token: str) -> None: await set_birth_date(token, session=session, lease=lease) await _grpc_call( NSFW_MGMT_URL, token, build_nsfw_mgmt_payload(), - label="enable_nsfw", origin="https://grok.com", referer="https://grok.com/?_s=data", + label="enable_nsfw", origin=GROK_ORIGIN, referer=f"{GROK_ORIGIN}/?_s=data", session=session, lease=lease, ) await proxy.feedback(lease, ProxyFeedback(kind=ProxyFeedbackKind.SUCCESS, status_code=200)) diff --git a/app/products/web/admin/batch.py b/app/products/web/admin/batch.py index dffd758f5..bb4711557 100644 --- a/app/products/web/admin/batch.py +++ b/app/products/web/admin/batch.py @@ -172,10 +172,11 @@ async def _one(token: str) -> None: # --------------------------------------------------------------------------- async def _nsfw_one(repo: "AccountRepository", token: str, enabled: bool) -> dict: - from app.dataplane.reverse.protocol.xai_auth import set_birth_date, set_nsfw + from app.dataplane.reverse.protocol.xai_auth import nsfw_sequence, set_nsfw if enabled: - await set_birth_date(token) - await set_nsfw(token, enabled) + await nsfw_sequence(token) + else: + await set_nsfw(token, enabled) patch = AccountPatch(token=token, add_tags=["nsfw"]) if enabled else AccountPatch(token=token, remove_tags=["nsfw"]) await repo.patch_accounts([patch]) return {"success": True, "tagged": enabled}