import asyncio
from datetime import datetime
from types import SimpleNamespace
from sqlalchemy import delete, func, select, text, update
from sqlalchemy.exc import SQLAlchemyError
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.models import Key, 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 invalidate_keys_list(tg_id: int) -> None:
await cache_delete(cache_key("keys_list", tg_id))
await cache_delete(cache_key("key_count", tg_id))
await invalidate_profile_cache(tg_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,
tg_id: 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,
):
"""Сохраняет или обновляет ключ подписки."""
try:
exists = await session.execute(select(Key).where(Key.tg_id == tg_id, 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,
}
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.tg_id == tg_id, Key.client_id == client_id).values(**values))
logger.info(f"[Store Key] Ключ обновлён: tg_id={tg_id}, 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(
tg_id=tg_id,
client_id=client_id,
email=email,
created_at=int(datetime.utcnow().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] Ключ создан: tg_id={tg_id}, client_id={client_id}, server_id={server_id}")
await session.commit()
invalidate_user_snapshot(tg_id)
await invalidate_keys_list(tg_id)
await invalidate_key_details(email)
except SQLAlchemyError as e:
logger.error(f"❌ Ошибка при сохранении ключа: {e}")
await session.rollback()
raise
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, tg_id: int):
ckey = cache_key("keys_list", tg_id)
cached = await cache_get(ckey)
if isinstance(cached, list):
return [SimpleNamespace(**d) for d in cached]
result = await session.execute(select(Key).where(Key.tg_id == tg_id))
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, tg_id: int, client_id: str):
stmt = select(Key).where(Key.tg_id == tg_id, Key.client_id == client_id)
result = await session.execute(stmt)
return result.scalar_one_or_none()
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.tg_id == User.tg_id).where(Key.email == email)
result = await session.execute(stmt)
row = result.first()
if not row:
return None
key, user = row
expiry_date = datetime.utcfromtimestamp(key.expiry_time / 1000)
current_date = datetime.utcnow()
time_left = expiry_date - current_date
if time_left.total_seconds() <= 0:
days_left_message = "Ключ истек."
elif time_left.days > 0:
days_left_message = f"Осталось дней: {time_left.days}"
else:
hours_left = time_left.seconds // 3600
days_left_message = f"Осталось часов: {hours_left}"
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, tg_id: int) -> int:
cached = await cache_get(cache_key("key_count", tg_id))
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.tg_id == tg_id))
count = result.scalar() or 0
await cache_set(cache_key("key_count", tg_id), count, KEY_COUNT_CACHE_TTL_SEC)
return count
async def delete_key(session: AsyncSession, identifier: int | str, commit: bool = True):
tg_id_for_cache = None
email_for_cache = None
if isinstance(identifier, str):
res = await session.execute(
select(Key.tg_id, Key.email).where(Key.client_id == identifier).limit(1)
)
row = res.first()
if row:
tg_id_for_cache, email_for_cache = row[0], row[1]
await cache_delete(cache_key("key_email", identifier))
else:
tg_id_for_cache = identifier
stmt = delete(Key).where(Key.tg_id == identifier if isinstance(identifier, int) else Key.client_id == identifier)
await session.execute(stmt)
if commit:
await session.commit()
if tg_id_for_cache is not None:
invalidate_user_snapshot(tg_id_for_cache)
await invalidate_keys_list(tg_id_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 session.commit()
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, tg_id: int, client_id: str):
await session.execute(update(Key).where(Key.tg_id == tg_id, Key.client_id == client_id).values(notified=True))
await session.commit()
await invalidate_keys_list(tg_id)
await invalidate_key_details_by_client_id(session, client_id)
async def mark_key_as_frozen(session: AsyncSession, tg_id: int, client_id: str, time_left: int):
await session.execute(
text(
"""
UPDATE keys
SET expiry_time = :expiry,
is_frozen = TRUE
WHERE tg_id = :tg_id
AND client_id = :client_id
"""
),
{"expiry": time_left, "tg_id": tg_id, "client_id": client_id},
)
await invalidate_keys_list(tg_id)
await invalidate_key_details_by_client_id(session, client_id)
async def mark_key_as_unfrozen(
session: AsyncSession,
tg_id: int,
client_id: str,
new_expiry_time: int,
):
await session.execute(
text(
"""
UPDATE keys
SET expiry_time = :expiry,
is_frozen = FALSE
WHERE tg_id = :tg_id
AND client_id = :client_id
"""
),
{"expiry": new_expiry_time, "tg_id": tg_id, "client_id": client_id},
)
await invalidate_keys_list(tg_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 session.commit()
await invalidate_key_details_by_client_id(session, client_id)
logger.info(f"Тариф ключа {client_id} обновлён на {tariff_id}")
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 session.commit()
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)
await session.commit()
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_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 session.commit()
await invalidate_key_details_by_client_id(session, client_id)
logger.info(f"Текущие лимиты ключа {client_id} сброшены к выбранным")