Only allow authorized users to use the bot

Signed-off-by: lzzy12 <jhashivam2020@gmail.com>

.gitignore

@@ -11,3 +11,4 @@ data*
 .idea
 *.json
 *.pickle
+authorized_chats.txt

bot/__init__.py

@@ -39,6 +39,13 @@ download_dict_lock = threading.Lock()
 status_reply_dict_lock = threading.Lock()
 status_reply_dict = {}
 download_dict = {}
+# Stores list of users and chats the bot is authorized to use in
+AUTHORIZED_CHATS = []
+with open('authorized_chats.txt', 'r+') as f:
+    lines = f.readlines()
+    for line in lines:
+        LOGGER.info(line.split())
+        AUTHORIZED_CHATS.append(line.split()[0])
 try:
     BOT_TOKEN = getConfig('BOT_TOKEN')
     parent_id = getConfig('GDRIVE_FOLDER_ID')
@@ -46,7 +53,7 @@ try:
     if DOWNLOAD_DIR[-1] != '/' or DOWNLOAD_DIR[-1] != '\\':
         DOWNLOAD_DIR = DOWNLOAD_DIR + '/'
     DOWNLOAD_STATUS_UPDATE_INTERVAL = int(getConfig('DOWNLOAD_STATUS_UPDATE_INTERVAL'))
-
+    OWNER_ID = int(getConfig('OWNER_ID'))
 except KeyError as e:
     LOGGER.error("One or more env variables missing! Exiting now")
     exit(1)

bot/__main__.py

@@ -1,12 +1,12 @@
 from telegram.ext import CommandHandler, run_async
 from bot import dispatcher, LOGGER, updater, aria2
-import bot.mirror, bot.list, bot.mirror_status, bot.cancel_mirror
+import bot.mirror, bot.list, bot.mirror_status, bot.cancel_mirror, bot.authorize
 from bot.helper import fs_utils
 import signal
 import time
 from bot.helper.message_utils import *
 import shutil
-
+from .helper.telegram_helper.filters import CustomFilters
 
 @run_async
 def disk_usage(update, context):
@@ -52,14 +52,19 @@ def bot_help(update, context):
 
 def main():
 
