diff --git a/.gitignore b/.gitignore index d80a558..99b50ea 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ config.json TBot/ logs/ -__pycache__/ \ No newline at end of file +__pycache__/ +.gitignore \ No newline at end of file diff --git a/databases/db_config.py b/databases/db_config.py index e7dd20c..baa3d55 100644 --- a/databases/db_config.py +++ b/databases/db_config.py @@ -1,32 +1,23 @@ +import os from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession from sqlalchemy.orm import sessionmaker from motor.motor_asyncio import AsyncIOMotorClient from databases.model import Base -from utils.LogCon import setup_logger, load_config -# Загружаем конфигурацию -config = load_config() - -# Настройки PostgreSQL -postgres_user = config['postgreSQL']['username'] -postgres_password = config['postgreSQL']['password_DB'] -postgres_host = "postgres" # Хост для PostgreSQL в Docker -POSTGRES_DSN = f"postgresql+asyncpg://{postgres_user}:{postgres_password}@{postgres_host}:5432/bot_db" +# Настройки PostgreSQL из переменных окружения +POSTGRES_DSN = os.getenv("POSTGRES_URL") # Создание движка для PostgreSQL postgres_engine = create_async_engine(POSTGRES_DSN, echo=False) AsyncSessionLocal = sessionmaker(bind=postgres_engine, class_=AsyncSession, expire_on_commit=False) -# Настройки MongoDB -mongodb_user = config['mongodb']['mongodb_username'] -mongodb_password = config['mongodb']['mongodb_password'] -mongodb_host = "mongodb" # Хост для MongoDB в Docker -mongodb_uri = f"mongodb://{mongodb_user}:{mongodb_password}@{mongodb_host}:27017" -database_name = config['mongodb']['database_name'] +# Настройки MongoDB из переменных окружения +MONGO_URI = os.getenv("MONGO_URL") +DATABASE_NAME = os.getenv("DB_NAME") # Создание клиента MongoDB -mongo_client = AsyncIOMotorClient(mongodb_uri) -mongo_db = mongo_client[database_name] +mongo_client = AsyncIOMotorClient(MONGO_URI) +mongo_db = mongo_client[DATABASE_NAME] # Инициализация PostgreSQL async def init_postgresql(): diff --git a/databases/mongodb.py b/databases/mongodb.py index e87580a..a95400f 100644 --- a/databases/mongodb.py +++ b/databases/mongodb.py @@ -1,53 +1,103 @@ -from pymongo import MongoClient -from utils.LogCon import setup_logger, load_config +import os +from motor.motor_asyncio import AsyncIOMotorClient +import logging + class MongoDBRepository: - def __init__(self, config_path="config.json"): - self.config = load_config() - self.client = MongoClient(config["mongodb_uri"]) - self.db = self.client[config["database_name"]] - self.collection = self.db["vpn_servers"] + def __init__(self): + # Настройки MongoDB из переменных окружения + mongo_uri = os.getenv("MONGO_URL") + database_name = os.getenv("DB_NAME") + server_collection = os.getenv("SERVER_COLLECTION", "servers") + plan_collection = os.getenv("PLAN_COLLECTION", "plans") - def add_server(self, server_data): - """Добавляет новый VPN сервер в коллекцию.""" - result = self.collection.insert_one(server_data) - print(f"VPN сервер добавлен с ID: {result.inserted_id}") + # Подключение к базе данных и коллекциям + self.client = AsyncIOMotorClient(mongo_uri) + self.db = self.client[database_name] + self.collection = self.db[server_collection] # Коллекция серверов + self.plans_collection = self.db[plan_collection] # Коллекция тарифных планов + self.logger = logging.getLogger(__name__) + + async def add_subscription_plan(self, plan_data): + """Добавляет новый тарифный план в коллекцию.""" + result = await self.plans_collection.insert_one(plan_data) + self.logger.debug(f"Тарифный план добавлен с ID: {result.inserted_id}") return result.inserted_id - def get_server(self, server_id): - """Получает сервер VPN по его ID.""" - server = self.collection.find_one({"_id": server_id}) - if server: - print(f"Найден VPN сервер: {server}") + async def get_subscription_plan(self, plan_id): + """Получает тарифный план по его имени.""" + plan = await self.plans_collection.find_one({"_id": plan_id}) + if plan: + self.logger.debug(f"Найден тарифный план: {plan}") else: - print(f"VPN сервер с ID {server_id} не найден.") + self.logger.error(f"Тарифный план {plan_id} не найден.") + return plan + + async def add_server(self, server_data): + """Добавляет новый VPN сервер в коллекцию.""" + result = await self.collection.insert_one(server_data) + self.logger.debug(f"VPN сервер добавлен с ID: {result.inserted_id}") + return result.inserted_id + + async def get_server(self, server_id): + """Получает сервер VPN по его ID.""" + server = await self.collection.find_one({"_id": server_id}) + if server: + self.logger.debug(f"Найден VPN сервер: {server}") + else: + self.logger.debug(f"VPN сервер с ID {server_id} не найден.") return server - def update_server(self, server_id, update_data): - """Обновляет данные VPN сервера.""" - result = self.collection.update_one({"_id": server_id}, {"$set": update_data}) - if result.matched_count > 0: - print(f"VPN сервер с ID {server_id} обновлен.") + async def get_server_with_least_clients(self): + """Возвращает сервер с наименьшим количеством подключенных клиентов.""" + pipeline = [ + { + "$addFields": { + "current_clients": {"$size": {"$ifNull": ["$clients", []]}} + } + }, + { + "$sort": {"current_clients": 1} + }, + { + "$limit": 1 + } + ] + + result = await self.collection.aggregate(pipeline).to_list(length=1) + if result: + server = result[0] + self.logger.debug(f"Найден сервер с наименьшим количеством клиентов: {server}") + return server else: - print(f"VPN сервер с ID {server_id} не найден.") + self.logger.debug("Не найдено серверов.") + return None + + async def update_server(self, server_id, update_data): + """Обновляет данные VPN сервера.""" + result = await self.collection.update_one({"_id": server_id}, {"$set": update_data}) + if result.matched_count > 0: + self.logger.debug(f"VPN сервер с ID {server_id} обновлен.") + else: + self.logger.debug(f"VPN сервер с ID {server_id} не найден.") return result.matched_count > 0 - def delete_server(self, server_id): + async def delete_server(self, server_id): """Удаляет VPN сервер по его ID.""" - result = self.collection.delete_one({"_id": server_id}) + result = await self.collection.delete_one({"_id": server_id}) if result.deleted_count > 0: - print(f"VPN сервер с ID {server_id} удален.") + self.logger.debug(f"VPN сервер с ID {server_id} удален.") else: - print(f"VPN сервер с ID {server_id} не найден.") + self.logger.debug(f"VPN сервер с ID {server_id} не найден.") return result.deleted_count > 0 - def list_servers(self): + async def list_servers(self): """Возвращает список всех VPN серверов.""" - servers = list(self.collection.find()) - print(f"Найдено {len(servers)} VPN серверов.") + servers = await self.collection.find().to_list(length=1000) # Получить до 1000 серверов (можно настроить) + self.logger.debug(f"Найдено {len(servers)} VPN серверов.") return servers - def close_connection(self): + async def close_connection(self): """Закрывает подключение к базе данных MongoDB.""" self.client.close() - print("Подключение к MongoDB закрыто.") + self.logger.debug("Подключение к MongoDB закрыто.") diff --git a/databases/postgresql.py b/databases/postgresql.py index 6b4835e..2da1641 100644 --- a/databases/postgresql.py +++ b/databases/postgresql.py @@ -2,10 +2,14 @@ from databases.model import User, Subscription, Transaction, Administrators from sqlalchemy.future import select from sqlalchemy.exc import SQLAlchemyError from sqlalchemy import desc +from dateutil.relativedelta import relativedelta +from datetime import datetime +from utils.panel import PanelInteraction +from databases.mongodb import MongoDBRepository import random import string import logging - +import asyncio class DatabaseManager: def __init__(self, session_generator): @@ -14,18 +18,19 @@ class DatabaseManager: """ self.session_generator = session_generator self.logger = logging.getLogger(__name__) + self.mongo_repo = MongoDBRepository() - async def create_user(self, telegram_id): + async def create_user(self, telegram_id: int): """ Создаёт нового пользователя, если его нет. """ async for session in self.session_generator(): try: username = self.generate_string(6) - result = await session.execute(select(User).where(User.telegram_id == telegram_id)) + result = await session.execute(select(User).where(User.telegram_id == int(telegram_id))) user = result.scalars().first() if not user: - new_user = User(telegram_id=telegram_id, username=username) + new_user = User(telegram_id=int(telegram_id), username=username) session.add(new_user) await session.commit() return new_user @@ -35,7 +40,6 @@ class DatabaseManager: await session.rollback() return "ERROR" - async def get_user_by_telegram_id(self, telegram_id: int): """ Возвращает пользователя по Telegram ID. @@ -112,6 +116,105 @@ class DatabaseManager: self.logger.error(f"Ошибка при получении транзакций пользователя {user_id}: {e}") return "ERROR" + async def buy_sub(self, telegram_id: str, plan_id: str): + async for session in self.session_generator(): + try: + result = await self.create_user(telegram_id) + if not result: + self.logger.error(f"Пользователь с Telegram ID {telegram_id} не найден.") + return "ERROR" + + # Получение тарифного плана из MongoDB + plan = await self.mongo_repo.get_subscription_plan(plan_id) + if not plan: + self.logger.error(f"Тарифный план {plan_id} не найден.") + return "ERROR" + + # Проверка достаточности средств для покупки подписки + cost = plan["cost"] + if result.balance >= cost: + result.balance -= cost + await session.commit() + + # Создание подписки для пользователя + expiry_date = datetime.now(datetime.timezone.utc) + relativedelta(months=plan["duration_months"]) + new_subscription = Subscription(user_id=result.id, vpn_server_id=None, plan=plan_id, expiry_date=expiry_date) + session.add(new_subscription) + await session.commit() + + self.logger.info(f"Подписка успешно оформлена для пользователя {telegram_id} на план {plan_id}.") + return "OK" + else: + self.logger.error(f"Недостаточно средств у пользователя {telegram_id} для покупки плана {plan_id}.") + return "INSUFFICIENT_FUNDS" + except SQLAlchemyError as e: + self.logger.error(f"Ошибка при покупке подписки {plan_id} для пользователя {telegram_id}: {e}") + await session.rollback() + return "ERROR" + + async def add_to_server(self, telegram_id: int): + """ + Метод для добавления пользователя на сервер. + """ + async for session in self.session_generator(): + try: + # Получаем подписку пользователя по telegram_id + result = await session.execute(select(Subscription).join(User).where(User.telegram_id == int(telegram_id))) + user_sub = result.scalars().first() + + if not user_sub: + self.logger.error(f"Не удалось найти подписку для пользователя с Telegram ID {telegram_id}.") + return "ERROR" + + # Получаем информацию о пользователе + user_result = await session.execute(select(User).where(User.telegram_id == telegram_id)) + user = user_result.scalars().first() + + # Получаем сервер с MongoDB + server = await self.mongo_repo.get_server(user_sub.vpn_server_id) + + if not server: + self.logger.error(f"Не удалось найти сервер с ID {user_sub.vpn_server_id}.") + return "ERROR" + + # Найдем клиента на сервере по telegram_id + client = None + for client_info in server['clients']: + if client_info['subscriptions']['tgId'] == telegram_id: + client = client_info + break + + if not client: + self.logger.error(f"Не удалось найти клиента с Telegram ID {telegram_id} на сервере.") + return "ERROR" + + # Доступ к данным сервера для добавления клиента + server_info = server['server'] + url_base = f"https://{server_info['ip']}:{server_info['port']}/{server_info['secretKey']}" + login_data = { + 'username': server_info['login'], + 'password': server_info['password'], + } + + # Инициализируем взаимодействие с панелью управления + panel = PanelInteraction(url_base, login_data, self.logger) + + # Добавляем клиента на сервер + client_id = client['id'] + expiry_date_iso = user_sub.expiry_date.isoformat() + response = await panel.add_client(client_id, expiry_date_iso, user.username) + + # Логируем результат + if response == "OK": + self.logger.info(f"Клиент {telegram_id} успешно добавлен на сервер.") + return "OK" + else: + self.logger.error(f"Ошибка при добавлении клиента {telegram_id} на сервер: {response}") + return "ERROR" + + except Exception as e: + self.logger.error(f"Ошибка при установке на сервер для пользователя {telegram_id}: {e}") + return "ERROR" @staticmethod def generate_string(length): diff --git a/handlers/handlers.py b/handlers/handlers.py index da37624..d1e5390 100644 --- a/handlers/handlers.py +++ b/handlers/handlers.py @@ -3,7 +3,7 @@ from aiogram.filters import Command from databases.postgresql import DatabaseManager from databases.model import User, Subscription, Transaction, Administrators from databases.db_config import get_postgres_session -from keyboard.keyboards import subhist_keyboard, popup_keyboard, main_keyboard,faq_keyboard,about_tarifs_keyboard, account_keyboard, buy_keyboard,balance_keyboard,guide_keyboard,tarif_Lark_keyboard,tarif_Lark_pro_keyboard,tranhist_keyboard +from keyboard.keyboards import subhist_keyboard,tarif_confirm_keyboard, popup_keyboard, main_keyboard,faq_keyboard,about_tarifs_keyboard, account_keyboard, buy_keyboard,balance_keyboard,guide_keyboard,tarif_Lark_keyboard,tarif_Lark_pro_keyboard,tranhist_keyboard # Инициализируем менеджер базы данных db_manager = DatabaseManager(get_postgres_session) @@ -219,6 +219,64 @@ async def faq_callback_handler(callback:types.CallbackQuery): reply_markup=faq_keyboard() ) +async def lark_tariff_callback_handler(callback: types.CallbackQuery): + """ + Обработчик для выбора тарифа Lark. + """ + data = callback.data.split(":") + tariff_name = data[0] + tariff_class = data[1] + tariff_time = int(data[2]) + + # Определение окончания для месяцев + if tariff_time == 1: + months = f"{tariff_time} месяц" + elif 2 <= tariff_time <= 4: + months = f"{tariff_time} месяца" + else: + months = f"{tariff_time} месяцев" + + text = f"Тариф {tariff_name} на {months}. Продолжите покупку..." + + # Рендеринг клавиатуры + keyboard = tarif_confirm_keyboard(tariff_name, tariff_time, tariff_class) + await callback.message.edit_text(text=text, reply_markup=keyboard) + +async def confirm_callback_handler(callback: types.CallbackQuery): + """ + Обработчик подтверждения покупки тарифа. + """ + tariff_info = callback.data.split(":")[1].split("_") + tariff_name = tariff_info[0] + tariff_class = tariff_info[1] + tariff_amount = int(tariff_info[2]) + + sub = await db_manager.buy_sub(callback.from_user.id, f"{tariff_name}_{tariff_class}_{tariff_amount}") + if sub == "ERROR": + await callback.message.answer( + "Произошла ошибка, попробуйте позже или свяжитесь с администрацией." + ) + await callback.answer() + return + elif sub == "INSUFFICIENT_FUNDS": + await callback.message.answer( + "Произошла ошибка, не достаточно средств на балансе." + ) + await callback.answer() + return + add_to_server = await db_manager.add_to_server(callback.from_user.id) + if add_to_server == "ERROR": + await callback.message.answer( + "Произошла ошибка, попробуйте позже или свяжитесь с администрацией." + ) + await callback.answer() + return + + # Текст подтверждения на основе тарифа + months_text = f"{tariff_amount} месяцев" if tariff_amount > 1 else f"{tariff_amount} месяц" + text = f"Вы успешно оформили тариф {tariff_name} на {months_text}. Спасибо за покупку!" + + await callback.message.edit_text(text=text) def register_handlers(dp: Dispatcher): """ @@ -236,4 +294,5 @@ def register_handlers(dp: Dispatcher): dp.callback_query.register(balance_callback_handler, lambda c: c.data == "balance") dp.callback_query.register(guide_callback_handler, lambda c: c.data == "guide") dp.callback_query.register(about_tarifs_callback_handler, lambda c: c.data == "about_tarifs") - + dp.callback_query.register(lark_tariff_callback_handler, lambda c: c.data.startswith("Lark:")) + dp.callback_query.register(confirm_callback_handler, lambda c: c.data.startswith("confirm:")) \ No newline at end of file diff --git a/keyboard/keyboards.py b/keyboard/keyboards.py index 6abce53..00d7d31 100644 --- a/keyboard/keyboards.py +++ b/keyboard/keyboards.py @@ -20,6 +20,7 @@ def account_keyboard(): builder.row(InlineKeyboardButton(text="Баланс", callback_data="balance")) builder.row(InlineKeyboardButton(text="Приобрести подписку", callback_data="buy_subscription")) builder.row(InlineKeyboardButton(text="Руководство по подключению", callback_data="guide")) + builder.row(InlineKeyboardButton(text="Назад", callback_data="base")) return builder.as_markup() @@ -66,9 +67,9 @@ def tarif_Lark_keyboard(): Тариф Lark """ builder = InlineKeyboardBuilder() - builder.row(InlineKeyboardButton(text="Тариф Lark 1 Месяц", callback_data="lark1")) - builder.row(InlineKeyboardButton(text="Тариф Lark 3 Месяц", callback_data="lark3")) - builder.row(InlineKeyboardButton(text="Тариф Lark 6 Месяц", callback_data="lark6")) + builder.row(InlineKeyboardButton(text="Тариф Lark 1 Месяц", callback_data="Lark:Standart:1")) + builder.row(InlineKeyboardButton(text="Тариф Lark 3 Месяц", callback_data="Lark:Standart:3")) + builder.row(InlineKeyboardButton(text="Тариф Lark 6 Месяц", callback_data="Lark:Standart:6")) builder.row(InlineKeyboardButton(text="Назад", callback_data="buy_subscription")) return builder.as_markup() @@ -77,9 +78,9 @@ def tarif_Lark_pro_keyboard(): Тариф Lark Pro """ builder = InlineKeyboardBuilder() - builder.row(InlineKeyboardButton(text="Тариф Lark Pro 1 Месяц", callback_data="lark1pro")) - builder.row(InlineKeyboardButton(text="Тариф Lark Pro 3 Месяц", callback_data="lark3pro")) - builder.row(InlineKeyboardButton(text="Тариф Lark Pro 6 Месяц", callback_data="lark6pro")) + builder.row(InlineKeyboardButton(text="Тариф Lark Pro 1 Месяц", callback_data="Lark:Pro:1")) + builder.row(InlineKeyboardButton(text="Тариф Lark Pro 3 Месяц", callback_data="Lark:Pro:3")) + builder.row(InlineKeyboardButton(text="Тариф Lark Pro 6 Месяц", callback_data="Lark:Pro:6")) builder.row(InlineKeyboardButton(text="Назад", callback_data="buy_subscription")) return builder.as_markup() @@ -117,4 +118,13 @@ def tranhist_keyboard(): """ builder = InlineKeyboardBuilder() builder.row(InlineKeyboardButton(text="Назад",callback_data="balance")) - return builder.as_markup() \ No newline at end of file + return builder.as_markup() + +def tarif_confirm_keyboard(name,amount,classif): + """ + Подтверждение покупки тарифа + """ + builder = InlineKeyboardBuilder() + builder.row(InlineKeyboardButton(text="Подтвердить", callback_data=f"confirm:{name}_{classif}_{amount}")) + builder.row(InlineKeyboardButton(text="Отменить",callback_data="buy_subscription")) + return builder.as_markup() diff --git a/main.py b/main.py index c9b4229..5d4b9a3 100644 --- a/main.py +++ b/main.py @@ -1,48 +1,72 @@ +import os import asyncio from aiogram import Bot, Dispatcher from databases.db_config import init_postgresql, init_mongodb, close_connections from aiogram.types import BotCommand -from utils.LogCon import setup_logger, load_config from Middleware.anti_spam_middleware import AntiSpamMiddleware import logging -setup_logger() -logger = logging.getLogger(__name__) +# Получение токена бота из переменных окружения +BOT_TOKEN = os.getenv("TOKEN") -BOT_TOKEN = load_config()['token'] +if not BOT_TOKEN: + raise ValueError("Не задан токен бота. Убедитесь, что переменная окружения 'TOKEN' установлена.") bot = Bot(token=BOT_TOKEN) dp = Dispatcher() +# Установка middleware для защиты от спама dp.message.middleware(AntiSpamMiddleware(rate_limit=1)) async def set_commands(): + """Устанавливает команды для бота.""" commands = [ BotCommand(command="/start", description="Запустить бота"), ] await bot.set_my_commands(commands) + async def on_startup(): + """Действия при запуске бота.""" + # Инициализация баз данных await init_mongodb() + await init_postgresql() + + # Установка команд бота await set_commands() - print("Бот запущен!") + + # Настройка логирования + logging.basicConfig(level=logging.INFO) + logger = logging.getLogger(__name__) + logger.info("Бот запущен!") + async def on_shutdown(): + """Действия при остановке бота.""" + # Закрытие подключений к базам данных await close_connections() + + # Закрытие сессии бота await bot.session.close() print("Бот остановлен.") + async def main(): + """Основной цикл работы бота.""" from handlers.handlers import register_handlers - register_handlers(dp) - await init_postgresql() # Убедитесь, что таблицы создаются здесь + register_handlers(dp) # Регистрация хендлеров + await on_startup() + try: + # Запуск polling await dp.start_polling(bot) finally: + # Действия при завершении работы await on_shutdown() + if __name__ == "__main__": asyncio.run(main())