remna/app/remnawave.py

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"),
        }