Modified authentication

README.md

@@ -103,6 +103,7 @@ sudo docker run mirror-bot
 ```
 
 ## Deploying on Heroku
+
 - Run the script to generate token file(token.pickle) for Google Drive:
 ```
 python3 generate_drive_token.py
@@ -124,6 +125,10 @@ heroku git:remote -a appname
 ```
 heroku stack:set container
 ```
+- Add heroku postgres database, connect to database and create a table:
+```
+CREATE TABLE users (uid bigint, sudo boolean DEFAULT FALSE)
+```
 - Add Private Credentials and Config Stuff:
 ```
 git add -f credentials.json token.pickle config.env heroku.yml

bot/__init__.py

@@ -8,6 +8,9 @@
 from dotenv import load_dotenv
 import socket
 
+import psycopg2
+from psycopg2 import Error
+
 socket.setdefaulttimeout(600)
 
 botStartTime = time.time()
@@ -58,14 +61,11 @@ def getConfig(name: str):
 download_dict = {}
 # Stores list of users and chats the bot is authorized to use in
 AUTHORIZED_CHATS = set()
-if os.path.exists('authorized_chats.txt'):
-    with open('authorized_chats.txt', 'r+') as f:
-        lines = f.readlines()
-        for line in lines:
-            #    LOGGER.info(line.split())
-            AUTHORIZED_CHATS.add(int(line.split()[0]))
+SUDO_USERS = set()
+
 try:
     BOT_TOKEN = getConfig('BOT_TOKEN')
+    DB_URI = getConfig('DATABASE_URL')
     parent_id = getConfig('GDRIVE_FOLDER_ID')
     DOWNLOAD_DIR = getConfig('DOWNLOAD_DIR')
     if DOWNLOAD_DIR[-1] != '/' or DOWNLOAD_DIR[-1] != '\\':
@@ -79,6 +79,27 @@ def getConfig(name: str):
 except KeyError as e:
     LOGGER.error("One or more env variables missing! Exiting now")
     exit(1)
+
+try:
+    conn = psycopg2.connect(DB_URI)
+    cur = conn.cursor()
+    sql = "SELECT * from users;"
+    cur.execute(sql)
+    rows = cur.fetchall()  #returns a list ==> (uid, sudo)
+    for row in rows:
+        AUTHORIZED_CHATS.add(row[0])
+        if row[1]:
+            SUDO_USERS.add(row[0])
+    print("Connected to DB")  
+except psycopg2.DatabaseError as error :
+    LOGGER.error(f"Error : {error}")
+    exit(1)
+finally:
+    #closing database connection.
+    if(conn):
+        cur.close()
+        conn.close()
+
 try:
     INDEX_URL = getConfig('INDEX_URL')
     if len(INDEX_URL) == 0:

bot/helper/ext_utils/db_handler.py

@@ -0,0 +1,78 @@
+import psycopg2
+from psycopg2 import Error
+from bot import AUTHORIZED_CHATS, SUDO_USERS, DB_URI, LOGGER
+
+class DbManger:
+    def __init__(self):
+        self.err = False
+
+    def connect(self):
+        try:
+            self.conn = psycopg2.connect(DB_URI)
+            self.cur = self.conn.cursor()
+        except psycopg2.DatabaseError as error :
+            LOGGER.error("Error in dbMang : ", error)
+            self.err = True
+
+    def disconnect(self):
+        self.cur.close()
+        self.conn.close()
+
+    def db_auth(self,chat_id: int):
+        self.connect()
+        if self.err :
+            return "There's some error check log for details"
+        else:
+            sql = 'INSERT INTO users VALUES ({});'.format(chat_id)
+            self.cur.execute(sql)
+            self.conn.commit()
+            self.disconnect()
+            AUTHORIZED_CHATS.add(chat_id)
+            return 'Authorized successfully'
+
+    def db_unauth(self,chat_id: int):
+        self.connect()
+        if self.err :
+            return "There's some error check log for details"
+        else:
+            sql = 'DELETE from users where uid = {};'.format(chat_id)
+            self.cur.execute(sql)
+            self.conn.commit()
+            self.disconnect()
+            AUTHORIZED_CHATS.remove(chat_id)
+            if chat_id in SUDO_USERS:
+                SUDO_USERS.remove(chat_id)
+            return 'Unauthorized successfully'
+
+    def db_addsudo(self,chat_id: int):
+        self.connect()
+        if self.err :
+            return "There's some error check log for details"
+        else:
+            if chat_id in AUTHORIZED_CHATS:
+                sql = 'UPDATE users SET sudo = TRUE where uid = {};'.format(chat_id)
+                self.cur.execute(sql)
+                self.conn.commit()
+                self.disconnect()
+                SUDO_USERS.add(chat_id)
+                return 'Successfully promoted as sudo'
+            else:
+                sql = 'INSERT INTO users VALUES ({},TRUE);'.format(chat_id)
+                self.cur.execute(sql)
+                self.conn.commit()
+                self.disconnect()
+                AUTHORIZED_CHATS.add(chat_id)
+                SUDO_USERS.add(chat_id)
+                return 'Successfully Authorized and promoted as sudo'
+
+    def db_rmsudo(self,chat_id: int):
+        self.connect()
+        if self.err :
+            return "There's some error check log for details"
+        else:
+            sql = 'UPDATE users SET sudo = FALSE where uid = {};'.format(chat_id)
+            self.cur.execute(sql)
+            self.conn.commit()
+            self.disconnect()
+            SUDO_USERS.remove(chat_id)
+            return 'Successfully removed from Sudo'

bot/helper/telegram_helper/bot_commands.py

@@ -8,8 +8,11 @@ def __init__(self):
         self.CancelAllCommand = 'cancelall'
         self.ListCommand = 'list'
         self.StatusCommand = 'status'
+        self.AuthorizedUsersCommand = 'users'
         self.AuthorizeCommand = 'auth'
         self.UnAuthorizeCommand = 'unauth'
+        self.AddSudoCommand = 'addsudo'
+        self.RmSudoCommand = 'rmsudo'
         self.PingCommand = 'ping'
         self.RestartCommand = 'restart'
         self.StatsCommand = 'stats'

bot/helper/telegram_helper/filters.py

@@ -1,6 +1,6 @@
 from telegram.ext import BaseFilter
 from telegram import Message
-from bot import AUTHORIZED_CHATS, OWNER_ID, download_dict, download_dict_lock
+from bot import AUTHORIZED_CHATS, SUDO_USERS, OWNER_ID, download_dict, download_dict_lock
 
 
 class CustomFilters:
@@ -23,6 +23,12 @@ def filter(self, message):
 
     authorized_chat = _AuthorizedChat()
 
+    class _SudoUser(BaseFilter):
+        def filter(self,message):
+            return bool(message.from_user.id in SUDO_USERS)
+
+    sudo_user = _SudoUser()
+
     class _MirrorOwner(BaseFilter):
         def filter(self, message: Message):
             user_id = message.from_user.id

bot/modules/authorize.py

@@ -1,69 +1,141 @@
 from bot.helper.telegram_helper.message_utils import sendMessage
 from telegram.ext import run_async
-from bot import AUTHORIZED_CHATS, dispatcher
+from bot import AUTHORIZED_CHATS, SUDO_USERS, dispatcher
 from telegram.ext import CommandHandler
 from bot.helper.telegram_helper.filters import CustomFilters
-from telegram.ext import Filters
 from telegram import Update
 from bot.helper.telegram_helper.bot_commands import BotCommands
-
+from bot.helper.ext_utils.db_handler import DbManger
 
 @run_async
 def authorize(update,context):
+    reply_message = None
+    message_ = None
     reply_message = update.message.reply_to_message
-    msg = ''
-    with open('authorized_chats.txt', 'a') as file:
+    message_ = update.message.text.split(' ')
+    if len(message_) == 2:
+        chat_id = int(message_[1])
+        if chat_id not in AUTHORIZED_CHATS:
+            msg = DbManger().db_auth(chat_id)
+        else:
+            msg = 'User already authorized'
+    else:
         if reply_message is None:
             # Trying to authorize a chat
             chat_id = update.effective_chat.id
             if chat_id not in AUTHORIZED_CHATS:
-                file.write(f'{chat_id}\n')
-                AUTHORIZED_CHATS.add(chat_id)
-                msg = 'Chat authorized'
+                msg = DbManger().db_auth(chat_id)
             else:
                 msg = 'Already authorized chat'
+
         else:
             # Trying to authorize someone in specific
             user_id = reply_message.from_user.id
             if user_id not in AUTHORIZED_CHATS:
-                file.write(f'{user_id}\n')
-                AUTHORIZED_CHATS.add(user_id)
-                msg = 'Person Authorized to use the bot!'
+                msg = DbManger().db_auth(user_id)
             else:
-                msg = 'Person already authorized'
-        sendMessage(msg, context.bot, update)
+                msg = 'User already authorized'
+    sendMessage(msg, context.bot, update)
 
 
 @run_async
 def unauthorize(update,context):
+    reply_message = None
+    message_ = None
     reply_message = update.message.reply_to_message
-    if reply_message is None:
-        # Trying to unauthorize a chat
-        chat_id = update.effective_chat.id
+    message_ = update.message.text.split(' ')
+    if len(message_) == 2:
+        chat_id = int(message_[1])
         if chat_id in AUTHORIZED_CHATS:
-            AUTHORIZED_CHATS.remove(chat_id)
-            msg = 'Chat unauthorized'
+            msg = DbManger().db_unauth(chat_id)
+        else:
+            msg = 'User already unauthorized'
+    else:
+        if reply_message is None:
+            # Trying to unauthorize a chat
+            chat_id = update.effective_chat.id
+            if chat_id in AUTHORIZED_CHATS:
+                msg = DbManger().db_unauth(chat_id)
+            else:
+                msg = 'Already unauthorized chat'
+        else:
+            # Trying to authorize someone in specific
+            user_id = reply_message.from_user.id
+            if user_id in AUTHORIZED_CHATS:
+                msg = DbManger().db_unauth(user_id)
+            else:
+                msg = 'User already unauthorized'
+    sendMessage(msg, context.bot, update)
+
+
+@run_async
+def addSudo(update,context):
+    reply_message = None
+    message_ = None
+    reply_message = update.message.reply_to_message
+    message_ = update.message.text.split(' ')
+    if len(message_) == 2:
+        chat_id = int(message_[1])
+        if chat_id not in SUDO_USERS:
+            msg = DbManger().db_addsudo(chat_id)
         else:
-            msg = 'Already unauthorized chat'
+            msg = 'Already Sudo'
     else:
-        # Trying to authorize someone in specific
-        user_id = reply_message.from_user.id
-        if user_id in AUTHORIZED_CHATS:
-            AUTHORIZED_CHATS.remove(user_id)
-            msg = 'Person unauthorized to use the bot!'
+        if reply_message is None:
+            msg = "Give ID or Reply To message of whom you want to Promote"
         else:
-            msg = 'Person already unauthorized!'
-    with open('authorized_chats.txt', 'a') as file:
-        file.truncate(0)
-        for i in AUTHORIZED_CHATS:
-            file.write(f'{i}\n')
+            # Trying to authorize someone in specific
+            user_id = reply_message.from_user.id
+            if user_id not in SUDO_USERS:
+                msg = DbManger().db_addsudo(user_id)
+            else:
+                msg = 'Already Sudo'
     sendMessage(msg, context.bot, update)
 
 
+@run_async
+def removeSudo(update,context):
+    reply_message = None
+    message_ = None
+    reply_message = update.message.reply_to_message
+    message_ = update.message.text.split(' ') 
+    if len(message_) == 2:
+        chat_id = int(message_[1])
+        if chat_id in SUDO_USERS:
+            msg = DbManger().db_rmsudo(chat_id)
+        else:
+            msg = 'Not a Sudo'
+    else:
+        if reply_message is None:
+            msg = "Give ID or Reply To message of whom you want to remove from Sudo"
+        else:
+            user_id = reply_message.from_user.id
+            if user_id in SUDO_USERS:
+                msg = DbManger().db_rmsudo(user_id)
+            else:
+                msg = 'Not a Sudo'
+    sendMessage(msg, context.bot, update)
+
+
+@run_async
+def sendAuthChats(update,context):
+    sendMessage(f'Authorized Chats are : {AUTHORIZED_CHATS.__str__()}\nSudo Users are : {SUDO_USERS}', context.bot, update)
+
+
+send_auth_handler = CommandHandler(command=BotCommands.AuthorizedUsersCommand, callback=sendAuthChats,
+                                    filters=CustomFilters.owner_filter | CustomFilters.sudo_user)
 authorize_handler = CommandHandler(command=BotCommands.AuthorizeCommand, callback=authorize,
-                                   filters=CustomFilters.owner_filter & Filters.group)
+                                   filters=CustomFilters.owner_filter | CustomFilters.sudo_user)
 unauthorize_handler = CommandHandler(command=BotCommands.UnAuthorizeCommand, callback=unauthorize,
-                                     filters=CustomFilters.owner_filter & Filters.group)
+                                     filters=CustomFilters.owner_filter | CustomFilters.sudo_user)
+addsudo_handler = CommandHandler(command=BotCommands.AddSudoCommand, callback=addSudo,
+                                    filters=CustomFilters.owner_filter)
+removesudo_handler = CommandHandler(command=BotCommands.RmSudoCommand, callback=removeSudo,
+                                    filters=CustomFilters.owner_filter)
+
+dispatcher.add_handler(send_auth_handler)
 dispatcher.add_handler(authorize_handler)
 dispatcher.add_handler(unauthorize_handler)
+dispatcher.add_handler(addsudo_handler)
+dispatcher.add_handler(removesudo_handler)
 

requirements.txt

@@ -12,3 +12,4 @@ beautifulsoup4>=4.8.2,<4.8.10
 Pyrogram>=0.16.0,<0.16.10
 TgCrypto>=1.1.1,<1.1.10
 git+git://github.com/lzzy12/youtube-dl@d7c2b43#youtube_dl
+psycopg2-binary