From 342c039ad8a2ea06035f13558b6a1111243cb8e8 Mon Sep 17 00:00:00 2001 From: Vlad Date: Mon, 16 Sep 2024 04:02:17 +0300 Subject: [PATCH] =?UTF-8?q?=D0=B4=D0=BE=D0=B1=D0=B0=D0=B2=D0=BB=D0=B5?= =?UTF-8?q?=D0=BD=D0=B0=20=D0=B4=D0=B0=D1=82=D0=B0=20=D0=BE=D0=BA=D0=BE?= =?UTF-8?q?=D0=BD=D1=87=D0=B0=D0=BD=D0=B8=D1=8F=20=D0=BF=D0=BE=D0=B4=D0=BF?= =?UTF-8?q?=D0=B8=D1=81=D0=BA=D0=B8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bot.py | 85 +++++++++++++++++++++++++++++++++-------------------- database.py | 18 ++++++++++-- 2 files changed, 68 insertions(+), 35 deletions(-) diff --git a/bot.py b/bot.py index e98a009e..808e1e51 100644 --- a/bot.py +++ b/bot.py @@ -4,19 +4,19 @@ from aiogram.types import InlineKeyboardButton, InlineKeyboardMarkup, Message from aiogram.fsm.storage.memory import MemoryStorage from aiogram.fsm.context import FSMContext from aiogram.fsm.state import State, StatesGroup -from auth import login_with_credentials, link, get_clients -from client import add_client, generate_client_id +from auth import login_with_credentials, link +from client import add_client from datetime import datetime, timedelta -from config import API_TOKEN, ADMIN_PASSWORD, ADMIN_USERNAME -from database import add_connection, DATABASE_PATH +from config import API_TOKEN, ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_PATH +from database import add_connection, has_active_key, get_active_key_email import uuid -import aiosqlite -import requests import re +import aiosqlite class Form(StatesGroup): waiting_for_key_name = State() waiting_for_statistics = State() + waiting_for_expiry_date = State() bot = Bot(token=API_TOKEN) storage = MemoryStorage() @@ -28,11 +28,12 @@ session = login_with_credentials(ADMIN_USERNAME, ADMIN_PASSWORD) @dp.message(Command("start")) async def start_command(message: Message): - welcome_text = "Добро пожаловать! Вы можете создать ключ для подключения VPN или просмотреть статистику использования." - + welcome_text = "Добро пожаловать! Вы можете создать ключ для подключения VPN, просмотреть статистику использования или узнать дату окончания ключа." + button_create_key = InlineKeyboardButton(text='Создать ключ', callback_data='create_key') button_view_stats = InlineKeyboardButton(text='Посмотреть статистику', callback_data='view_stats') - keyboard = InlineKeyboardMarkup(inline_keyboard=[[button_create_key], [button_view_stats]]) + button_view_expiry = InlineKeyboardButton(text='Дата окончания ключа', callback_data='view_expiry') + keyboard = InlineKeyboardMarkup(inline_keyboard=[[button_create_key], [button_view_stats], [button_view_expiry]]) await message.reply(welcome_text, reply_markup=keyboard) @@ -47,23 +48,20 @@ async def process_callback_view_stats(callback_query: types.CallbackQuery, state tg_id = callback_query.from_user.id try: - async with aiosqlite.connect(DATABASE_PATH) as db: - async with db.execute("SELECT email FROM connections WHERE tg_id = ? AND expiry_time > ?", (tg_id, int(datetime.utcnow().timestamp() * 1000))) as cursor: - record = await cursor.fetchone() - if record: - email = record[0] - connection_link = link(session, email) - - # Извлечение данных о загрузке и выгрузке из ссылки - up_match = re.search(r'up=(\d+)', connection_link) - down_match = re.search(r'down=(\d+)', connection_link) - - up = up_match.group(1) if up_match else "Неизвестно" - down = down_match.group(1) if down_match else "Неизвестно" - - statistics = f"Статистика вашего ключа:\nЗагрузка: {up} MB\nВыгрузка: {down} MB" - else: - statistics = "У вас нет активных ключей." + email = await get_active_key_email(tg_id) + if email: + connection_link = link(session, email) + + # Извлечение данных о загрузке и выгрузке из ссылки + up_match = re.search(r'up=(\d+)', connection_link) + down_match = re.search(r'down=(\d+)', connection_link) + + up = up_match.group(1) if up_match else "Неизвестно" + down = down_match.group(1) if down_match else "Неизвестно" + + statistics = f"Статистика вашего ключа:\nЗагрузка: {up} MB\nВыгрузка: {down} MB" + else: + statistics = "У вас нет активных ключей." except Exception as e: statistics = f"Ошибка при получении статистики: {e}" @@ -71,6 +69,32 @@ async def process_callback_view_stats(callback_query: types.CallbackQuery, state await callback_query.message.reply(f"Ваша статистика:\n{statistics}") await callback_query.answer() +@dp.callback_query(F.data == 'view_expiry') +async def process_callback_view_expiry(callback_query: types.CallbackQuery): + tg_id = callback_query.from_user.id + + try: + email = await get_active_key_email(tg_id) + if email: + async with aiosqlite.connect(DATABASE_PATH) as db: + async with db.execute("SELECT expiry_time FROM connections WHERE tg_id = ? AND expiry_time > ?", + (tg_id, int(datetime.utcnow().timestamp() * 1000))) as cursor: + record = await cursor.fetchone() + if record: + expiry_time = record[0] + expiry_date = datetime.utcfromtimestamp(expiry_time / 1000).strftime("%Y-%m-%d %H:%M:%S") + message_text = f"Дата окончания вашего ключа: {expiry_date}" + else: + message_text = "У вас нет активных ключей." + else: + message_text = "У вас нет активных ключей." + + except Exception as e: + message_text = f"Ошибка при получении даты окончания ключа: {e}" + + await callback_query.message.reply(message_text) + await callback_query.answer() + @dp.message() async def handle_text(message: types.Message, state: FSMContext): current_state = await state.get_state() @@ -80,12 +104,9 @@ async def handle_text(message: types.Message, state: FSMContext): tg_id = message.from_user.id # Проверка на наличие активного ключа - async with aiosqlite.connect(DATABASE_PATH) as db: - async with db.execute("SELECT COUNT(*) FROM connections WHERE tg_id = ? AND expiry_time > ?", (tg_id, int(datetime.utcnow().timestamp() * 1000))) as cursor: - count = await cursor.fetchone() - if count[0] > 0: - await message.reply("У вас уже есть активный ключ. Один клиент может иметь только один активный ключ.") - return + if await has_active_key(tg_id): + await message.reply("У вас уже есть активный ключ. Один клиент может иметь только один активный ключ.") + return # Создание уникального ID клиента client_id = str(uuid.uuid4()) diff --git a/database.py b/database.py index 3f02fd85..192bb1c8 100644 --- a/database.py +++ b/database.py @@ -1,7 +1,6 @@ import aiosqlite from datetime import datetime - -DATABASE_PATH = 'database.db' +from config import DATABASE_PATH async def init_db(): async with aiosqlite.connect(DATABASE_PATH) as db: @@ -18,10 +17,23 @@ async def init_db(): async def add_connection(tg_id: int, client_id: str, email: str, expiry_time: int): async with aiosqlite.connect(DATABASE_PATH) as db: - # Добавляем новый ключ await db.execute(''' INSERT INTO connections (tg_id, client_id, email, expiry_time) VALUES (?, ?, ?, ?) ''', (tg_id, client_id, email, expiry_time)) await db.commit() +async def has_active_key(tg_id: int) -> bool: + async with aiosqlite.connect(DATABASE_PATH) as db: + async with db.execute("SELECT COUNT(*) FROM connections WHERE tg_id = ? AND expiry_time > ?", + (tg_id, int(datetime.utcnow().timestamp() * 1000))) as cursor: + count = await cursor.fetchone() + return count[0] > 0 + +async def get_active_key_email(tg_id: int) -> str: + async with aiosqlite.connect(DATABASE_PATH) as db: + async with db.execute("SELECT email FROM connections WHERE tg_id = ? AND expiry_time > ?", + (tg_id, int(datetime.utcnow().timestamp() * 1000))) as cursor: + record = await cursor.fetchone() + return record[0] if record else None +