From 926774424e7ce3aaaa0fc96a5c22f61896b15316 Mon Sep 17 00:00:00 2001 From: Sergey Elpashev Date: Fri, 17 Jan 2025 13:19:13 +0300 Subject: [PATCH] Feat: added notifier --- main.py | 2 +- nwxraybot/fsm/__init__.py | 1 + nwxraybot/fsm/broadcast.py | 6 ++++ nwxraybot/handlers/admin.py | 56 +++++++++++++++++++++++++++++++++-- nwxraybot/meta/router.py | 4 +-- nwxraybot/middlewares/user.py | 2 +- 6 files changed, 64 insertions(+), 7 deletions(-) create mode 100644 nwxraybot/fsm/__init__.py create mode 100644 nwxraybot/fsm/broadcast.py diff --git a/main.py b/main.py index c445e37..9c81504 100644 --- a/main.py +++ b/main.py @@ -25,5 +25,5 @@ if __name__ == "__main__": # Start bot bot = NwXrayBot(config.bot_token.get_secret_value()) - bot.include_routers(HelloHandler(), MenuHandler(), AdminHandler()) + bot.include_routers(HelloHandler(), MenuHandler(), AdminHandler(bot.bot)) loop.run_until_complete(bot.start(skip_updates=True)) diff --git a/nwxraybot/fsm/__init__.py b/nwxraybot/fsm/__init__.py new file mode 100644 index 0000000..dca75b7 --- /dev/null +++ b/nwxraybot/fsm/__init__.py @@ -0,0 +1 @@ +from nwxraybot.fsm.broadcast import BroadcastStates diff --git a/nwxraybot/fsm/broadcast.py b/nwxraybot/fsm/broadcast.py new file mode 100644 index 0000000..70a7aca --- /dev/null +++ b/nwxraybot/fsm/broadcast.py @@ -0,0 +1,6 @@ +from aiogram.fsm.state import State, StatesGroup + + +class BroadcastStates(StatesGroup): + waiting_for_message = State() + confirming_message = State() diff --git a/nwxraybot/handlers/admin.py b/nwxraybot/handlers/admin.py index 5de4838..09f603b 100644 --- a/nwxraybot/handlers/admin.py +++ b/nwxraybot/handlers/admin.py @@ -1,19 +1,26 @@ +import asyncio +import logging import re from datetime import datetime +from typing import Optional +from aiogram import Bot, F from aiogram.enums import ParseMode from aiogram.filters import Command -from aiogram.types import Message +from aiogram.fsm.context import FSMContext +from aiogram.types import (CallbackQuery, InlineKeyboardButton, + InlineKeyboardMarkup, Message) from nwxraybot import get_code +from nwxraybot.fsm import BroadcastStates from nwxraybot.meta import Handler from nwxraybot.middlewares import AdminMiddleware from nwxraybot.models import User class AdminHandler(Handler): - def __init__(self) -> None: - super().__init__() + def __init__(self, bot: Optional[Bot] = None) -> None: + super().__init__(bot) self.router.message.middleware(AdminMiddleware()) @@ -51,3 +58,46 @@ class AdminHandler(Handler): User.name == user_dict['name']) query.execute() await message.answer('Информация о пользователе обновлена.') + + @self.router.message(Command('broadcast')) + async def start_broadcast(message: Message, state: FSMContext): + await message.answer('Отправьте для рассылки') + await state.set_state(BroadcastStates.waiting_for_message) + + @self.router.message(BroadcastStates.waiting_for_message) + async def broadcast_message(message: Message, state: FSMContext): + if message.media_group_id is not None: + await state.clear() + await message.answer('Пожалуйста, не отправляйте сообщение с несколькими изображениями.') + return + await state.set_state(BroadcastStates.confirming_message) + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [InlineKeyboardButton( + text="Подтвердить", callback_data='confirm_broadcast')], + [InlineKeyboardButton( + text="Отменить", callback_data='cancel_broadcast')] + ]) + await state.update_data(message_id=message.message_id) + await message.answer("Подтвердите отправку сообщения", reply_markup=keyboard) + + @self.router.callback_query(F.data == 'confirm_broadcast', BroadcastStates.confirming_message) + async def confirm_broadcast(callback: CallbackQuery, state: FSMContext): + await callback.answer() + data = await state.get_data() + message_id = data.get('message_id') + users = User.select().where((User.telegram_id != None)) + semaphore = asyncio.Semaphore(10) + for user in users: + async with semaphore: + try: + await self.bot.copy_message(user.telegram_id, callback.from_user.id, message_id) + except Exception as e: + logging.error(f"Error while broadcasting message to user { + user.telegram_id}: {e}") + await state.clear() + + @self.router.callback_query(F.data == 'cancel_broadcast', BroadcastStates.confirming_message) + async def cancel_broadcast(callback: CallbackQuery, state: FSMContext): + await callback.answer() + await state.clear() + await callback.message.answer('Рассылка отменена.') diff --git a/nwxraybot/meta/router.py b/nwxraybot/meta/router.py index 9646d7a..3095540 100644 --- a/nwxraybot/meta/router.py +++ b/nwxraybot/meta/router.py @@ -9,8 +9,8 @@ class Handler: def __init__(self, bot: Optional[Bot] = None) -> None: if bot: - assert isinstance(bot.bot, Optional[Bot]) - self.bot = bot.bot + assert isinstance(bot, Bot) + self.bot = bot self.router = Router() def __call__(self) -> Router: diff --git a/nwxraybot/middlewares/user.py b/nwxraybot/middlewares/user.py index fefb11b..3ab15c9 100644 --- a/nwxraybot/middlewares/user.py +++ b/nwxraybot/middlewares/user.py @@ -14,7 +14,7 @@ class UserMiddleware(BaseMiddleware): async def __call__(self, handler: Callable[[Message, Dict[str, Any]], Awaitable[Any]], event: Message, data: Dict[str, Any]) -> Any: if event.chat.type != ChatType.PRIVATE: return None - if event.text.startswith('/start'): + if event.text and event.text.startswith('/start'): return await handler(event, data) user: Optional[User] = User.select().where( User.telegram_id == event.from_user.id).first()