olgram/server/custom.py

240 lines
11 KiB
Python
Raw Normal View History

2021-09-10 21:32:06 +03:00
from aiogram import Bot as AioBot, Dispatcher
from aiogram.dispatcher.webhook import WebhookRequestHandler
from aiogram.dispatcher.webhook import SendMessage
2021-09-10 22:38:29 +03:00
from aiogram import exceptions
2021-09-10 21:32:06 +03:00
from aiogram import types
2021-09-10 22:07:02 +03:00
from contextvars import ContextVar
from aiohttp.web_exceptions import HTTPNotFound
2021-09-10 22:42:47 +03:00
from aioredis.commands import create_redis_pool
from aioredis import Redis
2022-02-18 21:47:40 +03:00
from tortoise.expressions import F
2021-09-17 11:07:55 +03:00
import logging
2021-09-10 22:38:29 +03:00
import typing as ty
from olgram.settings import ServerSettings
2022-01-18 23:28:03 +03:00
from olgram.models.models import Bot, GroupChat, BannedUser
2022-02-11 02:02:28 +03:00
from server.inlines import inline_handler
2022-03-13 17:04:15 +03:00
import socket
2021-09-17 11:07:55 +03:00
_logger = logging.getLogger(__name__)
2022-01-19 15:48:50 +03:00
_logger.setLevel(logging.INFO)
2021-09-17 11:07:55 +03:00
2021-09-10 22:07:02 +03:00
db_bot_instance: ContextVar[Bot] = ContextVar('db_bot_instance')
2021-09-10 22:42:47 +03:00
_redis: ty.Optional[Redis] = None
2021-09-10 22:38:29 +03:00
async def init_redis():
global _redis
2021-09-10 22:42:47 +03:00
_redis = await create_redis_pool(ServerSettings.redis_path())
2021-09-10 22:38:29 +03:00
def _message_unique_id(bot_id: int, message_id: int) -> str:
return f"{bot_id}_{message_id}"
2021-09-10 21:32:06 +03:00
2022-02-16 19:56:03 +03:00
def _thread_uniqie_id(bot_id: int, chat_id: int) -> str:
return f"thread_{bot_id}_{chat_id}"
2021-12-22 23:46:09 +03:00
async def message_handler(message: types.Message, *args, **kwargs):
2021-09-17 11:08:22 +03:00
_logger.info("message handler")
2021-09-10 22:38:29 +03:00
bot = db_bot_instance.get()
2022-01-18 23:28:03 +03:00
if message.text and message.text == "/start":
2021-09-10 21:32:06 +03:00
# На команду start нужно ответить, не пересылая сообщение никуда
2021-09-10 22:07:02 +03:00
return SendMessage(chat_id=message.chat.id,
2021-09-16 03:20:31 +03:00
text=bot.start_text + ServerSettings.append_text())
2022-03-13 17:04:15 +03:00
if message.text and message.text == "/gethostplz":
return SendMessage(chat_id=message.chat.id, text=str(socket.gethostname()))
2021-09-10 22:38:29 +03:00
super_chat_id = await bot.super_chat_id()
2022-02-16 19:56:03 +03:00
is_super_group = super_chat_id < 0
2021-09-10 22:38:29 +03:00
if message.chat.id != super_chat_id:
2022-01-18 23:28:03 +03:00
# Это обычный чат
# Проверить, не забанен ли пользователь
banned = await bot.banned_users.filter(telegram_id=message.chat.id)
if banned:
return SendMessage(chat_id=message.chat.id,
text="Вы заблокированы в этом боте")
2022-02-16 20:52:08 +03:00
# Пересылаем сообщение в супер-чат
2022-02-19 20:40:56 +03:00
if is_super_group and bot.enable_threads:
2022-02-16 19:56:03 +03:00
thread_first_message = await _redis.get(_thread_uniqie_id(bot.pk, message.chat.id))
if thread_first_message:
# переслать в супер-чат, отвечая на предыдущее сообщение
2022-02-19 02:41:59 +03:00
try:
new_message = await message.copy_to(super_chat_id, reply_to_message_id=int(thread_first_message))
except exceptions.BadRequest:
new_message = await message.forward(super_chat_id)
await _redis.set(_thread_uniqie_id(bot.pk, message.chat.id), new_message.message_id,
pexpire=ServerSettings.thread_timeout_ms())
2022-02-16 19:56:03 +03:00
else:
# переслать супер-чат
new_message = await message.forward(super_chat_id)
await _redis.set(_thread_uniqie_id(bot.pk, message.chat.id), new_message.message_id,
pexpire=ServerSettings.thread_timeout_ms())
2022-02-16 20:52:08 +03:00
else: # личные сообщения не поддерживают потоки сообщений: простой forward
2022-02-16 19:56:03 +03:00
new_message = await message.forward(super_chat_id)
2022-02-16 18:45:17 +03:00
await _redis.set(_message_unique_id(bot.pk, new_message.message_id), message.chat.id,
pexpire=ServerSettings.redis_timeout_ms())
bot.incoming_messages_count = F("incoming_messages_count") + 1
await bot.save(update_fields=["incoming_messages_count"])
# И отправить пользователю специальный текст, если он указан
if bot.second_text:
return SendMessage(chat_id=message.chat.id, text=bot.second_text)
2021-09-10 22:38:29 +03:00
else:
# Это супер-чат
2021-12-22 23:46:09 +03:00
2021-09-10 22:38:29 +03:00
if message.reply_to_message:
# В супер-чате кто-то ответил на сообщение пользователя, нужно переслать тому пользователю
2021-09-10 22:38:29 +03:00
chat_id = await _redis.get(_message_unique_id(bot.pk, message.reply_to_message.message_id))
if not chat_id:
chat_id = message.reply_to_message.forward_from_chat
if not chat_id:
2021-09-10 23:09:00 +03:00
return SendMessage(chat_id=message.chat.id,
2022-02-16 18:45:17 +03:00
text="<i>Невозможно переслать сообщение: автор не найден "
"(сообщение слишком старое?)</i>",
2021-09-11 15:58:51 +03:00
parse_mode="HTML")
2021-09-10 22:38:29 +03:00
chat_id = int(chat_id)
2022-01-18 23:28:03 +03:00
if message.text == "/ban":
2022-01-19 15:56:39 +03:00
user, _ = await BannedUser.get_or_create(telegram_id=chat_id, bot=bot)
2022-01-18 23:28:03 +03:00
await user.save()
return SendMessage(chat_id=message.chat.id, text="Пользователь заблокирован")
if message.text == "/unban":
banned_user = await bot.banned_users.filter(telegram_id=chat_id).first()
if not banned_user:
return SendMessage(chat_id=message.chat.id, text="Пользователь не был забанен")
else:
await banned_user.delete()
return SendMessage(chat_id=message.chat.id, text="Пользователь разбанен")
2021-09-10 22:38:29 +03:00
try:
await message.copy_to(chat_id)
2022-01-23 00:25:49 +03:00
except (exceptions.MessageError, exceptions.Unauthorized):
2021-09-11 15:58:51 +03:00
await message.reply("<i>Невозможно переслать сообщение (автор заблокировал бота?)</i>",
parse_mode="HTML")
2021-09-10 22:38:29 +03:00
return
2022-02-18 21:47:40 +03:00
bot.outgoing_messages_count = F("outgoing_messages_count") + 1
await bot.save(update_fields=["outgoing_messages_count"])
2022-02-11 04:42:19 +03:00
elif super_chat_id > 0:
# в супер-чате кто-то пишет сообщение сам себе, только для личных сообщений
2021-09-10 22:38:29 +03:00
await message.forward(super_chat_id)
# И отправить пользователю специальный текст, если он указан
if bot.second_text:
return SendMessage(chat_id=message.chat.id, text=bot.second_text)
2021-09-10 21:32:06 +03:00
2021-09-10 23:02:40 +03:00
async def receive_invite(message: types.Message):
bot = db_bot_instance.get()
for member in message.new_chat_members:
if member.id == message.bot.id:
chat, _ = await GroupChat.get_or_create(chat_id=message.chat.id,
defaults={"name": message.chat.full_name})
chat.name = message.chat.full_name
2021-09-17 11:49:48 +03:00
await chat.save()
2021-09-10 23:02:40 +03:00
if chat not in await bot.group_chats.all():
await bot.group_chats.add(chat)
await bot.save()
break
async def receive_group_create(message: types.Message):
bot = db_bot_instance.get()
chat, _ = await GroupChat.get_or_create(chat_id=message.chat.id,
defaults={"name": message.chat.full_name})
chat.name = message.chat.full_name
await chat.save()
if chat not in await bot.group_chats.all():
await bot.group_chats.add(chat)
await bot.save()
2021-09-10 23:02:40 +03:00
async def receive_left(message: types.Message):
bot = db_bot_instance.get()
if message.left_chat_member.id == message.bot.id:
chat = await bot.group_chats.filter(chat_id=message.chat.id).first()
if chat:
await bot.group_chats.remove(chat)
2021-09-17 11:16:18 +03:00
bot_group_chat = await bot.group_chat
if bot_group_chat == chat:
2021-09-10 23:02:40 +03:00
bot.group_chat = None
await bot.save()
2021-09-10 23:02:40 +03:00
2022-02-11 02:02:28 +03:00
async def receive_inline(inline_query):
2022-02-11 02:09:09 +03:00
_logger.info("inline handler")
2022-02-11 02:02:28 +03:00
bot = db_bot_instance.get()
return await inline_handler(inline_query, bot)
2021-12-14 23:55:19 +03:00
async def receive_migrate(message: types.Message):
bot = db_bot_instance.get()
from_id = message.chat.id
to_id = message.migrate_to_chat_id
chats = await bot.group_chats.filter(chat_id=from_id)
for chat in chats:
chat.chat_id = to_id
await chat.save(update_fields=["chat_id"])
2021-09-10 21:32:06 +03:00
class CustomRequestHandler(WebhookRequestHandler):
def __init__(self, *args, **kwargs):
self._dispatcher = None
super(CustomRequestHandler, self).__init__(*args, **kwargs)
async def _create_dispatcher(self):
key = self.request.url.path[1:]
bot = await Bot.filter(code=key).first()
if not bot:
return None
2021-09-10 22:07:02 +03:00
db_bot_instance.set(bot)
2021-09-26 20:36:05 +03:00
dp = Dispatcher(AioBot(bot.decrypted_token()))
2021-09-10 21:32:06 +03:00
dp.register_message_handler(message_handler, content_types=[types.ContentType.TEXT,
types.ContentType.CONTACT,
types.ContentType.ANIMATION,
types.ContentType.AUDIO,
types.ContentType.DOCUMENT,
types.ContentType.PHOTO,
types.ContentType.STICKER,
types.ContentType.VIDEO,
types.ContentType.VOICE])
2021-09-10 23:02:40 +03:00
dp.register_message_handler(receive_invite, content_types=[types.ContentType.NEW_CHAT_MEMBERS])
dp.register_message_handler(receive_left, content_types=[types.ContentType.LEFT_CHAT_MEMBER])
2021-12-14 23:55:19 +03:00
dp.register_message_handler(receive_migrate, content_types=[types.ContentType.MIGRATE_TO_CHAT_ID])
dp.register_message_handler(receive_group_create, content_types=[types.ContentType.GROUP_CHAT_CREATED])
2022-02-11 02:02:28 +03:00
dp.register_inline_handler(receive_inline)
2021-09-10 21:32:06 +03:00
return dp
async def post(self):
2021-09-10 21:57:17 +03:00
dispatcher = await self._create_dispatcher()
2021-09-10 22:07:02 +03:00
if not dispatcher:
raise HTTPNotFound()
2021-09-10 21:57:17 +03:00
Dispatcher.set_current(dispatcher)
AioBot.set_current(dispatcher.bot)
return await super(CustomRequestHandler, self).post()
2021-09-10 21:32:06 +03:00
def get_dispatcher(self):
"""
Get Dispatcher instance from environment
:return: :class:`aiogram.Dispatcher`
"""
2021-09-10 21:57:17 +03:00
return Dispatcher.get_current()