monitor-shared-utils/src/monitor_utils/sessions.py

"""Shared HTTP session helpers for all monitors.

Known gap: UA/TLS fingerprint does not always match headers/header order,
but targeted antibots do not currently seem to care.
"""

import asyncio
import json
import logging
import random
from contextlib import asynccontextmanager
from typing import AsyncGenerator

import aiohttp
import requests
from aioproxies import AioProxies
from curl_cffi.requests import AsyncSession
from curl_cffi.requests.errors import RequestsError
from curl_cffi.requests.exceptions import Timeout as CurlTimeout
from fake_useragent import FakeUserAgent
from tenacity import (
    after_log,
    before_log,
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_fixed,
)

from monitor_utils.constants import (
    CLIENT_IDENTIFIERS,
    CLIENT_TIMEOUT,
    RETRY_LIMIT,
    RETRY_WAIT,
    TIMEOUT_SECONDS,
)

LOGGER = logging.getLogger(__name__)
UA = FakeUserAgent()


@asynccontextmanager
async def get_base_session(
    headers: dict | None = None,
    cookies: dict | None = None,
    connector: aiohttp.TCPConnector | None = None,
    timeout: aiohttp.ClientTimeout = CLIENT_TIMEOUT,
    proxies: AioProxies | None = None,
) -> AsyncGenerator[aiohttp.ClientSession, None]:
    """Yield an aiohttp session with a random User-Agent and attached proxy pool.

    The proxy pool is stored on the session as proxy_pool for send_request to
    rotate per request. The session is always closed on exit.

    Args:
        headers: Headers applied to the session; User-Agent is defaulted if
            absent.
        cookies: Cookies applied to the session.
        connector: Optional TCP connector.
        timeout: Client timeout applied to the session.
        proxies: Proxy pool attached to the session.

    Yields:
        The configured aiohttp session.
    """

    headers = headers if headers else {}
    headers.setdefault("User-Agent", UA.random)
    session = aiohttp.ClientSession(
        headers=headers, cookies=cookies, connector=connector, timeout=timeout
    )
    session.proxy_pool = proxies
    try:
        yield session
    finally:
        await session.close()


@asynccontextmanager
async def get_tls_session(
    headers: dict | None = None,
    timeout: int = TIMEOUT_SECONDS,
    proxies: AioProxies | None = None,
) -> AsyncGenerator[AsyncSession, None]:
    """Yield a curl_cffi session impersonating a random client fingerprint.

    The proxy pool is stored on the session as proxy_pool for send_request to
    rotate per request.

    Args:
        headers: Headers applied to the session; a Chrome User-Agent is
            defaulted if absent.
        timeout: Request timeout in seconds.
        proxies: Proxy pool attached to the session.

    Yields:
        The configured curl_cffi session.
    """

    headers = headers if headers else {}
    headers.setdefault("User-Agent", UA.chrome)
    async with AsyncSession(
        impersonate=random.choice(CLIENT_IDENTIFIERS),
        timeout=timeout,
        headers=headers,
    ) as session:
        session.proxy_pool = proxies
        yield session


def build_session_details(
    session: aiohttp.ClientSession | AsyncSession, headers: dict = None
):
    """Merge headers into an existing session, defaulting a random User-Agent.

    Args:
        session: Session whose headers are updated in place.
        headers: Headers to merge; User-Agent is defaulted if absent.
    """

    headers = headers or {}
    headers.setdefault("User-Agent", UA.random)
    session.headers.update(headers)


@retry(
    stop=stop_after_attempt(RETRY_LIMIT),
    wait=wait_fixed(RETRY_WAIT),
    retry=retry_if_exception_type(
        (
            aiohttp.ClientError,
            asyncio.TimeoutError,
            RequestsError,
            requests.Timeout,
            CurlTimeout,
        )
    ),
    before=before_log(LOGGER, logging.DEBUG),
    after=after_log(LOGGER, logging.DEBUG),
    reraise=True,
)
async def send_request(
    session: aiohttp.ClientSession | AsyncSession,
    url: str,
    method: str,
    params: dict | None = None,
    json_data: dict | None = None,
    data: dict | None = None,
    reraise_json_error: bool = False,
    development_proxy: str | None = None,
) -> dict | str:
    """Send a request through either session type with retries and proxy rotation.

    Retries transport errors and timeouts up to RETRY_LIMIT with a fixed wait.
    When development_proxy is set it overrides the session's proxy pool and
    disables TLS verification for local traffic inspection; otherwise the next
    proxy is drawn from the session's proxy_pool when present. Responses are
    parsed as JSON when possible.

    Args:
        session: aiohttp or curl_cffi session to send through.
        url: Request URL.
        method: HTTP method.
        params: Query parameters.
        json_data: JSON body.
        data: Form body.
        reraise_json_error: Raise JSONDecodeError instead of returning raw
            text when the response is not valid JSON.
        development_proxy: Local proxy URL that overrides the pool and
            disables verification.

    Returns:
        Parsed JSON dict, or raw response text when parsing fails and
        reraise_json_error is False.

    Raises:
        json.JSONDecodeError: When parsing fails and reraise_json_error is
            True.
    """

    verify = True
    if development_proxy:
        proxy = development_proxy
        verify = False
    else:
        manager = getattr(session, "proxy_pool", None)
        proxy = manager.next().url() if manager else None

    if isinstance(session, aiohttp.ClientSession):
        async with session.request(
            method=method,
            url=url,
            params=params,
            json=json_data,
            data=data,
            proxy=proxy,
            ssl=verify,
        ) as response:
            response.raise_for_status()
            text = await response.text()
    else:
        response = await session.request(
            method=method.upper(),
            url=url,
            params=params,
            json=json_data,
            data=data,
            proxy=proxy,
            verify=verify,
        )
        response.raise_for_status()
        text = response.text

    try:
        return json.loads(text)
    except json.JSONDecodeError:
        if reraise_json_error:
            raise

        return text