feat(rw): remnawave integration
This commit is contained in:
480
external/rw.py
vendored
Normal file
480
external/rw.py
vendored
Normal file
@@ -0,0 +1,480 @@
|
||||
import logging
|
||||
import re
|
||||
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
|
||||
|
||||
from config import cfg
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RW_USERNAME_MAX_LENGTH = 36
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────
|
||||
# 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,
|
||||
)
|
||||
|
||||
|
||||
def _build_rw_username(*, telegram_id: int | None, username: str | None) -> str | None:
|
||||
if telegram_id is not None:
|
||||
return f"tg_{telegram_id}"
|
||||
|
||||
if not username:
|
||||
return None
|
||||
|
||||
sanitized_username = re.sub(r"[^a-zA-Z0-9_-]+", "_", username).strip("_")
|
||||
if not sanitized_username:
|
||||
return None
|
||||
|
||||
return f"web_{sanitized_username}"[:RW_USERNAME_MAX_LENGTH]
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────
|
||||
# Публичные функции — обёртки с перехватом ошибок
|
||||
# ─────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
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_user_by_username(
|
||||
sdk: RemnawaveSDK,
|
||||
username: str,
|
||||
) -> RWUserInfo | None:
|
||||
try:
|
||||
user = await sdk.users.get_user_by_username(username=username)
|
||||
if user is None:
|
||||
return None
|
||||
return _parse_user(user)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to get RW user by username=%s: %s",
|
||||
username,
|
||||
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 | None,
|
||||
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
|
||||
|
||||
|
||||
def get_sdk() -> RemnawaveSDK | None:
|
||||
if not cfg.remnawave_base_url or not cfg.remnawave_token:
|
||||
return None
|
||||
|
||||
return RemnawaveSDK(base_url=cfg.remnawave_base_url, token=cfg.remnawave_token)
|
||||
|
||||
|
||||
async def sync_subscription_by_telegram_id(
|
||||
*,
|
||||
expires_at: datetime,
|
||||
devices: int,
|
||||
telegram_id: int | None = None,
|
||||
username: str | None = None,
|
||||
) -> bool:
|
||||
sdk = get_sdk()
|
||||
if sdk is None:
|
||||
logger.info("Skipping RW subscription sync: RemnaWave is not configured")
|
||||
return False
|
||||
|
||||
rw_user = None
|
||||
if telegram_id is not None:
|
||||
rw_user = await get_user_by_telegram_id(sdk, telegram_id)
|
||||
|
||||
if rw_user is None:
|
||||
rw_username = _build_rw_username(telegram_id=telegram_id, username=username)
|
||||
if rw_username is None:
|
||||
logger.warning(
|
||||
"Skipping RW subscription sync: cannot build username for telegram_id=%s username=%s",
|
||||
telegram_id,
|
||||
username,
|
||||
)
|
||||
return False
|
||||
|
||||
rw_user = await get_user_by_username(sdk, rw_username)
|
||||
if rw_user is None:
|
||||
rw_user = await create_user(
|
||||
sdk=sdk,
|
||||
username=rw_username,
|
||||
telegram_id=telegram_id,
|
||||
expire_at=expires_at,
|
||||
hwid_device_limit=None if devices == 0 else devices,
|
||||
squad_uuids=cfg.remnawave_default_squads,
|
||||
)
|
||||
if rw_user is None:
|
||||
logger.warning(
|
||||
"Failed to create RW user for telegram_id=%s username=%s",
|
||||
telegram_id,
|
||||
username,
|
||||
)
|
||||
return False
|
||||
|
||||
expire_synced = await update_expire_at(sdk, rw_user.uuid, expires_at)
|
||||
devices_synced = await set_hwid_limit(sdk, rw_user.uuid, devices)
|
||||
return expire_synced and devices_synced
|
||||
Reference in New Issue
Block a user