97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import ipaddress
|
|
from typing import Any
|
|
from urllib.parse import urlparse
|
|
|
|
import httpx
|
|
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
|
|
def _normalize_ship_url(url: str) -> str:
|
|
url = url.rstrip("/")
|
|
if not url.startswith(("http://", "https://")):
|
|
url = f"https://{url}"
|
|
|
|
try:
|
|
parsed = urlparse(url)
|
|
hostname = parsed.hostname
|
|
if hostname:
|
|
addr = ipaddress.ip_address(hostname.strip("[]"))
|
|
if isinstance(addr, ipaddress.IPv6Address):
|
|
host_part = f"[{addr}]" if "[" not in parsed.netloc else parsed.hostname
|
|
port_part = f":{parsed.port}" if parsed.port else ""
|
|
url = f"{parsed.scheme}://{host_part}{port_part}{parsed.path}"
|
|
except (ValueError, ipaddress.AddressValueError):
|
|
pass
|
|
|
|
return url
|
|
|
|
|
|
class UrbitClient:
|
|
def __init__(
|
|
self,
|
|
ship_url: str,
|
|
ship_name: str,
|
|
timeout: float = 15.0,
|
|
):
|
|
self.ship_url = _normalize_ship_url(ship_url)
|
|
self.ship_name = ship_name.lstrip("~")
|
|
self._http: httpx.AsyncClient | None = None
|
|
self._timeout = timeout
|
|
self._session_cookie: str | None = None
|
|
|
|
@property
|
|
def http(self) -> httpx.AsyncClient:
|
|
if self._http is None:
|
|
raise RuntimeError("UrbitClient not started. Call start() first.")
|
|
return self._http
|
|
|
|
@property
|
|
def session_cookie(self) -> str | None:
|
|
return self._session_cookie
|
|
|
|
@property
|
|
def auth_cookie_name(self) -> str:
|
|
return f"urbauth-~{self.ship_name}"
|
|
|
|
async def start(self) -> None:
|
|
self._http = httpx.AsyncClient(
|
|
base_url=self.ship_url,
|
|
timeout=httpx.Timeout(self._timeout),
|
|
follow_redirects=True,
|
|
)
|
|
logger.debug(f"[Urbit] HTTP client started for {self.ship_url}")
|
|
|
|
async def stop(self) -> None:
|
|
if self._http:
|
|
await self._http.aclose()
|
|
self._http = None
|
|
self._session_cookie = None
|
|
logger.debug("[Urbit] HTTP client stopped")
|
|
|
|
def set_session_cookie(self, cookie: str) -> None:
|
|
self._session_cookie = cookie
|
|
|
|
def _cookies(self) -> dict[str, str] | None:
|
|
if self._session_cookie:
|
|
return {self.auth_cookie_name: self._session_cookie}
|
|
return None
|
|
|
|
async def post(self, path: str, json: dict[str, Any] | None = None, **kwargs) -> httpx.Response:
|
|
return await self.http.post(path, json=json, cookies=self._cookies(), **kwargs)
|
|
|
|
async def get(self, path: str, **kwargs) -> httpx.Response:
|
|
return await self.http.get(path, cookies=self._cookies(), **kwargs)
|
|
|
|
async def put(self, path: str, json: dict[str, Any] | None = None, **kwargs) -> httpx.Response:
|
|
return await self.http.put(path, json=json, cookies=self._cookies(), **kwargs)
|
|
|
|
async def delete(self, path: str, **kwargs) -> httpx.Response:
|
|
return await self.http.delete(path, cookies=self._cookies(), **kwargs)
|
|
|
|
async def stream_get(self, path: str, timeout: float = 300.0, **kwargs) -> httpx.Response:
|
|
req = self.http.build_request("GET", self.ship_url + path, cookies=self._cookies(), **kwargs)
|
|
return await self.http.send(req, stream=True, timeout=httpx.Timeout(timeout))
|