marzban-sub-bot/vpn_bot/marzban.py

from __future__ import annotations

from typing import Any

import httpx

from .utils import (
    DAYS_PER_MONTH,
    PAID_TRAFFIC_BYTES,
    TRIAL_DAYS,
    TRIAL_TRAFFIC_BYTES,
    add_days_from_base,
    build_marzban_note,
    calculate_reset_expire,
    format_marzban_username,
    now_ts,
)

DEFAULT_PROXIES = {"vless": {"flow": "xtls-rprx-vision"}}
DEFAULT_INBOUNDS = {"vless": ["VLESS TCP REALITY"]}
DATA_LIMIT_RESET_STRATEGY = "month"


class MarzbanError(RuntimeError):
    pass


class MarzbanClient:
    def __init__(self, *, base_url: str, username: str, password: str, timeout: float = 30.0):
        self.base_url = base_url.rstrip("/")
        self.username = username
        self.password = password
        self.timeout = timeout
        self._token: str | None = None
        self._client = httpx.AsyncClient(base_url=self.base_url, timeout=timeout)

    async def close(self) -> None:
        await self._client.aclose()

    @staticmethod
    def build_trial_user_payload(
        *,
        tg_user_id: int,
        first_name: str | None,
        last_name: str | None,
        tg_username: str | None,
        phone: str | None,
        now: int | None = None,
    ) -> dict[str, Any]:
        current = now_ts() if now is None else now
        return {
            "username": format_marzban_username(tg_user_id),
            "proxies": DEFAULT_PROXIES,
            "inbounds": DEFAULT_INBOUNDS,
            "expire": current + TRIAL_DAYS * 24 * 60 * 60,
            "data_limit": TRIAL_TRAFFIC_BYTES,
            "data_limit_reset_strategy": DATA_LIMIT_RESET_STRATEGY,
            "status": "active",
            "note": build_marzban_note(
                first_name=first_name,
                last_name=last_name,
                username=tg_username,
                phone=phone,
            ),
        }

    async def authenticate(self) -> str:
        response = await self._client.post(
            "/api/admin/token",
            data={"username": self.username, "password": self.password},
            headers={"Content-Type": "application/x-www-form-urlencoded"},
        )
        if response.status_code >= 400:
            raise MarzbanError(f"ошибка авторизации Marzban: {response.status_code} {response.text}")
        payload = response.json()
        token = payload.get("access_token")
        if not token:
            raise MarzbanError("ответ авторизации Marzban не содержит access_token")
        self._token = str(token)
        return self._token

    async def _request(self, method: str, path: str, **kwargs: Any) -> Any:
        if not self._token:
            await self.authenticate()
        headers = kwargs.pop("headers", {})
        headers["Authorization"] = f"Bearer {self._token}"
        response = await self._client.request(method, path, headers=headers, **kwargs)
        if response.status_code == 401:
            await self.authenticate()
            headers["Authorization"] = f"Bearer {self._token}"
            response = await self._client.request(method, path, headers=headers, **kwargs)
        if response.status_code >= 400:
            raise MarzbanError(f"ошибка API Marzban {method} {path}: {response.status_code} {response.text}")
        if not response.content:
            return None
        return response.json()

    async def create_trial_user(
        self,
        *,
        tg_user_id: int,
        first_name: str | None,
        last_name: str | None,
        tg_username: str | None,
        phone: str | None = None,
    ) -> dict[str, Any]:
        payload = self.build_trial_user_payload(
            tg_user_id=tg_user_id,
            first_name=first_name,
            last_name=last_name,
            tg_username=tg_username,
            phone=phone,
        )
        return await self._request("POST", "/api/user", json=payload)

    async def get_user(self, username: str) -> dict[str, Any]:
        return await self._request("GET", f"/api/user/{username}")

    async def modify_user(self, username: str, payload: dict[str, Any]) -> dict[str, Any]:
        return await self._request("PUT", f"/api/user/{username}", json=payload)

    async def reset_user_usage(self, username: str) -> dict[str, Any] | None:
        return await self._request("POST", f"/api/user/{username}/reset")

    async def extend_user(self, username: str, months: int) -> dict[str, Any]:
        user = await self.get_user(username)
        current_expire = user.get("expire")
        new_expire = add_days_from_base(current_expire=current_expire, days=months * DAYS_PER_MONTH)
        payload = {
            "expire": new_expire,
            "data_limit": PAID_TRAFFIC_BYTES,
            "data_limit_reset_strategy": DATA_LIMIT_RESET_STRATEGY,
            "status": "active",
        }
        await self.modify_user(username, payload)
        await self.reset_user_usage(username)
        return await self.get_user(username)

    async def reset_paid_traffic_with_time_penalty(self, username: str) -> tuple[dict[str, Any], int, int]:
        user = await self.get_user(username)
        new_expire, full_months = calculate_reset_expire(expire=user.get("expire"))
        payload = {
            "expire": new_expire,
            "data_limit": PAID_TRAFFIC_BYTES,
            "data_limit_reset_strategy": DATA_LIMIT_RESET_STRATEGY,
            "status": "active",
        }
        await self.modify_user(username, payload)
        await self.reset_user_usage(username)
        fresh = await self.get_user(username)
        return fresh, full_months, new_expire