From 00dbe52d8401139794cb86e30b4f5deddf8190e0 Mon Sep 17 00:00:00 2001 From: AlberLC Date: Tue, 26 Mar 2024 03:21:46 +0100 Subject: [PATCH] Add MultiBot._on_new_message_raw whitelist/blacklist --- flanabot/bots/flana_bot.py | 44 ++++++++++++++++++++------------- flanabot/bots/flana_tele_bot.py | 11 ++++++--- flanabot/bots/penalty_bot.py | 11 ++++++--- 3 files changed, 43 insertions(+), 23 deletions(-) diff --git a/flanabot/bots/flana_bot.py b/flanabot/bots/flana_bot.py index 2c93867..7f61c8d 100644 --- a/flanabot/bots/flana_bot.py +++ b/flanabot/bots/flana_bot.py @@ -12,7 +12,7 @@ import pymongo import pytz from flanaapis import InstagramLoginError, MediaNotFoundError, PlaceNotFoundError from flanautils import return_if_first_empty -from multibot import BadRoleError, MessagesFormat, MultiBot, Platform, RegisteredCallback, Role, User, bot_mentioned, constants as multibot_constants, group, inline, owner, reply +from multibot import BadRoleError, MessagesFormat, MultiBot, Platform, RegisteredCallback, Role, User, bot_mentioned, constants as multibot_constants, group, ignore_self_message, inline, owner, reply from flanabot import constants from flanabot.bots.connect_4_bot import Connect4Bot @@ -33,8 +33,8 @@ class FlanaBot(Connect4Bot, PenaltyBot, PollBot, ScraperBot, UberEatsBot, Weathe def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.tunnel_chat = None - self.help_calls = {} + self.tunnel_chat: Chat | None = None + self.help_calls: dict[int, datetime.timedelta] = {} # -------------------------------------------------------- # # ------------------- PROTECTED METHODS ------------------ # @@ -398,14 +398,6 @@ class FlanaBot(Connect4Bot, PenaltyBot, PollBot, ScraperBot, UberEatsBot, Weathe and message.author.id != self.owner_id and - ( - not message.replied_message - or - message.replied_message.author.id != self.id - or - not message.replied_message.medias - ) - and ( self.is_bot_mentioned(message) or @@ -413,15 +405,33 @@ class FlanaBot(Connect4Bot, PenaltyBot, PollBot, ScraperBot, UberEatsBot, Weathe and random.random() < constants.INSULT_PROBABILITY ) - and - ( - not self.tunnel_chat - or - self.tunnel_chat != message.chat - ) ): await self.send_insult(message) + @ignore_self_message + async def _on_new_message_raw( + self, + message: Message, + whitelist_callbacks: set[RegisteredCallback] | None = None, + blacklist_callbacks: set[RegisteredCallback] | None = None + ): + if ( + message.replied_message + and + message.replied_message.author.id == self.id + and + message.replied_message.medias + ): + whitelist_callbacks = (whitelist_callbacks or set()) | {self._on_song_info} + elif self.tunnel_chat and message.chat == await self.owner_chat: + whitelist_callbacks = (whitelist_callbacks or set()) | { + self._on_delete, + self._on_tunnel_message, + self._on_deactivate_tunnel + } + + await super()._on_new_message_raw(message, whitelist_callbacks, blacklist_callbacks) + async def _on_ready(self): await super()._on_ready() flanautils.do_every(multibot_constants.CHECK_OLD_DATABASE_MESSAGES_EVERY_SECONDS, self.check_old_database_actions) diff --git a/flanabot/bots/flana_tele_bot.py b/flanabot/bots/flana_tele_bot.py index 666e392..7b4a146 100644 --- a/flanabot/bots/flana_tele_bot.py +++ b/flanabot/bots/flana_tele_bot.py @@ -8,7 +8,7 @@ from typing import Callable import telethon.tl.functions from flanautils import Media, OrderedSet -from multibot import TelegramBot, find_message, use_user_client, user_client +from multibot import RegisteredCallback, TelegramBot, find_message, use_user_client, user_client from flanabot import constants from flanabot.bots.flana_bot import FlanaBot @@ -74,8 +74,13 @@ class FlanaTeleBot(TelegramBot, FlanaBot): await super()._on_inline_query_raw(message) @whitelisted - async def _on_new_message_raw(self, message: Message): - await super()._on_new_message_raw(message) + async def _on_new_message_raw( + self, + message: Message, + whitelist_callbacks: set[RegisteredCallback] | None = None, + blacklist_callbacks: set[RegisteredCallback] | None = None + ): + await super()._on_new_message_raw(message, whitelist_callbacks, blacklist_callbacks) async def _on_ready(self): await super()._on_ready() diff --git a/flanabot/bots/penalty_bot.py b/flanabot/bots/penalty_bot.py index 1667b81..a75236a 100644 --- a/flanabot/bots/penalty_bot.py +++ b/flanabot/bots/penalty_bot.py @@ -6,7 +6,7 @@ from abc import ABC import flanautils from flanautils import TimeUnits -from multibot import MultiBot, User, admin, bot_mentioned, constants as multibot_constants, group, ignore_self_message +from multibot import MultiBot, RegisteredCallback, User, admin, bot_mentioned, constants as multibot_constants, group, ignore_self_message from flanabot import constants from flanabot.models import Chat, Message, Punishment @@ -100,8 +100,13 @@ class PenaltyBot(MultiBot, ABC): await self.mute(user, message, flanautils.text_to_time(await self.filter_mention_ids(message.text, message)), message) @ignore_self_message - async def _on_new_message_raw(self, message: Message): - await super()._on_new_message_raw(message) + async def _on_new_message_raw( + self, + message: Message, + whitelist_callbacks: set[RegisteredCallback] | None = None, + blacklist_callbacks: set[RegisteredCallback] | None = None + ): + await super()._on_new_message_raw(message, whitelist_callbacks, blacklist_callbacks) if message.chat.config['check_flood'] and message.chat.config['punish'] and not message.is_inline: async with self.lock: await self._check_message_flood(message)