Optional Database and minor changes

-force remove tasks from aria2
-purge completed,removed and failed downloads to free memory

Signed-off-by: anas <e.anastayyar@gmail.com>

bot/__init__.py

@@ -80,6 +80,16 @@ download_dict = {}
 # Stores list of users and chats the bot is authorized to use in
 AUTHORIZED_CHATS = set()
 SUDO_USERS = set()
+if os.path.exists('authorized_chats.txt'):
+    with open('authorized_chats.txt', 'r+') as f:
+        lines = f.readlines()
+        for line in lines:
+            AUTHORIZED_CHATS.add(int(line.split()[0]))
+if os.path.exists('sudo_users.txt'):
+    with open('sudo_users.txt', 'r+') as f:
+        lines = f.readlines()
+        for line in lines:
+            SUDO_USERS.add(int(line.split()[0]))
 try:
     achats = getConfig('AUTHORIZED_CHATS')
     achats = achats.split(" ")
@@ -87,10 +97,15 @@ try:
         AUTHORIZED_CHATS.add(int(chats))
 except:
     pass
-
+try:
+    schats = getConfig('SUDO_USERS')
+    schats = achats.split(" ")
+    for chats in schats:
+        SUDO_USERS.add(int(chats))
+except:
+    pass
 try:
     BOT_TOKEN = getConfig('BOT_TOKEN')
-    DB_URI = getConfig('DATABASE_URL')
     parent_id = getConfig('GDRIVE_FOLDER_ID')
     DOWNLOAD_DIR = getConfig('DOWNLOAD_DIR')
     if not DOWNLOAD_DIR.endswith("/"):
@@ -105,26 +120,33 @@ try:
 except KeyError as e:
     LOGGER.error("One or more env variables missing! Exiting now")
     exit(1)
-
 try:
-    conn = psycopg2.connect(DB_URI)
+    DB_URI = getConfig('DATABASE_URL')
-    cur = conn.cursor()
+    if len(DB_URI) == 0:
-    sql = "SELECT * from users;"
+        raise KeyError
-    cur.execute(sql)
+except KeyError:
-    rows = cur.fetchall()  #returns a list ==> (uid, sudo)
+    logging.warning('Database not provided!')
-    for row in rows:
+    DB_URI = None
-        AUTHORIZED_CHATS.add(row[0])
+if DB_URI is not None:
-        if row[1]:
+    try:
-            SUDO_USERS.add(row[0])
+        conn = psycopg2.connect(DB_URI)
-except Error as e:
+        cur = conn.cursor()
-    if 'relation "users" does not exist' in str(e):
+        sql = "SELECT * from users;"
-        mktable()
+        cur.execute(sql)
-    else:
+        rows = cur.fetchall()  #returns a list ==> (uid, sudo)
-        LOGGER.error(e)
+        for row in rows:
-        exit(1)
+            AUTHORIZED_CHATS.add(row[0])
-finally:
+            if row[1]:
-    cur.close()
+                SUDO_USERS.add(row[0])
-    conn.close()
+    except Error as e:
+        if 'relation "users" does not exist' in str(e):
+            mktable()
+        else:
+            LOGGER.error(e)
+            exit(1)
+    finally:
+        cur.close()
+        conn.close()
 
 LOGGER.info("Generating USER_SESSION_STRING")
 app = Client(':memory:', api_id=int(TELEGRAM_API), api_hash=TELEGRAM_HASH, bot_token=BOT_TOKEN)

bot/helper/mirror_utils/download_utils/aria2_download.py

@@ -15,16 +15,17 @@ class AriaDownloadHelper(DownloadHelper):
         super().__init__()
 
     @new_thread
-    def __onDownloadStarted(self, api, gid):
+    def __onDownloadStarted(self, api: API, gid):
         if STOP_DUPLICATE_MIRROR or TORRENT_DIRECT_LIMIT is not None or TAR_UNZIP_LIMIT is not None:
-            sleep(0.5)
+            sleep(1)
             dl = getDownloadByGid(gid)
             download = api.get_download(gid)
             
             if STOP_DUPLICATE_MIRROR:
                 LOGGER.info(f"Checking File/Folder if already in Drive...")
-                self.name = download.name
+                sleep(1)
-                sname = download.name
+                self.name = aria2.get_download(gid).name
+                sname = self.name
                 if self.listener.isTar:
                     sname = sname + ".tar"
                 if self.listener.extract:
@@ -33,8 +34,9 @@ class AriaDownloadHelper(DownloadHelper):
                     gdrive = GoogleDriveHelper(None)
                     smsg, button = gdrive.drive_list(sname)
                 if smsg:
