From 0867ebdd61be0d448b26ed46df4b9266e1d2e5eb Mon Sep 17 00:00:00 2001 From: Zakhar Izmaylov Date: Sun, 24 Nov 2024 01:19:00 +0300 Subject: [PATCH] Fix typing --- database.py | 75 ++++++++++---- handlers/admin/admin_panel.py | 2 +- handlers/admin/admin_user_editor.py | 4 +- handlers/coupons.py | 2 +- handlers/donate.py | 4 +- handlers/instructions/instructions.py | 6 +- handlers/keys/key_management.py | 65 +++++------- handlers/keys/keys.py | 140 ++++++++++++-------------- handlers/keys/trial_key.py | 120 ++++++---------------- handlers/pay.py | 2 +- handlers/payments/utils.py | 2 +- handlers/payments/yookassa_pay.py | 2 +- handlers/profile.py | 4 +- handlers/start.py | 20 ++-- middlewares/database.py | 13 ++- 15 files changed, 210 insertions(+), 251 deletions(-) diff --git a/database.py b/database.py index bc3cf86e..a372c3fa 100644 --- a/database.py +++ b/database.py @@ -1,4 +1,5 @@ from datetime import datetime +from typing import Any import asyncpg @@ -177,19 +178,64 @@ async def delete_coupon_from_db(coupon_code: str): await conn.close() -async def restore_trial(tg_id: int): - conn = await asyncpg.connect(DATABASE_URL) +async def restore_trial(tg_id: int, session: Any): + """ + Восстанавливает возможность использования триального периода для пользователя. + + Args: + tg_id (int): Telegram ID пользователя + session (Any): Сессия базы данных + + Returns: + bool: True, если триал успешно восстановлен, False в случае ошибки + """ try: - await conn.execute("UPDATE connections SET trial = 0 WHERE tg_id = $1", tg_id) + response = await session.execute( + """ + INSERT INTO connections (tg_id, trial) + VALUES ($1, 0) + ON CONFLICT (tg_id) + DO UPDATE SET trial = 0 + """, + tg_id, + ) + logger.info(response) + return True except Exception as e: logger.error(f"Ошибка при установке значения триала: {e}") - finally: - await conn.close() + return False -async def add_connection(tg_id: int, balance: float = 0.0, trial: int = 0): - conn = await asyncpg.connect(DATABASE_URL) - await conn.execute( +async def use_trial(tg_id: int, session: Any): + """ + Отмечает использование триального периода для пользователя. + + Args: + tg_id (int): Telegram ID пользователя + session (Any): Сессия базы данных + + Returns: + bool: True, если триал успешно использован, False в случае ошибки + """ + try: + response = await session.execute( + """ + INSERT INTO connections (tg_id, trial) + VALUES ($1, 1) + ON CONFLICT (tg_id) + DO UPDATE SET trial = 1 + """, + tg_id, + ) + logger.info(response) + return True + except Exception as e: + logger.error(f"Ошибка при использовании триала: {e}") + return False + + +async def add_connection(tg_id: int, balance: float = 0.0, trial: int = 0, session: Any = None): + await session.execute( """ INSERT INTO connections (tg_id, balance, trial) VALUES ($1, $2, $3) @@ -198,7 +244,6 @@ async def add_connection(tg_id: int, balance: float = 0.0, trial: int = 0): balance, trial, ) - await conn.close() async def check_connection_exists(tg_id: int): @@ -298,10 +343,8 @@ async def update_balance(tg_id: int, amount: float): await conn.close() -async def get_trial(tg_id: int) -> int: - conn = await asyncpg.connect(DATABASE_URL) - trial = await conn.fetchval("SELECT trial FROM connections WHERE tg_id = $1", tg_id) - await conn.close() +async def get_trial(tg_id: int, session: Any) -> int: + trial = await session.fetchval("SELECT trial FROM connections WHERE tg_id = $1", tg_id) return trial if trial is not None else 0 @@ -316,9 +359,8 @@ async def get_all_users(conn): return await conn.fetch("SELECT tg_id FROM connections") -async def add_referral(referred_tg_id: int, referrer_tg_id: int): - conn = await asyncpg.connect(DATABASE_URL) - await conn.execute( +async def add_referral(referred_tg_id: int, referrer_tg_id: int, session: Any): + await session.execute( """ INSERT INTO referrals (referred_tg_id, referrer_tg_id) VALUES ($1, $2) @@ -326,7 +368,6 @@ async def add_referral(referred_tg_id: int, referrer_tg_id: int): referred_tg_id, referrer_tg_id, ) - await conn.close() async def handle_referral_on_balance_update(tg_id: int, amount: float): diff --git a/handlers/admin/admin_panel.py b/handlers/admin/admin_panel.py index 6520ffd9..ce7466f6 100644 --- a/handlers/admin/admin_panel.py +++ b/handlers/admin/admin_panel.py @@ -40,7 +40,7 @@ async def handle_admin_message(message: types.Message, state: FSMContext): builder.row(InlineKeyboardButton(text="📢 Массовая рассылка", callback_data="send_to_alls")) builder.row(InlineKeyboardButton(text="💾 Создать резервную копию", callback_data="backups")) builder.row(InlineKeyboardButton(text="🔄 Перезагрузить бота", callback_data="restart_bot")) - builder.row(InlineKeyboardButton(text="⬅️ Вернуться в профиль", callback_data="profile")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) await message.answer("🤖 Панель администратора", reply_markup=builder.as_markup()) diff --git a/handlers/admin/admin_user_editor.py b/handlers/admin/admin_user_editor.py index 6f1ea996..878834f2 100644 --- a/handlers/admin/admin_user_editor.py +++ b/handlers/admin/admin_user_editor.py @@ -140,10 +140,10 @@ async def handle_tg_id_input(message: types.Message, state: FSMContext, session: @router.callback_query(F.data.startswith("restore_trial_"), IsAdminFilter()) -async def handle_restore_trial(callback_query: types.CallbackQuery): +async def handle_restore_trial(callback_query: types.CallbackQuery, session: Any): tg_id = int(callback_query.data.split("_")[2]) - await restore_trial(tg_id) + await restore_trial(tg_id, session) builder = InlineKeyboardBuilder() builder.row(InlineKeyboardButton(text="🔙 Назад в меню администратора", callback_data="admin")) diff --git a/handlers/coupons.py b/handlers/coupons.py index 7d9a5ced..9d79076a 100644 --- a/handlers/coupons.py +++ b/handlers/coupons.py @@ -20,7 +20,7 @@ router = Router() @router.callback_query(F.data == "activate_coupon") async def handle_activate_coupon(callback_query: types.CallbackQuery, state: FSMContext): builder = InlineKeyboardBuilder() - builder.row(InlineKeyboardButton(text="⬅️ Вернуться в профиль", callback_data="profile")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) await callback_query.message.answer( "🎫 Введите код купона:\n\n" diff --git a/handlers/donate.py b/handlers/donate.py index e6160205..d441463a 100644 --- a/handlers/donate.py +++ b/handlers/donate.py @@ -29,7 +29,7 @@ async def process_donate(callback_query: types.CallbackQuery, state: FSMContext) callback_data="enter_custom_donate_amount", ) ) - builder.row(InlineKeyboardButton(text="⬅️ Назад", callback_data="profile")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) await callback_query.message.answer( text="🌟 Поддержите наш проект! 💪\n\n" @@ -89,7 +89,7 @@ async def on_successful_donate(message: types.Message, state: FSMContext): try: amount = float(message.successful_payment.invoice_payload.split("_")[0]) builder = InlineKeyboardBuilder() - builder.row(InlineKeyboardButton(text="Вернуться в профиль", callback_data="profile")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) await message.answer( text=f"🙏 Спасибо за донат {amount} рублей! Ваша поддержка очень важна для нас. 💖", reply_markup=builder.as_markup(), diff --git a/handlers/instructions/instructions.py b/handlers/instructions/instructions.py index 16d9094d..ce82d45a 100644 --- a/handlers/instructions/instructions.py +++ b/handlers/instructions/instructions.py @@ -22,11 +22,11 @@ async def send_instructions(callback_query: types.CallbackQuery): builder = InlineKeyboardBuilder() builder.row(InlineKeyboardButton(text="💬 Поддержка", url=SUPPORT_CHAT_URL)) builder.row( - InlineKeyboardButton(text="⬅️ Вернуться в профиль", callback_data="profile"), + InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile"), ) with open(image_path, "rb") as image_from_buffer: - await callback_query.answer_photo( + await callback_query.message.answer_photo( BufferedInputFile(image_from_buffer.read(), filename="instructions.jpg"), caption=instructions_message, reply_markup=builder.as_markup(), @@ -59,6 +59,6 @@ async def process_connect_pc(callback_query: types.CallbackQuery, session: Any): builder = InlineKeyboardBuilder() builder.row(InlineKeyboardButton(text="💻 Подключить Windows", url=f"{CONNECT_WINDOWS}{key}")) builder.row(InlineKeyboardButton(text="🆘 Поддержка", url=f"{SUPPORT_CHAT_URL}")) - builder.row(InlineKeyboardButton(text="🔙 Назад в профиль", callback_data="profile")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) await callback_query.message.answer(instruction_message, reply_markup=builder.as_markup()) diff --git a/handlers/keys/key_management.py b/handlers/keys/key_management.py index 4869f884..7085b2cc 100644 --- a/handlers/keys/key_management.py +++ b/handlers/keys/key_management.py @@ -6,7 +6,8 @@ import uuid from aiogram import F, Router from aiogram.fsm.context import FSMContext from aiogram.fsm.state import State, StatesGroup -from aiogram.types import CallbackQuery, InlineKeyboardButton, InlineKeyboardMarkup, Message +from aiogram.types import CallbackQuery, InlineKeyboardButton, Message +from aiogram.utils.keyboard import InlineKeyboardBuilder from config import CONNECT_ANDROID, CONNECT_IOS, DOWNLOAD_ANDROID, DOWNLOAD_IOS, PUBLIC_LINK, SUPPORT_CHAT_URL from database import add_connection, get_balance, store_key, update_balance @@ -41,19 +42,15 @@ async def select_server(callback_query: CallbackQuery, state: FSMContext, sessio trial_status = existing_connection["trial"] if existing_connection else 0 if trial_status == 1: + builder = InlineKeyboardBuilder() + builder.row( + InlineKeyboardButton(text="✅ Да, подключить новое устройство", callback_data="confirm_create_new_key") + ) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) + await callback_query.message.answer( text=KEY, - reply_markup=InlineKeyboardMarkup( - inline_keyboard=[ - [ - InlineKeyboardButton( - text="✅ Да, подключить новое устройство", - callback_data="confirm_create_new_key", - ) - ], - [InlineKeyboardButton(text="↩️ Назад", callback_data="profile")], - ] - ), + reply_markup=builder.as_markup(), ) await state.update_data(creating_new_key=True) else: @@ -69,9 +66,9 @@ async def confirm_create_new_key(callback_query: CallbackQuery, state: FSMContex balance = await get_balance(tg_id) if balance < RENEWAL_PLANS["1"]["price"]: - replenish_button = InlineKeyboardButton(text="Перейти в профиль", callback_data="profile") - keyboard = InlineKeyboardMarkup(inline_keyboard=[[replenish_button]]) - await callback_query.message.edit_text(NULL_BALANCE, reply_markup=keyboard) + builder = InlineKeyboardBuilder() + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) + await callback_query.message.edit_text(NULL_BALANCE, reply_markup=builder.as_markup()) await state.clear() return @@ -125,11 +122,11 @@ async def handle_key_name_input(message: Message, state: FSMContext, session: An else: balance = await get_balance(tg_id) if balance < RENEWAL_PLANS["1"]["price"]: - replenish_button = InlineKeyboardButton(text="Перейти в профиль", callback_data="profile") - keyboard = InlineKeyboardMarkup(inline_keyboard=[[replenish_button]]) + builder = InlineKeyboardBuilder() + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) await message.answer( "💳 Недостаточно средств для создания подписки на новое устройство. Пополните баланс в личном кабинете.", - reply_markup=keyboard, + reply_markup=builder.as_markup(), ) logger.warning(f"User {tg_id} has insufficient funds for key creation.") await state.clear() @@ -144,29 +141,17 @@ async def handle_key_name_input(message: Message, state: FSMContext, session: An logger.info(f"Generated public link for the key: {public_link}") - button_support = InlineKeyboardButton(text="💬 Поддержка", url=SUPPORT_CHAT_URL) - - button_profile = InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile") - button_iphone = InlineKeyboardButton(text="🍏 Подключить", url=f"{CONNECT_IOS}{public_link}") - button_android = InlineKeyboardButton( - text="🤖 Подключить", - url=f"{CONNECT_ANDROID}{public_link}", + builder = InlineKeyboardBuilder() + builder.row(InlineKeyboardButton(text="💬 Поддержка", url=SUPPORT_CHAT_URL)) + builder.row( + InlineKeyboardButton(text="🍏 Скачать для iOS", url=DOWNLOAD_IOS), + InlineKeyboardButton(text="🤖 Скачать для Android", url=DOWNLOAD_ANDROID), ) - - button_download_ios = InlineKeyboardButton(text="🍏 Скачать", url=DOWNLOAD_IOS) - button_download_android = InlineKeyboardButton( - text="🤖 Скачать", - url=DOWNLOAD_ANDROID, - ) - - keyboard = InlineKeyboardMarkup( - inline_keyboard=[ - [button_support], - [button_download_ios, button_download_android], - [button_iphone, button_android], - [button_profile], - ] + builder.row( + InlineKeyboardButton(text="🍏 Подключить на iOS", url=f"{CONNECT_IOS}{public_link}"), + InlineKeyboardButton(text="🤖 Подключить на Android", url=f"{CONNECT_ANDROID}{public_link}"), ) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) remaining_time = expiry_time - current_time days = remaining_time.days @@ -174,7 +159,7 @@ async def handle_key_name_input(message: Message, state: FSMContext, session: An logger.info(f"Sending key message to user {tg_id} with the public link.") - await message.answer(key_message, reply_markup=keyboard) + await message.answer(key_message, reply_markup=builder.as_markup()) try: least_loaded_cluster = await get_least_loaded_cluster() diff --git a/handlers/keys/keys.py b/handlers/keys/keys.py index 7c0626dd..9a668893 100644 --- a/handlers/keys/keys.py +++ b/handlers/keys/keys.py @@ -5,7 +5,8 @@ import os from typing import Any from aiogram import F, Router, types -from aiogram.types import BufferedInputFile +from aiogram.types import BufferedInputFile, InlineKeyboardButton +from aiogram.utils.keyboard import InlineKeyboardBuilder from config import CLUSTERS, CONNECT_ANDROID, CONNECT_IOS, DOWNLOAD_ANDROID, DOWNLOAD_IOS, PUBLIC_LINK, TOTAL_GB from database import delete_key, get_balance, store_key, update_balance, update_key_expiry @@ -44,20 +45,14 @@ async def process_callback_view_keys(callback_query: types.CallbackQuery, sessio ) if records: - buttons = [] + builder = InlineKeyboardBuilder() for record in records: key_name = record["email"] - button = types.InlineKeyboardButton( - text=f"🔑 {key_name}", - callback_data=f"view_key|{key_name}", - ) + builder.row(InlineKeyboardButton(text=f"🔑 {key_name}", callback_data=f"view_key|{key_name}")) - buttons.append([button]) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) - back_button = types.InlineKeyboardButton(text="🔙 Назад", callback_data="profile") - buttons.append([back_button]) - - inline_keyboard = types.InlineKeyboardMarkup(inline_keyboard=buttons) + inline_keyboard = builder.as_markup() response_message = ( "🔑 Список ваших устройств\n\n" "👇 Выберите устройство для управления подпиской:" ) @@ -78,10 +73,11 @@ async def process_callback_view_keys(callback_query: types.CallbackQuery, sessio else: response_message = NO_KEYS - create_key_button = types.InlineKeyboardButton(text="➕ Создать подписку", callback_data="create_key") - back_button = types.InlineKeyboardButton(text="🔙 Назад", callback_data="profile") + builder = InlineKeyboardBuilder() + builder.row(InlineKeyboardButton(text="➕ Создать подписку", callback_data="create_key")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) - keyboard = types.InlineKeyboardMarkup(inline_keyboard=[[create_key_button], [back_button]]) + keyboard = builder.as_markup() image_path = os.path.join("img", "pic_keys.jpg") @@ -135,38 +131,33 @@ async def process_callback_view_key(callback_query: types.CallbackQuery, session formatted_expiry_date = expiry_date.strftime("%d %B %Y года") response_message = key_message(key, formatted_expiry_date, days_left_message, server_name) - download_android_button = types.InlineKeyboardButton(text="🤖 Скачать", url=DOWNLOAD_ANDROID) - download_iphone_button = types.InlineKeyboardButton(text="🍏 Скачать", url=DOWNLOAD_IOS) + builder = InlineKeyboardBuilder() - connect_iphone_button = types.InlineKeyboardButton(text="🍏 Подключить", url=f"{CONNECT_IOS}{key}") - connect_android_button = types.InlineKeyboardButton(text="🤖 Подключить", url=f"{CONNECT_ANDROID}{key}") - - connect_pc_button = types.InlineKeyboardButton( - text="💻 Windows/Linux", - callback_data=f"connect_pc|{key_name}", + builder.row( + InlineKeyboardButton(text="🍏 Скачать для iOS", url=DOWNLOAD_IOS), + InlineKeyboardButton(text="🤖 Скачать для Android", url=DOWNLOAD_ANDROID), ) - renew_button = types.InlineKeyboardButton(text="⏳ Продлить", callback_data=f"renew_key|{key_name}") - delete_button = types.InlineKeyboardButton(text="❌ Удалить", callback_data=f"delete_key|{key_name}") - back_button = types.InlineKeyboardButton(text="🔙 Назад в профиль", callback_data="profile") + builder.row( + InlineKeyboardButton(text="🍏 Подключить на iOS", url=f"{CONNECT_IOS}{key}"), + InlineKeyboardButton(text="🤖 Подключить на Android", url=f"{CONNECT_ANDROID}{key}"), + ) - inline_keyboard = [ - [download_iphone_button, download_android_button], - [connect_iphone_button, connect_android_button], - [connect_pc_button], - [renew_button, delete_button], - ] + builder.row(InlineKeyboardButton(text="💻 Windows/Linux", callback_data=f"connect_pc|{key_name}")) + + builder.row( + InlineKeyboardButton(text="⏳ Продлить", callback_data=f"renew_key|{key_name}"), + InlineKeyboardButton(text="❌ Удалить", callback_data=f"delete_key|{key_name}"), + ) if not key.startswith(PUBLIC_LINK): - update_subscription_button = types.InlineKeyboardButton( - text="🔄 Обновить подписку", - callback_data=f"update_subscription|{key_name}", + builder.row( + InlineKeyboardButton(text="🔄 Обновить подписку", callback_data=f"update_subscription|{key_name}") ) - inline_keyboard.append([update_subscription_button]) - inline_keyboard.append([back_button]) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) - keyboard = types.InlineKeyboardMarkup(inline_keyboard=inline_keyboard) + keyboard = builder.as_markup() image_path = os.path.join("img", "pic_view.jpg") @@ -251,12 +242,12 @@ async def process_callback_update_subscription(callback_query: types.CallbackQue server_id=least_loaded_cluster_id, ) response_message = f"Ваша подписка {email} обновлена!" - back_button = types.InlineKeyboardButton(text="🔙 Назад в профиль", callback_data="profile") - keyboard = types.InlineKeyboardMarkup(inline_keyboard=[[back_button]]) + builder = InlineKeyboardBuilder() + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) await callback_query.message.answer( response_message, - reply_markup=keyboard, + reply_markup=builder.as_markup(), ) else: await callback_query.message.answer( @@ -309,36 +300,31 @@ async def process_callback_renew_key(callback_query: types.CallbackQuery, sessio client_id = record["client_id"] expiry_time = record["expiry_time"] - keyboard = types.InlineKeyboardMarkup( - inline_keyboard=[ - [ - types.InlineKeyboardButton( - text=f'📅 1 месяц ({RENEWAL_PLANS["1"]["price"]} руб.)', - callback_data=f"renew_plan|1|{client_id}", - ) - ], - [ - types.InlineKeyboardButton( - text=f'📅 3 месяца ({RENEWAL_PLANS["3"]["price"]} руб.)', - callback_data=f"renew_plan|3|{client_id}", - ) - ], - [ - types.InlineKeyboardButton( - text=f'📅 6 месяцев ({RENEWAL_PLANS["6"]["price"]} руб.)', - callback_data=f"renew_plan|6|{client_id}", - ) - ], - [ - types.InlineKeyboardButton( - text=f'📅 12 месяцев ({RENEWAL_PLANS["12"]["price"]} руб.)', - callback_data=f"renew_plan|12|{client_id}", - ) - ], - [types.InlineKeyboardButton(text="🔙 Назад", callback_data="profile")], - ] + builder = InlineKeyboardBuilder() + builder.row( + InlineKeyboardButton( + text=f'📅 1 месяц ({RENEWAL_PLANS["1"]["price"]} руб.)', + callback_data=f"renew_plan|1|{client_id}", + ) + ) + builder.row( + InlineKeyboardButton( + text=f'📅 3 месяца ({RENEWAL_PLANS["3"]["price"]} руб.)', + callback_data=f"renew_plan|3|{client_id}", + ) + ) + builder.row( + InlineKeyboardButton( + text=f'📅 6 месяцев ({RENEWAL_PLANS["6"]["price"]} руб.)', + callback_data=f"renew_plan|6|{client_id}", + ) + ) + builder.row( + InlineKeyboardButton( + text=f'📅 12 месяцев ({RENEWAL_PLANS["12"]["price"]} руб.)', + callback_data=f"renew_plan|12|{client_id}", + ) ) - balance = await get_balance(tg_id) response_message = PLAN_SELECTION_MSG.format( @@ -348,7 +334,7 @@ async def process_callback_renew_key(callback_query: types.CallbackQuery, sessio await callback_query.message.answer( text=response_message, - reply_markup=keyboard, + reply_markup=builder.as_markup(), ) else: await callback_query.message.answer("Ключ не найден.") @@ -434,21 +420,21 @@ async def process_callback_renew_plan(callback_query: types.CallbackQuery, sessi balance = await get_balance(tg_id) if balance < cost: - replenish_button = types.InlineKeyboardButton(text="Пополнить баланс", callback_data="pay") - view_profile = types.InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile") - keyboard = types.InlineKeyboardMarkup(inline_keyboard=[[replenish_button], [view_profile]]) + builder = InlineKeyboardBuilder() + builder.row(InlineKeyboardButton(text="Пополнить баланс", callback_data="pay")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) await callback_query.message.answer( INSUFFICIENT_FUNDS_MSG, - reply_markup=keyboard, + reply_markup=builder.as_markup(), ) return response_message = SUCCESS_RENEWAL_MSG.format(months=RENEWAL_PLANS[plan]["months"]) - back_button = types.InlineKeyboardButton(text="Назад", callback_data="profile") - keyboard = types.InlineKeyboardMarkup(inline_keyboard=[[back_button]]) + builder = InlineKeyboardBuilder() + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) - await callback_query.message.answer(response_message, reply_markup=keyboard) + await callback_query.message.answer(response_message, reply_markup=builder.as_markup()) async def renew_key_on_servers(): tasks = [] diff --git a/handlers/keys/trial_key.py b/handlers/keys/trial_key.py index 85f02c33..13c74362 100644 --- a/handlers/keys/trial_key.py +++ b/handlers/keys/trial_key.py @@ -1,107 +1,53 @@ -import asyncio from datetime import datetime, timedelta +from typing import Any import uuid -import asyncpg from py3xui import AsyncApi from client import add_client -from config import ADMIN_PASSWORD, ADMIN_USERNAME, CLUSTERS, DATABASE_URL, PUBLIC_LINK, TRIAL_TIME -from database import store_key +from config import ADMIN_PASSWORD, ADMIN_USERNAME, CLUSTERS, PUBLIC_LINK, TRIAL_TIME +from database import store_key, use_trial from handlers.texts import INSTRUCTIONS from handlers.utils import generate_random_email, get_least_loaded_cluster -async def create_trial_key(tg_id: int): - conn = await asyncpg.connect(DATABASE_URL) - try: - client_id = str(uuid.uuid4()) - email = generate_random_email() +async def create_trial_key(tg_id: int, session: Any): + client_id = str(uuid.uuid4()) + email = generate_random_email() + public_link = f"{PUBLIC_LINK}{email}/{tg_id}" + instructions = INSTRUCTIONS + result = {"key": public_link, "instructions": instructions} + current_time = datetime.utcnow() + expiry_time = current_time + timedelta(days=TRIAL_TIME, hours=3) + expiry_timestamp = int(expiry_time.timestamp() * 1000) - public_link = f"{PUBLIC_LINK}{email}/{tg_id}" - instructions = INSTRUCTIONS + least_loaded_cluster = await get_least_loaded_cluster() + for server_id, server in CLUSTERS[least_loaded_cluster].items(): + xui = AsyncApi( + CLUSTERS[least_loaded_cluster][server_id]["API_URL"], + username=ADMIN_USERNAME, + password=ADMIN_PASSWORD, + ) - result = {"key": public_link, "instructions": instructions} - - asyncio.create_task(generate_and_store_keys(tg_id, client_id, email, public_link)) - - return result - - finally: - await conn.close() - - -async def generate_and_store_keys(tg_id: int, client_id: str, email: str, public_link: str): - conn = await asyncpg.connect(DATABASE_URL) - try: - current_time = datetime.utcnow() - expiry_time = current_time + timedelta(days=TRIAL_TIME, hours=3) - expiry_timestamp = int(expiry_time.timestamp() * 1000) - - least_loaded_cluster = await get_least_loaded_cluster() - - tasks = [] - for server_id, server in CLUSTERS[least_loaded_cluster].items(): - task = create_key_on_server( - least_loaded_cluster, - server_id, - client_id, - email, - tg_id, - expiry_timestamp, - ) - tasks.append(task) - - await asyncio.gather(*tasks) - - await store_key( - tg_id, + await add_client( + xui, client_id, email, - expiry_timestamp, - public_link, - server_id=least_loaded_cluster, - ) - - await conn.execute( - """ - INSERT INTO connections (tg_id, trial) - VALUES ($1, 1) - ON CONFLICT (tg_id) - DO UPDATE SET trial = 1 - """, tg_id, + limit_ip=1, + total_gb=0, + expiry_time=expiry_timestamp, + enable=True, + flow="xtls-rprx-vision", ) - finally: - await conn.close() - -async def create_key_on_server( - cluster_id: str, - server_id: str, - client_id: str, - email: str, - tg_id: int, - expiry_timestamp: int, -): - """Создает ключ на сервере в указанном кластере и возвращает результат.""" - - xui = AsyncApi( - CLUSTERS[cluster_id][server_id]["API_URL"], - username=ADMIN_USERNAME, - password=ADMIN_PASSWORD, - ) - - response = await add_client( - xui, + await store_key( + tg_id, client_id, email, - tg_id, - limit_ip=1, - total_gb=0, - expiry_time=expiry_timestamp, - enable=True, - flow="xtls-rprx-vision", + expiry_timestamp, + public_link, + server_id=least_loaded_cluster, ) - - return response + await use_trial(tg_id, session) + return result diff --git a/handlers/pay.py b/handlers/pay.py index 06ff1dc3..9c588072 100644 --- a/handlers/pay.py +++ b/handlers/pay.py @@ -48,7 +48,7 @@ async def handle_pay(callback_query: CallbackQuery): ) builder.row(InlineKeyboardButton(text="🎟️ Активировать купон", callback_data="activate_coupon")) - builder.row(InlineKeyboardButton(text="⬅️ Вернуться в профиль", callback_data="profile")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) await callback_query.message.answer( "💸 Выберите удобный способ пополнения баланса:\n\n" diff --git a/handlers/payments/utils.py b/handlers/payments/utils.py index 6f807f53..21015090 100644 --- a/handlers/payments/utils.py +++ b/handlers/payments/utils.py @@ -8,7 +8,7 @@ from logger import logger async def send_payment_success_notification(user_id: int, amount: float): try: builder = InlineKeyboardBuilder() - builder.row(InlineKeyboardButton(text="Перейти в профиль", callback_data="profile")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) await bot.send_message( chat_id=user_id, text=f"Ваш баланс успешно пополнен на {amount} рублей. Спасибо за оплату!", diff --git a/handlers/payments/yookassa_pay.py b/handlers/payments/yookassa_pay.py index 5cec58b3..0278fa0f 100644 --- a/handlers/payments/yookassa_pay.py +++ b/handlers/payments/yookassa_pay.py @@ -60,7 +60,7 @@ async def process_callback_pay_yookassa(callback_query: types.CallbackQuery, sta callback_data="enter_custom_amount_yookassa", ) ) - builder.row(InlineKeyboardButton(text="⬅️ Назад", callback_data="profile")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) key_count = await get_key_count(tg_id) diff --git a/handlers/profile.py b/handlers/profile.py index 9b1f52cd..2ab7cffe 100644 --- a/handlers/profile.py +++ b/handlers/profile.py @@ -66,7 +66,7 @@ async def process_callback_view_profile(callback_query: types.CallbackQuery, sta @router.callback_query(F.data == "view_tariffs") async def view_tariffs_handler(callback_query: types.CallbackQuery): builder = InlineKeyboardBuilder() - builder.row(InlineKeyboardButton(text="⬅️ Вернуться в профиль", callback_data="profile")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) await callback_query.message.answer( "🚀 Доступные тарифы VPN:\n\n" @@ -94,7 +94,7 @@ async def invite_handler(callback_query: types.CallbackQuery): image_path = os.path.join("img", "pic_invite.jpg") builder = InlineKeyboardBuilder() - builder.row(InlineKeyboardButton(text="⬅️ Вернуться в профиль", callback_data="profile")) + builder.row(InlineKeyboardButton(text="👤 Личный кабинет", callback_data="profile")) if os.path.isfile(image_path): with open(image_path, "rb") as image_file: await callback_query.message.answer_photo( diff --git a/handlers/start.py b/handlers/start.py index 7b8ecdf5..8737f2cc 100644 --- a/handlers/start.py +++ b/handlers/start.py @@ -8,30 +8,32 @@ from aiogram.types import BufferedInputFile, CallbackQuery, InlineKeyboardButton from aiogram.utils.keyboard import InlineKeyboardBuilder from config import CHANNEL_URL, CONNECT_ANDROID, CONNECT_IOS, DOWNLOAD_ANDROID, DOWNLOAD_IOS, SUPPORT_CHAT_URL -from database import add_connection, add_referral, check_connection_exists, get_trial, restore_trial +from database import add_connection, add_referral, check_connection_exists, get_trial, use_trial from handlers.keys.trial_key import create_trial_key from handlers.texts import INSTRUCTIONS_TRIAL, WELCOME_TEXT, get_about_vpn +from logger import logger router = Router() @router.callback_query(F.data == "start") -async def handle_start_callback_query(callback_query: CallbackQuery, state: FSMContext): - await start_command(callback_query.message, state) +async def handle_start_callback_query(callback_query: CallbackQuery, state: FSMContext, session: Any, admin: bool): + await start_command(callback_query.message, state, session, admin) @router.message(Command("start")) -async def start_command(message: Message, admin: bool): +async def start_command(message: Message, state: FSMContext, session: Any, admin: bool): if message.text: try: referrer_tg_id = int(message.text.split("referral_")[1]) - await add_referral(message.from_user.id, referrer_tg_id) - except (ValueError,IndexError): + await add_referral(message.from_user.id, referrer_tg_id, session) + except (ValueError, IndexError): pass connection_exists = await check_connection_exists(message.from_user.id) if not connection_exists: - await add_connection(message.from_user.id) - trial_status = await get_trial(message.from_user.id) + await add_connection(message.from_user.id, session) + trial_status = await get_trial(message.from_user.id, session) + logger.info(f'trial_status {trial_status}') image_path = os.path.join("img", "pic.jpg") builder = InlineKeyboardBuilder() @@ -71,7 +73,7 @@ async def handle_connect_vpn(callback_query: CallbackQuery, session: Any): if "error" in trial_key_info: await callback_query.message.answer(trial_key_info["error"]) else: - await restore_trial(user_id) + await use_trial(user_id, session) key_message = ( f"🔑 Ваш персональный ключ доступа:\n" diff --git a/middlewares/database.py b/middlewares/database.py index d529dad4..b1687aa6 100644 --- a/middlewares/database.py +++ b/middlewares/database.py @@ -14,10 +14,9 @@ class DatabaseMiddleware(BaseMiddleware): event: TelegramObject, data: Dict[str, Any], ) -> Any: - async with await asyncpg.create_pool(DATABASE_URL) as pool: - async with pool.acquire() as session: - data["session"] = session - try: - return await handler(event, data) - finally: - await pool.release(session) + conn = await asyncpg.connect(DATABASE_URL) + try: + data["session"] = conn + return await handler(event, data) + finally: + await conn.close()