Files
Solo_bot/database/keys.py
T

734 lines
27 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import asyncio
from datetime import UTC, datetime
from types import SimpleNamespace
from sqlalchemy import delete, func, select, text, update
from sqlalchemy.ext.asyncio import AsyncSession
from core.cache_config import (
KEY_COUNT_CACHE_TTL_SEC,
KEY_DETAILS_CACHE_TTL_SEC,
KEYS_LIST_CACHE_TTL_SEC,
)
from core.redis_cache import cache_delete, cache_get, cache_key, cache_set
from database.access.resolution import resolve_user_optional
from database.models import Key, Tariff, User
from database.users import invalidate_profile_cache, invalidate_user_snapshot
from logger import logger
async def invalidate_key_details(email: str) -> None:
await cache_delete(cache_key("key_details", email))
async def invalidate_key_email(client_id: str) -> None:
await cache_delete(cache_key("key_email", client_id))
async def _purge_keys_cache_ids(*ids: int) -> None:
for i in ids:
await cache_delete(cache_key("keys_list", i))
await cache_delete(cache_key("key_count", i))
await invalidate_profile_cache(i)
async def invalidate_keys_list(session: AsyncSession, legacy_user_ref: int) -> None:
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
await _purge_keys_cache_ids(legacy_user_ref)
return
if u.tg_id is not None:
await _purge_keys_cache_ids(u.id, u.tg_id)
else:
await _purge_keys_cache_ids(u.id)
async def invalidate_key_details_by_client_id(session: AsyncSession, client_id: str) -> None:
email = await cache_get(cache_key("key_email", client_id))
await cache_delete(cache_key("key_email", client_id))
if email:
await invalidate_key_details(str(email))
else:
res = await session.execute(select(Key.email).where(Key.client_id == client_id).limit(1))
row = res.scalar_one_or_none()
if row is not None:
await invalidate_key_details(str(row))
async def store_key(
session: AsyncSession,
legacy_user_ref: int,
client_id: str,
email: str,
expiry_time: int,
key: str,
server_id: str,
remnawave_link: str = None,
tariff_id: int | None = None,
alias: str | None = None,
selected_device_limit: int | None = None,
selected_traffic_limit: int | None = None,
selected_price_rub: int | None = None,
current_device_limit: int | None = None,
current_traffic_limit: int | None = None,
):
"""Сохраняет или обновляет ключ подписки."""
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
raise ValueError(f"Пользователь не найден для ключа: {legacy_user_ref}")
uid = u.id
exists = await session.execute(select(Key).where(Key.user_id == uid, Key.client_id == client_id))
existing_key = exists.scalar_one_or_none()
if existing_key:
values: dict = {
"email": email,
"expiry_time": expiry_time,
"key": key,
"server_id": server_id,
"remnawave_link": remnawave_link,
"tariff_id": tariff_id,
"alias": alias,
"tg_id": u.tg_id,
}
if selected_device_limit is not None:
values["selected_device_limit"] = selected_device_limit
if selected_traffic_limit is not None:
values["selected_traffic_limit"] = selected_traffic_limit
if selected_price_rub is not None:
values["selected_price_rub"] = selected_price_rub
if current_device_limit is not None:
values["current_device_limit"] = current_device_limit
if current_traffic_limit is not None:
values["current_traffic_limit"] = current_traffic_limit
await session.execute(update(Key).where(Key.user_id == uid, Key.client_id == client_id).values(**values))
logger.info(f"[Store Key] Ключ обновлён: user_id={uid}, client_id={client_id}, server_id={server_id}")
else:
if current_device_limit is None:
current_device_limit = selected_device_limit
if current_traffic_limit is None:
current_traffic_limit = selected_traffic_limit
new_key = Key(
user_id=uid,
tg_id=u.tg_id,
client_id=client_id,
email=email,
created_at=int(datetime.now(UTC).timestamp() * 1000),
expiry_time=expiry_time,
key=key,
server_id=server_id,
remnawave_link=remnawave_link,
tariff_id=tariff_id,
alias=alias,
selected_device_limit=selected_device_limit,
selected_traffic_limit=selected_traffic_limit,
selected_price_rub=selected_price_rub,
current_device_limit=current_device_limit,
current_traffic_limit=current_traffic_limit,
)
add_result = session.add(new_key)
if asyncio.iscoroutine(add_result):
await add_result
logger.info(f"[Store Key] Ключ создан: user_id={uid}, client_id={client_id}, server_id={server_id}")
invalidate_user_snapshot(uid)
if u.tg_id is not None:
invalidate_user_snapshot(u.tg_id)
await invalidate_keys_list(session, uid)
await invalidate_key_details(email)
def _key_to_cache_dict(k: Key) -> dict:
return {
"email": k.email,
"alias": k.alias,
"client_id": k.client_id,
"expiry_time": int(k.expiry_time) if k.expiry_time is not None else 0,
"created_at": int(k.created_at) if k.created_at is not None else 0,
"tariff_id": k.tariff_id,
"server_id": k.server_id,
"key": k.key,
"remnawave_link": k.remnawave_link,
"is_frozen": bool(k.is_frozen) if k.is_frozen is not None else False,
}
async def get_keys(session: AsyncSession, legacy_user_ref: int):
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
return []
uid = u.id
ckey = cache_key("keys_list", uid)
cached = await cache_get(ckey)
if isinstance(cached, list):
return [SimpleNamespace(**d) for d in cached]
result = await session.execute(select(Key).where(Key.user_id == uid))
rows = result.scalars().all()
serialized = [_key_to_cache_dict(k) for k in rows]
await cache_set(ckey, serialized, KEYS_LIST_CACHE_TTL_SEC)
return rows
async def get_all_keys(session: AsyncSession):
result = await session.execute(select(Key))
return result.scalars().all()
async def get_key_by_server(session: AsyncSession, legacy_user_ref: int, client_id: str):
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
return None
stmt = select(Key).where(Key.user_id == u.id, Key.client_id == client_id)
result = await session.execute(stmt)
return result.scalar_one_or_none()
async def get_key_by_email(session: AsyncSession, email: str, legacy_user_ref: int | None = None) -> Key | None:
stmt = select(Key).where(Key.email == email)
if legacy_user_ref is not None:
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
return None
stmt = stmt.where(Key.user_id == u.id)
result = await session.execute(stmt.limit(1))
return result.scalar_one_or_none()
async def get_key_by_client_id(session: AsyncSession, client_id: str, legacy_user_ref: int | None = None) -> Key | None:
stmt = select(Key).where(Key.client_id == client_id)
if legacy_user_ref is not None:
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
return None
stmt = stmt.where(Key.user_id == u.id)
result = await session.execute(stmt.limit(1))
return result.scalar_one_or_none()
async def get_key_expiry_presets(session: AsyncSession, email: str) -> tuple[str | None, list[int]]:
key_obj = await get_key_by_email(session, email)
if not key_obj:
return None, []
if not key_obj.tariff_id:
return key_obj.client_id, []
tariff = await session.execute(select(Tariff.group_code).where(Tariff.id == key_obj.tariff_id))
group_code = tariff.scalar_one_or_none()
if not group_code:
return key_obj.client_id, []
result = await session.execute(
select(Tariff.duration_days)
.where(Tariff.group_code == group_code, Tariff.is_active.is_(True))
.order_by(Tariff.duration_days)
)
unique_durations: list[int] = []
seen: set[int] = set()
for (days,) in result.all():
if days is None or days < 1 or days in seen:
continue
seen.add(int(days))
unique_durations.append(int(days))
return key_obj.client_id, unique_durations
async def get_key_details(session: AsyncSession, email: str) -> dict | None:
"""Возвращает подробную информацию о ключе по email. Горячие данные кэшируются в Redis."""
ckey = cache_key("key_details", email)
cached = await cache_get(ckey)
if isinstance(cached, dict):
return cached
stmt = select(Key, User).join(User, Key.user_id == User.id).where(Key.email == email)
result = await session.execute(stmt)
row = result.first()
if not row:
return None
key, user = row
expiry_date = datetime.fromtimestamp(key.expiry_time / 1000, UTC)
current_date = datetime.now(UTC)
time_left = expiry_date - current_date
if time_left.total_seconds() <= 0:
days_left_message = "<b>Ключ истек.</b>"
elif time_left.days > 0:
days_left_message = f"Осталось дней: <b>{time_left.days}</b>"
else:
hours_left = time_left.seconds // 3600
days_left_message = f"Осталось часов: <b>{hours_left}</b>"
out = {
"key": key.key,
"remnawave_link": key.remnawave_link,
"server_id": key.server_id,
"created_at": key.created_at,
"expiry_time": key.expiry_time,
"client_id": key.client_id,
"tg_id": user.tg_id,
"email": key.email,
"is_frozen": key.is_frozen,
"balance": user.balance,
"alias": key.alias,
"expiry_date": expiry_date.strftime("%d %B %Y года %H:%M"),
"days_left_message": days_left_message,
"link": key.key or key.remnawave_link,
"cluster_name": key.server_id,
"location_name": key.server_id,
"tariff_id": key.tariff_id,
"selected_device_limit": key.selected_device_limit,
"selected_traffic_limit": key.selected_traffic_limit,
"selected_price_rub": key.selected_price_rub,
"current_device_limit": key.current_device_limit,
"current_traffic_limit": key.current_traffic_limit,
}
await cache_set(ckey, out, KEY_DETAILS_CACHE_TTL_SEC)
if key.client_id:
await cache_set(cache_key("key_email", key.client_id), email, KEY_DETAILS_CACHE_TTL_SEC)
return out
async def get_key_count(session: AsyncSession, legacy_user_ref: int) -> int:
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
return 0
uid = u.id
cached = await cache_get(cache_key("key_count", uid))
if cached is not None:
try:
return int(cached)
except (TypeError, ValueError):
pass
result = await session.execute(select(func.count()).select_from(Key).where(Key.user_id == uid))
count = result.scalar() or 0
await cache_set(cache_key("key_count", uid), count, KEY_COUNT_CACHE_TTL_SEC)
return count
async def get_key_by_user_and_email(session: AsyncSession, user_id: int, email: str) -> Key | None:
"""Возвращает ORM-объект Key по паре (users.id, email) или None."""
result = await session.execute(
select(Key).where(Key.user_id == int(user_id), Key.email == email)
)
return result.scalar_one_or_none()
async def delete_key_by_user_and_email(session: AsyncSession, user_id: int, email: str) -> None:
"""Удаляет ключ по паре (users.id, email). Commit — ответственность caller'а."""
await session.execute(
delete(Key).where(Key.user_id == int(user_id), Key.email == email)
)
async def get_user_keys_with_servers_by_email(
session: AsyncSession, user_id: int, email: str
) -> list[tuple[str, str, dict]]:
"""Возвращает ключи пользователя + инфо о серверах (join Key × Server).
Каждый элемент — ``(client_id, server_id, server_info_dict)``. Join
делается по (Key.server_id == Server.server_name OR Server.cluster_name),
чтобы поддержать и country-mode (server_id = cluster), и cluster-mode
(server_id = server_name).
Используется в ``services.operations.traffic.get_user_traffic``.
"""
from sqlalchemy import or_
from database.models import Server
join_cond = or_(
Key.server_id == Server.server_name,
Key.server_id == Server.cluster_name,
)
result = await session.execute(
select(Key.client_id, Key.server_id, Server)
.select_from(Key)
.join(Server, join_cond)
.where(Server.enabled.is_(True), Key.user_id == int(user_id), Key.email == email)
)
rows = []
for client_id, server_id, server in result.all():
rows.append((
client_id,
server_id,
{
"server_name": server.server_name,
"cluster_name": server.cluster_name,
"api_url": server.api_url,
"panel_type": server.panel_type,
},
))
return rows
async def get_key_client_id_by_email_and_server(
session: AsyncSession, email: str, server_id: str
) -> str | None:
"""Возвращает ``client_id`` первого ключа для пары (email, server_id).
Используется для remnawave traffic reset, где нам нужен только client_id,
без остальных полей ключа.
"""
result = await session.execute(
select(Key.client_id)
.where(Key.email == email, Key.server_id == server_id)
.limit(1)
)
return result.scalar()
async def count_keys_by_server_id(session: AsyncSession, server_id: str) -> int:
"""Сколько всего ключей привязано к указанному server_id (кластеру или серверу).
Используется для проверки max_keys лимита. ``server_id`` — строка
(у ``keys.server_id`` колонка типа String, содержит либо cluster_name,
либо server_name в зависимости от страны/кластера).
"""
result = await session.execute(
select(func.count()).select_from(Key).where(Key.server_id == server_id)
)
return int(result.scalar() or 0)
async def get_all_key_server_ids(session: AsyncSession) -> list[str]:
"""Список всех ``server_id`` из таблицы keys (с повторениями).
Используется в ``services.clusters.select_cluster`` для подсчёта загрузки
кластеров. Возвращаем только server_id строки без подгрузки остальных
полей, чтобы не тянуть сотни мегабайт для огромных deployments.
"""
result = await session.execute(select(Key.server_id))
return [row[0] for row in result.all() if row[0] is not None]
async def count_active_keys_for_user(session: AsyncSession, user_id: int) -> int:
"""Количество незамороженных ключей у пользователя (по internal users.id).
Отличается от `get_key_count`: не кэшируется и явно исключает замороженные.
Используется в проверке "новый пользователь" для купонных правил.
"""
result = await session.execute(
select(func.count())
.select_from(Key)
.where(Key.user_id == int(user_id), Key.is_frozen.is_(False))
)
return int(result.scalar() or 0)
async def delete_key(session: AsyncSession, identifier: int | str):
legacy_for_cache = None
email_for_cache = None
if isinstance(identifier, str):
res = await session.execute(
select(Key.user_id, Key.email).where(Key.client_id == identifier).limit(1)
)
row = res.first()
if row:
legacy_for_cache, email_for_cache = row[0], row[1]
await cache_delete(cache_key("key_email", identifier))
await session.execute(delete(Key).where(Key.client_id == identifier))
else:
u = await resolve_user_optional(session, identifier)
if u is None:
logger.info(f"Ключ не удалён: пользователь {identifier} не найден")
return
legacy_for_cache = u.id
await session.execute(delete(Key).where(Key.user_id == u.id))
if legacy_for_cache is not None:
invalidate_user_snapshot(legacy_for_cache)
await invalidate_keys_list(session, legacy_for_cache)
if email_for_cache is not None:
await invalidate_key_details(str(email_for_cache))
logger.info(f"Ключ с идентификатором {identifier} удалён")
async def update_key_expiry(session: AsyncSession, client_id: str, new_expiry_time: int):
await session.execute(update(Key).where(Key.client_id == client_id).values(expiry_time=new_expiry_time))
await invalidate_key_details_by_client_id(session, client_id)
logger.info(f"Срок действия ключа {client_id} обновлён до {new_expiry_time}")
async def get_client_id_by_email(session: AsyncSession, email: str):
result = await session.execute(select(Key.client_id).where(Key.email == email))
return result.scalar_one_or_none()
async def update_key_notified(session: AsyncSession, legacy_user_ref: int, client_id: str):
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
return
await session.execute(update(Key).where(Key.user_id == u.id, Key.client_id == client_id).values(notified=True))
await invalidate_keys_list(session, u.id)
await invalidate_key_details_by_client_id(session, client_id)
async def mark_key_as_frozen(session: AsyncSession, legacy_user_ref: int, client_id: str, time_left: int):
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
return
await session.execute(
text(
"""
UPDATE keys
SET expiry_time = :expiry,
is_frozen = TRUE
WHERE user_id = :user_id
AND client_id = :client_id
"""
),
{"expiry": time_left, "user_id": u.id, "client_id": client_id},
)
await invalidate_keys_list(session, u.id)
await invalidate_key_details_by_client_id(session, client_id)
async def mark_key_as_unfrozen(
session: AsyncSession,
legacy_user_ref: int,
client_id: str,
new_expiry_time: int,
):
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
return
await session.execute(
text(
"""
UPDATE keys
SET expiry_time = :expiry,
is_frozen = FALSE
WHERE user_id = :user_id
AND client_id = :client_id
"""
),
{"expiry": new_expiry_time, "user_id": u.id, "client_id": client_id},
)
await invalidate_keys_list(session, u.id)
await invalidate_key_details_by_client_id(session, client_id)
async def update_key_tariff(session: AsyncSession, client_id: str, tariff_id: int):
await session.execute(update(Key).where(Key.client_id == client_id).values(tariff_id=tariff_id))
await invalidate_key_details_by_client_id(session, client_id)
logger.info(f"Тариф ключа {client_id} обновлён на {tariff_id}")
async def update_key_renewal_snapshot(
session: AsyncSession,
email: str,
*,
tariff_id: int,
selected_device_limit: int | None = None,
current_device_limit: int | None = None,
selected_traffic_limit: int | None = None,
current_traffic_limit: int | None = None,
apply_limits: bool = True,
) -> None:
"""Обновляет tariff_id и (опционально) лимиты ключа после продления.
``apply_limits=True`` — выставить все четыре лимита (для non-configurable
тарифов). ``apply_limits=False`` — обновить только ``tariff_id``, лимиты
не трогать (configurable-тарифы обновляют их через `save_key_config_with_mode`).
"""
values: dict = {"tariff_id": tariff_id}
if apply_limits:
values["selected_device_limit"] = selected_device_limit
values["current_device_limit"] = current_device_limit
values["selected_traffic_limit"] = selected_traffic_limit
values["current_traffic_limit"] = current_traffic_limit
await session.execute(update(Key).where(Key.email == email).values(**values))
await invalidate_key_details(email)
async def update_key_post_creation_snapshot(
session: AsyncSession,
*,
user_id: int,
email: str,
selected_device_limit: int | None,
selected_traffic_limit: int | None,
selected_price_rub: int | None,
) -> None:
"""Дозаписывает выбранные пользователем параметры ключа сразу после создания.
Используется из `services.keys.create_vpn_key_headless` — тариф/лимиты не
всегда известны на момент `create_key_on_cluster`, поэтому после него
идёт snapshot-апдейт для полей, которые нужны для отображения в UI.
"""
await session.execute(
update(Key)
.where(Key.user_id == int(user_id), Key.email == email)
.values(
selected_device_limit=selected_device_limit,
selected_traffic_limit=selected_traffic_limit,
selected_price_rub=selected_price_rub,
)
)
await invalidate_key_details(email)
async def get_subscription_link(session: AsyncSession, email: str) -> str | None:
result = await session.execute(select(func.coalesce(Key.key, Key.remnawave_link)).where(Key.email == email))
return result.scalar_one_or_none()
async def update_key_client_id(session: AsyncSession, email: str, new_client_id: str):
await session.execute(update(Key).where(Key.email == email).values(client_id=new_client_id))
await invalidate_key_details(email)
logger.info(f"client_id обновлён для {email} -> {new_client_id}")
async def update_key_link(session: AsyncSession, email: str, link: str) -> bool:
q = update(Key).where(Key.email == email).values(key=link).returning(Key.client_id)
res = await session.execute(q)
ok = res.scalar_one_or_none() is not None
if ok:
await invalidate_key_details(email)
return ok
async def update_key_subscription_links(session: AsyncSession, email: str, link: str) -> bool:
stmt = (
update(Key)
.where(Key.email == email)
.values(
key=link,
remnawave_link=link,
)
.returning(Key.client_id)
)
res = await session.execute(stmt)
ok = res.scalar_one_or_none() is not None
if ok:
await invalidate_key_details(email)
return ok
async def save_key_config_with_mode(
session: AsyncSession,
email: str,
selected_devices: int | None,
selected_traffic_gb: int | None,
total_price: int,
has_device_choice: bool,
has_traffic_choice: bool,
config_mode: str,
) -> None:
values: dict = {}
if config_mode == "pack":
if has_device_choice and selected_devices is not None:
values["current_device_limit"] = int(selected_devices)
if has_traffic_choice and selected_traffic_gb is not None:
values["current_traffic_limit"] = int(selected_traffic_gb)
else:
device_val = int(selected_devices) if selected_devices is not None and has_device_choice else None
traffic_val = int(selected_traffic_gb) if selected_traffic_gb is not None and has_traffic_choice else None
values["selected_device_limit"] = device_val
values["selected_traffic_limit"] = traffic_val
values["selected_price_rub"] = int(total_price)
values["current_device_limit"] = device_val
values["current_traffic_limit"] = traffic_val
if not values:
return
await session.execute(update(Key).where(Key.email == email).values(**values))
await invalidate_key_details(email)
async def reset_key_tariff_state(session: AsyncSession, legacy_user_ref: int, email: str, tariff_id: int) -> None:
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
return
await session.execute(
update(Key)
.where(Key.user_id == u.id, Key.email == email)
.values(
tariff_id=tariff_id,
selected_device_limit=None,
current_device_limit=None,
selected_traffic_limit=None,
current_traffic_limit=None,
selected_price_rub=None,
)
)
await invalidate_keys_list(session, u.id)
await invalidate_key_details(email)
async def save_key_tariff_selection(
session: AsyncSession,
legacy_user_ref: int,
email: str,
tariff_id: int,
selected_devices: int | None,
selected_traffic_gb: int | None,
) -> None:
u = await resolve_user_optional(session, legacy_user_ref)
if u is None:
return
selected_devices_val = int(selected_devices) if selected_devices is not None else None
selected_traffic_val = int(selected_traffic_gb) if selected_traffic_gb is not None and int(selected_traffic_gb) > 0 else None
await session.execute(
update(Key)
.where(Key.user_id == u.id, Key.email == email)
.values(
tariff_id=tariff_id,
selected_device_limit=selected_devices_val,
current_device_limit=selected_devices_val,
selected_traffic_limit=selected_traffic_val,
current_traffic_limit=selected_traffic_val,
selected_price_rub=None,
)
)
await invalidate_keys_list(session, u.id)
await invalidate_key_details(email)
async def save_admin_key_config(
session: AsyncSession,
email: str,
base_devices: int,
total_devices: int,
base_traffic: int | None,
total_traffic: int | None,
selected_price: int | None,
) -> None:
await session.execute(
update(Key)
.where(Key.email == email)
.values(
selected_device_limit=base_devices,
current_device_limit=total_devices,
selected_traffic_limit=base_traffic,
current_traffic_limit=total_traffic,
selected_price_rub=selected_price,
)
)
await invalidate_key_details(email)
async def reset_key_current_limits_to_selected(session: AsyncSession, client_id: str):
"""Сбрасывает текущие лимиты к выбранным для ключа."""
await session.execute(
text(
"""
UPDATE keys
SET current_device_limit = selected_device_limit,
current_traffic_limit = selected_traffic_limit
WHERE client_id = :client_id
"""
),
{"client_id": client_id},
)
await invalidate_key_details_by_client_id(session, client_id)
logger.info(f"Текущие лимиты ключа {client_id} сброшены к выбранным")