-    start_handler = CommandHandler('start', start)
-    ping_handler = CommandHandler('ping', ping)
-    help_handler = CommandHandler('help', bot_help)
-    disk_usage_handler = CommandHandler('disk', disk_usage)
+    start_handler = CommandHandler('start', start,
+                                   filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
+    ping_handler = CommandHandler('ping', ping,
+                                  filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
+    help_handler = CommandHandler('help',
+                                  bot_help, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
+    disk_handler = CommandHandler('disk',
+                                  disk_usage, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
+
     dispatcher.add_handler(start_handler)
     dispatcher.add_handler(ping_handler)
     dispatcher.add_handler(help_handler)
-    dispatcher.add_handler(disk_usage_handler)
+    dispatcher.add_handler(disk_handler)
     updater.start_polling()
     LOGGER.info("Bot Started!")
     signal.signal(signal.SIGINT, fs_utils.exit_clean_up)

bot/authorize.py

@@ -0,0 +1,32 @@
+from .helper.message_utils import sendMessage
+from telegram.ext import run_async
+from bot import AUTHORIZED_CHATS, dispatcher
+from telegram.ext import CommandHandler
+from .helper.telegram_helper.filters import CustomFilters
+from telegram.ext import Filters
+from telegram import Update
+
+
+@run_async
+def authorize(update: Update, context):
+    reply_message = update.message.reply_to_message
+    with open('authorized_chats.txt', 'a') as file:
+        if reply_message is None:
+            chat_id = update.effective_chat.id
+            if chat_id not in AUTHORIZED_CHATS:
+                file.write(f'{chat_id}\n')
+                AUTHORIZED_CHATS.append(chat_id)
+                sendMessage('Chat authorized', context, update)
+            else:
+                sendMessage('Already authorized chat', context, update)
+        else:
+            user_id = reply_message.from_user.id
+            if user_id not in AUTHORIZED_CHATS:
+                file.write(f'{user_id}\n')
+                AUTHORIZED_CHATS.append(user_id)
+
+
+authorize_handler = CommandHandler(command='authorize', callback=authorize,
+                                   filters=CustomFilters.owner_filter & Filters.group)
+dispatcher.add_handler(authorize_handler)
+

bot/cancel_mirror.py

@@ -1,7 +1,7 @@
 from telegram.ext import CommandHandler, run_async
 from bot.helper.message_utils import *
 from bot import download_dict, aria2, dispatcher, download_dict_lock
-
+from .helper.telegram_helper.filters import CustomFilters
 
 @run_async
 def cancel_mirror(update: Update, context):
@@ -21,5 +21,6 @@ def cancel_mirror(update: Update, context):
     sendMessage("Download canceled", context, update)
 
 
-cancel_mirror_handler = CommandHandler('cancel', cancel_mirror)
-dispatcher.add_handler(cancel_mirror_handler)
+cancel_mirror_handler = CommandHandler('cancel', cancel_mirror,
+                                       filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
+dispatcher.add_handler(cancel_mirror_handler)

bot/config_sample.ini

@@ -6,4 +6,5 @@ _____REMOVE_THIS_LINE_____=True
 BOT_TOKEN =
 GDRIVE_FOLDER_ID = 
 DOWNLOAD_DIR = ./download
-DOWNLOAD_STATUS_UPDATE_INTERVAL = 5
+DOWNLOAD_STATUS_UPDATE_INTERVAL = 5
+OWNER_ID =

bot/helper/telegram_helper/filters.py

@@ -0,0 +1,24 @@
+from telegram.ext import BaseFilter
+from telegram import Update
+from bot import AUTHORIZED_CHATS, OWNER_ID
+
+
+class CustomFilters:
+    class _OwnerFilter(BaseFilter):
+        def filter(self, message):
+            return bool(message.from_user.id == OWNER_ID)
+
+    owner_filter = _OwnerFilter()
+
+    class _AuthorizedUserFilter(BaseFilter):
+        def filter(self, message):
+            id = message.from_user.id
+            return bool(id in AUTHORIZED_CHATS or id == OWNER_ID)
+
+    authorized_user = _AuthorizedUserFilter()
+
+    class _AuthorizedChat(BaseFilter):
+        def filter(self, message):
+            return bool(message.chat.id in AUTHORIZED_CHATS)
+
+    authorized_chat = _AuthorizedChat()

bot/mirror.py

@@ -7,7 +7,7 @@ from bot import download_dict, status_reply_dict, status_reply_dict_lock, downlo
 from bot.helper.message_utils import *
 from bot.helper.bot_utils import get_readable_message, MirrorStatus
 from bot.helper.exceptions import KillThreadException
-
+from .helper.telegram_helper.filters import CustomFilters
 
 class MirrorListener(listeners.MirrorListeners):
     def __init__(self, context, update, reply_message):
@@ -98,5 +98,5 @@ def mirror(update, context):
     aria.add_download(link)
 
 
-mirror_handler = CommandHandler('mirror', mirror)
+mirror_handler = CommandHandler('mirror', mirror, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
 dispatcher.add_handler(mirror_handler)

bot/mirror_status.py

@@ -4,7 +4,7 @@ from bot.helper.message_utils import *
 from time import sleep
 from bot.helper.bot_utils import get_readable_message
 from telegram.error import BadRequest
-
+from .helper.telegram_helper.filters import CustomFilters
 
 @run_async
 def mirror_status(update: Update, context):
@@ -42,5 +42,6 @@ def mirror_status(update: Update, context):
         sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL)
 
 
-mirror_status_handler = CommandHandler('status', mirror_status)
+mirror_status_handler = CommandHandler('status', mirror_status,
+                                       filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
 dispatcher.add_handler(mirror_status_handler)