diff --git a/database.py b/database.py index 6c3f1958..8dba5f4a 100644 --- a/database.py +++ b/database.py @@ -3,6 +3,7 @@ from datetime import datetime from typing import Any import asyncpg +import pytz from config import DATABASE_URL, REFERRAL_BONUS_PERCENTAGES from logger import logger @@ -57,7 +58,7 @@ async def remove_blocked_user(tg_id: int | list[int], conn: asyncpg.Connection): async def init_db(file_path: str = "assets/schema.sql"): - with open(file_path, mode="r") as file: + with open(file_path) as file: sql_content = file.read() statements = [stmt.strip() for stmt in sql_content.split(";") if stmt.strip()] @@ -843,39 +844,6 @@ async def update_key_expiry(client_id: str, new_expiry_time: int): logger.info("Закрытие подключения к базе данных") -async def delete_key(client_id: str): - """ - Удаление ключа из базы данных для указанного клиента. - - Args: - client_id (str): Уникальный идентификатор клиента, ключ которого будет удален - - Raises: - Exception: В случае ошибки при подключении к базе данных или удалении ключа - """ - conn = None - try: - conn = await asyncpg.connect(DATABASE_URL) - logger.info(f"Установлено подключение к базе данных для удаления ключа клиента {client_id}") - - await conn.execute( - """ - DELETE FROM keys - WHERE client_id = $1 - """, - client_id, - ) - logger.info(f"Успешно удален ключ для клиента {client_id}") - - except Exception as e: - logger.error(f"Ошибка при удалении ключа для клиента {client_id}: {e}") - raise - finally: - if conn: - await conn.close() - logger.info("Закрытие подключения к базе данных") - - async def add_balance_to_client(client_id: str, amount: float): """ Добавление баланса клиенту по его идентификатору Telegram. @@ -1202,7 +1170,7 @@ async def delete_user_data(session: Any, tg_id: int): await session.execute("DELETE FROM payments WHERE tg_id = $1", tg_id) await session.execute("DELETE FROM users WHERE tg_id = $1", tg_id) await session.execute("DELETE FROM connections WHERE tg_id = $1", tg_id) - await session.execute("DELETE FROM keys WHERE tg_id = $1", tg_id) + await delete_key(tg_id, session) await session.execute("DELETE FROM referrals WHERE referrer_tg_id = $1", tg_id) @@ -1261,3 +1229,69 @@ async def store_gift_link( finally: if conn is not None and session is None: await conn.close() + + +async def get_key_details(email, session): + record = await session.fetchrow( + """ + SELECT k.key, k.expiry_time, k.server_id,k, k.client_id, k.created_at, c.tg_id, c.balance + FROM keys k + JOIN connections c ON k.tg_id = c.tg_id + WHERE k.email = $1 + """, + email, + ) + + if not record: + return None + + cluster_name = record["server_id"] + + moscow_tz = pytz.timezone("Europe/Moscow") + expiry_date = datetime.fromtimestamp(record["expiry_time"] / 1000, tz=moscow_tz) + current_date = datetime.now(moscow_tz) + 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}" + + return { + "key": record["key"], + "created_at": record["created_at"], + "expiry_time": record["expiry_time"], + "client_id": record["client_id"], + "expiry_date": expiry_date.strftime("%d %B %Y года %H:%M"), + "days_left_message": days_left_message, + "server_name": cluster_name, + "balance": record["balance"], + "tg_id": record["tg_id"], + } + + +async def delete_key(identifier, session): + """ + Удаляет ключ из базы данных по client_id или tg_id + + Args: + identifier (str): client_id или tg_id для удаления + session: Сессия базы данных + + Raises: + Exception: В случае ошибки при удалении ключа + """ + try: + # Проверяем, является ли идентификатор числом (tg_id) или строкой (client_id) + if identifier.isdigit(): + query = "DELETE FROM keys WHERE tg_id = $1" + else: + query = "DELETE FROM keys WHERE client_id = $1" + + await session.execute(query, identifier) + logger.info(f"Ключ с идентификатором {identifier} успешно удалён") + except Exception as e: + logger.error(f"Ошибка при удалении ключа с идентификатором {identifier} из базы данных: {e}") diff --git a/handlers/admin/admin_user_editor.py b/handlers/admin/admin_user_editor.py index 94e02f33..64c7b6e1 100644 --- a/handlers/admin/admin_user_editor.py +++ b/handlers/admin/admin_user_editor.py @@ -13,6 +13,7 @@ from config import TOTAL_GB from database import ( delete_user_data, get_client_id_by_email, + get_key_details, get_servers_from_db, restore_trial, update_key_expiry, @@ -257,45 +258,6 @@ async def handle_new_balance_input(message: types.Message, state: FSMContext, se await state.clear() -async def get_key_details(email, session): - record = await session.fetchrow( - """ - SELECT k.key, k.expiry_time, k.server_id, c.tg_id, c.balance - FROM keys k - JOIN connections c ON k.tg_id = c.tg_id - WHERE k.email = $1 - """, - email, - ) - - if not record: - return None - - cluster_name = record["server_id"] - - moscow_tz = pytz.timezone("Europe/Moscow") - expiry_date = datetime.fromtimestamp(record["expiry_time"] / 1000, tz=moscow_tz) - current_date = datetime.now(moscow_tz) - 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}" - - return { - "key": record["key"], - "expiry_date": expiry_date.strftime("%d %B %Y года %H:%M"), - "days_left_message": days_left_message, - "server_name": cluster_name, - "balance": record["balance"], - "tg_id": record["tg_id"], - } - - @router.callback_query(F.data.startswith("edit_key_"), IsAdminFilter()) async def process_key_edit(callback_query: CallbackQuery, session: Any): email = callback_query.data.split("_", 2)[2] @@ -492,9 +454,9 @@ async def handle_expiry_time_input(message: types.Message, state: FSMContext, se @router.callback_query(F.data.startswith("delete_key_admin|"), IsAdminFilter()) async def process_callback_delete_key(callback_query: types.CallbackQuery, session: Any): email = callback_query.data.split("|")[1] - client_id = await session.fetchval("SELECT client_id FROM keys WHERE email = $1", email) + key_details = await get_key_details(email, session) - if client_id is None: + if key_details is None: builder = InlineKeyboardBuilder() builder.row(InlineKeyboardButton(text="🔙 Назад", callback_data="user_editor")) await callback_query.message.answer("🔍 Ключ не найден. 🚫", reply_markup=builder.as_markup()) @@ -504,7 +466,7 @@ async def process_callback_delete_key(callback_query: types.CallbackQuery, sessi builder.row( types.InlineKeyboardButton( text="✅ Да, удалить", - callback_data=f"confirm_delete_admin|{client_id}", + callback_data=f"confirm_delete_admin|{key_details['client_id']}", ) ) builder.row(types.InlineKeyboardButton(text="❌ Нет, отменить", callback_data="user_editor")) diff --git a/handlers/instructions/instructions.py b/handlers/instructions/instructions.py index 9663531d..dd373ccb 100644 --- a/handlers/instructions/instructions.py +++ b/handlers/instructions/instructions.py @@ -2,12 +2,12 @@ import os from typing import Any import aiofiles -import asyncpg from aiogram import F, Router, types from aiogram.types import BufferedInputFile, InlineKeyboardButton from aiogram.utils.keyboard import InlineKeyboardBuilder -from config import CONNECT_MACOS, CONNECT_WINDOWS, DATABASE_URL, SUPPORT_CHAT_URL +from config import CONNECT_MACOS, CONNECT_WINDOWS, SUPPORT_CHAT_URL +from database import get_key_details from handlers.texts import ( CONNECT_TV_TEXT, INSTRUCTION_PC, @@ -60,15 +60,7 @@ async def process_connect_pc(callback_query: types.CallbackQuery, session: Any): tg_id = callback_query.message.chat.id key_name = callback_query.data.split("|")[1] - record = await session.fetchrow( - """ - SELECT k.key - FROM keys k - WHERE k.tg_id = $1 AND k.email = $2 - """, - tg_id, - key_name, - ) + record = await get_key_details(key_name, session) if not record: await callback_query.message.answer("❌ Ключ не найден. Проверьте имя ключа. 🔍") @@ -107,27 +99,15 @@ async def process_connect_tv(callback_query: types.CallbackQuery): @router.callback_query(F.data.startswith("continue_tv|")) -async def process_continue_tv(callback_query: types.CallbackQuery): +async def process_continue_tv(callback_query: types.CallbackQuery, session: Any): key_name = callback_query.data.split("|")[1] tg_id = callback_query.from_user.id logger.info(f"tg_id: {tg_id}, key_name: {key_name}") - conn = await asyncpg.connect(DATABASE_URL) - try: - record = await conn.fetchrow( - """ - SELECT k.key - FROM keys k - WHERE k.tg_id = $1 AND k.email = $2 - """, - tg_id, - key_name, - ) + record = await get_key_details(key_name, session) - logger.info(f"Query result: {record}") - finally: - await conn.close() + logger.info(f"Query result: {record}") subscription_link = record["key"] diff --git a/handlers/keys/key_management.py b/handlers/keys/key_management.py index 3ab75926..a06f004f 100644 --- a/handlers/keys/key_management.py +++ b/handlers/keys/key_management.py @@ -25,6 +25,7 @@ from config import ( ) from database import ( get_balance, + get_key_details, get_trial, save_temporary_data, store_key, @@ -352,11 +353,7 @@ async def finalize_key_creation( key_name = generate_random_email() logger.info(f"Generated random key name for user {tg_id}: {key_name}") - existing_key = await session.fetchrow( - "SELECT * FROM keys WHERE email = $1 AND tg_id = $2", - key_name, - tg_id, - ) + existing_key = await get_key_details(key_name, session) if not existing_key: break logger.warning(f"Key name '{key_name}' already exists for user {tg_id}. Generating a new one.") diff --git a/handlers/keys/key_utils.py b/handlers/keys/key_utils.py index f59b3072..3657f77f 100644 --- a/handlers/keys/key_utils.py +++ b/handlers/keys/key_utils.py @@ -142,13 +142,6 @@ async def renew_key_in_cluster(cluster_id, email, client_id, new_expiry_time, to raise e -async def delete_key_from_db(client_id, session): - try: - await session.execute("DELETE FROM keys WHERE client_id = $1", client_id) - except Exception as e: - logger.error(f"Ошибка при удалении ключа {client_id} из базы данных: {e}") - - async def delete_key_from_cluster(cluster_id, email, client_id): """Удаление ключа с серверов в кластере""" try: diff --git a/handlers/keys/keys.py b/handlers/keys/keys.py index a0b4c6b5..20c036fb 100644 --- a/handlers/keys/keys.py +++ b/handlers/keys/keys.py @@ -29,6 +29,7 @@ from config import ( from database import ( delete_key, get_balance, + get_key_details, get_servers_from_db, save_temporary_data, store_key, @@ -163,15 +164,7 @@ async def process_callback_view_key(callback_query: types.CallbackQuery, session tg_id = callback_query.message.chat.id key_name = callback_query.data.split("|")[1] try: - record = await session.fetchrow( - """ - SELECT k.expiry_time, k.server_id, k.key - FROM keys k - WHERE k.tg_id = $1 AND k.email = $2 - """, - tg_id, - key_name, - ) + record = await get_key_details(key_name, session) if record: key = record["key"] @@ -264,15 +257,7 @@ async def process_callback_update_subscription(callback_query: types.CallbackQue tg_id = callback_query.message.chat.id email = callback_query.data.split("|")[1] try: - record = await session.fetchrow( - """ - SELECT k.key, k.expiry_time, k.email, k.server_id, k.client_id - FROM keys k - WHERE k.tg_id = $1 AND k.email = $2 - """, - tg_id, - email, - ) + record = await get_key_details(email, session) if record: expiry_time = record["expiry_time"] @@ -353,14 +338,7 @@ async def process_callback_renew_key(callback_query: types.CallbackQuery, sessio tg_id = callback_query.message.chat.id key_name = callback_query.data.split("|")[1] try: - record = await session.fetchrow( - """ - SELECT client_id, expiry_time - FROM keys - WHERE email = $1 - """, - key_name, - ) + record = await get_key_details(key_name, session) if record: client_id = record["client_id"] @@ -405,7 +383,7 @@ async def process_callback_renew_key(callback_query: types.CallbackQuery, sessio async def process_callback_confirm_delete(callback_query: types.CallbackQuery, session: Any): email = callback_query.data.split("|")[1] try: - record = await session.fetchrow("SELECT client_id FROM keys WHERE email = $1", email) + record = await get_key_details(email, session) if record: client_id = record["client_id"] @@ -413,7 +391,7 @@ async def process_callback_confirm_delete(callback_query: types.CallbackQuery, s back_button = types.InlineKeyboardButton(text="Назад", callback_data="view_keys") keyboard = types.InlineKeyboardMarkup(inline_keyboard=[[back_button]]) - await delete_key(client_id) + await delete_key(client_id, session) await callback_query.message.answer( response_message, reply_markup=keyboard, diff --git a/handlers/keys/subscriptions.py b/handlers/keys/subscriptions.py index 4f21ec14..d3f6d36e 100644 --- a/handlers/keys/subscriptions.py +++ b/handlers/keys/subscriptions.py @@ -7,7 +7,7 @@ import asyncpg from aiohttp import web from config import DATABASE_URL, PROJECT_NAME, SUB_MESSAGE, SUPERNODE, TRANSITION_DATE_STR -from database import get_servers_from_db +from database import get_key_details, get_servers_from_db from logger import logger # Глобальная переменная для пула соединений @@ -92,7 +92,7 @@ async def handle_old_subscription(request): await init_db_pool() async with db_pool.acquire() as conn: - key_info = await conn.fetchrow("SELECT created_at, server_id FROM keys WHERE email = $1", email) + key_info = await get_key_details(email, conn) if not key_info: logger.warning(f"Клиент с email {email} не найден в базе.") @@ -161,7 +161,7 @@ async def handle_new_subscription(request): await init_db_pool() async with db_pool.acquire() as conn: - client_data = await conn.fetchrow("SELECT tg_id, server_id FROM keys WHERE email = $1", email) + client_data = await get_key_details(email, conn) if not client_data: logger.warning(f"Клиент с email {email} не найден в базе.") diff --git a/handlers/keys/trial_key.py b/handlers/keys/trial_key.py index 9ec15090..00048577 100644 --- a/handlers/keys/trial_key.py +++ b/handlers/keys/trial_key.py @@ -8,7 +8,7 @@ from py3xui import AsyncApi from client import add_client from config import ADMIN_PASSWORD, ADMIN_USERNAME, LIMIT_IP, PUBLIC_LINK, SUPERNODE, TOTAL_GB, TRIAL_TIME -from database import get_servers_from_db, store_key, use_trial +from database import get_servers_from_db, get_trial, store_key, use_trial from handlers.texts import INSTRUCTIONS from handlers.utils import generate_random_email, get_least_loaded_cluster from logger import logger @@ -16,7 +16,7 @@ from logger import logger async def create_trial_key(tg_id: int, session: Any): try: - trial_status = await session.fetchval("SELECT trial FROM connections WHERE tg_id = $1", tg_id) + trial_status = await get_trial(tg_id, session) if trial_status == 1: return {"error": "Вы уже использовали пробную версию."} except Exception as e: diff --git a/handlers/notifications.py b/handlers/notifications.py index c01bb371..369e8e48 100644 --- a/handlers/notifications.py +++ b/handlers/notifications.py @@ -401,7 +401,7 @@ async def handle_expired_keys(bot: Bot, conn: asyncpg.Connection, current_time: for record in expired_keys: try: await delete_key_from_cluster(record["server_id"], record["email"], record["email"]) - await conn.execute("DELETE FROM keys WHERE client_id = $1", record["client_id"]) + await delete_key(record["client_id"], conn) logger.info(f"Удалён истёкший ключ {record['client_id']} пользователя {record['tg_id']}.") except Exception as e: logger.error(f"Ошибка при удалении истёкшего ключа {record['client_id']}: {e}") @@ -497,7 +497,7 @@ async def process_key(record, bot, conn): logger.error(f"Ошибка при удалении клиента {client_id} из кластера {cluster_id}: {e}") try: - await delete_key(client_id) + await delete_key(client_id, conn) logger.info(f"Ключ {client_id} удалён из базы данных.") except Exception as e: logger.error(f"Ошибка при удалении ключа {client_id} из базы данных: {e}") diff --git a/handlers/start.py b/handlers/start.py index a48f02d6..cf7bec02 100644 --- a/handlers/start.py +++ b/handlers/start.py @@ -101,14 +101,7 @@ async def start_command(message: Message, state: FSMContext, session: Any, admin await message.answer("❌ Вы не можете получить подарок от самого себя.") return await show_start_menu(message, admin, session) - await session.execute( - """ - INSERT INTO connections (tg_id, balance, trial) - VALUES ($1, 0, 1) - ON CONFLICT (tg_id) DO UPDATE SET trial = 1 - """, - recipient_tg_id, - ) + await add_connection(tg_id=recipient_tg_id, session=session) selected_months = gift_info["selected_months"] expiry_time = gift_info["expiry_time"]