Add MultiBot._on_new_message_raw whitelist/blacklist

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)

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()

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)