diff --git a/backup.py b/backup.py index 8550e619..c7a81649 100644 --- a/backup.py +++ b/backup.py @@ -1,14 +1,24 @@ import os import subprocess -from datetime import datetime +from datetime import datetime, timedelta +from pathlib import Path +from typing import Optional, Tuple, Union +import aiofiles from aiogram.types import BufferedInputFile - from config import ADMIN_ID, BACK_DIR, DB_NAME, DB_PASSWORD, DB_USER, PG_HOST, PG_PORT + +from bot import bot from logger import logger async def backup_database() -> Exception | None: + """ + Создает резервную копию базы данных и отправляет ее администраторам. + + Returns: + Optional[Exception]: Исключение в случае ошибки или None при успешном выполнении + """ backup_file_path, exception = _create_database_backup() if exception: @@ -30,16 +40,25 @@ async def backup_database() -> Exception | None: def _create_database_backup() -> tuple[str | None, Exception | None]: + """ + Создает резервную копию базы данных PostgreSQL. + + Returns: + Tuple[Optional[str], Optional[Exception]]: Путь к файлу бэкапа и исключение (если произошла ошибка) + """ date_formatted = datetime.now().strftime("%Y-%m-%d-%H%M%S") - if not os.path.exists(BACK_DIR): - os.makedirs(BACK_DIR) + # Создаем директорию для бэкапов, если она не существует + backup_dir = Path(BACK_DIR) + backup_dir.mkdir(parents=True, exist_ok=True) - filename = os.path.join(BACK_DIR, f"{DB_NAME}-backup-{date_formatted}.sql") + filename = backup_dir / f"{DB_NAME}-backup-{date_formatted}.sql" try: + # Устанавливаем пароль PostgreSQL через переменную окружения os.environ["PGPASSWORD"] = DB_PASSWORD + # Запускаем pg_dump для создания бэкапа subprocess.run( [ "pg_dump", @@ -52,61 +71,94 @@ def _create_database_backup() -> tuple[str | None, Exception | None]: "-F", "c", "-f", - filename, + str(filename), DB_NAME, ], check=True, + capture_output=True, + text=True, ) logger.info(f"Бэкап базы данных создан: {filename}") - return filename, None + return str(filename), None except subprocess.CalledProcessError as e: + logger.error(f"Ошибка при выполнении pg_dump: {e.stderr}") + return None, e + except Exception as e: + logger.error(f"Непредвиденная ошибка при создании бэкапа: {e}") return None, e finally: - del os.environ["PGPASSWORD"] + # Удаляем переменную окружения с паролем + if "PGPASSWORD" in os.environ: + del os.environ["PGPASSWORD"] -def _cleanup_old_backups() -> None | Exception: +def _cleanup_old_backups() -> Exception | None: + """ + Удаляет бэкапы старше 3 дней. + + Returns: + Optional[Exception]: Исключение в случае ошибки или None при успешном выполнении + """ try: - subprocess.run( - [ - "find", - BACK_DIR, - "-type", - "f", - "-name", - "*.sql", - "-mtime", - "+3", - "-exec", - "rm", - "{}", - ";", - ], - check=True, - ) - logger.info("Старые бэкапы удалены.") + backup_dir = Path(BACK_DIR) + if not backup_dir.exists(): + return None + + # Вычисляем дату, старше которой нужно удалить файлы + cutoff_date = datetime.now() - timedelta(days=3) + + # Находим и удаляем старые файлы бэкапов + for backup_file in backup_dir.glob("*.sql"): + if backup_file.is_file(): + file_mtime = datetime.fromtimestamp(backup_file.stat().st_mtime) + if file_mtime < cutoff_date: + backup_file.unlink() + logger.info(f"Удален старый бэкап: {backup_file}") + + logger.info("Очистка старых бэкапов завершена") return None - except subprocess.CalledProcessError as e: + except Exception as e: + logger.error(f"Ошибка при удалении старых бэкапов: {e}") return e -async def create_backup_and_send_to_admins(xui) -> None: - await xui.login() - await xui.database.export() +async def create_backup_and_send_to_admins(client) -> None: + """ + Создает бэкап и отправляет администраторам через переданный клиент. + + Args: + client: Клиент для работы с базой данных + """ + await client.login() + await client.database.export() async def _send_backup_to_admins(backup_file_path: str) -> None: + """ + Отправляет файл бэкапа всем администраторам через Telegram. + + Args: + backup_file_path: Путь к файлу бэкапа + + Raises: + Exception: При ошибке отправки файла + """ + if not backup_file_path or not os.path.exists(backup_file_path): + raise FileNotFoundError(f"Файл бэкапа не найден: {backup_file_path}") + try: - import aiofiles - - from bot import bot - async with aiofiles.open(backup_file_path, "rb") as backup_file: backup_data = await backup_file.read() - backup_input_file = BufferedInputFile(file=backup_data, filename=os.path.basename(backup_file_path)) - admin_ids = ADMIN_ID if isinstance(ADMIN_ID, list) else [ADMIN_ID] - for admin_id in admin_ids: - await bot.send_document(chat_id=admin_id, document=backup_input_file) - logger.info(f"Бэкап базы данных отправлен админу: {admin_id}") + filename = os.path.basename(backup_file_path) + backup_input_file = BufferedInputFile(file=backup_data, filename=filename) + + # Отправляем файл каждому администратору + for admin_id in ADMIN_ID: + try: + await bot.send_document(chat_id=admin_id, document=backup_input_file) + logger.info(f"Бэкап базы данных отправлен админу: {admin_id}") + except Exception as e: + logger.error(f"Не удалось отправить бэкап админу {admin_id}: {e}") except Exception as e: logger.error(f"Ошибка при отправке бэкапа в Telegram: {e}") + raise diff --git a/bot.py b/bot.py index d53f1188..ce55c505 100644 --- a/bot.py +++ b/bot.py @@ -8,8 +8,8 @@ from aiogram.filters import ExceptionTypeFilter from aiogram.fsm.storage.memory import MemoryStorage from aiogram.types import BufferedInputFile, ErrorEvent from aiogram.utils.markdown import hbold - from config import ADMIN_ID, API_TOKEN + from filters.private import IsPrivateFilter from logger import logger from middlewares import register_middleware diff --git a/client.py b/client.py index 750d5520..ef0643c3 100644 --- a/client.py +++ b/client.py @@ -2,8 +2,8 @@ from dataclasses import dataclass from typing import Any import py3xui - from config import LIMIT_IP, SUPERNODE + from logger import logger diff --git a/database.py b/database.py index 45e56bbd..1490e0b7 100644 --- a/database.py +++ b/database.py @@ -4,8 +4,8 @@ from typing import Any import asyncpg import pytz - from config import CASHBACK, CHECK_REFERRAL_REWARD_ISSUED, DATABASE_URL, REFERRAL_BONUS_PERCENTAGES + from logger import logger @@ -842,7 +842,9 @@ async def get_total_referral_bonus(conn, referrer_tg_id: int, max_levels: int) - ORDER BY tg_id, created_at ) """ - bonus_query = bonus_cte + f""" + bonus_query = ( + bonus_cte + + f""" SELECT COALESCE(SUM( CASE @@ -859,6 +861,7 @@ async def get_total_referral_bonus(conn, referrer_tg_id: int, max_levels: int) - JOIN earliest_payments ep ON rl.referred_tg_id = ep.tg_id WHERE rl.level <= {max_levels} """ + ) else: bonus_cte = f""" WITH RECURSIVE @@ -881,7 +884,9 @@ async def get_total_referral_bonus(conn, referrer_tg_id: int, max_levels: int) - WHERE rl.level < {max_levels} ) """ - bonus_query = bonus_cte + f""" + bonus_query = ( + bonus_cte + + f""" SELECT COALESCE(SUM( CASE @@ -898,6 +903,7 @@ async def get_total_referral_bonus(conn, referrer_tg_id: int, max_levels: int) - JOIN payments p ON rl.referred_tg_id = p.tg_id WHERE p.status = 'success' AND rl.level <= {max_levels} """ + ) total_bonus = await conn.fetchval(bonus_query, referrer_tg_id) logger.debug(f"Получена общая сумма бонусов от рефералов: {total_bonus}") return total_bonus @@ -1080,7 +1086,8 @@ async def upsert_user( last_name: str = None, language_code: str = None, is_bot: bool = False, -): + session: Any = None, +) -> dict: """ Обновляет или вставляет информацию о пользователе в базу данных. @@ -1091,16 +1098,29 @@ async def upsert_user( last_name (str, optional): Фамилия пользователя language_code (str, optional): Код языка пользователя is_bot (bool, optional): Флаг, указывающий является ли пользователь ботом + session (Any, optional): Существующая сессия базы данных + + Returns: + dict: Словарь с информацией о пользователе после обновления/вставки Raises: Exception: В случае ошибки при работе с базой данных """ conn = None - try: - conn = await asyncpg.connect(DATABASE_URL) - logger.info(f"Установлено подключение к базе данных для обновления пользователя {tg_id}") + close_conn = False - await conn.execute( + try: + # Используем переданную сессию или создаем новое подключение + if session: + conn = session + logger.debug(f"Используем существующую сессию для обновления пользователя {tg_id}") + else: + conn = await asyncpg.connect(DATABASE_URL) + close_conn = True + logger.info(f"Установлено новое подключение к базе данных для обновления пользователя {tg_id}") + + # Выполняем вставку/обновление и сразу получаем обновленные данные + user_data = await conn.fetchrow( """ INSERT INTO users (tg_id, username, first_name, last_name, language_code, is_bot, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) @@ -1112,6 +1132,9 @@ async def upsert_user( language_code = COALESCE(EXCLUDED.language_code, users.language_code), is_bot = EXCLUDED.is_bot, updated_at = CURRENT_TIMESTAMP + RETURNING + id, tg_id, username, first_name, last_name, language_code, + is_bot, created_at, updated_at """, tg_id, username, @@ -1120,14 +1143,19 @@ async def upsert_user( language_code, is_bot, ) - logger.info(f"Успешно обновлена информация о пользователе {tg_id}") + + logger.debug(f"Успешно обновлена информация о пользователе {tg_id}") + + # Преобразуем результат в словарь + return dict(user_data) except Exception as e: logger.error(f"Ошибка при обновлении информации о пользователе {tg_id}: {e}") raise finally: - if conn: + # Закрываем соединение только если мы его создали + if conn and close_conn: await conn.close() - logger.info("Закрытие подключения к базе данных") + logger.debug("Закрытие подключения к базе данных") async def add_payment(tg_id: int, amount: float, payment_system: str): diff --git a/filters/admin.py b/filters/admin.py index 5a392217..d04fa5a3 100644 --- a/filters/admin.py +++ b/filters/admin.py @@ -1,14 +1,11 @@ from aiogram.filters import BaseFilter from aiogram.types import Message - from config import ADMIN_ID class IsAdminFilter(BaseFilter): async def __call__(self, message: Message) -> bool: try: - admin_ids: int | list[int] = ADMIN_ID - if isinstance(admin_ids, list): - return message.from_user.id in admin_ids + return message.from_user.id in ADMIN_ID except Exception: return False diff --git a/handlers/admin/admin_coupons.py b/handlers/admin/admin_coupons.py index a2ba1ca1..3dfc16f6 100644 --- a/handlers/admin/admin_coupons.py +++ b/handlers/admin/admin_coupons.py @@ -4,8 +4,8 @@ from aiogram import F, Router, types from aiogram.fsm.context import FSMContext from aiogram.fsm.state import State, StatesGroup from aiogram.types import CallbackQuery, Message - from config import USERNAME_BOT + from database import create_coupon, delete_coupon, get_all_coupons from filters.admin import IsAdminFilter from keyboards.admin.coupons_kb import AdminCouponDeleteCallback, build_coupons_kb, build_coupons_list_kb diff --git a/handlers/admin/admin_servers.py b/handlers/admin/admin_servers.py index 4ddae1a9..73cf9c65 100644 --- a/handlers/admin/admin_servers.py +++ b/handlers/admin/admin_servers.py @@ -6,10 +6,10 @@ from aiogram import F, Router, types from aiogram.fsm.context import FSMContext from aiogram.fsm.state import State, StatesGroup from aiogram.types import CallbackQuery, Message +from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL from py3xui import AsyncApi from backup import create_backup_and_send_to_admins -from config import ADMIN_PASSWORD, ADMIN_USERNAME, DATABASE_URL from database import check_unique_server_name, delete_server, get_servers from filters.admin import IsAdminFilter from handlers.keys.key_utils import create_key_on_cluster diff --git a/handlers/admin/admin_stats.py b/handlers/admin/admin_stats.py index bb5697a8..187f3106 100644 --- a/handlers/admin/admin_stats.py +++ b/handlers/admin/admin_stats.py @@ -1,10 +1,10 @@ from datetime import datetime from typing import Any +import pytz from aiogram import F, Router from aiogram.exceptions import TelegramBadRequest from aiogram.types import CallbackQuery -import pytz from filters.admin import IsAdminFilter from keyboards.admin.panel_kb import AdminPanelCallback, build_admin_back_kb diff --git a/handlers/admin/admin_users.py b/handlers/admin/admin_users.py index 343ab510..c219cf11 100644 --- a/handlers/admin/admin_users.py +++ b/handlers/admin/admin_users.py @@ -8,8 +8,8 @@ from aiogram.exceptions import TelegramBadRequest from aiogram.fsm.context import FSMContext from aiogram.fsm.state import State, StatesGroup from aiogram.types import CallbackQuery, Message - from config import TOTAL_GB + from database import ( delete_key, delete_user_data, diff --git a/handlers/captcha.py b/handlers/captcha.py index 66643304..4f38075c 100644 --- a/handlers/captcha.py +++ b/handlers/captcha.py @@ -6,8 +6,8 @@ from aiogram import F, Router from aiogram.fsm.context import FSMContext from aiogram.types import CallbackQuery, Message from aiogram.utils.keyboard import InlineKeyboardBuilder - from config import CAPTCHA_EMOJIS + from logger import logger router = Router() diff --git a/handlers/coupons.py b/handlers/coupons.py index ed8a3126..fc230c49 100644 --- a/handlers/coupons.py +++ b/handlers/coupons.py @@ -38,9 +38,9 @@ async def handle_activate_coupon(callback_query_or_message: Message | CallbackQu await edit_or_send_message( target_message=target_message, text="🎫 Введите код купона:\n\n" - "📝 Пожалуйста, введите действующий код купона, который вы хотите активировать. 🔑", + "📝 Пожалуйста, введите действующий код купона, который вы хотите активировать. 🔑", reply_markup=builder.as_markup(), - media_path=None + media_path=None, ) await state.set_state(CouponActivationState.waiting_for_coupon_code) diff --git a/handlers/donate.py b/handlers/donate.py index 73d6f76e..66c1c392 100644 --- a/handlers/donate.py +++ b/handlers/donate.py @@ -3,8 +3,8 @@ from aiogram.fsm.context import FSMContext from aiogram.fsm.state import State, StatesGroup from aiogram.types import CallbackQuery, InlineKeyboardButton, LabeledPrice, Message, PreCheckoutQuery from aiogram.utils.keyboard import InlineKeyboardBuilder - from config import RUB_TO_XTR + from logger import logger diff --git a/handlers/instructions/instructions.py b/handlers/instructions/instructions.py index 0f1ad70e..e8019b72 100644 --- a/handlers/instructions/instructions.py +++ b/handlers/instructions/instructions.py @@ -8,8 +8,8 @@ from aiogram.types import ( Message, ) from aiogram.utils.keyboard import InlineKeyboardBuilder - from config import CONNECT_MACOS, CONNECT_WINDOWS, SUPPORT_CHAT_URL + from database import get_key_details from handlers.texts import ( CONNECT_TV_TEXT, diff --git a/handlers/keys/key_management.py b/handlers/keys/key_management.py index 3902c216..5afb5f8c 100644 --- a/handlers/keys/key_management.py +++ b/handlers/keys/key_management.py @@ -8,10 +8,6 @@ from aiogram import F, Router from aiogram.fsm.context import FSMContext from aiogram.types import CallbackQuery, InlineKeyboardButton, Message from aiogram.utils.keyboard import InlineKeyboardBuilder -from py3xui import AsyncApi - -from bot import bot -from client import delete_client from config import ( ADMIN_PASSWORD, ADMIN_USERNAME, @@ -29,6 +25,10 @@ from config import ( USE_COUNTRY_SELECTION, USE_NEW_PAYMENT_FLOW, ) +from py3xui import AsyncApi + +from bot import bot +from client import delete_client from database import ( add_connection, check_connection_exists, diff --git a/handlers/keys/key_utils.py b/handlers/keys/key_utils.py index 197d5ef7..834de969 100644 --- a/handlers/keys/key_utils.py +++ b/handlers/keys/key_utils.py @@ -1,10 +1,10 @@ import asyncio from typing import Any +from config import ADMIN_PASSWORD, ADMIN_USERNAME, LIMIT_IP, PUBLIC_LINK, SUPERNODE, TOTAL_GB, USE_COUNTRY_SELECTION from py3xui import AsyncApi from client import ClientConfig, add_client, delete_client, extend_client_key, get_client_traffic, toggle_client -from config import ADMIN_PASSWORD, ADMIN_USERNAME, LIMIT_IP, PUBLIC_LINK, SUPERNODE, TOTAL_GB, USE_COUNTRY_SELECTION from database import get_servers, store_key from handlers.utils import get_least_loaded_cluster from logger import logger diff --git a/handlers/keys/keys.py b/handlers/keys/keys.py index 23d4e0c1..1a155bca 100644 --- a/handlers/keys/keys.py +++ b/handlers/keys/keys.py @@ -9,8 +9,6 @@ import pytz from aiogram import F, Router, types from aiogram.types import CallbackQuery, InlineKeyboardButton, Message from aiogram.utils.keyboard import InlineKeyboardBuilder - -from bot import bot from config import ( CONNECT_ANDROID, CONNECT_IOS, @@ -26,6 +24,8 @@ from config import ( USE_COUNTRY_SELECTION, USE_NEW_PAYMENT_FLOW, ) + +from bot import bot from database import ( check_server_name_by_cluster, create_temporary_data, diff --git a/handlers/keys/subscriptions.py b/handlers/keys/subscriptions.py index 231e0847..970858ac 100644 --- a/handlers/keys/subscriptions.py +++ b/handlers/keys/subscriptions.py @@ -5,11 +5,12 @@ import re import time import urllib.parse from datetime import datetime +from typing import Dict, List, Optional, Union import aiohttp import asyncpg +import pytz from aiohttp import web - from config import ( DATABASE_URL, PROJECT_NAME, @@ -21,75 +22,91 @@ from config import ( USE_COUNTRY_SELECTION, USERNAME_BOT, ) + from database import get_key_details, get_servers from handlers.utils import convert_to_bytes from logger import logger -db_pool = None +# Функции для работы с URL и подписками +async def fetch_url_content(url: str, identifier: str) -> list[str]: + """ + Получает содержимое подписки по URL и декодирует его. -async def init_db_pool(): - """Инициализация пула соединений, если он ещё не создан.""" - global db_pool - if not db_pool: - db_pool = await asyncpg.create_pool(dsn=DATABASE_URL, min_size=5, max_size=20) + Args: + url: URL для получения содержимого + identifier: Идентификатор пользователя (tg_id или email) - -async def fetch_url_content(url, tg_id): - """Получает содержимое подписки по URL и декодирует его.""" + Returns: + Список строк из подписки + """ try: - logger.info(f"Получение URL: {url} для tg_id: {tg_id}") + logger.info(f"Получение URL: {url} для идентификатора: {identifier}") timeout = aiohttp.ClientTimeout(total=5) async with aiohttp.ClientSession(timeout=timeout) as session: async with session.get(url, ssl=False) as response: if response.status == 200: content = await response.text() - logger.info(f"Успешно получен контент с {url} для tg_id: {tg_id}") + logger.info(f"Успешно получен контент с {url} для идентификатора: {identifier}") return base64.b64decode(content).decode("utf-8").split("\n") else: - logger.error(f"Не удалось получить {url} для tg_id: {tg_id}, статус: {response.status}") + logger.error( + f"Не удалось получить {url} для идентификатора: {identifier}, статус: {response.status}" + ) return [] except TimeoutError: - logger.error(f"Таймаут при получении {url} для tg_id: {tg_id}") + logger.error(f"Таймаут при получении {url} для идентификатора: {identifier}") return [] except Exception as e: - logger.error(f"Ошибка при получении {url} для tg_id: {tg_id}: {e}") + logger.error(f"Ошибка при получении {url} для идентификатора: {identifier}: {e}") return [] -async def combine_unique_lines(urls, tg_id, query_string): - """Объединяет строки подписки, удаляя дубликаты.""" +async def combine_unique_lines(urls: list[str], identifier: str, query_string: str) -> list[str]: + """ + Объединяет строки подписки из нескольких URL, удаляя дубликаты. + + Args: + urls: Список URL для получения подписок + identifier: Идентификатор пользователя (tg_id или email) + query_string: Строка запроса для добавления к URL + + Returns: + Список уникальных строк из всех подписок + """ if SUPERNODE: - logger.info(f"Режим SUPERNODE активен. Возвращаем первую ссылку для tg_id: {tg_id}") + logger.info(f"Режим SUPERNODE активен. Возвращаем первую ссылку для идентификатора: {identifier}") if not urls: return [] url_with_query = f"{urls[0]}?{query_string}" if query_string else urls[0] - return await fetch_url_content(url_with_query, tg_id) + return await fetch_url_content(url_with_query, identifier) - logger.info(f"Начинаем объединение подписок для tg_id: {tg_id}, запрос: {query_string}") + logger.info(f"Начинаем объединение подписок для идентификатора: {identifier}, запрос: {query_string}") urls_with_query = [f"{url}?{query_string}" if query_string else url for url in urls] logger.info(f"Составлены URL-адреса: {urls_with_query}") - tasks = [fetch_url_content(url, tg_id) for url in urls_with_query] + tasks = [fetch_url_content(url, identifier) for url in urls_with_query] results = await asyncio.gather(*tasks) all_lines = set() for lines in results: all_lines.update(filter(None, lines)) - logger.info(f"Объединено {len(all_lines)} строк после фильтрации и удаления дубликатов для tg_id: {tg_id}") + logger.info( + f"Объединено {len(all_lines)} строк после фильтрации и удаления дубликатов для идентификатора: {identifier}" + ) return list(all_lines) -transition_date = datetime.strptime(TRANSITION_DATE_STR, "%Y-%m-%d %H:%M:%S") -transition_timestamp_ms = int(transition_date.timestamp() * 1000) -transition_timestamp_ms_adjusted = transition_timestamp_ms - (3 * 60 * 60 * 1000) -logger.info(f"Время перехода (с поправкой на часовой пояс): {transition_timestamp_ms_adjusted}") - - -async def get_subscription_urls(server_id: str, email: str, conn) -> list: +async def get_subscription_urls(server_id: str, email: str, conn) -> list[str]: """ - Универсальная функция, которая в зависимости от флага USE_COUNTRY_SELECTION - получает список URL-адресов для подписки. Возвращает пустой список, если - нужные данные не найдены. + Получает список URL-адресов для подписки в зависимости от режима выбора страны. + + Args: + server_id: Идентификатор сервера или кластера + email: Email пользователя + conn: Соединение с базой данных + + Returns: + Список URL-адресов для подписки """ if USE_COUNTRY_SELECTION: logger.info(f"Режим выбора страны активен. Ищем сервер {server_id} в БД.") @@ -102,7 +119,7 @@ async def get_subscription_urls(server_id: str, email: str, conn) -> list: logger.info(f"Используем подписку {urls[0]}") return urls - servers = await get_servers() + servers = await get_servers(conn) logger.info(f"Режим выбора страны отключен. Используем кластер {server_id}.") cluster_servers = servers.get(server_id, []) if not cluster_servers: @@ -113,20 +130,49 @@ async def get_subscription_urls(server_id: str, email: str, conn) -> list: return urls -def calculate_traffic(cleaned_subscriptions, expiry_time_ms): +def get_transition_timestamp() -> int: + """ + Получает временную метку перехода с учетом часового пояса Москвы. + + Returns: + Временная метка перехода в миллисекундах + """ + moscow_tz = pytz.timezone("Europe/Moscow") + transition_date_naive = datetime.strptime(TRANSITION_DATE_STR, "%Y-%m-%d %H:%M:%S") + transition_date = moscow_tz.localize(transition_date_naive) + transition_timestamp_ms = int(transition_date.timestamp() * 1000) + return transition_timestamp_ms + + +# Функции для обработки и форматирования данных +def calculate_traffic(cleaned_subscriptions: list[str], expiry_time_ms: int | None) -> str: + """ + Рассчитывает информацию о трафике на основе подписок. + + Args: + cleaned_subscriptions: Список строк подписки + expiry_time_ms: Время истечения подписки в миллисекундах + + Returns: + Строка с информацией о трафике + """ expire_timestamp = int(expiry_time_ms / 1000) if expiry_time_ms else 0 + if TOTAL_GB != 0: country_remaining = {} for line in cleaned_subscriptions: if "#" not in line: continue + try: _, meta = line.split("#", 1) except ValueError: continue + parts = meta.split("-") country = parts[0].strip() remaining_str = parts[1].strip() if len(parts) == 2 else "" + if remaining_str: remaining_str = remaining_str.replace(",", ".") m_total = re.search(r"([\d\.]+)\s*([GMKTB]B)", remaining_str, re.IGNORECASE) @@ -135,10 +181,12 @@ def calculate_traffic(cleaned_subscriptions, expiry_time_ms): unit = m_total.group(2).upper() remaining_bytes = convert_to_bytes(value, unit) country_remaining[country] = remaining_bytes + num_countries = len(country_remaining) issued_per_country = TOTAL_GB total_traffic_bytes = issued_per_country * num_countries consumed_traffic_bytes = total_traffic_bytes - sum(country_remaining.values()) + if consumed_traffic_bytes < 0: consumed_traffic_bytes = 0 else: @@ -148,8 +196,116 @@ def calculate_traffic(cleaned_subscriptions, expiry_time_ms): return f"upload=0; download={consumed_traffic_bytes}; total={total_traffic_bytes}; expire={expire_timestamp}" -async def handle_subscription(request, old_subscription=False): - """Обрабатывает запрос на подписку (старую или новую).""" +def clean_subscription_line(line: str) -> str: + """ + Очищает строку подписки, оставляя только нужную информацию. + + Args: + line: Исходная строка подписки + + Returns: + Очищенная строка подписки + """ + if "#" not in line: + return line + + try: + base, meta = line.split("#", 1) + except ValueError: + return line + + parts = meta.split("-") + country = parts[0].strip() if parts else "" + traffic = "" + + for part in parts[1:]: + part_decoded = urllib.parse.unquote(part).strip() + if re.search(r"\d+(?:[.,]\d+)?\s*(?:GB|MB|KB|TB)", part_decoded, re.IGNORECASE): + traffic = part_decoded + break + + meta_clean = f"{country} - {traffic}" if traffic else country + return base + "#" + meta_clean + + +def format_time_left(expiry_time_ms: int | None) -> str: + """ + Форматирует оставшееся время подписки. + + Args: + expiry_time_ms: Время истечения подписки в миллисекундах + + Returns: + Отформатированная строка с оставшимся временем + """ + if not expiry_time_ms: + return "N/A" + + now_ms = int(time.time() * 1000) + remaining_sec = max((expiry_time_ms - now_ms) / 1000, 0) + days = int(remaining_sec // 86400) + hours = int((remaining_sec % 86400) // 3600) + + return f"{days}D,{hours}H ⏳" if days else f"{hours}H ⏳" + + +def prepare_headers( + user_agent: str, project_name: str, subscription_info: str, subscription_userinfo: str +) -> dict[str, str]: + """ + Подготавливает заголовки ответа в зависимости от User-Agent клиента. + + Args: + user_agent: User-Agent клиента + project_name: Название проекта + subscription_info: Информация о подписке + subscription_userinfo: Информация о трафике + + Returns: + Словарь с заголовками ответа + """ + if "Happ" in user_agent: + encoded_project_name = f"{project_name}" + announce_str = f"↖️Бот | {subscription_info} | Поддержка↗️" + return { + "Content-Type": "text/plain; charset=utf-8", + "Content-Disposition": "inline", + "profile-update-interval": "3", + "profile-title": "base64:" + base64.b64encode(encoded_project_name.encode("utf-8")).decode("utf-8"), + "support-url": SUPPORT_CHAT_URL, + "announce": "base64:" + base64.b64encode(announce_str.encode("utf-8")).decode("utf-8"), + "profile-web-page-url": f"https://t.me/{USERNAME_BOT}", + "subscription-userinfo": subscription_userinfo, + } + elif "Hiddify" in user_agent: + encoded_project_name = f"{project_name}\n📄 Подписка: {subscription_info.split(' - ')[0].split(': ')[1]}" + return { + "profile-update-interval": "3", + "profile-title": "base64:" + base64.b64encode(encoded_project_name.encode("utf-8")).decode("utf-8"), + "subscription-userinfo": subscription_userinfo, + } + else: + encoded_project_name = f"{project_name}\n{subscription_info}" + return { + "Content-Type": "text/plain; charset=utf-8", + "Content-Disposition": "inline", + "profile-update-interval": "3", + "profile-title": "base64:" + base64.b64encode(encoded_project_name.encode("utf-8")).decode("utf-8"), + } + + +# Основные обработчики запросов +async def handle_subscription(request: web.Request, old_subscription: bool = False) -> web.Response: + """ + Обрабатывает запрос на подписку (старую или новую). + + Args: + request: Объект запроса + old_subscription: Флаг, указывающий на тип подписки (старая или новая) + + Returns: + Ответ с подпиской или сообщением об ошибке + """ email = request.match_info.get("email") tg_id = request.match_info.get("tg_id") if not old_subscription else None @@ -160,9 +316,9 @@ async def handle_subscription(request, old_subscription=False): logger.info( f"Обработка запроса для {'старого' if old_subscription else 'нового'} клиента: email={email}, tg_id={tg_id}" ) - await init_db_pool() - async with db_pool.acquire() as conn: + async with asyncpg.connect(DATABASE_URL) as conn: + # Получение данных клиента client_data = await get_key_details(email, conn) if not client_data: logger.warning(f"Клиент с email {email} не найден в базе.") @@ -171,103 +327,61 @@ async def handle_subscription(request, old_subscription=False): stored_tg_id = client_data.get("tg_id") server_id = client_data["server_id"] + # Проверка tg_id для новых клиентов if not old_subscription and str(tg_id) != str(stored_tg_id): logger.warning(f"Неверный tg_id для клиента с email {email}.") return web.Response(text="❌ Неверные данные. Получите свой ключ в боте.", status=403) + # Проверка даты создания для старых клиентов if old_subscription: created_at_ms = client_data["created_at"] + # Используем pytz для корректного сравнения времени created_at_datetime = datetime.utcfromtimestamp(created_at_ms / 1000) logger.info(f"created_at для {email}: {created_at_datetime}, server_id: {server_id}") - if created_at_ms >= transition_timestamp_ms_adjusted: + + # Получаем временную метку перехода + transition_timestamp_ms = get_transition_timestamp() + logger.info(f"Время перехода (с учетом часового пояса Москвы): {transition_timestamp_ms}") + + if created_at_ms >= transition_timestamp_ms: logger.info(f"Клиент с email {email} является новым.") return web.Response(text="❌ Эта ссылка устарела. Пожалуйста, обновите ссылку.", status=400) + # Форматирование оставшегося времени expiry_time_ms = client_data.get("expiry_time") - if expiry_time_ms: - now_ms = int(time.time() * 1000) - remaining_sec = max((expiry_time_ms - now_ms) / 1000, 0) - days = int(remaining_sec // 86400) - hours = int((remaining_sec % 86400) // 3600) - time_left = f"{days}D,{hours}H ⏳" if days else f"{hours}H ⏳" - else: - time_left = "N/A" + time_left = format_time_left(expiry_time_ms) + # Получение URL-адресов подписки urls = await get_subscription_urls(server_id, email, conn) if not urls: return web.Response(text="❌ Сервер не найден.", status=404) + # Получение и обработка подписок query_string = request.query_string if not old_subscription else "" combined_subscriptions = await combine_unique_lines(urls, tg_id or email, query_string) random.shuffle(combined_subscriptions) - cleaned_subscriptions = [] - for line in combined_subscriptions: - if "#" in line: - try: - base, meta = line.split("#", 1) - except ValueError: - cleaned_line = line - else: - parts = meta.split("-") - country = parts[0].strip() if parts else "" - traffic = "" - for part in parts[1:]: - part_decoded = urllib.parse.unquote(part).strip() - if re.search(r"\d+(?:[.,]\d+)?\s*(?:GB|MB|KB|TB)", part_decoded, re.IGNORECASE): - traffic = part_decoded - break - meta_clean = f"{country} - {traffic}" if traffic else country - cleaned_line = base + "#" + meta_clean - else: - cleaned_line = line - cleaned_subscriptions.append(cleaned_line) + # Очистка строк подписки + cleaned_subscriptions = [clean_subscription_line(line) for line in combined_subscriptions] - final_subscriptions = cleaned_subscriptions - base64_encoded = base64.b64encode("\n".join(final_subscriptions).encode("utf-8")).decode("utf-8") + # Кодирование подписки в base64 + base64_encoded = base64.b64encode("\n".join(cleaned_subscriptions).encode("utf-8")).decode("utf-8") subscription_info = f"📄 Подписка: {email} - {time_left}" + # Подготовка заголовков ответа user_agent = request.headers.get("User-Agent", "") subscription_userinfo = calculate_traffic(cleaned_subscriptions, expiry_time_ms) - if "Happ" in user_agent: - encoded_project_name = f"{PROJECT_NAME}" - support_username = SUPPORT_CHAT_URL.split("https://t.me/")[-1] - announce_str = f"↖️Бот | {subscription_info} | Поддержка↗️" - headers = { - "Content-Type": "text/plain; charset=utf-8", - "Content-Disposition": "inline", - "profile-update-interval": "3", - "profile-title": "base64:" + base64.b64encode(encoded_project_name.encode("utf-8")).decode("utf-8"), - "support-url": SUPPORT_CHAT_URL, - "announce": "base64:" + base64.b64encode(announce_str.encode("utf-8")).decode("utf-8"), - "profile-web-page-url": f"https://t.me/{USERNAME_BOT}", - "subscription-userinfo": subscription_userinfo, - } - elif "Hiddify" in user_agent: - encoded_project_name = f"{PROJECT_NAME}\n📄 Подписка: {email}" - headers = { - "profile-update-interval": "3", - "profile-title": "base64:" + base64.b64encode(encoded_project_name.encode("utf-8")).decode("utf-8"), - "subscription-userinfo": subscription_userinfo, - } - else: - encoded_project_name = f"{PROJECT_NAME}\n{subscription_info}" - headers = { - "Content-Type": "text/plain; charset=utf-8", - "Content-Disposition": "inline", - "profile-update-interval": "3", - "profile-title": "base64:" + base64.b64encode(encoded_project_name.encode("utf-8")).decode("utf-8"), - } + headers = prepare_headers(user_agent, PROJECT_NAME, subscription_info, subscription_userinfo) logger.info(f"Возвращаем объединенные подписки для email: {email}") return web.Response(text=base64_encoded, headers=headers) -async def handle_old_subscription(request): +async def handle_old_subscription(request: web.Request) -> web.Response: """Обработка запроса для старых клиентов.""" return await handle_subscription(request, old_subscription=True) -async def handle_new_subscription(request): +async def handle_new_subscription(request: web.Request) -> web.Response: """Обработка запроса для новых клиентов.""" return await handle_subscription(request, old_subscription=False) diff --git a/handlers/notifications/general_notifications.py b/handlers/notifications/general_notifications.py index d9cf2152..75af5057 100644 --- a/handlers/notifications/general_notifications.py +++ b/handlers/notifications/general_notifications.py @@ -4,7 +4,6 @@ from datetime import datetime, timedelta import asyncpg import pytz from aiogram import Bot, Router - from config import ( DATABASE_URL, NOTIFICATION_TIME, @@ -17,6 +16,7 @@ from config import ( TOTAL_GB, TRIAL_TIME_DISABLE, ) + from database import ( add_notification, check_notification_time, diff --git a/handlers/notifications/special_notifications.py b/handlers/notifications/special_notifications.py index ced6f1df..cd48c2fa 100644 --- a/handlers/notifications/special_notifications.py +++ b/handlers/notifications/special_notifications.py @@ -6,8 +6,8 @@ import pytz from aiogram import Bot, Router, types from aiogram.exceptions import TelegramForbiddenError from aiogram.utils.keyboard import InlineKeyboardBuilder - from config import NOTIFY_EXTRA_DAYS, NOTIFY_INACTIVE, NOTIFY_INACTIVE_TRAFFIC, SUPPORT_CHAT_URL, TRIAL_TIME + from database import ( add_notification, check_notification_time, diff --git a/handlers/pay.py b/handlers/pay.py index 3f7604cd..353b5fde 100644 --- a/handlers/pay.py +++ b/handlers/pay.py @@ -1,7 +1,6 @@ from aiogram import F, Router from aiogram.types import CallbackQuery, InlineKeyboardButton from aiogram.utils.keyboard import InlineKeyboardBuilder - from config import ( CRYPTO_BOT_ENABLE, DONATIONS_ENABLE, diff --git a/handlers/payments/__init__.py b/handlers/payments/__init__.py index e1ccb3fc..ba0f040b 100644 --- a/handlers/payments/__init__.py +++ b/handlers/payments/__init__.py @@ -1,7 +1,6 @@ __all__ = ("router",) from aiogram import Router - from config import ( CRYPTO_BOT_ENABLE, ROBOKASSA_ENABLE, diff --git a/handlers/payments/robokassa_pay.py b/handlers/payments/robokassa_pay.py index 93e8b747..442f707a 100644 --- a/handlers/payments/robokassa_pay.py +++ b/handlers/payments/robokassa_pay.py @@ -8,8 +8,6 @@ from aiogram.fsm.state import State, StatesGroup from aiogram.types import InlineKeyboardButton, InlineKeyboardMarkup from aiogram.utils.keyboard import InlineKeyboardBuilder from aiohttp import web -from robokassa import HashAlgorithm, Robokassa - from config import ( DATABASE_URL, ROBOKASSA_ENABLE, @@ -18,6 +16,8 @@ from config import ( ROBOKASSA_PASSWORD2, ROBOKASSA_TEST_MODE, ) +from robokassa import HashAlgorithm, Robokassa + from database import ( add_connection, add_payment, @@ -28,9 +28,8 @@ from database import ( ) from handlers.payments.utils import send_payment_success_notification from handlers.texts import PAYMENT_OPTIONS -from logger import logger - from handlers.utils import edit_or_send_message +from logger import logger router = Router() @@ -54,9 +53,7 @@ if ROBOKASSA_ENABLE: def generate_payment_link(amount, inv_id, description, tg_id): """Генерация ссылки на оплату.""" - logger.debug( - f"Generating payment link for amount: {amount}, inv_id: {inv_id}, description: {description}" - ) + logger.debug(f"Generating payment link for amount: {amount}, inv_id: {inv_id}, description: {description}") payment_link = robokassa._payment.link.generate_by_script( out_sum=amount, inv_id=inv_id, @@ -68,9 +65,7 @@ def generate_payment_link(amount, inv_id, description, tg_id): @router.callback_query(F.data == "pay_robokassa") -async def process_callback_pay_robokassa( - callback_query: types.CallbackQuery, state: FSMContext, session: Any -): +async def process_callback_pay_robokassa(callback_query: types.CallbackQuery, state: FSMContext, session: Any): tg_id = callback_query.message.chat.id logger.info(f"User {tg_id} initiated Robokassa payment.") @@ -108,16 +103,14 @@ async def process_callback_pay_robokassa( target_message=callback_query.message, text="Выберите сумму пополнения:", reply_markup=builder.as_markup(), - force_text=True + force_text=True, ) await state.set_state(ReplenishBalanceState.choosing_amount_robokassa) logger.info(f"Displayed amount selection for user {tg_id}.") @router.callback_query(F.data.startswith("robokassa_amount|")) -async def process_amount_selection( - callback_query: types.CallbackQuery, state: FSMContext -): +async def process_amount_selection(callback_query: types.CallbackQuery, state: FSMContext): logger.info(f"Получены данные callback_data: {callback_query.data}") data = callback_query.data.split("|") @@ -127,7 +120,7 @@ async def process_amount_selection( target_message=callback_query.message, text="Ошибка: данные повреждены.", reply_markup=types.InlineKeyboardMarkup(), - force_text=True + force_text=True, ) return @@ -142,7 +135,7 @@ async def process_amount_selection( target_message=callback_query.message, text="Некорректная сумма.", reply_markup=types.InlineKeyboardMarkup(), - force_text=True + force_text=True, ) return @@ -166,7 +159,7 @@ async def process_amount_selection( target_message=callback_query.message, text=f"Вы выбрали пополнение на {amount} рублей. Для оплаты перейдите по ссылке ниже:", reply_markup=confirm_keyboard, - force_text=True + force_text=True, ) logger.info(f"Payment link sent to user {callback_query.message.chat.id}.") @@ -183,9 +176,7 @@ async def robokassa_webhook(request): shp_id = params.get("shp_id") signature_value = params.get("SignatureValue") - logger.info( - f"OutSum: {amount}, InvId: {inv_id}, shp_id: {shp_id}, SignatureValue: {signature_value}" - ) + logger.info(f"OutSum: {amount}, InvId: {inv_id}, shp_id: {shp_id}, SignatureValue: {signature_value}") if not check_payment_signature(params): logger.error("Неверная подпись или данные запроса.") @@ -224,9 +215,7 @@ def check_payment_signature(params): logger.info(f"Signature string before hashing: {signature_string}") - expected_signature = ( - hashlib.md5(signature_string.encode("utf-8")).hexdigest().upper() - ) + expected_signature = hashlib.md5(signature_string.encode("utf-8")).hexdigest().upper() logger.info(f"Expected signature: {expected_signature}") logger.info(f"Received signature: {signature_value}") @@ -235,9 +224,7 @@ def check_payment_signature(params): @router.callback_query(F.data == "enter_custom_amount_robokassa") -async def process_custom_amount_selection( - callback_query: types.CallbackQuery, state: FSMContext -): +async def process_custom_amount_selection(callback_query: types.CallbackQuery, state: FSMContext): tg_id = callback_query.message.chat.id logger.info(f"User {tg_id} chose to enter a custom amount.") @@ -248,16 +235,16 @@ async def process_custom_amount_selection( target_message=callback_query.message, text="Пожалуйста, введите сумму пополнения.", reply_markup=builder.as_markup(), - force_text=True + force_text=True, ) - await state.set_state( - ReplenishBalanceState.waiting_for_payment_confirmation_robokassa - ) + await state.set_state(ReplenishBalanceState.waiting_for_payment_confirmation_robokassa) @router.message(ReplenishBalanceState.waiting_for_payment_confirmation_robokassa) -async def handle_custom_amount_input(message: types.Message | types.CallbackQuery, state: FSMContext = None, session: Any = None): +async def handle_custom_amount_input( + message: types.Message | types.CallbackQuery, state: FSMContext = None, session: Any = None +): if isinstance(message, types.CallbackQuery): tg_id = message.message.chat.id target_message = message.message @@ -277,7 +264,7 @@ async def handle_custom_amount_input(message: types.Message | types.CallbackQuer await edit_or_send_message( target_message=target_message, text="Данные для оплаты не найдены. Попробуйте снова.", - reply_markup=types.InlineKeyboardMarkup() + reply_markup=types.InlineKeyboardMarkup(), ) return @@ -288,7 +275,7 @@ async def handle_custom_amount_input(message: types.Message | types.CallbackQuer await edit_or_send_message( target_message=target_message, text="Недостаточная сумма для пополнения.", - reply_markup=types.InlineKeyboardMarkup() + reply_markup=types.InlineKeyboardMarkup(), ) return @@ -300,22 +287,22 @@ async def handle_custom_amount_input(message: types.Message | types.CallbackQuer builder.row(InlineKeyboardButton(text="⬅️ Назад", callback_data="pay_robokassa")) if state_type == "waiting_for_payment": - message_text = f"Вы выбрали пополнение на {amount} рублей для создания нового ключа. Перейдите по ссылке для оплаты:" + message_text = ( + f"Вы выбрали пополнение на {amount} рублей для создания нового ключа. Перейдите по ссылке для оплаты:" + ) elif state_type == "waiting_for_renewal_payment": - message_text = f"Вы выбрали пополнение на {amount} рублей для продления ключа. Перейдите по ссылке для оплаты:" + message_text = ( + f"Вы выбрали пополнение на {amount} рублей для продления ключа. Перейдите по ссылке для оплаты:" + ) else: await edit_or_send_message( target_message=target_message, text="Некорректное состояние данных. Попробуйте снова.", - reply_markup=types.InlineKeyboardMarkup() + reply_markup=types.InlineKeyboardMarkup(), ) return - await edit_or_send_message( - target_message=target_message, - text=message_text, - reply_markup=builder.as_markup() - ) + await edit_or_send_message(target_message=target_message, text=message_text, reply_markup=builder.as_markup()) if isinstance(state, FSMContext): await state.clear() @@ -325,5 +312,5 @@ async def handle_custom_amount_input(message: types.Message | types.CallbackQuer await edit_or_send_message( target_message=target_message, text="Произошла ошибка при создании платежа. Попробуйте позже.", - reply_markup=types.InlineKeyboardMarkup() + reply_markup=types.InlineKeyboardMarkup(), ) diff --git a/handlers/profile.py b/handlers/profile.py index 4805363a..fe1bff39 100644 --- a/handlers/profile.py +++ b/handlers/profile.py @@ -14,7 +14,6 @@ from aiogram.types import ( Message, ) from aiogram.utils.keyboard import InlineKeyboardBuilder - from config import ( DATABASE_URL, INLINE_MODE, @@ -25,6 +24,7 @@ from config import ( TRIAL_TIME, USERNAME_BOT, ) + from database import get_balance, get_key_count, get_last_payments, get_referral_stats, get_trial from handlers.buttons.profile import ( ADD_SUB, diff --git a/handlers/start.py b/handlers/start.py index b8793223..9e32037c 100644 --- a/handlers/start.py +++ b/handlers/start.py @@ -10,8 +10,6 @@ from aiogram.types import ( Message, ) from aiogram.utils.keyboard import InlineKeyboardBuilder - -from bot import bot from config import ( CAPTCHA_ENABLE, CHANNEL_EXISTS, @@ -21,6 +19,8 @@ from config import ( DONATIONS_ENABLE, SUPPORT_CHAT_URL, ) + +from bot import bot from database import ( add_connection, add_referral, diff --git a/handlers/utils.py b/handlers/utils.py index ffd17160..19d2d972 100644 --- a/handlers/utils.py +++ b/handlers/utils.py @@ -8,9 +8,9 @@ import aiofiles import aiohttp import asyncpg from aiogram.types import BufferedInputFile, InlineKeyboardMarkup, InputMediaPhoto, Message +from config import DATABASE_URL from bot import bot -from config import DATABASE_URL from database import get_all_keys, get_servers from logger import logger diff --git a/keyboards/admin/servers_kb.py b/keyboards/admin/servers_kb.py index 10118eba..e7d148e8 100644 --- a/keyboards/admin/servers_kb.py +++ b/keyboards/admin/servers_kb.py @@ -1,5 +1,5 @@ from aiogram.filters.callback_data import CallbackData -from aiogram.types import InlineKeyboardMarkup, InlineKeyboardButton +from aiogram.types import InlineKeyboardButton, InlineKeyboardMarkup from aiogram.utils.keyboard import InlineKeyboardBuilder from keyboards.admin.panel_kb import AdminPanelCallback, build_admin_back_btn @@ -46,7 +46,7 @@ def build_manage_cluster_kb(cluster_servers, cluster_name) -> InlineKeyboardMark InlineKeyboardButton( text="🔄 Синхронизация", callback_data=AdminServerEditorCallback(action="clusters_sync", data=cluster_name).pack(), - ) + ), ) builder.button( text="💾 Создать бэкап кластера", diff --git a/keyboards/admin/users_kb.py b/keyboards/admin/users_kb.py index be179bc1..52931cf9 100644 --- a/keyboards/admin/users_kb.py +++ b/keyboards/admin/users_kb.py @@ -3,8 +3,8 @@ from datetime import datetime, timezone from aiogram.filters.callback_data import CallbackData from aiogram.types import InlineKeyboardButton, InlineKeyboardMarkup from aiogram.utils.keyboard import InlineKeyboardBuilder - from config import RENEWAL_PRICES + from keyboards.admin.panel_kb import build_admin_back_btn diff --git a/middlewares/__init__.py b/middlewares/__init__.py index 889a6af3..f2b1d9d9 100644 --- a/middlewares/__init__.py +++ b/middlewares/__init__.py @@ -1,6 +1,8 @@ -from typing import Any +from collections.abc import Iterable +from typing import Any, Optional, Type, Union from aiogram import Dispatcher +from aiogram.dispatcher.middlewares.base import BaseMiddleware from .admin import AdminMiddleware from .loggings import LoggingMiddleware @@ -9,16 +11,48 @@ from .throttling import ThrottlingMiddleware from .user import UserMiddleware -def register_middleware(dispatcher: Dispatcher) -> None: - middlewares = [ - AdminMiddleware(), - SessionMiddleware(), - LoggingMiddleware(), - ThrottlingMiddleware(), - UserMiddleware(), +def register_middleware( + dispatcher: Dispatcher, + middlewares: Iterable[BaseMiddleware | type[BaseMiddleware]] | None = None, + exclude: Iterable[str] | None = None, +) -> None: + """Регистрирует middleware в диспетчере. + + Args: + dispatcher: Экземпляр диспетчера Aiogram + middlewares: Опциональный список middleware для регистрации. + Если не указан, регистрируются все стандартные middleware. + exclude: Опциональный список имен middleware, которые нужно исключить из регистрации. + Применяется только если middlewares не указан. + """ + # Если middleware не указаны, используем стандартный набор + if middlewares is None: + # Словарь всех доступных middleware + available_middlewares = { + "admin": AdminMiddleware(), + "session": SessionMiddleware(), + "logging": LoggingMiddleware(), + "throttling": ThrottlingMiddleware(), + "user": UserMiddleware(), + } + + # Фильтруем middleware по списку исключений + exclude_set = set(exclude or []) + middlewares = [middleware for name, middleware in available_middlewares.items() if name not in exclude_set] + + # Регистрируем middleware для всех типов обработчиков + handlers = [ + dispatcher.message, + dispatcher.callback_query, + dispatcher.inline_query, + # Можно добавить другие типы обработчиков при необходимости ] + # Регистрируем каждый middleware для каждого типа обработчика for middleware in middlewares: - dispatcher.message.outer_middleware(middleware) - dispatcher.callback_query.outer_middleware(middleware) - dispatcher.inline_query.outer_middleware(middleware) + # Если передан класс, а не экземпляр, создаем экземпляр + if isinstance(middleware, type): + middleware = middleware() + + for handler in handlers: + handler.outer_middleware(middleware) diff --git a/middlewares/admin.py b/middlewares/admin.py index bc21afe4..afc2129a 100644 --- a/middlewares/admin.py +++ b/middlewares/admin.py @@ -1,27 +1,57 @@ from collections.abc import Awaitable, Callable -from typing import Any +from typing import Any, Dict, Set, cast from aiogram import BaseMiddleware -from aiogram.types import TelegramObject - +from aiogram.types import CallbackQuery, Message, TelegramObject from config import ADMIN_ID class AdminMiddleware(BaseMiddleware): + """Middleware для проверки прав администратора. + + Добавляет в data['admin'] = True/False в зависимости от того, + является ли пользователь администратором. + """ + + # Кэш ID администраторов для быстрого доступа + _admin_ids: set[int] = set(ADMIN_ID) if isinstance(ADMIN_ID, (list, tuple)) else {ADMIN_ID} + async def __call__( self, handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]], event: TelegramObject, data: dict[str, Any], ) -> Any: + """Обрабатывает событие и добавляет флаг администратора в data. + + Args: + handler: Обработчик события + event: Событие Telegram + data: Словарь с данными события + + Returns: + Результат выполнения обработчика + """ data["admin"] = self._check_admin_access(event) return await handler(event, data) def _check_admin_access(self, event: TelegramObject) -> bool: + """Проверяет, имеет ли пользователь права администратора. + + Args: + event: Событие Telegram + + Returns: + True, если пользователь администратор, иначе False + """ try: - admin_ids: int | list[int] = ADMIN_ID - if isinstance(admin_ids, list): - return event.from_user.id in admin_ids - return event.from_user.id == admin_ids + if isinstance(event, Message): + return event.from_user and event.from_user.id in self._admin_ids + elif isinstance(event, CallbackQuery): + return event.from_user and event.from_user.id in self._admin_ids + + # Для других типов событий пытаемся получить from_user + user_id = getattr(getattr(event, "from_user", None), "id", None) + return user_id in self._admin_ids if user_id else False except Exception: return False diff --git a/middlewares/loggings.py b/middlewares/loggings.py index e99a60d6..42035ecc 100644 --- a/middlewares/loggings.py +++ b/middlewares/loggings.py @@ -1,13 +1,21 @@ from collections.abc import Awaitable, Callable -from typing import Any +from typing import Any, Dict, Optional, TypedDict, Union from aiogram import BaseMiddleware -from aiogram.types import CallbackQuery, InlineQuery, Message, TelegramObject +from aiogram.types import CallbackQuery, InlineQuery, Message, TelegramObject, User from logger import logger +class UserInfo(TypedDict): + user_id: int | None + username: str | None + action: str | None + + class LoggingMiddleware(BaseMiddleware): + """Middleware для логирования действий пользователя.""" + async def __call__( self, handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]], @@ -16,33 +24,38 @@ class LoggingMiddleware(BaseMiddleware): ) -> Any: user_info = self._extract_user_info(event) - logger.info( - f"Активность пользователя - " - f"ID пользователя: {user_info['user_id']}, " - f"Имя пользователя: {user_info['username']}, " - f"Действие: {user_info['action']}" - ) + if user_info["user_id"]: + logger.info( + f"Активность пользователя - " + f"ID пользователя: {user_info['user_id']}, " + f"Имя пользователя: {user_info['username'] or 'Не указано'}, " + f"Действие: {user_info['action'] or 'Неизвестно'}" + ) + return await handler(event, data) - def _extract_user_info(self, event: TelegramObject) -> dict[str, str | None]: - user_id = None - username = None - action = None + def _extract_user_info(self, event: TelegramObject) -> UserInfo: + """Извлекает информацию о пользователе из различных типов событий. - if isinstance(event, Message): - user = event.from_user - user_id = user.id - username = user.username - action = f"Сообщение: {event.text}" - elif isinstance(event, CallbackQuery): - user = event.from_user - user_id = user.id - username = user.username - action = f"Обратный вызов: {event.data}" - elif isinstance(event, InlineQuery): - user = event.from_user - user_id = user.id - username = user.username - action = f"Inline запрос: {event.query}" + Args: + event: Событие Telegram - return {"user_id": user_id, "username": username, "action": action} + Returns: + Словарь с информацией о пользователе + """ + result: UserInfo = {"user_id": None, "username": None, "action": None} + + # Обработка общих атрибутов для разных типов событий + if hasattr(event, "from_user") and isinstance(event.from_user, User): + result["user_id"] = event.from_user.id + result["username"] = event.from_user.username + + # Определение типа действия в зависимости от типа события + if isinstance(event, Message): + result["action"] = f"Сообщение: {event.text}" + elif isinstance(event, CallbackQuery): + result["action"] = f"Обратный вызов: {event.data}" + elif isinstance(event, InlineQuery): + result["action"] = f"Inline запрос: {event.query}" + + return result diff --git a/middlewares/session.py b/middlewares/session.py index a5d34441..c6cfdd55 100644 --- a/middlewares/session.py +++ b/middlewares/session.py @@ -1,23 +1,33 @@ from collections.abc import Awaitable, Callable -from typing import Any +from typing import Any, Optional import asyncpg from aiogram import BaseMiddleware from aiogram.types import TelegramObject - from config import DATABASE_URL class SessionMiddleware(BaseMiddleware): + pool: asyncpg.Pool | None = None + async def __call__( self, handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]], event: TelegramObject, data: dict[str, Any], ) -> Any: - conn = await asyncpg.connect(DATABASE_URL) - try: + # Создаем пул соединений при первом вызове + if self.pool is None: + self.pool = await asyncpg.create_pool(DATABASE_URL, min_size=5, max_size=20) + + # Получаем соединение из пула + async with self.pool.acquire() as conn: data["session"] = conn return await handler(event, data) - finally: - await conn.close() + + @classmethod + async def close(cls) -> None: + """Закрыть пул соединений при завершении работы приложения.""" + if cls.pool is not None: + await cls.pool.close() + cls.pool = None diff --git a/middlewares/throttling.py b/middlewares/throttling.py index d1085dba..a55b2989 100644 --- a/middlewares/throttling.py +++ b/middlewares/throttling.py @@ -1,5 +1,5 @@ from collections.abc import Awaitable, Callable, MutableMapping -from typing import Any +from typing import Any, Dict, Optional, Union from aiogram import BaseMiddleware from aiogram.dispatcher.flags import get_flag @@ -10,22 +10,61 @@ from logger import logger class ThrottlingMiddleware(BaseMiddleware): + """ + Middleware для ограничения частоты запросов от пользователей. + Позволяет настраивать разные временные интервалы для разных типов запросов. + """ + def __init__( self, *, default_key: str | None = "default", default_ttl: float = 0.5, + cache_size: int = 10_000, **ttl_map: float, ) -> None: + """ + Инициализация middleware для ограничения частоты запросов. + + Args: + default_key: Ключ по умолчанию для ограничения + default_ttl: Время ограничения по умолчанию в секундах + cache_size: Максимальный размер кэша для каждого ключа + **ttl_map: Словарь с ключами и временем ограничения + """ + # Добавляем ключ по умолчанию в карту TTL, если он указан if default_key: ttl_map[default_key] = default_ttl self.default_key = default_key self.caches: dict[str, MutableMapping[int, None]] = {} + self.cache_size = cache_size + # Инициализация кэшей для каждого ключа for name, ttl in ttl_map.items(): - self.caches[name] = TTLCache(maxsize=10_000, ttl=ttl) - logger.debug("ThrottlingMiddleware initialized.") + self.caches[name] = TTLCache(maxsize=self.cache_size, ttl=ttl) + + logger.debug(f"ThrottlingMiddleware initialized with {len(self.caches)} throttling keys") + + def _should_skip_throttling(self, event: Update) -> bool: + """ + Проверяет, нужно ли пропустить ограничение для данного события. + + Args: + event: Событие Telegram + + Returns: + True, если ограничение следует пропустить + """ + # Пропускаем предварительные запросы на оплату + if event.pre_checkout_query: + return True + + # Пропускаем уведомления об успешной оплате + if event.message and event.message.successful_payment: + return True + + return False async def __call__( self, @@ -33,32 +72,38 @@ class ThrottlingMiddleware(BaseMiddleware): event: TelegramObject, data: dict[str, Any], ) -> Any: + # Проверяем, что событие является экземпляром Update if not isinstance(event, Update): - logger.debug(f"Received event of type {type(event)}, skipping throttling.") + logger.debug(f"Skipping throttling for non-Update event: {type(event).__name__}") return await handler(event, data) - if event.pre_checkout_query: - logger.debug("Pre-checkout query event, skipping throttling.") + # Проверяем, нужно ли пропустить ограничение + if self._should_skip_throttling(event): + logger.debug("Skipping throttling for special event type") return await handler(event, data) - if event.message and event.message.successful_payment: - logger.debug("Successful payment event, skipping throttling.") + # Получаем пользователя из данных события + user: User | None = data.get("event_from_user") + + if user is None: + logger.debug("No user found in event data, proceeding without throttle") return await handler(event, data) - user: User | None = data.get("event_from_user", None) + # Получаем ключ ограничения из флагов или используем ключ по умолчанию + key = get_flag(data, "throttling_key", default=self.default_key) - if user is not None: - key = get_flag(data, "throttling_key", default=self.default_key) + if not key: + logger.debug(f"No throttling key provided for user {user.id}, proceeding without throttle") + return await handler(event, data) - if key: - if user.id in self.caches[key]: - logger.warning(f"User {user.id} is being throttled with key: {key}") - return None - logger.debug( - f"User {user.id} is allowed to proceed, adding to cache with key: {key}", - ) - self.caches[key][user.id] = None - else: - logger.debug(f"No throttling key provided for user {user.id}, proceeding without throttle.") + # Проверяем, находится ли пользователь в кэше (т.е. ограничен) + if user.id in self.caches[key]: + logger.warning(f"User {user.id} is throttled with key: {key}") + return None + # Добавляем пользователя в кэш + self.caches[key][user.id] = None + logger.debug(f"User {user.id} allowed to proceed with key: {key}") + + # Продолжаем обработку события return await handler(event, data) diff --git a/middlewares/user.py b/middlewares/user.py index 6d276e75..e9a65750 100644 --- a/middlewares/user.py +++ b/middlewares/user.py @@ -5,25 +5,60 @@ from aiogram import BaseMiddleware from aiogram.types import TelegramObject, User from database import upsert_user +from logger import logger class UserMiddleware(BaseMiddleware): + """ + Middleware для обработки информации о пользователе. + Сохраняет или обновляет данные пользователя в базе данных. + """ + async def __call__( self, handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]], event: TelegramObject, data: dict[str, Any], ) -> Any: - if user := data.get("event_from_user"): - await self._process_user(user) + try: + # Получаем пользователя из данных события + if user := data.get("event_from_user"): + # Получаем сессию из контекста, если она есть + session = data.get("session") + # Обрабатываем пользователя и получаем его данные из БД + db_user = await self._process_user(user, session) + # Добавляем данные пользователя из БД в контекст + if db_user: + data["user"] = db_user + except Exception as e: + # Логируем ошибку, но не прерываем обработку события + logger.error(f"Ошибка при обработке пользователя: {e}") + + # Продолжаем обработку события в любом случае return await handler(event, data) - async def _process_user(self, user: User) -> None: - await upsert_user( + async def _process_user(self, user: User, session: Any = None) -> dict: + """ + Обрабатывает информацию о пользователе и сохраняет её в базу данных. + + Args: + user (User): Объект пользователя Telegram + session (Any, optional): Сессия базы данных, если доступна + + Returns: + dict: Словарь с информацией о пользователе из базы данных + """ + logger.debug(f"Обработка пользователя: {user.id}") + # Получаем данные пользователя из БД после вставки/обновления + user_data = await upsert_user( tg_id=user.id, username=user.username, first_name=user.first_name, last_name=user.last_name, language_code=user.language_code, is_bot=user.is_bot, + session=session, # Передаем сессию, если она есть ) + + logger.debug(f"Получены данные пользователя из БД: {user.id}") + return user_data diff --git a/servers.py b/servers.py index caa0921f..63ecc5c1 100644 --- a/servers.py +++ b/servers.py @@ -4,10 +4,10 @@ from datetime import datetime, timedelta from aiogram.types import InlineKeyboardButton from aiogram.utils.keyboard import InlineKeyboardBuilder +from config import ADMIN_ID, DATABASE_URL, PING_TIME from ping3 import ping from bot import bot -from config import ADMIN_ID, DATABASE_URL, PING_TIME from database import get_servers from keyboards.admin.servers_kb import AdminServerEditorCallback from logger import logger