387 lines
13 KiB
Python
387 lines
13 KiB
Python
import logging
|
||
from dataclasses import dataclass
|
||
from datetime import UTC, datetime, timedelta
|
||
from uuid import UUID
|
||
|
||
from remnawave import RemnawaveSDK
|
||
from remnawave.models import (
|
||
CreateUserRequestDto,
|
||
DeleteUserHwidDeviceResponseDto,
|
||
HWIDDeleteRequest,
|
||
UpdateUserRequestDto,
|
||
)
|
||
from remnawave.models.hwid import HwidDeviceDto
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────
|
||
# Dataclass — унифицированное представление пользователя RW
|
||
# ─────────────────────────────────────────────────────────────────
|
||
|
||
|
||
@dataclass
|
||
class RWUserInfo:
|
||
uuid: str
|
||
username: str
|
||
status: str # "ACTIVE", "DISABLED", "EXPIRED", "ON_HOLD", "LIMITED"
|
||
expire_at: datetime
|
||
used_traffic_bytes: float
|
||
traffic_limit_bytes: int # 0 = unlimited
|
||
hwid_device_limit: int | None # None = unlimited
|
||
active_squads: list[dict] # [{"uuid": "...", "name": "..."}]
|
||
description: str | None
|
||
short_uuid: str
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────
|
||
# Внутренние хелперы
|
||
# ─────────────────────────────────────────────────────────────────
|
||
|
||
|
||
def _parse_user(user_dto) -> RWUserInfo:
|
||
"""Преобразовать UserResponseDto → RWUserInfo."""
|
||
active_squads: list[dict] = []
|
||
for squad in user_dto.active_internal_squads or []:
|
||
active_squads.append({"uuid": str(squad.uuid), "name": squad.name})
|
||
|
||
# UserStatus может быть enum-объектом или строкой — нормализуем
|
||
status_raw = user_dto.status
|
||
if hasattr(status_raw, "value"):
|
||
status_str = str(status_raw.value).upper()
|
||
else:
|
||
status_str = str(status_raw).upper()
|
||
|
||
# Трафик лежит в user_traffic.used_traffic_bytes
|
||
used_traffic: float = 0.0
|
||
if user_dto.user_traffic is not None:
|
||
used_traffic = float(user_dto.user_traffic.used_traffic_bytes or 0)
|
||
|
||
return RWUserInfo(
|
||
uuid=str(user_dto.uuid),
|
||
username=user_dto.username,
|
||
status=status_str,
|
||
expire_at=user_dto.expire_at,
|
||
used_traffic_bytes=used_traffic,
|
||
traffic_limit_bytes=int(user_dto.traffic_limit_bytes or 0),
|
||
hwid_device_limit=user_dto.hwid_device_limit,
|
||
active_squads=active_squads,
|
||
description=user_dto.description,
|
||
short_uuid=user_dto.short_uuid,
|
||
)
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────
|
||
# Публичные функции — обёртки с перехватом ошибок
|
||
# ─────────────────────────────────────────────────────────────────
|
||
|
||
|
||
async def get_user_by_telegram_id(
|
||
sdk: RemnawaveSDK,
|
||
telegram_id: int,
|
||
) -> RWUserInfo | None:
|
||
"""
|
||
Получить первого пользователя Remnawave, привязанного к Telegram ID.
|
||
Возвращает None, если пользователь не найден или Remnawave недоступен.
|
||
"""
|
||
try:
|
||
users_list = await sdk.users.get_users_by_telegram_id(str(telegram_id))
|
||
if not users_list:
|
||
return None
|
||
return _parse_user(users_list[0])
|
||
except Exception as exc:
|
||
logger.warning(
|
||
"Failed to get RW user by telegram_id=%s: %s",
|
||
telegram_id,
|
||
exc,
|
||
)
|
||
return None
|
||
|
||
|
||
async def get_all_squads(sdk: RemnawaveSDK) -> list[dict]:
|
||
"""
|
||
Получить все Internal Squads из Remnawave.
|
||
Возвращает список [{"uuid": "...", "name": "..."}].
|
||
При ошибке возвращает пустой список.
|
||
"""
|
||
try:
|
||
response = await sdk.internal_squads.get_internal_squads()
|
||
result: list[dict] = []
|
||
# GetAllInternalSquadsResponseDto имеет поле .internal_squads
|
||
for squad in response.internal_squads:
|
||
result.append({"uuid": str(squad.uuid), "name": squad.name})
|
||
return result
|
||
except Exception as exc:
|
||
logger.warning("Failed to get all internal squads: %s", exc)
|
||
return []
|
||
|
||
|
||
async def add_days(
|
||
sdk: RemnawaveSDK,
|
||
user_uuid: str,
|
||
days: int,
|
||
) -> datetime | None:
|
||
"""
|
||
Добавить N дней к expire_at пользователя.
|
||
Возвращает новую дату окончания или None при ошибке.
|
||
"""
|
||
try:
|
||
# Получаем актуальный expire_at — нельзя слепо использовать
|
||
# закешированное значение, чтобы не потерять изменения других агентов
|
||
user_dto = await sdk.users.get_user_by_uuid(uuid=user_uuid)
|
||
if user_dto is None or user_dto.expire_at is None:
|
||
logger.warning("Cannot add days: user %s has no expire_at", user_uuid)
|
||
return None
|
||
|
||
# Обеспечиваем timezone-aware datetime
|
||
current_expire = user_dto.expire_at
|
||
if current_expire.tzinfo is None:
|
||
current_expire = current_expire.replace(tzinfo=UTC)
|
||
|
||
new_expire = current_expire + timedelta(days=days)
|
||
|
||
result = await sdk.users.update_user(
|
||
UpdateUserRequestDto(
|
||
uuid=UUID(user_uuid),
|
||
expire_at=new_expire,
|
||
)
|
||
)
|
||
if result is None:
|
||
return None
|
||
|
||
new_dt = result.expire_at
|
||
if new_dt is not None and new_dt.tzinfo is None:
|
||
new_dt = new_dt.replace(tzinfo=UTC)
|
||
return new_dt
|
||
except Exception as exc:
|
||
logger.warning("Failed to add %d days for user %s: %s", days, user_uuid, exc)
|
||
return None
|
||
|
||
|
||
async def remove_days(
|
||
sdk: RemnawaveSDK,
|
||
user_uuid: str,
|
||
days: int,
|
||
) -> datetime | None:
|
||
"""
|
||
Вычесть N дней из expire_at пользователя.
|
||
Возвращает новую дату окончания или None при ошибке.
|
||
"""
|
||
try:
|
||
user_dto = await sdk.users.get_user_by_uuid(uuid=user_uuid)
|
||
if user_dto is None or user_dto.expire_at is None:
|
||
logger.warning("Cannot remove days: user %s has no expire_at", user_uuid)
|
||
return None
|
||
|
||
current_expire = user_dto.expire_at
|
||
if current_expire.tzinfo is None:
|
||
current_expire = current_expire.replace(tzinfo=UTC)
|
||
|
||
new_expire = current_expire - timedelta(days=days)
|
||
|
||
result = await sdk.users.update_user(
|
||
UpdateUserRequestDto(
|
||
uuid=UUID(user_uuid),
|
||
expire_at=new_expire,
|
||
)
|
||
)
|
||
if result is None:
|
||
return None
|
||
|
||
new_dt = result.expire_at
|
||
if new_dt is not None and new_dt.tzinfo is None:
|
||
new_dt = new_dt.replace(tzinfo=UTC)
|
||
return new_dt
|
||
except Exception as exc:
|
||
logger.warning("Failed to remove %d days for user %s: %s", days, user_uuid, exc)
|
||
return None
|
||
|
||
|
||
async def set_hwid_limit(
|
||
sdk: RemnawaveSDK,
|
||
user_uuid: str,
|
||
limit: int,
|
||
) -> bool:
|
||
"""
|
||
Установить лимит HWID-устройств.
|
||
limit == 0 → убрать ограничение (None).
|
||
Возвращает True при успехе, False при ошибке.
|
||
"""
|
||
try:
|
||
hwid_value: int | None = None if limit == 0 else limit
|
||
await sdk.users.update_user(
|
||
UpdateUserRequestDto(
|
||
uuid=UUID(user_uuid),
|
||
hwid_device_limit=hwid_value,
|
||
)
|
||
)
|
||
return True
|
||
except Exception as exc:
|
||
logger.warning(
|
||
"Failed to set HWID limit=%s for user %s: %s",
|
||
limit,
|
||
user_uuid,
|
||
exc,
|
||
)
|
||
return False
|
||
|
||
|
||
async def set_description(
|
||
sdk: RemnawaveSDK,
|
||
user_uuid: str,
|
||
note: str,
|
||
) -> bool:
|
||
"""
|
||
Установить заметку (description) пользователя в Remnawave.
|
||
Возвращает True при успехе, False при ошибке.
|
||
"""
|
||
try:
|
||
await sdk.users.update_user(
|
||
UpdateUserRequestDto(
|
||
uuid=UUID(user_uuid),
|
||
description=note,
|
||
)
|
||
)
|
||
return True
|
||
except Exception as exc:
|
||
logger.warning(
|
||
"Failed to set description for user %s: %s",
|
||
user_uuid,
|
||
exc,
|
||
)
|
||
return False
|
||
|
||
|
||
async def reset_traffic(
|
||
sdk: RemnawaveSDK,
|
||
user_uuid: str,
|
||
) -> bool:
|
||
"""
|
||
Сбросить счётчик использованного трафика пользователя.
|
||
Возвращает True при успехе, False при ошибке.
|
||
"""
|
||
try:
|
||
await sdk.users.reset_user_traffic(uuid=user_uuid)
|
||
return True
|
||
except Exception as exc:
|
||
logger.warning("Failed to reset traffic for user %s: %s", user_uuid, exc)
|
||
return False
|
||
|
||
|
||
async def disable_user(
|
||
sdk: RemnawaveSDK,
|
||
user_uuid: str,
|
||
) -> bool:
|
||
"""
|
||
Отключить пользователя (статус → DISABLED).
|
||
Возвращает True при успехе, False при ошибке.
|
||
"""
|
||
try:
|
||
await sdk.users.disable_user(uuid=user_uuid)
|
||
return True
|
||
except Exception as exc:
|
||
logger.warning("Failed to disable user %s: %s", user_uuid, exc)
|
||
return False
|
||
|
||
|
||
async def enable_user(
|
||
sdk: RemnawaveSDK,
|
||
user_uuid: str,
|
||
) -> bool:
|
||
"""
|
||
Включить пользователя (статус → ACTIVE).
|
||
Возвращает True при успехе, False при ошибке.
|
||
"""
|
||
try:
|
||
await sdk.users.enable_user(uuid=user_uuid)
|
||
return True
|
||
except Exception as exc:
|
||
logger.warning("Failed to enable user %s: %s", user_uuid, exc)
|
||
return False
|
||
|
||
|
||
async def update_squads(
|
||
sdk: RemnawaveSDK,
|
||
user_uuid: str,
|
||
squad_uuid_list: list[str],
|
||
) -> bool:
|
||
"""
|
||
Полностью заменить список Internal Squads пользователя.
|
||
squad_uuid_list — строковые UUID всех squad'ов (полный новый список).
|
||
Возвращает True при успехе, False при ошибке.
|
||
"""
|
||
try:
|
||
await sdk.users.update_user(
|
||
UpdateUserRequestDto(
|
||
uuid=UUID(user_uuid),
|
||
active_internal_squads=[UUID(s) for s in squad_uuid_list],
|
||
)
|
||
)
|
||
return True
|
||
except Exception as exc:
|
||
logger.warning(
|
||
"Failed to update squads for user %s: %s",
|
||
user_uuid,
|
||
exc,
|
||
)
|
||
return False
|
||
|
||
|
||
async def create_user(
|
||
sdk: RemnawaveSDK,
|
||
username: str,
|
||
telegram_id: int,
|
||
expire_at: datetime,
|
||
hwid_device_limit: int | None,
|
||
squad_uuids: list[str],
|
||
) -> RWUserInfo | None:
|
||
"""
|
||
Создать нового пользователя в Remnawave.
|
||
hwid_device_limit=None → без ограничений.
|
||
Возвращает RWUserInfo при успехе, None при ошибке.
|
||
"""
|
||
try:
|
||
dto = CreateUserRequestDto(
|
||
username=username,
|
||
telegram_id=telegram_id,
|
||
expire_at=expire_at,
|
||
hwid_device_limit=hwid_device_limit,
|
||
active_internal_squads=[UUID(s) for s in squad_uuids] if squad_uuids else None,
|
||
)
|
||
result = await sdk.users.create_user(dto)
|
||
if result is None:
|
||
return None
|
||
return _parse_user(result)
|
||
except Exception as exc:
|
||
logger.warning("Failed to create user username=%s: %s", username, exc)
|
||
return None
|
||
|
||
|
||
async def get_hwid_list(sdk: RemnawaveSDK, user_uuid: str) -> list[HwidDeviceDto] | None:
|
||
try:
|
||
resp = await sdk.hwid.get_hwid_user(user_uuid)
|
||
return resp.devices
|
||
except Exception:
|
||
logger.exception("failed to fetch hwid list for user=%s", user_uuid)
|
||
return
|
||
|
||
|
||
async def delete_hwid(sdk: RemnawaveSDK, user_uuid: str, hwid: str) -> bool:
|
||
body = HWIDDeleteRequest(user_uuid=user_uuid, hwid=hwid)
|
||
try:
|
||
resp = await sdk.hwid.delete_hwid_to_user(body)
|
||
return isinstance(resp, DeleteUserHwidDeviceResponseDto)
|
||
except Exception:
|
||
logger.exception("failed to delete hwid=%s for user=%s", hwid, user_uuid)
|
||
return False
|
||
|
||
|
||
async def update_expire_at(sdk: RemnawaveSDK, user_uuid: str, expire_at: datetime):
|
||
dto = UpdateUserRequestDto(uuid=user_uuid, expire_at=expire_at) # type: ignore
|
||
try:
|
||
await sdk.users.update_user(body=dto)
|
||
return True
|
||
except Exception:
|
||
logger.exception("failed to update expire_at=%s for user=%s", str(expire_at), user_uuid)
|
||
return False
|