diff --git a/bot.py b/bot.py index 7d65eb8..53ca92b 100644 --- a/bot.py +++ b/bot.py @@ -1,29 +1,65 @@ import asyncio -from aiogram import Bot, Dispatcher, executor + +import aiogram.types +import tortoise.transactions +from aiogram import Bot as AioBot, Dispatcher from aiogram.contrib.fsm_storage.memory import MemoryStorage -from settings import BotSettings +from olgram.settings import BotSettings from olgram.bot.bots import router as bots_router from olgram.bot.start import router as start_router from olgram.bot.bot import router as bot_router from olgram.utils.database import init_database +from olgram.models.models import Bot, GroupChat + +from instance.bot import BotInstance + +import typing as ty + + +async def invite_callback(identify: int, message: aiogram.types.Message): + bot = await Bot.get(id=identify) + chat, _ = await GroupChat.get_or_create(chat_id=message.chat.id, + defaults={"name": message.chat.full_name}) + if chat not in await bot.group_chats.all(): + await bot.group_chats.add(chat) + + +def run_bot(bot: BotInstance, loop: ty.Optional[asyncio.AbstractEventLoop] = None): + loop = loop or asyncio.get_event_loop() + loop.create_task(bot.start_polling()) + + +async def run_all_bots(loop: asyncio.AbstractEventLoop): + bots = await Bot.all() + for bot in bots: + run_bot(BotInstance(bot.token, + bot.super_chat_id, + bot.start_text, + invite_callback=invite_callback, + identify=bot.id), loop) + def main(): """ Classic polling """ - asyncio.get_event_loop().run_until_complete(init_database()) + loop = asyncio.get_event_loop() + loop.run_until_complete(init_database()) - bot = Bot(BotSettings.token()) + bot = AioBot(BotSettings.token()) dp = Dispatcher(bot, storage=MemoryStorage()) start_router.setup(dp) bots_router.setup(dp) bot_router.setup(dp) - executor.start_polling(dp, skip_updates=True) + loop.run_until_complete(run_all_bots(loop)) + loop.create_task(dp.start_polling()) + + loop.run_forever() if __name__ == '__main__': diff --git a/instance/bot.py b/instance/bot.py index 5ad70c8..32db39e 100644 --- a/instance/bot.py +++ b/instance/bot.py @@ -1,77 +1,100 @@ import asyncio -import aioredis import typing as ty -from aiogram import Bot, Dispatcher, executor, types, exceptions +import aiogram +import aioredis +from aiogram import Dispatcher, types, exceptions from aiogram.contrib.fsm_storage.memory import MemoryStorage -from olgram.utils.router import Router -token = "(token)" -bot_id = token.split(":")[0] -start_text = 'Здравствуйте! Напишите тут что-то' -super_chat_id = -1 # ID чата здесь +class BotInstance: + def __init__(self, token: str, super_chat_id: int, start_text: str, + invite_callback: ty.Optional[ty.Callable] = None, + left_callback: ty.Optional[ty.Callable] = None, + identify: ty.Optional[int] = None): + self._token = token + self._bot_id = self._token.split(":")[0] + self._super_chat_id = super_chat_id + self._start_text = start_text + self._redis: aioredis.Redis = None + self._dp: aiogram.Dispatcher = None + self._identify = identify -router = Router() -redis: ty.Optional[aioredis.Redis] = None + self._invite_callback = invite_callback + self._left_callback = left_callback + def stop_polling(self): + self._dp.stop_polling() -def message_unique_id(message_id) -> str: - return bot_id + "-" + str(message_id) + async def start_polling(self): + self._redis = await aioredis.create_redis_pool('redis://localhost:6370') + bot = aiogram.Bot(self._token) + self._dp = Dispatcher(bot, storage=MemoryStorage()) -@router.message_handler(content_types=[types.ContentType.ANY]) -async def receive_text(message: types.Message): - """ - Some text received - :param message: - :return: - """ - if message.text and message.text.startswith("/start"): - await message.answer(start_text) - return + self._dp.register_message_handler(self._receive_text, content_types=[types.ContentType.TEXT]) + self._dp.register_message_handler(self._receive_invite, content_types=[types.ContentType.NEW_CHAT_MEMBERS]) + self._dp.register_message_handler(self._receive_left, content_types=[types.ContentType.LEFT_CHAT_MEMBER]) - if message.chat.id != super_chat_id: - # Это обычный чат - new_message = await message.forward(super_chat_id) - await redis.set(message_unique_id(new_message.message_id), message.chat.id) - else: - # Это чат, в который бот должен пересылать сообщения - if message.reply_to_message: - chat_id = await redis.get(message_unique_id(message.reply_to_message.message_id)) - if not chat_id: - chat_id = message.reply_to_message.forward_from_chat - if not chat_id: - await message.reply("Невозможно ответить, автор сообщения не найден") - return - chat_id = int(chat_id) - try: - await message.copy_to(chat_id) - except exceptions.MessageError: - await message.reply("Невозможно отправить сообщение пользователю: возможно, он заблокировал бота") - return + await self._dp.start_polling() + def _message_unique_id(self, message_id) -> str: + return self._bot_id + "-" + str(message_id) + + async def _receive_invite(self, message: types.Message): + if not self._invite_callback: + return + + for member in message.new_chat_members: + if member.id == message.bot.id: + await self._invite_callback(self._identify, message) + + async def _receive_left(self, message: types.Message): + if not self._left_callback: + return + + if message.left_chat_member.id == message.bot.id: + await self._left_callback(self._identify, message) + + async def _receive_text(self, message: types.Message): + """ + Some text received + :param message: + :return: + """ + if message.text and message.text.startswith("/start"): + await message.answer(self._start_text) + return + + if message.chat.id != self._super_chat_id: + # Это обычный чат + new_message = await message.forward(self._super_chat_id) + await self._redis.set(self._message_unique_id(new_message.message_id), message.chat.id) else: - await message.forward(super_chat_id) + # Это чат, в который бот должен пересылать сообщения + if message.reply_to_message: + chat_id = await self._redis.get(self._message_unique_id(message.reply_to_message.message_id)) + if not chat_id: + chat_id = message.reply_to_message.forward_from_chat + if not chat_id: + await message.reply("Невозможно ответить, автор сообщения не найден") + return + chat_id = int(chat_id) + try: + await message.copy_to(chat_id) + except exceptions.MessageError: + await message.reply("Невозможно отправить сообщение пользователю: возможно, он заблокировал бота") + return - -async def init_redis(): - global redis - redis = await aioredis.create_redis_pool('redis://localhost:6370') - - -def main(): - """ - Classic polling - """ - - asyncio.get_event_loop().run_until_complete(init_redis()) - - bot = Bot(token) - dp = Dispatcher(bot, storage=MemoryStorage()) - router.setup(dp) - - executor.start_polling(dp, skip_updates=True) + else: + await message.forward(self._super_chat_id) if __name__ == '__main__': - main() + # Single instance mode + import os + bot = BotInstance( + os.getenv("TOKEN"), + int(os.getenv("CHAT_ID")), + os.getenv("START_TEXT") + ) + asyncio.get_event_loop().run_until_complete(bot.start_polling()) diff --git a/olgram/bot/bot.py b/olgram/bot/bot.py index 5cad383..c4f71e8 100644 --- a/olgram/bot/bot.py +++ b/olgram/bot/bot.py @@ -12,10 +12,12 @@ router = Router() # Пользователь выбрал бота select_bot = CallbackData('bot_select', 'bot_id') # Пользователь выбрал, что хочет сделать со своим ботом -bot_operation = CallbackData('bot_operation', 'operation') +bot_operation = CallbackData('bot_operation', 'bot_id', 'operation') +# Пользователь выбрал чат +select_bot_chat = CallbackData('chat_select', 'bot_id', 'chat_id') -@router.callback_query_handler(select_bot.filter()) +@router.callback_query_handler(select_bot.filter(), state="*") async def select_bot_callback(call: types.CallbackQuery, callback_data: dict, state: FSMContext): """ Пользователь выбрал бота для редактирования @@ -23,20 +25,20 @@ async def select_bot_callback(call: types.CallbackQuery, callback_data: dict, st bot_id = callback_data["bot_id"] bot = await Bot.get_or_none(id=bot_id) if not bot or (await bot.owner).telegram_id != call.from_user.id: - await call.answer("Такого бота нет", show_alert=True) + await call.answer("Такого бота нет, либо он принадлежит не вам", show_alert=True) return - async with state.proxy() as proxy: - proxy["bot"] = bot - await try_delete_message(call.message) keyboard = types.InlineKeyboardMarkup(row_width=2) - keyboard.insert(types.InlineKeyboardButton(text="Текст", callback_data=bot_operation.new(operation="text"))) - keyboard.insert(types.InlineKeyboardButton(text="Чат", callback_data=bot_operation.new(operation="chat"))) - keyboard.insert(types.InlineKeyboardButton(text="Удалить бот", callback_data=bot_operation.new(operation="delete"))) + keyboard.insert(types.InlineKeyboardButton(text="Текст", + callback_data=bot_operation.new(bot_id=bot_id, operation="text"))) + keyboard.insert(types.InlineKeyboardButton(text="Чат", + callback_data=bot_operation.new(bot_id=bot_id, operation="chat"))) + keyboard.insert(types.InlineKeyboardButton(text="Удалить бот", + callback_data=bot_operation.new(bot_id=bot_id, operation="delete"))) keyboard.insert(types.InlineKeyboardButton(text="<<Вернуться к списку ботов", - callback_data=bot_operation.new(operation="back"))) + callback_data=bot_operation.new(bot_id=bot_id, operation="back"))) await AioBot.get_current().send_message(call.message.chat.id, dedent(f""" Управление ботом @{bot.name}. @@ -44,3 +46,45 @@ async def select_bot_callback(call: types.CallbackQuery, callback_data: dict, st Если у вас возникли вопросы по настройке бота, то посмотрите нашу справку /help. """), reply_markup=keyboard) + +@router.callback_query_handler(bot_operation.filter(operation="delete"), state="*") +async def delete_bot_callback(call: types.CallbackQuery, callback_data: dict, state: FSMContext): + bot_id = callback_data["bot_id"] + bot = await Bot.get_or_none(id=bot_id) + if not bot or (await bot.owner).telegram_id != call.from_user.id: + await call.answer("Такого бота нет, либо он принадлежит не вам", show_alert=True) + return + + await bot.delete() + await call.answer("Бот удалён") + await try_delete_message(call.message) + + +@router.callback_query_handler(bot_operation.filter(operation="chat"), state="*") +async def chats_bot_callback(call: types.CallbackQuery, callback_data: dict, state: FSMContext): + bot_id = callback_data["bot_id"] + bot = await Bot.get_or_none(id=bot_id) + if not bot or (await bot.owner).telegram_id != call.from_user.id: + await call.answer("Такого бота нет, либо он принадлежит не вам", show_alert=True) + return + + await try_delete_message(call.message) + + keyboard = types.InlineKeyboardMarkup(row_width=2) + + chats = await bot.group_chats.all() + + if not chats: + return await AioBot.get_current().send_message(call.message.chat.id, dedent(f""" + Этот бот не добавлен в чаты, поэтому все сообщения будут приходить вам в бот. + Чтобы подключить чат — просто добавьте бот @{bot.name} в чат. + """), reply_markup=keyboard) + + for chat in chats: + keyboard.insert(types.InlineKeyboardButton(text=chat.name, + callback_data=select_bot_chat.new(bot_id=bot_id, chat_id=chat.id))) + + await AioBot.get_current().send_message(call.message.chat.id, dedent(f""" + В этом разделе вы можете привязать бота @{bot.name} к чату. + Выберите чат, куда бот будет пересылать сообщения. + """), reply_markup=keyboard) diff --git a/olgram/bot/bots.py b/olgram/bot/bots.py index 2d70dcb..9f2c6eb 100644 --- a/olgram/bot/bots.py +++ b/olgram/bot/bots.py @@ -1,6 +1,7 @@ from aiogram import types, Bot as AioBot from aiogram.dispatcher import FSMContext from aiogram.utils.exceptions import Unauthorized, TelegramAPIError +from tortoise.exceptions import IntegrityError import re from textwrap import dedent @@ -40,9 +41,9 @@ async def add_bot(message: types.Message, state: FSMContext): """ Команда /addbot (добавить бота) """ - bot_count = await Bot.filter(user__telegram_id=message.from_user.id).count() - if bot_count > OlgramSettings.max_bots_per_user(): - await message.answer("У вас уже слишком много ботов") + bot_count = await Bot.filter(owner__telegram_id=message.from_user.id).count() + if bot_count >= OlgramSettings.max_bots_per_user(): + await message.answer("У вас уже слишком много ботов.") return await message.answer(dedent(""" @@ -83,6 +84,11 @@ async def bot_added(message: types.Message, state: FSMContext): Не удалось запустить этого бота: непредвиденная ошибка """)) + async def on_duplication_bot(): + await message.answer(dedent(""" + Такой бот уже есть в базе данных + """)) + if not token: return await on_invalid_token() @@ -100,7 +106,11 @@ async def bot_added(message: types.Message, state: FSMContext): return await on_unknown_error() user, _ = await User.get_or_create(telegram_id=message.from_user.id) - bot = Bot(token=token, owner=user, name=test_bot_info.username) - await bot.save() + bot = Bot(token=token, owner=user, name=test_bot_info.username, super_chat_id=message.from_user.id) + try: + await bot.save() + except IntegrityError: + return await on_duplication_bot() await message.answer("Бот добавлен!") + await state.reset_state() diff --git a/olgram/models/models.py b/olgram/models/models.py index dee8349..935710a 100644 --- a/olgram/models/models.py +++ b/olgram/models/models.py @@ -1,12 +1,21 @@ from tortoise.models import Model from tortoise import fields +from textwrap import dedent + class Bot(Model): id = fields.IntField(pk=True) token = fields.CharField(max_length=50, unique=True) owner = fields.ForeignKeyField("models.User", related_name="bots") name = fields.CharField(max_length=33) + start_text = fields.TextField(default=dedent(""" + Здравствуйте! + Напишите ваш вопрос и мы ответим Вам в ближайшее время. + """)) + + super_chat_id = fields.IntField() + group_chats = fields.ManyToManyField("models.GroupChat", related_name="bots", on_delete=fields.relational.SET_NULL) class Meta: table = 'bot' @@ -18,3 +27,12 @@ class User(Model): class Meta: table = 'user' + + +class GroupChat(Model): + id = fields.IntField(pk=True) + chat_id = fields.IntField(index=True, unique=True) + name = fields.CharField(max_length=50) + + class Meta: + table = 'group_chat' diff --git a/olgram/utils/database.py b/olgram/utils/database.py index 4c81d85..6b24a69 100644 --- a/olgram/utils/database.py +++ b/olgram/utils/database.py @@ -1,5 +1,5 @@ from tortoise import Tortoise -from settings import DatabaseSettings +from olgram.settings import DatabaseSettings async def init_database():