Fix mentions at private chats

AlberLC · Jun 23, 2022 · a9ccb0a · a9ccb0a
1 parent f98c8eb
commit a9ccb0a
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 7 deletions.
diff --git a/multibot/bots/discord_bot.py b/multibot/bots/discord_bot.py
@@ -137,15 +137,26 @@ async def _get_mentions(self, original_message: constants.DISCORD_MESSAGE) -> li
         mentions = OrderedSet(self._create_user_from_discord_user(user) for user in original_message.mentions)
 
         chat = await self._get_chat(original_message)
-        if chat.original_object.guild:
-            text = await self._get_text(original_message)
-            text = flanautils.remove_symbols(text, replace_with=' ')
-            words = text.lower().split()
+        text = await self._get_text(original_message)
+        text = flanautils.remove_symbols(text, ignore=('#',), replace_with=' ')
+        words = text.lower().split()
+
+        if chat.is_group:
             for member in chat.original_object.guild.members:
                 user_name = f'{member.name}#{member.discriminator}'.lower()
                 short_user_name = member.name.lower()
                 if user_name in words or short_user_name in words:
                     mentions.add(self._create_user_from_discord_user(member))
+        else:
+            user_name = f'{chat.original_object.recipient.name}#{chat.original_object.recipient.discriminator}'.lower()
+            short_user_name = chat.original_object.recipient.name.lower()
+            if user_name in words or short_user_name in words:
+                mentions.add(self._create_user_from_discord_user(chat.original_object.recipient))
+
+            bot_name = f'{chat.original_object.me.name}#{chat.original_object.me.discriminator}'.lower()
+            short_bot_name = chat.original_object.me.name.lower()
+            if bot_name in words or short_bot_name in words:
+                mentions.add(self._create_user_from_discord_user(chat.original_object.me))
 
         return list(mentions - None)
 

diff --git a/multibot/bots/telegram_bot.py b/multibot/bots/telegram_bot.py
@@ -128,7 +128,7 @@ async def _create_user_from_telegram_user(self, original_user: constants.TELEGRA
         return User(
             platform=self.platform.value,
             id=original_user.id,
-            name=self._get_name_from_entity(original_user).strip(' @'),
+            name=self._get_name_from_entity(original_user),
             is_admin=is_admin,
             is_bot=original_user.bot,
             original_object=original_user
@@ -172,10 +172,14 @@ async def _get_mentions(self, original_message: constants.TELEGRAM_MESSAGE) -> l
 
         text = flanautils.remove_symbols(text, replace_with=' ')
         words = text.lower().split()
+
         for participant in await self.client.get_participants(chat.original_object):
-            user_name = self._get_name_from_entity(participant).strip(' @').lower()
+            user_name = self._get_name_from_entity(participant).lower()
             if user_name in words:
                 mentions.add(await self._create_user_from_telegram_user(participant, chat))
+        if chat.is_private:
+            if self.name.lower() in words:
+                mentions.add(await self.get_me())
 
         return list(mentions - None)
 
@@ -186,7 +190,7 @@ async def _get_message_id(self, original_message: constants.TELEGRAM_MESSAGE) ->
     @return_if_first_empty('', exclude_self_types='TelegramBot', globals_=globals())
     def _get_name_from_entity(self, entity: telethon.hints.EntityLike) -> str:
         if isinstance(entity, telethon.types.User):
-            return f'@{entity.username}' if entity.username else entity.first_name
+            return entity.username or entity.first_name
         elif isinstance(entity, (telethon.types.Channel, telethon.types.Chat)):
             return entity.title