ForcePilot/backend/package/yuxi/channels/adapters/bluebubbles/client.py

from __future__ import annotations

import hashlib
import ipaddress
from typing import Any, Literal
from urllib.parse import urlparse

import httpx

from yuxi.channels.adapters.bluebubbles.exceptions import (
    BlueBubblesAuthenticationError,
    BlueBubblesHTTPError,
)
from yuxi.channels.infra.circuit_breaker import CircuitBreaker
from yuxi.utils.logging_config import logger

_PRIVATE_IP_RANGES = [
    ipaddress.ip_network("10.0.0.0/8"),
    ipaddress.ip_network("172.16.0.0/12"),
    ipaddress.ip_network("192.168.0.0/16"),
    ipaddress.ip_network("127.0.0.0/8"),
    ipaddress.ip_network("169.254.0.0/16"),
    ipaddress.ip_network("::1/128"),
    ipaddress.ip_network("fc00::/7"),
    ipaddress.ip_network("fe80::/10"),
]

SSRFMode = Literal["global_allow", "hostname_whitelist", "default_deny"]
AuthStrategy = Literal["header", "query", "both"]


def _is_private_host(host: str) -> bool:
    try:
        addr = ipaddress.ip_address(host)
    except ValueError:
        return False
    return any(addr in net for net in _PRIVATE_IP_RANGES)


def _resolve_hostname(host: str) -> set[str]:
    import socket

    try:
        addrs = socket.getaddrinfo(host, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
        return {a[4][0] for a in addrs}
    except socket.gaierror:
        return set()


def _build_client_cache_key(server_url: str, account_id: str = "", credential: str = "") -> str:
    raw = f"{server_url}|{account_id}|{credential}"
    return hashlib.sha256(raw.encode()).hexdigest()


class BlueBubblesClient:
    def __init__(
        self,
        server_url: str,
        password: str,
        timeout: float = 30.0,
        max_retries: int = 3,
        allow_private_network: bool = False,
        ssrf_mode: SSRFMode = "default_deny",
        hostname_whitelist: set[str] | None = None,
        auth_strategy: AuthStrategy = "header",
        account_id: str = "",
        connect_timeout: float | None = None,
        send_timeout: float | None = None,
        auth_header_name: str = "X-BB-Password",
    ):
        self._base_url = server_url.rstrip("/")
        self._password = password
        self._timeout = timeout
        self._max_retries = max_retries
        self._allow_private_network = allow_private_network
        self._ssrf_mode: SSRFMode = ssrf_mode
        self._hostname_whitelist: set[str] = hostname_whitelist or set()
        self._auth_strategy: AuthStrategy = auth_strategy
        self._auth_header_name = auth_header_name
        self._account_id = account_id
        self._client_cache_key = _build_client_cache_key(server_url, account_id, password)
        self._client: httpx.AsyncClient | None = None
        self._breaker = CircuitBreaker(failure_threshold=5, recovery_timeout=30.0)
        self._connect_timeout = connect_timeout or timeout
        self._send_timeout = send_timeout or timeout

    @property
    def base_url(self) -> str:
        return self._base_url

    @property
    def client_cache_key(self) -> str:
        return self._client_cache_key

    async def __aenter__(self) -> BlueBubblesClient:
        self._client = httpx.AsyncClient(
            timeout=httpx.Timeout(
                connect=self._connect_timeout,
                read=self._send_timeout,
                write=self._send_timeout,
                pool=self._connect_timeout,
            )
        )
        return self

    async def __aexit__(self, *args) -> None:
        if self._client:
            await self._client.aclose()
            self._client = None

    async def get(self, path: str, **kwargs) -> dict[str, Any]:
        return await self._breaker.call(lambda: self._request("GET", path, **kwargs))

    async def post(self, path: str, **kwargs) -> dict[str, Any]:
        return await self._breaker.call(lambda: self._request("POST", path, **kwargs))

    async def put(self, path: str, **kwargs) -> dict[str, Any]:
        return await self._breaker.call(lambda: self._request("PUT", path, **kwargs))

    async def delete(self, path: str, **kwargs) -> dict[str, Any]:
        return await self._breaker.call(lambda: self._request("DELETE", path, **kwargs))

    def _check_ssrf(self, url: str) -> None:
        if self._ssrf_mode == "global_allow":
            return

        parsed = urlparse(url)
        host = parsed.hostname or ""
        if not host:
            return

        if self._ssrf_mode == "hostname_whitelist":
            if host in self._hostname_whitelist:
                return
            resolved_ips = _resolve_hostname(host)
            for ip_addr in resolved_ips:
                if _is_private_host(ip_addr) and not self._allow_private_network:
                    raise BlueBubblesHTTPError(0, f"SSRF blocked: hostname '{host}' resolves to private IP '{ip_addr}'")
            return

        if self._allow_private_network:
            return

        for resolved_ip in _resolve_hostname(host):
            if _is_private_host(resolved_ip):
                raise BlueBubblesHTTPError(
                    0, f"SSRF blocked: private network host '{host}' resolves to '{resolved_ip}'"
                )

        if _is_private_host(host):
            raise BlueBubblesHTTPError(0, f"SSRF blocked: private network host '{host}'")

    def _build_auth_headers(self) -> dict[str, str]:
        headers: dict[str, str] = {}
        if self._auth_strategy in ("header", "both"):
            headers[self._auth_header_name] = self._password
        return headers

    def _build_url(self, path: str) -> str:
        base = f"{self._base_url}{path}"
        if self._auth_strategy in ("query", "both") and "?" not in path:
            import urllib.parse

            password_encoded = urllib.parse.quote(self._password, safe="")
            sep = "&" if "?" in path else "?"
            return f"{base}{sep}password={password_encoded}"
        return base

    async def _request(self, method: str, path: str, **kwargs) -> dict[str, Any]:
        if self._client is None:
            raise BlueBubblesHTTPError(0, "Client not initialized")

        url = self._build_url(path)
        self._check_ssrf(url)

        headers = kwargs.pop("headers", {})
        headers.update(self._build_auth_headers())

        logger.debug(f"[BlueBubbles] {method} {self._base_url}{path}")

        resp = await self._client.request(method, url, headers=headers, **kwargs)

        if resp.status_code == 401:
            raise BlueBubblesAuthenticationError("Invalid server password")

        if resp.status_code >= 400:
            raise BlueBubblesHTTPError(
                resp.status_code,
                message=resp.text[:500],
                headers=dict(resp.headers),
            )

        return resp.json()

    async def close(self) -> None:
        if self._client:
            await self._client.aclose()
            self._client = None

    async def download(self, file_path: str) -> bytes:
        if self._client is None:
            raise BlueBubblesHTTPError(0, "Client not initialized")

        url = self._build_url(file_path)
        self._check_ssrf(url)

        headers = self._build_auth_headers()
        resp = await self._breaker.call(lambda: self._client.get(url, headers=headers))
        return resp.content