From 6f8b39cda329eab849787d35dd17925836a9aceb Mon Sep 17 00:00:00 2001 From: anas Date: Tue, 13 Jul 2021 17:01:09 +0300 Subject: [PATCH] Optional Database and minor changes -force remove tasks from aria2 -purge completed,removed and failed downloads to free memory Signed-off-by: anas --- bot/__init__.py | 64 ++++++---- .../download_utils/aria2_download.py | 28 +++-- .../download_utils/mega_downloader.py | 2 +- .../status_utils/aria_download_status.py | 8 +- .../mirror_utils/upload_utils/gdriveTools.py | 4 +- bot/modules/authorize.py | 115 +++++++++++++----- config_sample.env | 5 +- 7 files changed, 158 insertions(+), 68 deletions(-) diff --git a/bot/__init__.py b/bot/__init__.py index 3371b9b..678b288 100644 --- a/bot/__init__.py +++ b/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) - 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]) -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() + DB_URI = getConfig('DATABASE_URL') + if len(DB_URI) == 0: + raise KeyError +except KeyError: + logging.warning('Database not provided!') + DB_URI = None +if DB_URI is not None: + 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]) + 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) diff --git a/bot/helper/mirror_utils/download_utils/aria2_download.py b/bot/helper/mirror_utils/download_utils/aria2_download.py index e9aaba3..1fcb701 100644 --- a/bot/helper/mirror_utils/download_utils/aria2_download.py +++ b/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 - sname = download.name + sleep(1) + 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') - aria2.remove([download]) + dl.getListener().onDownloadError(f'File/Folder already available in Drive.\n\n') + 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, diff --git a/bot/helper/mirror_utils/download_utils/mega_downloader.py b/bot/helper/mirror_utils/download_utils/mega_downloader.py index 3ddd730..b807418 100644 --- a/bot/helper/mirror_utils/download_utils/mega_downloader.py +++ b/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: diff --git a/bot/helper/mirror_utils/status_utils/aria_download_status.py b/bot/helper/mirror_utils/status_utils/aria_download_status.py index 1d0c00f..394319d 100644 --- a/bot/helper/mirror_utils/status_utils/aria_download_status.py +++ b/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() diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 6a2ce04..b2fab5d 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/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: diff --git a/bot/modules/authorize.py b/bot/modules/authorize.py index be1bd50..c20cc93 100644 --- a/bot/modules/authorize.py +++ b/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]) - if chat_id not in AUTHORIZED_CHATS: - msg = DbManger().db_auth(chat_id) + user_id = int(message_[1]) + if user_id not in AUTHORIZED_CHATS: + 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]) - if chat_id in AUTHORIZED_CHATS: - msg = DbManger().db_unauth(chat_id) + user_id = int(message_[1]) + if user_id in AUTHORIZED_CHATS: + 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]) - if chat_id not in SUDO_USERS: - msg = DbManger().db_addsudo(chat_id) + user_id = int(message_[1]) + if user_id not in SUDO_USERS: + 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]) - if chat_id in SUDO_USERS: - msg = DbManger().db_rmsudo(chat_id) + user_id = int(message_[1]) + if user_id in SUDO_USERS: + 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) diff --git a/config_sample.env b/config_sample.env index d0061cb..58eaf8b 100644 --- a/config_sample.env +++ b/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 = ""