first com2

2026-04-08 17:10:36 +07:00
parent 8a597f880f
commit fc45816eba
34 changed files with 1310 additions and 29 deletions
--- a/services/rw.py
+++ b/services/rw.py
@@ -0,0 +1,320 @@
+import logging
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from typing import Optional
+from uuid import UUID
+
+from remnawave import RemnawaveSDK
+from remnawave.models import UpdateUserRequestDto
+
+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: Optional[int]  # None = unlimited
+    active_squads: list[dict]  # [{"uuid": "...", "name": "..."}]
+    description: Optional[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,
+    )
+
+
+# ─────────────────────────────────────────────────────────────────
+#  Публичные функции — обёртки с перехватом ошибок
+# ─────────────────────────────────────────────────────────────────
+
+
+async def get_user_by_telegram_id(
+    sdk: RemnawaveSDK,
+    telegram_id: int,
+) -> Optional[RWUserInfo]:
+    """
+    Получить первого пользователя 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,
+) -> Optional[datetime]:
+    """
+    Добавить 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=timezone.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=timezone.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,
+) -> Optional[datetime]:
+    """
+    Вычесть 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=timezone.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=timezone.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: Optional[int] = 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