-                    dl.getListener().onDownloadError(f'File/Folder is already available in Drive.\n\n')
+                    dl.getListener().onDownloadError(f'File/Folder already available in Drive.\n\n')
-                    aria2.remove([download])
+                    aria2.remove([download], force = True)
+                    aria2.purge()
                     sendMarkup("Here are the search results:", dl.getListener().bot, dl.getListener().update, button)
                     return
 
@@ -44,7 +46,7 @@ class AriaDownloadHelper(DownloadHelper):
                     LOGGER.info(f"Checking File/Folder Size...")
                     limit = TAR_UNZIP_LIMIT
                     mssg = f'Tar/Unzip limit is {TAR_UNZIP_LIMIT}'
-                elif TORRENT_DIRECT_LIMIT is not None and limit is None:
+                if TORRENT_DIRECT_LIMIT is not None and limit is None:
                     LOGGER.info(f"Checking File/Folder Size...")
                     limit = TORRENT_DIRECT_LIMIT
                     mssg = f'Torrent/Direct limit is {TORRENT_DIRECT_LIMIT}'
@@ -56,12 +58,14 @@ class AriaDownloadHelper(DownloadHelper):
                     if 'G' in limit[1] or 'g' in limit[1]:
                         if size > limitint * 1024**3:
                             dl.getListener().onDownloadError(f'{mssg}.\nYour File/Folder size is {get_readable_file_size(size)}')
-                            aria2.remove([download])
+                            aria2.remove([download], force = True)
+                            aria2.purge()
                             return
                     elif 'T' in limit[1] or 't' in limit[1]:
                         if size > limitint * 1024**4:
                             dl.getListener().onDownloadError(f'{mssg}.\nYour File/Folder size is {get_readable_file_size(size)}')
-                            aria2.remove([download])
+                            aria2.remove([download], force = True)
+                            aria2.purge()
                             return
         update_all_messages()
 
@@ -73,7 +77,6 @@ class AriaDownloadHelper(DownloadHelper):
             new_gid = download.followed_by_ids[0]
             new_download = api.get_download(new_gid)
             with download_dict_lock:
-                sleep(0.5)
                 download_dict[dl.uid()] = AriaDownloadStatus(new_gid, dl.getListener())
                 if new_download.is_torrent:
                     download_dict[dl.uid()].is_torrent = True
@@ -82,13 +85,15 @@ class AriaDownloadHelper(DownloadHelper):
         else:
             if dl:
                 threading.Thread(target=dl.getListener().onDownloadComplete).start()
+                aria2.purge()
 
     @new_thread
     def __onDownloadStopped(self, api, gid):
-        sleep(1)
+        sleep(5)
         dl = getDownloadByGid(gid)
         if dl: 
             dl.getListener().onDownloadError('Dead torrent!')
+        aria2.purge()
 
     @new_thread
     def __onDownloadError(self, api, gid):
@@ -99,6 +104,7 @@ class AriaDownloadHelper(DownloadHelper):
         LOGGER.info(f"Download Error: {error}")
         if dl: 
             dl.getListener().onDownloadError(error)
