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