ForcePilot/backend/package/yuxi/channel/extensions/bluebubbles/client.py

import hashlib
from urllib.parse import urlparse

import httpx


class BlueBubblesAuthStrategy:
    def apply(self, request_kwargs: dict, password: str) -> dict:
        raise NotImplementedError


class QueryStringAuth(BlueBubblesAuthStrategy):
    def apply(self, request_kwargs: dict, password: str) -> dict:
        params = request_kwargs.get("params", {})
        if isinstance(params, dict):
            params["password"] = password
        else:
            params = {"password": password}
        request_kwargs["params"] = params
        return request_kwargs


class HeaderAuth(BlueBubblesAuthStrategy):
    def __init__(self, header_name: str = "X-BB-Password"):
        self.header_name = header_name

    def apply(self, request_kwargs: dict, password: str) -> dict:
        headers = request_kwargs.get("headers", {})
        if isinstance(headers, dict):
            headers[self.header_name] = password
        else:
            headers = {self.header_name: password}
        request_kwargs["headers"] = headers
        return request_kwargs


class BlueBubblesClient:
    def __init__(
        self,
        server_url: str,
        password: str,
        account_id: str = "default",
        timeout_ms: int = 30000,
        auth_strategy: BlueBubblesAuthStrategy | None = None,
        allow_private_network: bool = False,
    ):
        self.server_url = server_url.rstrip("/")
        self.password = password
        self.account_id = account_id
        self.timeout = httpx.Timeout(timeout_ms / 1000.0)
        self.auth_strategy = auth_strategy or QueryStringAuth()
        self._trusted_hostname = self._resolve_hostname(server_url)
        self._allow_private_network = allow_private_network
        self._client: httpx.AsyncClient | None = None

    @staticmethod
    def _resolve_hostname(url: str) -> str | None:
        try:
            return urlparse(url).hostname
        except Exception:
            return None

    def _build_transport(self) -> httpx.AsyncHTTPTransport | None:
        if self._allow_private_network:
            return None
        return None

    async def _get_client(self) -> httpx.AsyncClient:
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                base_url=self.server_url,
                timeout=self.timeout,
                transport=self._build_transport(),
            )
        return self._client

    def _fingerprint(self) -> str:
        raw = f"{self.account_id}:{self.password}:{type(self.auth_strategy).__name__}"
        return hashlib.sha256(raw.encode()).hexdigest()[:8]

    async def request(self, method: str, path: str, **kwargs) -> httpx.Response:
        kwargs = self.auth_strategy.apply(kwargs, self.password)
        client = await self._get_client()
        return await client.request(method, path, **kwargs)

    async def get(self, path: str, **kwargs) -> httpx.Response:
        return await self.request("GET", path, **kwargs)

    async def post(self, path: str, **kwargs) -> httpx.Response:
        return await self.request("POST", path, **kwargs)

    async def put(self, path: str, **kwargs) -> httpx.Response:
        return await self.request("PUT", path, **kwargs)

    async def delete(self, path: str, **kwargs) -> httpx.Response:
        return await self.request("DELETE", path, **kwargs)

    async def ping(self) -> bool:
        try:
            resp = await self.get("/")
            return resp.status_code < 500
        except Exception:
            return False

    async def close(self):
        if self._client and not self._client.is_closed:
            await self._client.aclose()
            self._client = None


_client_cache: dict[str, BlueBubblesClient] = {}


def get_or_create_client(
    server_url: str,
    password: str,
    account_id: str = "default",
    **kwargs,
) -> BlueBubblesClient:
    fingerprint = hashlib.sha256(f"{account_id}:{password}:{server_url}".encode()).hexdigest()[:12]
    if fingerprint not in _client_cache:
        _client_cache[fingerprint] = BlueBubblesClient(server_url, password, account_id, **kwargs)
    return _client_cache[fingerprint]


def clear_client_cache():
    _client_cache.clear()