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 = "", ): 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._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) @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(self._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["X-BlueBubbles-Password"] = 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