SurfSense/surfsense_backend/app/proprietary/platforms/youtube/innertube.py

"""Proxy-aware fetch seam + InnerTube payload builder for the YouTube scraper.

All network I/O for the scraper flows through :func:`fetch_html` and
:func:`post_innertube`, and always egresses through the residential proxy (never
a direct connection — that would expose and risk-block the server IP). Within a
continuation chain, :func:`proxy_session` opens ONE reused keep-alive session so
sequential pages share a connection and a sticky exit IP (roughly halving warm
latency vs a fresh rotating IP per request). HTML pages fall back to a
``StealthyFetcher`` browser tier for anti-bot walls. Parsers and the orchestrator
never see which tier served the bytes.

The ``AsyncFetcher.post(..., json=...)`` + ``page.json()`` / ``page.html_content``
pattern is already proven in ``app/routes/youtube_routes.py``.
"""

from __future__ import annotations

import asyncio
import logging
import time
from contextlib import asynccontextmanager, suppress
from contextvars import ContextVar
from typing import Any

from scrapling.fetchers import AsyncFetcher, FetcherSession

from app.utils.proxy import get_proxy_url

try:  # browser tier is optional (needs patchright browsers installed)
    from scrapling.fetchers import StealthyFetcher
except Exception:  # pragma: no cover - import guard
    StealthyFetcher = None  # type: ignore[assignment]

logger = logging.getLogger(__name__)

# Per-flow proxy session (set by ``proxy_session`` around one continuation
# chain). Reusing one keep-alive connection halves warm latency AND pins a
# single sticky exit IP on the rotating gateway, instead of paying a fresh
# TCP+TLS handshake and drawing a new (often slow) residential node per request.
# A ContextVar keeps each concurrent fan-out flow on its own session/IP without
# threading a param through every parser call.
_current_session: ContextVar[_RotatingSession | None] = ContextVar(
    "yt_proxy_session", default=None
)

# Statuses that mean "this IP is throttled/blocked" -> rotate to a fresh one.
# A probe of 120 sequential requests on one sticky IP saw zero blocks, so we
# stay sticky for speed and rotate only *reactively*, up to this many times per
# request. ponytail upgrade path: also rotate proactively every N requests if a
# future run ever trips these.
_BLOCK_STATUSES = frozenset({403, 429})
_MAX_ROTATIONS = 3


class _RotatingSession:
    """Owns one live ``FetcherSession`` (sticky IP) and can swap it for a fresh one.

    ``rotate()`` closes the current keep-alive connection and opens a new one, so
    the rotating gateway hands out a different residential exit IP. Used
    sequentially within a single flow (never shared across concurrent tasks), so
    no locking is needed. ``session`` is ``None`` only when no proxy is configured.
    """

    def __init__(self) -> None:
        self._cm: Any | None = None
        self.session: Any | None = None
        self.rotations = 0

    async def _open(self) -> None:
        proxy = get_proxy_url()
        if proxy is None:
            self._cm = self.session = None
            return
        self._cm = FetcherSession(
            proxy=proxy, stealthy_headers=True, impersonate="chrome"
        )
        self.session = await self._cm.__aenter__()

    async def close(self) -> None:
        if self._cm is not None:
            with suppress(Exception):  # best-effort teardown
                await self._cm.__aexit__(None, None, None)
        self._cm = self.session = None

    async def rotate(self) -> Any | None:
        """Drop the current IP and connect through a fresh one. Returns new session."""
        await self.close()
        self.rotations += 1
        await self._open()
        logger.info("[youtube] rotated proxy session (rotation #%d)", self.rotations)
        return self.session


async def open_proxy_holder() -> _RotatingSession:
    """Open a warm rotate-on-block session holder (caller owns ``close()``)."""
    holder = _RotatingSession()
    await holder._open()
    return holder