+        aria2.purge()
 
     def start_listener(self):
         aria2.listen_to_notifications(threaded=True, on_download_start=self.__onDownloadStarted,

bot/helper/mirror_utils/download_utils/mega_downloader.py

@@ -184,7 +184,7 @@ class MegaDownloadHelper:
             if TAR_UNZIP_LIMIT is not None and (listener.isTar or listener.extract):
                 limit = TAR_UNZIP_LIMIT
                 msg3 = f'Failed, Tar/Unzip limit is {TAR_UNZIP_LIMIT}.\nYour File/Folder size is {get_readable_file_size(api.getSize(node))}.'
-            elif MEGA_LIMIT is not None and limit is None:
+            if MEGA_LIMIT is not None and limit is None:
                 limit = MEGA_LIMIT
                 msg3 = f'Failed, Mega limit is {MEGA_LIMIT}.\nYour File/Folder size is {get_readable_file_size(api.getSize(node))}.'
             if limit is not None:

bot/helper/mirror_utils/status_utils/aria_download_status.py

@@ -95,11 +95,13 @@ class AriaDownloadStatus(Status):
         LOGGER.info(f"Cancelling Download: {self.name()}")
         download = self.aria_download()
         if download.is_waiting:
-            aria2.remove([download])
+            aria2.remove([download], force = True)
+            aria2.purge()
             self.__listener.onDownloadError("Cancelled by user")
             return
         if len(download.followed_by_ids) != 0:
             downloads = aria2.get_downloads(download.followed_by_ids)
-            aria2.remove(downloads)
+            aria2.remove(downloads, force = True)
         self.__listener.onDownloadError("Download stopped by user!")
-        aria2.remove([download])
+        aria2.remove([download], force = True)
+        aria2.purge()

bot/helper/mirror_utils/upload_utils/gdriveTools.py

@@ -298,7 +298,7 @@ class GoogleDriveHelper:
                 reason = json.loads(err.content).get('error').get('errors')[0].get('reason')
                 if reason == 'userRateLimitExceeded' or reason == 'dailyLimitExceeded':
                     if USE_SERVICE_ACCOUNTS:
-                        if self.sa_count > self.service_account_count:
+                        if self.sa_count == self.service_account_count:
                             self.is_cancelled = True
                             raise err
                         else:
@@ -811,7 +811,7 @@ class GoogleDriveHelper:
                     reason = json.loads(err.content).get('error').get('errors')[0].get('reason')
                     if reason == 'downloadQuotaExceeded' or reason == 'dailyLimitExceeded':
                         if USE_SERVICE_ACCOUNTS:
-                            if self.sa_count > self.service_account_count:
+                            if self.sa_count == self.service_account_count:
                                 self.is_cancelled = True
                                 raise err
                             else:

bot/modules/authorize.py

@@ -1,5 +1,5 @@
 from bot.helper.telegram_helper.message_utils import sendMessage
-from bot import AUTHORIZED_CHATS, SUDO_USERS, dispatcher
+from bot import AUTHORIZED_CHATS, SUDO_USERS, dispatcher, DB_URI
 from telegram.ext import CommandHandler
 from bot.helper.telegram_helper.filters import CustomFilters
 from telegram.ext import Filters
@@ -14,27 +14,45 @@ def authorize(update, context):
     reply_message = update.message.reply_to_message
     message_ = update.message.text.split(' ')
     if len(message_) == 2:
-        chat_id = int(message_[1])
+        user_id = int(message_[1])
-        if chat_id not in AUTHORIZED_CHATS:
+        if user_id not in AUTHORIZED_CHATS:
-            msg = DbManger().db_auth(chat_id)
+            if DB_URI is not None:
+                msg = DbManger().db_auth(user_id)
+            else:
+                with open('authorized_chats.txt', 'a') as file:
+                    file.write(f'{user_id}\n')
+                    AUTHORIZED_CHATS.add(user_id)
+                    msg = 'User Authorized'
         else:
-            msg = 'User already authorized'
+            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:
-                msg = DbManger().db_auth(chat_id)
+                if DB_URI is not None:
+                    msg = DbManger().db_auth(chat_id)
+                else:
+                    with open('authorized_chats.txt', 'a') as file:
+                        file.write(f'{chat_id}\n')
+                        AUTHORIZED_CHATS.add(chat_id)
+                        msg = 'Chat Authorized'
             else:
-                msg = 'Already authorized chat'
+                msg = 'Chat Already Authorized'
 
         else:
-            # Trying to authorize someone in specific
+            # Trying to authorize someone by replying
             user_id = reply_message.from_user.id
             if user_id not in AUTHORIZED_CHATS:
-                msg = DbManger().db_auth(user_id)
+                if DB_URI is not None:
+                    msg = DbManger().db_auth(user_id)
+                else:
+                    with open('authorized_chats.txt', 'a') as file:
+                        file.write(f'{user_id}\n')
+                        AUTHORIZED_CHATS.add(user_id)
+                        msg = 'User Authorized'
             else:
-                msg = 'User already authorized'
+                msg = 'User Already Authorized'
     sendMessage(msg, context.bot, update)
 
 
@@ -44,26 +62,42 @@ def unauthorize(update, context):
     reply_message = update.message.reply_to_message
     message_ = update.message.text.split(' ')
     if len(message_) == 2:
-        chat_id = int(message_[1])
+        user_id = int(message_[1])
-        if chat_id in AUTHORIZED_CHATS:
+        if user_id in AUTHORIZED_CHATS:
-            msg = DbManger().db_unauth(chat_id)
+            if DB_URI is not None:
+                msg = DbManger().db_unauth(user_id)
+            else:
+                AUTHORIZED_CHATS.remove(user_id)
+                msg = 'User Unauthorized'
         else:
-            msg = 'User already unauthorized'
+            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)
+                if DB_URI is not None:
+                    msg = DbManger().db_unauth(chat_id)
+                else:
+                    AUTHORIZED_CHATS.remove(chat_id)
+                    msg = 'Chat Unauthorized'
             else:
