From 0119698e813e27db0f1b5bdd96e4428f968c302e Mon Sep 17 00:00:00 2001
From: Nick80835 <nick80835@gmail.com>
Date: Sat, 21 Oct 2023 20:35:06 -0400
Subject: [PATCH] improve/fix moderation mentions

---
 ubot/custom.py             |  7 ++++++-
 ubot/fixes/utils.py        |  2 +-
 ubot/modules/moderation.py | 12 ++++++++----
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/ubot/custom.py b/ubot/custom.py
index 77bff46..a761321 100644
--- a/ubot/custom.py
+++ b/ubot/custom.py
@@ -6,7 +6,8 @@ from telethon.events.inlinequery import InlineQuery
 from telethon.events.newmessage import NewMessage
 from telethon.tl.types import (DocumentAttributeFilename,
                                DocumentAttributeImageSize,
-                               DocumentAttributeSticker)
+                               DocumentAttributeSticker,
+                               MessageEntityMentionName)
 
 from ubot.command import (CallbackQueryCommand, Command, InlineArticleCommand,
                           InlinePhotoCommand)
@@ -24,6 +25,10 @@ class ExtendedNewMessage(NewMessage.Event):
     other_args: tuple # any groups between the args group and the command itself
     nsfw_disabled: bool # only set if pass_nsfw is True; this value is the opposite of nsfw_enabled in chat_db
 
+    @property
+    def has_user_entities(self) -> bool:
+        return any([i for i in self.entities if isinstance(i, MessageEntityMentionName)]) if self.entities else False
+
     async def get_text(self, return_msg=False, default=""):
         if self.args:
             if return_msg:
diff --git a/ubot/fixes/utils.py b/ubot/fixes/utils.py
index a7ddf84..04a3628 100644
--- a/ubot/fixes/utils.py
+++ b/ubot/fixes/utils.py
@@ -4,7 +4,7 @@ from telethon.tl.types import Channel, Chat, MessageEntityMentionName
 
 
 async def get_user(event, allow_channel=False):
-    if mention_entities := [i for i in event.get_entities_text() if isinstance(i, MessageEntityMentionName)]:
+    if mention_entities := [i for i in event.entities if isinstance(i, MessageEntityMentionName)] if event.entities else False:
         if len(mention_entities) > 1:
             await event.reply("You provided too many arguments!")
             return
diff --git a/ubot/modules/moderation.py b/ubot/modules/moderation.py
index f853b99..292b9b6 100644
--- a/ubot/modules/moderation.py
+++ b/ubot/modules/moderation.py
@@ -18,7 +18,7 @@ async def kick_user(event):
         await event.reply("I can't kick users in this chat.")
         return
 
-    if time_regex.sub("", event.args).strip().lower() == "me":
+    if time_regex.sub("", event.args).strip().lower() == "me" and not event.has_user_entities:
         self_harm = True
         user_to_kick = await event.get_sender()
     else:
@@ -63,7 +63,7 @@ async def ban_user(event):
     if time_match := time_regex.search(event.args):
         event.args = time_regex.sub("", event.args).strip()
 
-    if time_regex.sub("", event.args).strip().lower() == "me":
+    if time_regex.sub("", event.args).strip().lower() == "me" and not event.has_user_entities:
         await event.reply("I don't think I should do that…")
         return
 
@@ -105,6 +105,10 @@ async def ban_user(event):
 @ldr.add("unban", moderation=True, help="Unban a user.")
 @ldr.add(f"{bot_name}(,|) unban", moderation=True, simple_pattern=True, hide_help=True)
 async def unban_user(event):
+    if event.args.lower() == "me" and not event.has_user_entities:
+        await event.reply("You probably aren't banned.")
+        return
+
     if not (await event.client.get_permissions(event.chat, "me")).ban_users:
         await event.reply("I can't unban users in this chat.")
         return
@@ -135,7 +139,7 @@ async def mute_user(event):
     if time_match := time_regex.search(event.args):
         event.args = time_regex.sub("", event.args).strip()
 
-    if time_regex.sub("", event.args).strip().lower() == "me":
+    if time_regex.sub("", event.args).strip().lower() == "me" and not event.has_user_entities:
         self_harm = True
         user_to_mute = await event.get_sender()
     else:
@@ -186,7 +190,7 @@ async def mute_user(event):
 @ldr.add("unmute", moderation=True, help="Unmute a user.")
 @ldr.add(f"{bot_name}(,|) unmute", moderation=True, simple_pattern=True, hide_help=True)
 async def unmute_user(event):
-    if event.args.lower() == "me":
+    if event.args.lower() == "me" and not event.has_user_entities:
         await event.reply("You probably aren't muted.")
         return