@asynccontextmanager
async def bind_proxy_holder(holder: _RotatingSession):
    """Route this task's fetches through ``holder`` for the enclosed block.

    Does NOT close the holder — enables pooling warm sessions across sequential
    jobs so each job skips the ~2s proxy TCP+TLS handshake.
    """
    token = _current_session.set(holder)
    try:
        yield holder
    finally:
        _current_session.reset(token)


@asynccontextmanager
async def proxy_session():
    """Open one reused, rotate-on-block proxy session for a continuation chain."""
    holder = await open_proxy_holder()
    try:
        async with bind_proxy_holder(holder):
            yield holder
    finally:
        await holder.close()


# Consent cookies to dodge the EU consent interstitial that otherwise returns a
# page with no ``ytInitialData``. Mirrors app/routes/youtube_routes.py.
CONSENT_COOKIES = {
    "CONSENT": "PENDING+999",
    "SOCS": "CAISNQgDEitib3FfaWRlbnRpdHlmcm9udGVuZHVpc2VydmVyXzIwMjMwODI5LjA3X3AxGgJlbiADGgYIgOa_pgY",
}

# Public InnerTube "WEB" client block. The web API key below is YouTube's
# long-standing public key; ``search``/``next`` may reject a keyless POST where
# ``browse`` accepts it, so the builder can attach both key and visitorData.
INNERTUBE_PUBLIC_API_KEY = "AIzaSyAO_FJ2SlqU8Q4STEHLGCilw_Y9_11qcW8"

INNERTUBE_BROWSE_URL = "https://www.youtube.com/youtubei/v1/browse"
INNERTUBE_SEARCH_URL = "https://www.youtube.com/youtubei/v1/search"
INNERTUBE_NEXT_URL = "https://www.youtube.com/youtubei/v1/next"


def build_innertube_payload(
    *,
    continuation_token: str | None = None,
    browse_id: str | None = None,
    search_query: str | None = None,
    search_params: str | None = None,
    visitor_data: str | None = None,
    video_id: str | None = None,
    hl: str | None = None,
) -> dict[str, Any]:
    """Build the standard InnerTube ``context.client`` payload.

    Ported from the Scrapfly reference ``call_youtube_api`` but I/O-free so it is
    unit-testable and reusable across browse/search/next/player endpoints.
    ``hl`` selects the response language: a ``/next`` call with a non-default
    ``hl`` returns the video's creator-localized title/description.
    """
    client: dict[str, Any] = {
        "hl": hl or "en",
        "gl": "US",
        "clientName": "WEB",
        "clientVersion": "2.20241111.07.00",
        "platform": "DESKTOP",
        "userInterfaceTheme": "USER_INTERFACE_THEME_DARK",
    }
    if visitor_data:
        client["visitorData"] = visitor_data

    payload: dict[str, Any] = {
        "context": {
            "client": client,
            "user": {"lockedSafetyMode": False},
            "request": {"useSsl": True},
        }
    }
    if browse_id is not None:
        payload["browseId"] = browse_id
    if video_id is not None:
        payload["videoId"] = video_id
    if search_query is not None:
        payload["query"] = search_query
        if search_params is not None:
            payload["params"] = search_params
    if continuation_token is not None:
        payload["continuation"] = continuation_token
    return payload