-                msg = 'Already unauthorized chat'
+                msg = 'Chat Already Unauthorized'
         else:
-            # Trying to authorize someone in specific
+            # Trying to authorize someone by replying
             user_id = reply_message.from_user.id
             if user_id in AUTHORIZED_CHATS:
-                msg = DbManger().db_unauth(user_id)
+                if DB_URI is not None:
+                    msg = DbManger().db_unauth(user_id)
+                else:
+                    AUTHORIZED_CHATS.remove(user_id)
+                    msg = 'User Unauthorized'
             else:
-                msg = 'User already unauthorized'
+                msg = 'User Already Unauthorized'
+    with open('authorized_chats.txt', 'a') as file:
+        file.truncate(0)
+        for i in AUTHORIZED_CHATS:
+            file.write(f'{i}\n')
     sendMessage(msg, context.bot, update)
 
 
@@ -73,19 +107,31 @@ def addSudo(update, context):
     reply_message = update.message.reply_to_message
     message_ = update.message.text.split(' ')
     if len(message_) == 2:
-        chat_id = int(message_[1])
+        user_id = int(message_[1])
-        if chat_id not in SUDO_USERS:
+        if user_id not in SUDO_USERS:
-            msg = DbManger().db_addsudo(chat_id)
+            if DB_URI is not None:
+                msg = DbManger().db_addsudo(user_id)
+            else:
+                with open('sudo_users.txt', 'a') as file:
+                    file.write(f'{user_id}\n')
+                    SUDO_USERS.add(user_id)
+                    msg = 'Promoted as Sudo'
         else:
             msg = 'Already Sudo'
     else:
         if reply_message is None:
             msg = "Give ID or Reply To message of whom you want to Promote"
         else:
-            # Trying to authorize someone in specific
+            # Trying to authorize someone by replying
             user_id = reply_message.from_user.id
             if user_id not in SUDO_USERS:
-                msg = DbManger().db_addsudo(user_id)
+                if DB_URI is not None:
+                    msg = DbManger().db_addsudo(user_id)
+                else:
+                    with open('sudo_users.txt', 'a') as file:
+                        file.write(f'{user_id}\n')
+                        SUDO_USERS.add(user_id)
+                        msg = 'Promoted as Sudo'
             else:
                 msg = 'Already Sudo'
     sendMessage(msg, context.bot, update)
@@ -97,9 +143,13 @@ def removeSudo(update, context):
     reply_message = update.message.reply_to_message
     message_ = update.message.text.split(' ') 
     if len(message_) == 2:
-        chat_id = int(message_[1])
+        user_id = int(message_[1])
-        if chat_id in SUDO_USERS:
+        if user_id in SUDO_USERS:
-            msg = DbManger().db_rmsudo(chat_id)
+            if DB_URI is not None:
+                msg = DbManger().db_rmsudo(user_id)
+            else:
+                SUDO_USERS.remove(user_id)
+                msg = 'Demoted'
         else:
             msg = 'Not a Sudo'
     else:
@@ -108,9 +158,18 @@ def removeSudo(update, context):
         else:
             user_id = reply_message.from_user.id
             if user_id in SUDO_USERS:
-                msg = DbManger().db_rmsudo(user_id)
+                if DB_URI is not None:
+                    msg = DbManger().db_rmsudo(user_id)
+                else:
+                    SUDO_USERS.remove(user_id)
+                    msg = 'Demoted'
             else:
                 msg = 'Not a Sudo'
+    if DB_URI is None:
+        with open('sudo_users.txt', 'a') as file:
+            file.truncate(0)
+            for i in SUDO_USERS:
+                file.write(f'{i}\n')
     sendMessage(msg, context.bot, update)
 
 

config_sample.env

@@ -11,11 +11,12 @@ AUTO_DELETE_MESSAGE_DURATION = 20
 IS_TEAM_DRIVE = ""
 TELEGRAM_API = 
 TELEGRAM_HASH = ""
-DATABASE_URL = ""
 UPSTREAM_REPO = "https://github.com/breakdowns/slam-mirrorbot"
 UPSTREAM_BRANCH = "master"
 # OPTIONAL CONFIG
-AUTHORIZED_CHATS = ""
+DATABASE_URL = ""
+AUTHORIZED_CHATS = ""    #split by space
+SUDO_USERS = ""    #split by space
 IGNORE_PENDING_REQUESTS = ""
 USE_SERVICE_ACCOUNTS = ""
 INDEX_URL = ""