app/dataplane/reverse/transport/websocket.py

"""WebSocket transport with proxy and SOCKS support."""

import ssl
from typing import Any, Awaitable, Callable, Mapping, Optional
from urllib.parse import urlparse

import aiohttp
import certifi
from aiohttp_socks import ProxyConnector

from app.platform.logging.logger import logger
from app.platform.config.snapshot import get_config
from app.control.proxy.models import ProxyLease


def _ssl_ctx() -> ssl.SSLContext:
    ctx = ssl.create_default_context()
    ctx.load_verify_locations(certifi.where())
    return ctx


def _normalize_socks(proxy_url: str) -> tuple[str, Optional[bool]]:
    scheme = urlparse(proxy_url).scheme.lower()
    rdns: Optional[bool] = None
    base = scheme
    if scheme == "socks":
        base, rdns = "socks5", True
    elif scheme == "socks5h":
        base, rdns = "socks5", True
    elif scheme == "socks4a":
        base, rdns = "socks4", True
    if base != scheme:
        proxy_url = proxy_url.replace(f"{scheme}://", f"{base}://", 1)
    return proxy_url, rdns


def _build_connector(
    proxy_url: str,
    ssl_ctx:   ssl.SSLContext,
) -> tuple[aiohttp.BaseConnector, str | None]:
    if not proxy_url:
        return aiohttp.TCPConnector(ssl=ssl_ctx), None
    scheme = urlparse(proxy_url).scheme.lower()
    if scheme.startswith("socks"):
        normalized, rdns = _normalize_socks(proxy_url)
        kwargs: dict = {"ssl": ssl_ctx}
        if rdns is not None:
            kwargs["rdns"] = rdns
        logger.debug("websocket connector selected: proxy_type=socks proxy_url={}", proxy_url)
        return ProxyConnector.from_url(normalized, **kwargs), None
    logger.debug("websocket connector selected: proxy_type=http proxy_url={}", proxy_url)
    return aiohttp.TCPConnector(ssl=ssl_ctx), proxy_url


class WebSocketConnection:
    """Wraps aiohttp WebSocketResponse with lifecycle cleanup."""

    def __init__(
        self,
        session: aiohttp.ClientSession,
        ws:      aiohttp.ClientWebSocketResponse,
        *,
        on_close: Callable[[], Awaitable[None]] | None = None,
    ) -> None:
        self.session  = session
        self.ws       = ws
        self._on_close = on_close

    async def close(self) -> None:
        if not self.ws.closed:
            await self.ws.close()
        await self.session.close()
        if self._on_close:
            await self._on_close()
            self._on_close = None

    async def __aenter__(self) -> aiohttp.ClientWebSocketResponse:
        return self.ws

    async def __aexit__(self, *_: Any) -> None:
        await self.close()


class WebSocketClient:
    """Establish WebSocket connections through optional proxy."""

    def __init__(self, proxy: str | None = None) -> None:
        self._proxy_override = proxy
        self._ssl = _ssl_ctx()

    async def connect(
        self,
        url:       str,
        headers:   Optional[Mapping[str, str]] = None,
        timeout:   Optional[float]             = None,
        ws_kwargs: Optional[Mapping[str, Any]] = None,
        *,
        lease:    ProxyLease | None                      = None,
        on_close: Callable[[], Awaitable[None]] | None   = None,
    ) -> WebSocketConnection:
        proxy_url  = self._proxy_override or (lease.proxy_url if lease else "")
        connector, http_proxy = _build_connector(proxy_url, self._ssl)

        cfg           = get_config()
        total_s       = float(timeout) if timeout is not None else cfg.get_float("voice.timeout", 120.0)
        client_timeout = aiohttp.ClientTimeout(total=total_s)
        session        = aiohttp.ClientSession(connector=connector, timeout=client_timeout)

        try:
            extra: dict[str, Any] = dict(ws_kwargs or {})
            skip_ssl = cfg.get_bool("proxy.egress.skip_ssl_verify", False) and bool(proxy_url)

            if skip_ssl and urlparse(proxy_url).scheme.lower() == "https":
                proxy_ssl = ssl.create_default_context()
                proxy_ssl.check_hostname = False
                proxy_ssl.verify_mode    = ssl.CERT_NONE
                try:
                    ws = await session.ws_connect(
                        url, headers=headers, proxy=http_proxy,
                        ssl=self._ssl, proxy_ssl=proxy_ssl, **extra,
                    )
                except TypeError:
                    ws = await session.ws_connect(
                        url, headers=headers, proxy=http_proxy,
                        ssl=self._ssl, **extra,
                    )
            else:
                ws = await session.ws_connect(
                    url, headers=headers, proxy=http_proxy,
                    ssl=self._ssl, **extra,
                )

            return WebSocketConnection(session, ws, on_close=on_close)
        except Exception:
            await session.close()
            raise


__all__ = ["WebSocketClient", "WebSocketConnection"]