From 3b5e2d43ef97689e9cf24226e02e1bc1de64a0d0 Mon Sep 17 00:00:00 2001
From: Nick80835 <nick80835@gmail.com>
Date: Fri, 12 Jun 2020 13:49:57 -0400
Subject: [PATCH] add a bot-wide user blacklist

---
 ubot/command_handler.py | 14 +++++++++
 ubot/modules/system.py  | 64 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/ubot/command_handler.py b/ubot/command_handler.py
index cf322db..a3393b0 100644
--- a/ubot/command_handler.py
+++ b/ubot/command_handler.py
@@ -20,6 +20,10 @@ class CommandHandler():
         client.add_event_handler(self.handle_inline, events.InlineQuery())
 
     async def handle_incoming(self, event):
+        if self.is_blacklisted(event):
+            print(f"Attempted command ({event.text}) from blacklisted ID {event.from_id}")
+            return
+
         prefix = escape(self.settings.get_config("cmd_prefix") or '.')
 
         for key, value in self.incoming_commands.items():
@@ -49,6 +53,10 @@ class CommandHandler():
                     raise exception
 
     async def handle_inline(self, event):
+        if self.is_blacklisted(event):
+            print(f"Attempted command ({event.text}) from blacklisted ID {event.from_id}")
+            return
+
         for key, value in self.inline_photo_commands.items():
             pattern_match = search(self.inline_pattern_template.format(key), event.text)
 
@@ -136,3 +144,9 @@ class CommandHandler():
                 return True
 
         return False
+
+    def is_blacklisted(self, event):
+        if str(event.from_id) in self.settings.get_list("blacklisted_users"):
+            return True
+        else:
+            return False
diff --git a/ubot/modules/system.py b/ubot/modules/system.py
index 858ad78..7c7faaa 100644
--- a/ubot/modules/system.py
+++ b/ubot/modules/system.py
@@ -109,3 +109,67 @@ async def nsfw_toggle(event):
     elif event.args == "off":
         ldr.settings.add_to_list("nsfw_blacklist", event.chat.id)
         await event.reply("NSFW commands disabled for this chat!")
+
+
+@ldr.add("blacklist", sudo=True)
+async def add_blacklist(event):
+    if event.args:
+        try:
+            user_entity = await event.client.get_entity(event.args)
+        except (ValueError, TypeError):
+            await event.reply("The ID or username you provided was invalid!")
+            return
+    elif event.is_reply:
+        reply = await event.get_reply_message()
+        reply_id = reply.from_id
+
+        if reply_id:
+            try:
+                user_entity = await event.client.get_entity(reply_id)
+            except (ValueError, TypeError):
+                await event.reply("There was an error getting the user's ID!")
+                return
+        else:
+            await event.reply("Blacklisting failed!")
+            return
+    else:
+        await event.reply("Give me a user ID, username or reply!")
+        return
+
+    userid = user_entity.id
+    userfullname = f"{user_entity.first_name} {user_entity.last_name or ''}"
+
+    ldr.settings.add_to_list("blacklisted_users", userid)
+    await event.reply(f"Successfully blacklisted **{userfullname}** `({userid})`")
+
+
+@ldr.add("unblacklist", sudo=True)
+async def rem_blacklist(event):
+    if event.args:
+        try:
+            user_entity = await event.client.get_entity(event.args)
+        except (ValueError, TypeError):
+            await event.reply("The ID or username you provided was invalid!")
+            return
+    elif event.is_reply:
+        reply = await event.get_reply_message()
+        reply_id = reply.from_id
+
+        if reply_id:
+            try:
+                user_entity = await event.client.get_entity(reply_id)
+            except (ValueError, TypeError):
+                await event.reply("There was an error getting the user's ID!")
+                return
+        else:
+            await event.reply("Blacklisting failed!")
+            return
+    else:
+        await event.reply("Give me a user ID, username or reply!")
+        return
+
+    userid = user_entity.id
+    userfullname = f"{user_entity.first_name} {user_entity.last_name or ''}"
+
+    ldr.settings.remove_from_list("blacklisted_users", userid)
+    await event.reply(f"Successfully unblacklisted **{userfullname}** `({userid})`")