Instance в режиме Polling

This commit is contained in:
mihalin 2021-07-01 19:51:59 +03:00
parent d738da2fa9
commit 60bb00bcc9
6 changed files with 212 additions and 81 deletions

46
bot.py
View File

@ -1,29 +1,65 @@
import asyncio 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 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.bots import router as bots_router
from olgram.bot.start import router as start_router from olgram.bot.start import router as start_router
from olgram.bot.bot import router as bot_router from olgram.bot.bot import router as bot_router
from olgram.utils.database import init_database 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(): def main():
""" """
Classic polling 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()) dp = Dispatcher(bot, storage=MemoryStorage())
start_router.setup(dp) start_router.setup(dp)
bots_router.setup(dp) bots_router.setup(dp)
bot_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__': if __name__ == '__main__':

View File

@ -1,43 +1,78 @@
import asyncio import asyncio
import aioredis
import typing as ty 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 aiogram.contrib.fsm_storage.memory import MemoryStorage
from olgram.utils.router import Router
token = "(token)" class BotInstance:
bot_id = token.split(":")[0] def __init__(self, token: str, super_chat_id: int, start_text: str,
start_text = 'Здравствуйте! Напишите тут что-то' invite_callback: ty.Optional[ty.Callable] = None,
super_chat_id = -1 # ID чата здесь 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() self._invite_callback = invite_callback
redis: ty.Optional[aioredis.Redis] = None self._left_callback = left_callback
def stop_polling(self):
self._dp.stop_polling()
def message_unique_id(message_id) -> str: async def start_polling(self):
return bot_id + "-" + str(message_id) 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]) self._dp.register_message_handler(self._receive_text, content_types=[types.ContentType.TEXT])
async def receive_text(message: types.Message): 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])
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 Some text received
:param message: :param message:
:return: :return:
""" """
if message.text and message.text.startswith("/start"): if message.text and message.text.startswith("/start"):
await message.answer(start_text) await message.answer(self._start_text)
return return
if message.chat.id != super_chat_id: if message.chat.id != self._super_chat_id:
# Это обычный чат # Это обычный чат
new_message = await message.forward(super_chat_id) new_message = await message.forward(self._super_chat_id)
await redis.set(message_unique_id(new_message.message_id), message.chat.id) await self._redis.set(self._message_unique_id(new_message.message_id), message.chat.id)
else: else:
# Это чат, в который бот должен пересылать сообщения # Это чат, в который бот должен пересылать сообщения
if message.reply_to_message: if message.reply_to_message:
chat_id = await redis.get(message_unique_id(message.reply_to_message.message_id)) chat_id = await self._redis.get(self._message_unique_id(message.reply_to_message.message_id))
if not chat_id: if not chat_id:
chat_id = message.reply_to_message.forward_from_chat chat_id = message.reply_to_message.forward_from_chat
if not chat_id: if not chat_id:
@ -51,27 +86,15 @@ async def receive_text(message: types.Message):
return return
else: else:
await message.forward(super_chat_id) await message.forward(self._super_chat_id)
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)
if __name__ == '__main__': 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())

View File

@ -12,10 +12,12 @@ router = Router()
# Пользователь выбрал бота # Пользователь выбрал бота
select_bot = CallbackData('bot_select', 'bot_id') 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): 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_id = callback_data["bot_id"]
bot = await Bot.get_or_none(id=bot_id) bot = await Bot.get_or_none(id=bot_id)
if not bot or (await bot.owner).telegram_id != call.from_user.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 return
async with state.proxy() as proxy:
proxy["bot"] = bot
await try_delete_message(call.message) await try_delete_message(call.message)
keyboard = types.InlineKeyboardMarkup(row_width=2) keyboard = types.InlineKeyboardMarkup(row_width=2)
keyboard.insert(types.InlineKeyboardButton(text="Текст", callback_data=bot_operation.new(operation="text"))) keyboard.insert(types.InlineKeyboardButton(text="Текст",
keyboard.insert(types.InlineKeyboardButton(text="Чат", callback_data=bot_operation.new(operation="chat"))) callback_data=bot_operation.new(bot_id=bot_id, operation="text")))
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="chat")))
keyboard.insert(types.InlineKeyboardButton(text="Удалить бот",
callback_data=bot_operation.new(bot_id=bot_id, operation="delete")))
keyboard.insert(types.InlineKeyboardButton(text="<<Вернуться к списку ботов", 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""" await AioBot.get_current().send_message(call.message.chat.id, dedent(f"""
Управление ботом @{bot.name}. Управление ботом @{bot.name}.
@ -44,3 +46,45 @@ async def select_bot_callback(call: types.CallbackQuery, callback_data: dict, st
Если у вас возникли вопросы по настройке бота, то посмотрите нашу справку /help. Если у вас возникли вопросы по настройке бота, то посмотрите нашу справку /help.
"""), reply_markup=keyboard) """), 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)

View File

@ -1,6 +1,7 @@
from aiogram import types, Bot as AioBot from aiogram import types, Bot as AioBot
from aiogram.dispatcher import FSMContext from aiogram.dispatcher import FSMContext
from aiogram.utils.exceptions import Unauthorized, TelegramAPIError from aiogram.utils.exceptions import Unauthorized, TelegramAPIError
from tortoise.exceptions import IntegrityError
import re import re
from textwrap import dedent from textwrap import dedent
@ -40,9 +41,9 @@ async def add_bot(message: types.Message, state: FSMContext):
""" """
Команда /addbot (добавить бота) Команда /addbot (добавить бота)
""" """
bot_count = await Bot.filter(user__telegram_id=message.from_user.id).count() bot_count = await Bot.filter(owner__telegram_id=message.from_user.id).count()
if bot_count > OlgramSettings.max_bots_per_user(): if bot_count >= OlgramSettings.max_bots_per_user():
await message.answer("У вас уже слишком много ботов") await message.answer("У вас уже слишком много ботов.")
return return
await message.answer(dedent(""" 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: if not token:
return await on_invalid_token() return await on_invalid_token()
@ -100,7 +106,11 @@ async def bot_added(message: types.Message, state: FSMContext):
return await on_unknown_error() return await on_unknown_error()
user, _ = await User.get_or_create(telegram_id=message.from_user.id) user, _ = await User.get_or_create(telegram_id=message.from_user.id)
bot = Bot(token=token, owner=user, name=test_bot_info.username) bot = Bot(token=token, owner=user, name=test_bot_info.username, super_chat_id=message.from_user.id)
try:
await bot.save() await bot.save()
except IntegrityError:
return await on_duplication_bot()
await message.answer("Бот добавлен!") await message.answer("Бот добавлен!")
await state.reset_state()

View File

@ -1,12 +1,21 @@
from tortoise.models import Model from tortoise.models import Model
from tortoise import fields from tortoise import fields
from textwrap import dedent
class Bot(Model): class Bot(Model):
id = fields.IntField(pk=True) id = fields.IntField(pk=True)
token = fields.CharField(max_length=50, unique=True) token = fields.CharField(max_length=50, unique=True)
owner = fields.ForeignKeyField("models.User", related_name="bots") owner = fields.ForeignKeyField("models.User", related_name="bots")
name = fields.CharField(max_length=33) 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: class Meta:
table = 'bot' table = 'bot'
@ -18,3 +27,12 @@ class User(Model):
class Meta: class Meta:
table = 'user' 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'

View File

@ -1,5 +1,5 @@
from tortoise import Tortoise from tortoise import Tortoise
from settings import DatabaseSettings from olgram.settings import DatabaseSettings
async def init_database(): async def init_database():