from __future__ import annotations from datetime import datetime, timedelta, timezone from typing import Any from urllib.parse import quote import httpx from config import get_settings class RemnawaveError(Exception): def __init__(self, message: str, status_code: int | None = None): super().__init__(message) self.status_code = status_code class RemnawaveClient: def __init__(self) -> None: settings = get_settings() self.base_url = settings.remnawave_base_url.rstrip("/") self.token = settings.remnawave_token def _headers(self) -> dict[str, str]: return { "Authorization": f"Bearer {self.token}", "Content-Type": "application/json", "Accept": "application/json", } async def _request(self, method: str, path: str, **kwargs: Any) -> Any: if not self.token: raise RemnawaveError("REMNAWAVE_TOKEN не задан в .env") url = f"{self.base_url}{path}" async with httpx.AsyncClient(timeout=30.0) as client: response = await client.request( method, url, headers=self._headers(), **kwargs ) if response.status_code >= 400: detail = response.text[:500] raise RemnawaveError( f"Remnawave API error {response.status_code}: {detail}", response.status_code, ) if not response.content: return None data = response.json() if isinstance(data, dict) and "response" in data: return data["response"] return data async def find_user(self, identifier: str) -> dict[str, Any]: identifier = identifier.strip() if not identifier: raise RemnawaveError("Укажите username или UUID") encoded = quote(identifier, safe="") for resolve_body in ( {"username": identifier}, {"shortUuid": identifier}, {"uuid": identifier}, ): try: resolved = await self._request( "POST", "/api/users/resolve", json=resolve_body ) uuid = None if isinstance(resolved, dict): uuid = resolved.get("uuid") if uuid: user = await self._request("GET", f"/api/users/{uuid}") if user: return user except RemnawaveError as exc: if exc.status_code not in (404, 400, 422): raise attempts = [ f"/api/users/by-username/{encoded}", f"/api/users/by-short-uuid/{encoded}", f"/api/users/{encoded}", ] last_error: RemnawaveError | None = None for path in attempts: try: user = await self._request("GET", path) if user: return user except RemnawaveError as exc: last_error = exc if exc.status_code not in (404, 400, 422): raise raise RemnawaveError( "Пользователь не найден. Проверьте username или short UUID." ) from last_error async def update_user(self, payload: dict[str, Any]) -> dict[str, Any]: return await self._request("PATCH", "/api/users", json=payload) @staticmethod def _parse_expire_at(value: str | None) -> datetime: if not value: return datetime.now(timezone.utc) raw = str(value).replace("Z", "+00:00") dt = datetime.fromisoformat(raw) if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc) return dt async def extend_subscription( self, identifier: str, days: int, plan_traffic_gb: int, ) -> dict[str, Any]: """Продление + фиксация лимита тарифа + сброс использованного трафика.""" user = await self.find_user(identifier) now = datetime.now(timezone.utc) current = self._parse_expire_at(user.get("expireAt") or user.get("expire_at")) base = current if current > now else now new_expire = base + timedelta(days=days) old_limit = int( user.get("trafficLimitBytes") or user.get("traffic_limit_bytes") or 0 ) new_limit = int(plan_traffic_gb) * 1024**3 updated = await self.update_user( { "uuid": user["uuid"], "expireAt": new_expire.strftime("%Y-%m-%dT%H:%M:%S.000Z"), "trafficLimitBytes": new_limit, } ) # Сброс счётчика использованного трафика — новый период тарифа try: await self._request( "POST", f"/api/users/{user['uuid']}/actions/reset-traffic" ) except RemnawaveError: # лимит уже выставлен; сброс — желателен, но не откатываем продление pass return { "user": updated or user, "old_expire": current.isoformat(), "new_expire": new_expire.isoformat(), "days": days, "old_limit_gb": round(old_limit / 1024**3, 2), "new_limit_gb": plan_traffic_gb, "plan_traffic_gb": plan_traffic_gb, "username": user.get("username"), "uuid": user.get("uuid"), } async def add_traffic(self, identifier: str, gigabytes: int) -> dict[str, Any]: user = await self.find_user(identifier) current_limit = int( user.get("trafficLimitBytes") or user.get("traffic_limit_bytes") or 0 ) add_bytes = gigabytes * 1024**3 new_limit = current_limit + add_bytes updated = await self.update_user( { "uuid": user["uuid"], "trafficLimitBytes": new_limit, } ) return { "user": updated or user, "old_limit_bytes": current_limit, "new_limit_bytes": new_limit, "old_limit_gb": round(current_limit / 1024**3, 2), "new_limit_gb": round(new_limit / 1024**3, 2), "added_gb": gigabytes, "bytes_added": add_bytes, "username": user.get("username"), "uuid": user.get("uuid"), } async def remove_traffic_bytes(self, user_uuid: str, bytes_to_remove: int) -> dict[str, Any]: user = await self._request("GET", f"/api/users/{user_uuid}") if not user: raise RemnawaveError("Пользователь не найден при откате трафика", 404) current_limit = int( user.get("trafficLimitBytes") or user.get("traffic_limit_bytes") or 0 ) new_limit = max(0, current_limit - int(bytes_to_remove)) updated = await self.update_user( { "uuid": user["uuid"], "trafficLimitBytes": new_limit, } ) return { "user": updated or user, "old_limit_bytes": current_limit, "new_limit_bytes": new_limit, "username": user.get("username"), "uuid": user.get("uuid"), }