From 83d589fff5e1a42bf5a4faec0eec978d903f7707 Mon Sep 17 00:00:00 2001 From: Vladless Date: Sun, 27 Oct 2024 01:03:00 +0300 Subject: [PATCH] =?UTF-8?q?=D0=9F=D0=BE=D1=80=D1=8F=D0=B4=D0=BE=D0=BA.=20?= =?UTF-8?q?=D0=98=D1=81=D0=BF=D1=80=D0=B0=D0=B2=D0=BB=D0=B5=D0=BD=D0=B8?= =?UTF-8?q?=D1=8F=20=D0=BC=D0=B5=D0=BB=D0=BA=D0=B8=D1=85=20=D0=BD=D0=B5?= =?UTF-8?q?=D0=B4=D0=BE=D1=87=D0=B5=D1=82=D0=BE=D0=B2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bot.py | 3 +- handlers/admin.py | 22 +----- handlers/backup_handler.py | 9 ++- handlers/commands.py | 135 +++++++++++++++++++++++++++++++++++++ handlers/trial_key.py | 25 ++----- handlers/utils.py | 23 ++++++- 6 files changed, 171 insertions(+), 46 deletions(-) create mode 100644 handlers/commands.py diff --git a/bot.py b/bot.py index 11d67ae7..2c685583 100644 --- a/bot.py +++ b/bot.py @@ -9,8 +9,9 @@ dp = Dispatcher(bot=bot, storage=storage) router = Router() from handlers import (backup_handler, key_management, keys, notifications, pay, - profile, start, admin) + profile, start, admin, commands) +dp.include_router(commands.router) dp.include_router(start.router) dp.include_router(profile.router) dp.include_router(keys.router) diff --git a/handlers/admin.py b/handlers/admin.py index 1edb3b4a..4bfafe10 100644 --- a/handlers/admin.py +++ b/handlers/admin.py @@ -24,24 +24,4 @@ async def cmd_add_balance(message: types.Message): except ValueError: await message.reply("Пожалуйста, используйте формат: /add_balance ") except Exception as e: - await message.reply(f"Произошла ошибка: {e}") - -@router.message(Command('check_balance')) -async def cmd_check_balance(message: types.Message): - if message.from_user.id != ADMIN_ID: - await message.reply("У вас нет доступа к этой команде.") - return - - try: - _, client_id = message.text.split() - - if not await check_connection_exists(int(client_id)): - await message.reply(f"Клиент с ID {client_id} не найден.") - return - - balance = await get_balance(int(client_id)) - await message.reply(f"Баланс клиента {client_id}: {balance} у.е.") - except ValueError: - await message.reply("Пожалуйста, используйте формат: /check_balance ") - except Exception as e: - await message.reply(f"Произошла ошибка: {e}") + await message.reply(f"Произошла ошибка: {e}") \ No newline at end of file diff --git a/handlers/backup_handler.py b/handlers/backup_handler.py index 6ad684a3..0dcbe436 100644 --- a/handlers/backup_handler.py +++ b/handlers/backup_handler.py @@ -1,12 +1,17 @@ from aiogram import Router from aiogram.filters import Command from aiogram.types import Message +from config import ADMIN_ID router = Router() @router.message(Command('backup')) async def backup_command(message: Message): - from backup import backup_database + if message.from_user.id != ADMIN_ID: + await message.answer("У вас нет прав для выполнения этой команды.") + return + + from backup import backup_database await message.answer("Запускаю бэкап базы данных...") await backup_database() - await message.answer("Бэкап завершен и отправлен админу.") + await message.answer("Бэкап завершен и отправлен админу.") \ No newline at end of file diff --git a/handlers/commands.py b/handlers/commands.py new file mode 100644 index 00000000..603febe9 --- /dev/null +++ b/handlers/commands.py @@ -0,0 +1,135 @@ +from aiogram import F, Router, types +from aiogram.filters import Command +from aiogram.fsm.context import FSMContext +from aiogram.fsm.state import State, StatesGroup +import asyncpg + +from bot import bot +from config import ADMIN_ID, DATABASE_URL +from handlers.backup_handler import backup_command +from handlers.pay import ReplenishBalanceState, process_custom_amount_input +from handlers.profile import process_callback_view_profile +from handlers.start import start_command +from handlers.texts import TRIAL +from handlers.admin import cmd_add_balance +from handlers.key_management import handle_key_name_input + +router = Router() + +class Form(StatesGroup): + waiting_for_server_selection = State() + waiting_for_key_name = State() + viewing_profile = State() + waiting_for_message = State() + +@router.message(Command('start')) +async def handle_start(message: types.Message, state: FSMContext): + await start_command(message) + +@router.message(Command('add_balance')) +async def handle_add_balance(message: types.Message, state: FSMContext): + await cmd_add_balance(message) + +@router.message(Command('menu')) +async def handle_menu(message: types.Message, state: FSMContext): + await start_command(message) + +@router.message(Command('send_trial')) +async def handle_send_trial_command(message: types.Message, state: FSMContext): + try: + conn = await asyncpg.connect(DATABASE_URL) + try: + records = await conn.fetch(''' + SELECT tg_id FROM connections WHERE trial = 0 + ''') + + if records: + for record in records: + tg_id = record['tg_id'] + trial_message = ( + TRIAL + ) + try: + await bot.send_message(chat_id=tg_id, text=trial_message) + except Exception as e: + if "Forbidden: bot was blocked by the user" in str(e): + print(f"Бот заблокирован пользователем с tg_id: {tg_id}") + else: + print(f"Ошибка при отправке сообщения пользователю {tg_id}: {e}") + + await message.answer("Сообщения о пробном периоде отправлены всем пользователям с не использованным ключом.") + else: + await message.answer("Нет пользователей с не использованными пробными ключами.") + + finally: + await conn.close() + + except Exception as e: + await message.answer(f"Ошибка при отправке сообщений: {e}") + +@router.message(Command('send_to_all')) +async def send_message_to_all_clients(message: types.Message, state: FSMContext): + if message.from_user.id != ADMIN_ID: + await message.answer("У вас нет прав для выполнения этой команды.") + return + + await message.answer("Введите текст сообщения, который вы хотите отправить всем клиентам:") + await state.set_state(Form.waiting_for_message) + +@router.message(Form.waiting_for_message) +async def process_message_to_all(message: types.Message, state: FSMContext): + text_message = message.text + + try: + conn = await asyncpg.connect(DATABASE_URL) + tg_ids = await conn.fetch('SELECT tg_id FROM connections') + + for record in tg_ids: + tg_id = record['tg_id'] + try: + await bot.send_message(chat_id=tg_id, text=text_message) + except Exception as e: + print(f"Ошибка при отправке сообщения пользователю {tg_id}: {e}. Пропускаем этого пользователя.") + + await message.answer("Сообщение было отправлено всем клиентам.") + except Exception as e: + print(f"Ошибка при подключении к базе данных: {e}") + await message.answer("Произошла ошибка при отправке сообщения.") + finally: + await conn.close() + + await state.clear() + +@router.message() +async def handle_text(message: types.Message, state: FSMContext): + current_state = await state.get_state() + + if message.text in ["/send_to_all"]: + await send_message_to_all_clients(message, state) + return + + if message.text == "Мой профиль": + callback_query = types.CallbackQuery( + id="1", + from_user=message.from_user, + chat_instance='', + data='view_profile', + message=message + ) + await process_callback_view_profile(callback_query, state) + return + + if current_state == ReplenishBalanceState.entering_custom_amount.state: + await process_custom_amount_input(message, state) + return + + if current_state == Form.waiting_for_key_name.state: + await handle_key_name_input(message, state) + return + + if message.text == "/backup": + await backup_command(message) + return + + elif current_state is None: + await start_command(message) diff --git a/handlers/trial_key.py b/handlers/trial_key.py index db3388d2..1d57a599 100644 --- a/handlers/trial_key.py +++ b/handlers/trial_key.py @@ -1,18 +1,14 @@ import random import asyncpg -import time -import uuid # Импортируем модуль для генерации UUID +import uuid from config import DATABASE_URL, SERVERS, ADMIN_USERNAME, ADMIN_PASSWORD -from auth import login_with_credentials, link # Импортируйте необходимые функции +from auth import login_with_credentials, link from client import add_client from database import store_key, add_connection from handlers.texts import INSTRUCTIONS from datetime import datetime, timedelta +from handlers.utils import generate_random_email, get_least_loaded_server -def generate_random_email(): - """Генерирует случайный набор символов.""" - random_string = ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=6)) - return random_string async def create_trial_key(tg_id: int): conn = await asyncpg.connect(DATABASE_URL) @@ -21,7 +17,7 @@ async def create_trial_key(tg_id: int): session = await login_with_credentials(server_id, ADMIN_USERNAME, ADMIN_PASSWORD) current_time = datetime.utcnow() - expiry_time = current_time + timedelta(days=30, hours=3) + expiry_time = current_time + timedelta(days=1, hours=3) expiry_timestamp = int(expiry_time.timestamp() * 1000) client_id = str(uuid.uuid4()) @@ -52,16 +48,3 @@ async def create_trial_key(tg_id: int): return {'error': 'Не удалось добавить клиента на панель'} finally: await conn.close() - -async def get_least_loaded_server(conn): - least_loaded_server_id = None - min_load_percentage = float('inf') - - for server_id, server in SERVERS.items(): - count = await conn.fetchval('SELECT COUNT(*) FROM keys WHERE server_id = $1', server_id) - percent_full = (count / 60) * 100 if count <= 60 else 100 - if percent_full < min_load_percentage: - min_load_percentage = percent_full - least_loaded_server_id = server_id - - return least_loaded_server_id diff --git a/handlers/utils.py b/handlers/utils.py index 5c2436f8..06b13e24 100644 --- a/handlers/utils.py +++ b/handlers/utils.py @@ -1,3 +1,24 @@ import re +import random +from config import SERVERS def sanitize_key_name(key_name: str) -> str: - return re.sub(r'[^a-z0-9@._-]', '', key_name.lower()) \ No newline at end of file + return re.sub(r'[^a-z0-9@._-]', '', key_name.lower()) + +def generate_random_email(): + """Генерирует случайный набор символов.""" + random_string = ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=6)) + return random_string + +async def get_least_loaded_server(conn): + """Находит сервер с наименьшей загрузкой.""" + least_loaded_server_id = None + min_load_percentage = float('inf') + + for server_id, server in SERVERS.items(): + count = await conn.fetchval('SELECT COUNT(*) FROM keys WHERE server_id = $1', server_id) + percent_full = (count / 60) * 100 if count <= 60 else 100 + if percent_full < min_load_percentage: + min_load_percentage = percent_full + least_loaded_server_id = server_id + + return least_loaded_server_id \ No newline at end of file