async def post_innertube(
    base_url: str,
    payload: dict[str, Any],
    *,
    api_key: str | None = None,
) -> dict[str, Any] | None:
    """POST an InnerTube payload and return parsed JSON (or ``None`` on failure).

    Always egresses through the proxy (a reused per-flow session when one is
    open, else a one-shot ``AsyncFetcher``). Never connects directly — a direct
    hit would expose and risk-block the server IP.
    """
    url = f"{base_url}?key={api_key}" if api_key else base_url
    holder = _current_session.get()
    for attempt in range(_MAX_ROTATIONS + 1):
        session = holder.session if holder else None
        try:
            started = time.perf_counter()
            if session is not None:
                page = await session.post(url, json=payload)
            else:
                page = await AsyncFetcher.post(
                    url, json=payload, proxy=get_proxy_url(), stealthy_headers=True
                )
            fetch_ms = (time.perf_counter() - started) * 1000
            logger.info(
                "[youtube][perf] source=innertube reused=%s url=%s status=%s fetch_ms=%.1f",
                session is not None,
                base_url,
                page.status,
                fetch_ms,
            )
            if page.status == 200:
                return page.json()
            logger.warning("InnerTube POST %s returned %s", base_url, page.status)
            if not (
                holder and page.status in _BLOCK_STATUSES and attempt < _MAX_ROTATIONS
            ):
                return None
        except Exception as e:
            logger.warning("InnerTube POST %s failed: %s", base_url, e)
            if not (holder and attempt < _MAX_ROTATIONS):
                return None
        await holder.rotate()  # blocked/errored on a proxy session: try a fresh IP
    return None


async def fetch_html(url: str, *, cookies: dict[str, str] | None = None) -> str | None:
    """GET a YouTube page and return raw HTML (or ``None`` on failure).

    Fetches through the proxy (reused per-flow session when open, else a one-shot
    ``AsyncFetcher``); falls back to the ``StealthyFetcher`` browser tier for
    anti-bot walls. Sets consent cookies by default so EU-egress fetches don't
    hit the consent wall and return a page without ``ytInitialData``.
    """
    merged = {**CONSENT_COOKIES, **(cookies or {})}
    headers = {"Accept-Language": "en-US,en;q=0.9"}
    holder = _current_session.get()
    for attempt in range(_MAX_ROTATIONS + 1):
        session = holder.session if holder else None
        try:
            started = time.perf_counter()
            if session is not None:
                page = await session.get(url, headers=headers, cookies=merged)
            else:
                page = await AsyncFetcher.get(
                    url,
                    headers=headers,
                    cookies=merged,
                    proxy=get_proxy_url(),
                    stealthy_headers=True,
                )
            fetch_ms = (time.perf_counter() - started) * 1000
            logger.info(
                "[youtube][perf] source=html reused=%s url=%s status=%s fetch_ms=%.1f",
                session is not None,
                url,
                page.status,
                fetch_ms,
            )
            if page.status == 200:
                return page.html_content
            logger.warning("HTML GET %s returned %s", url, page.status)
            if not (
                holder and page.status in _BLOCK_STATUSES and attempt < _MAX_ROTATIONS
            ):
                break
        except Exception as e:
            logger.warning("HTML GET %s failed: %s", url, e)
            if not (holder and attempt < _MAX_ROTATIONS):
                break
        await holder.rotate()  # blocked/errored on a proxy session: try a fresh IP
    # All proxy attempts blocked/failed -> last-resort browser tier.
    return await _fetch_html_stealthy(url)


async def _fetch_html_stealthy(url: str) -> str | None:
    """Last-resort browser fetch for anti-bot walls (mirrors the crawler tier).

    ``StealthyFetcher.fetch`` is synchronous, so it runs in a worker thread to
    keep the event loop free. Returns ``None`` when the browser tier is
    unavailable or still blocked.
    """
    if StealthyFetcher is None:
        return None
    try:
        started = time.perf_counter()
        page = await asyncio.to_thread(
            StealthyFetcher.fetch,
            url,
            headless=True,
            network_idle=True,
            solve_cloudflare=True,
            proxy=get_proxy_url(),
        )
        fetch_ms = (time.perf_counter() - started) * 1000
        logger.info(
            "[youtube][perf] source=html tier=stealthy url=%s status=%s fetch_ms=%.1f",
            url,
            page.status,
            fetch_ms,
        )
        if page.status == 200:
            return page.html_content
        logger.warning("HTML GET %s tier=stealthy returned %s", url, page.status)
    except Exception as e:
        logger.warning("HTML GET %s tier=stealthy failed: %s", url, e)
    return None