From 4f0a312632bd4b914c5b45f4a15675f336056812 Mon Sep 17 00:00:00 2001 From: Levi <57452819+l3v11@users.noreply.github.com> Date: Sun, 31 Jul 2022 22:37:57 +0600 Subject: [PATCH] Add archival and extraction progress in compress module - This adds the ability to show the progress of archival and extraction tasks - Also adds the ability to cancel archival or extraction task - Tidy up --- bot/__init__.py | 4 +- bot/__main__.py | 1 + bot/helper/download_utils/__init__.py | 1 + bot/helper/drive_utils/gdriveTools.py | 72 ++++++--------- bot/helper/ext_utils/bot_utils.py | 41 +++++---- bot/helper/ext_utils/database.py | 3 +- bot/helper/ext_utils/fs_utils.py | 84 +++-------------- bot/helper/status_utils/__init__.py | 1 + bot/helper/status_utils/archive_status.py | 52 ++++++++--- bot/helper/status_utils/clone_status.py | 6 +- bot/helper/status_utils/download_status.py | 11 +-- bot/helper/status_utils/extract_status.py | 52 ++++++++--- bot/helper/status_utils/upload_status.py | 5 -- bot/helper/telegram_helper/filters.py | 8 +- bot/helper/telegram_helper/message_utils.py | 65 ++++++++++---- bot/modules/__init__.py | 1 + bot/modules/auth.py | 60 ++++++------- bot/modules/cancel.py | 19 ++-- bot/modules/clone.py | 12 +-- bot/modules/compress.py | 99 +++++++++++---------- bot/modules/count.py | 9 +- bot/modules/delete.py | 9 +- bot/modules/eval.py | 2 +- bot/modules/list.py | 6 +- bot/modules/permission.py | 14 +-- bot/modules/shell.py | 16 ++-- bot/modules/status.py | 35 +++++--- 27 files changed, 347 insertions(+), 341 deletions(-) create mode 100644 bot/helper/download_utils/__init__.py create mode 100644 bot/helper/status_utils/__init__.py create mode 100644 bot/modules/__init__.py diff --git a/bot/__init__.py b/bot/__init__.py index 66e97f66..cfa201d6 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -69,9 +69,9 @@ def get_config(name: str): try: users = get_config('AUTHORIZED_CHATS') - users = users.split(" ") + users = users.split() for user in users: - AUTHORIZED_CHATS.add(int(user)) + AUTHORIZED_CHATS.add(int(user.strip())) except: pass diff --git a/bot/__main__.py b/bot/__main__.py index 00a0eef3..4ac260e2 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -67,6 +67,7 @@ def restart(update, context): restart_message = sendMessage("Restart in progress...", context.bot, update.message) if Interval: Interval[0].cancel() + Interval.clear() clean_all() with open(".restartmsg", "w") as f: f.truncate(0) diff --git a/bot/helper/download_utils/__init__.py b/bot/helper/download_utils/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/bot/helper/download_utils/__init__.py @@ -0,0 +1 @@ + diff --git a/bot/helper/drive_utils/gdriveTools.py b/bot/helper/drive_utils/gdriveTools.py index 1b3647fe..d33437ef 100644 --- a/bot/helper/drive_utils/gdriveTools.py +++ b/bot/helper/drive_utils/gdriveTools.py @@ -9,7 +9,7 @@ from urllib.parse import parse_qs, urlparse from random import randrange from tenacity import retry, wait_exponential, stop_after_attempt, \ - retry_if_exception_type, before_log, RetryError + retry_if_exception_type, RetryError from telegram import InlineKeyboardMarkup from telegraph.exceptions import RetryAfterError @@ -18,7 +18,7 @@ from google.oauth2 import service_account from google.oauth2.credentials import Credentials from googleapiclient.discovery import build -from googleapiclient.errors import HttpError +from googleapiclient.errors import Error as GCError, HttpError from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload from bot import LOGGER, DOWNLOAD_DIR, DRIVE_NAMES, DRIVE_IDS, INDEX_URLS, PARENT_ID, \ @@ -80,7 +80,7 @@ def speed(self): """ try: return self.uploaded_bytes / self.total_time - except ZeroDivisionError: + except: return 0 def dspeed(self): @@ -90,7 +90,7 @@ def dspeed(self): """ try: return self.downloaded_bytes / self.dtotal_time - except ZeroDivisionError: + except: return 0 def cspeed(self): @@ -100,7 +100,7 @@ def cspeed(self): """ try: return self.transferred_size / int(time.time() - self.start_time) - except ZeroDivisionError: + except: return 0 def authorize(self): @@ -202,7 +202,7 @@ def __set_permission_email(self, file_id, email): body=permissions, sendNotificationEmail=False).execute() - def setPerm(self, link, access): + def setPermission(self, link, access): try: file_id = self.getIdFromUrl(link) except (KeyError, IndexError): @@ -226,7 +226,7 @@ def setPerm(self, link, access): token_service = self.alt_authorize() if token_service is not None: self.__service = token_service - return self.setPerm(link, access) + return self.setPermission(link, access) msg = "Insufficient file permissions" else: msg = str(err) @@ -234,8 +234,7 @@ def setPerm(self, link, access): @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(3), - retry=retry_if_exception_type(HttpError), - before=before_log(LOGGER, logging.DEBUG)) + retry=retry_if_exception_type(GCError)) def copyFile(self, file_id, dest_id): body = { 'parents': [dest_id] @@ -266,8 +265,7 @@ def copyFile(self, file_id, dest_id): @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(3), - retry=retry_if_exception_type(HttpError), - before=before_log(LOGGER, logging.DEBUG)) + retry=retry_if_exception_type(GCError)) def getFileMetadata(self, file_id): return self.__service.files().get( supportsAllDrives=True, @@ -276,8 +274,7 @@ def getFileMetadata(self, file_id): @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(3), - retry=retry_if_exception_type(HttpError), - before=before_log(LOGGER, logging.DEBUG)) + retry=retry_if_exception_type(GCError)) def getFilesByFolderId(self, folder_id): page_token = None query = f"'{folder_id}' in parents and trashed = false" @@ -357,10 +354,9 @@ def clone(self, link, key): LOGGER.info(f"Total attempts: {err.last_attempt.attempt_number}") err = err.last_attempt.exception() err = str(err).replace('>', '').replace('<', '') - LOGGER.error(err) - if "User rate limit exceeded" in str(err): + if "User rate limit exceeded" in err: msg = "User rate limit exceeded" - elif "File not found" in str(err): + elif "File not found" in err: token_service = self.alt_authorize() if token_service is not None: self.__service = token_service @@ -389,8 +385,7 @@ def cloneFolder(self, name, local_path, folder_id, parent_id): @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(3), - retry=retry_if_exception_type(HttpError), - before=before_log(LOGGER, logging.DEBUG)) + retry=retry_if_exception_type(GCError)) def create_directory(self, directory_name, parent_id): file_metadata = { "name": directory_name, @@ -437,8 +432,7 @@ def count(self, link): LOGGER.info(f"Total attempts: {err.last_attempt.attempt_number}") err = err.last_attempt.exception() err = str(err).replace('>', '').replace('<', '') - LOGGER.error(err) - if "File not found" in str(err): + if "File not found" in err: token_service = self.alt_authorize() if token_service is not None: self.__service = token_service @@ -493,8 +487,7 @@ def helper(self, link): LOGGER.info(f"Total attempts: {err.last_attempt.attempt_number}") err = err.last_attempt.exception() err = str(err).replace('>', '').replace('<', '') - LOGGER.error(err) - if "File not found" in str(err): + if "File not found" in err: token_service = self.alt_authorize() if token_service is not None: self.__service = token_service @@ -687,13 +680,10 @@ def upload(self, file_name: str): link = f"https://drive.google.com/folderview?id={dir_id}" if self.is_cancelled: return - except Exception as e: - if isinstance(e, RetryError): - LOGGER.info(f"Total attempts: {e.last_attempt.attempt_number}") - err = e.last_attempt.exception() - else: - err = e - LOGGER.error(err) + except Exception as err: + if isinstance(err, RetryError): + LOGGER.info(f"Total attempts: {err.last_attempt.attempt_number}") + err = err.last_attempt.exception() self.__listener.onUploadError(str(err)) self.is_errored = True finally: @@ -732,8 +722,7 @@ def upload_dir(self, input_directory, parent_id): @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(3), - retry=(retry_if_exception_type(HttpError) | retry_if_exception_type(IOError)), - before=before_log(LOGGER, logging.DEBUG)) + retry=(retry_if_exception_type(GCError) | retry_if_exception_type(IOError))) def upload_file(self, file_path, file_name, mime_type, parent_id): file_metadata = { 'name': file_name, @@ -788,10 +777,6 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): download_url = self.__G_DRIVE_BASE_DOWNLOAD_URL.format(drive_file.get('id')) return download_url - @retry(wait=wait_exponential(multiplier=2, min=3, max=6), - stop=stop_after_attempt(3), - retry=retry_if_exception_type(HttpError), - before=before_log(LOGGER, logging.DEBUG)) def _on_upload_progress(self): if self.status is not None: chunk_size = self.status.total_size * self.status.progress() - self._file_uploaded_bytes @@ -816,10 +801,9 @@ def download(self, link): LOGGER.info(f"Total attempts: {err.last_attempt.attempt_number}") err = err.last_attempt.exception() err = str(err).replace('>', '').replace('<', '') - LOGGER.error(err) - if "downloadQuotaExceeded" in str(err): + if "downloadQuotaExceeded" in err: err = "Download quota exceeded." - elif "File not found" in str(err): + elif "File not found" in err: token_service = self.alt_authorize() if token_service is not None: self.__service = token_service @@ -860,11 +844,15 @@ def download_folder(self, folder_id, path, folder_name): @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(3), - retry=(retry_if_exception_type(HttpError) | retry_if_exception_type(IOError)), - before=before_log(LOGGER, logging.DEBUG)) + retry=(retry_if_exception_type(GCError) | retry_if_exception_type(IOError))) def download_file(self, file_id, path, filename, mime_type): request = self.__service.files().get_media(fileId=file_id) filename = filename.replace('/', '') + if len(filename.encode()) > 255: + ext = os.path.splitext(filename)[1] + filename = filename[:245] + ext + if self.name.endswith(ext): + self.name = filename fh = FileIO('{}{}'.format(path, filename), 'wb') downloader = MediaIoBaseDownload(fh, request, chunksize=50 * 1024 * 1024) done = False @@ -891,10 +879,6 @@ def download_file(self, file_id, path, filename, mime_type): raise err self._file_downloaded_bytes = 0 - @retry(wait=wait_exponential(multiplier=2, min=3, max=6), - stop=stop_after_attempt(3), - retry=retry_if_exception_type(HttpError), - before=before_log(LOGGER, logging.DEBUG)) def _on_download_progress(self): if self.dstatus is not None: chunk_size = self.dstatus.total_size * self.dstatus.progress() - self._file_downloaded_bytes diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index 4457f2d1..fcae93f8 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/bot/helper/ext_utils/bot_utils.py @@ -13,9 +13,9 @@ URL_REGEX = r'(?:(?:https?|ftp):\/\/)?[\w/\-?=%.]+\.[\w/\-?=%.]+' class TaskStatus: - STATUS_UPLOADING = "Uploading...📤" - STATUS_DOWNLOADING = "Downloading...📥" STATUS_CLONING = "Cloning...♻️" + STATUS_DOWNLOADING = "Downloading...📥" + STATUS_UPLOADING = "Uploading...📤" STATUS_ARCHIVING = "Archiving...🔐" STATUS_EXTRACTING = "Extracting...📂" @@ -39,11 +39,8 @@ def cancel(self): def getDownloadByGid(gid): with download_dict_lock: for dl in list(download_dict.values()): - status = dl.status() - if status not in [TaskStatus.STATUS_ARCHIVING, - TaskStatus.STATUS_EXTRACTING]: - if dl.gid() == gid: - return dl + if dl.gid() == gid: + return dl return None def get_progress_bar_string(status): @@ -63,21 +60,23 @@ def get_readable_message(): for download in list(download_dict.values()): msg += f"Name: {escape(str(download.name()))}" msg += f"\nStatus: {download.status()}" - if download.status() not in [TaskStatus.STATUS_ARCHIVING, - TaskStatus.STATUS_EXTRACTING]: - msg += f"\n{get_progress_bar_string(download)} {download.progress()}" - if download.status() == TaskStatus.STATUS_CLONING: - msg += f"\nCloned: {get_readable_file_size(download.processed_bytes())} / {download.size()}" - msg += f"\nTransfers: {download.processed_files()} / {download.files()}" - elif download.status() == TaskStatus.STATUS_UPLOADING: - msg += f"\nUploaded: {get_readable_file_size(download.processed_bytes())} / {download.size()}" - else: - msg += f"\nDownloaded: {get_readable_file_size(download.processed_bytes())} / {download.size()}" - msg += f"\nSpeed: {download.speed()} | ETA: {download.eta()}" - msg += f"\n/{BotCommands.CancelCommand} {download.gid()}" - else: - msg += f"\nSize: {download.size()}" + msg += f"\n{get_progress_bar_string(download)} {download.progress()}" + if download.status() == TaskStatus.STATUS_CLONING: + msg += f"\nCloned: {get_readable_file_size(download.processed_bytes())} / {download.size()}" + msg += f"\nTransfers: {download.processed_files()} / {download.files()}" + elif download.status() == TaskStatus.STATUS_DOWNLOADING: + msg += f"\nDownloaded: {get_readable_file_size(download.processed_bytes())} / {download.size()}" + elif download.status() == TaskStatus.STATUS_UPLOADING: + msg += f"\nUploaded: {get_readable_file_size(download.processed_bytes())} / {download.size()}" + elif download.status() == TaskStatus.STATUS_ARCHIVING: + msg += f"\nArchived: {get_readable_file_size(download.processed_bytes())} / {download.size()}" + elif download.status() == TaskStatus.STATUS_EXTRACTING: + msg += f"\nExtracted: {get_readable_file_size(download.processed_bytes())} / {download.size()}" + msg += f"\nSpeed: {download.speed()} | ETA: {download.eta()}" + msg += f"\n/{BotCommands.CancelCommand} {download.gid()}" msg += "\n\n" + if len(msg) == 0: + return None cpu = cpu_percent(interval=0.5) ram = virtual_memory().percent disk = disk_usage('/').percent diff --git a/bot/helper/ext_utils/database.py b/bot/helper/ext_utils/database.py index a753af34..215fa4df 100644 --- a/bot/helper/ext_utils/database.py +++ b/bot/helper/ext_utils/database.py @@ -24,5 +24,4 @@ def load_users(self): AUTHORIZED_CHATS.add(user['user_id']) if DATABASE_URL is not None: - db = DatabaseHelper() - db.load_users() + DatabaseHelper().load_users() diff --git a/bot/helper/ext_utils/fs_utils.py b/bot/helper/ext_utils/fs_utils.py index d5afe062..80e77b84 100644 --- a/bot/helper/ext_utils/fs_utils.py +++ b/bot/helper/ext_utils/fs_utils.py @@ -1,11 +1,17 @@ import magic import os +import re import shutil import sys from bot import LOGGER, DOWNLOAD_DIR from bot.helper.ext_utils.exceptions import CompressExceptionHandler +ARCH_EXT = [".tar.bz2", ".tar.gz", ".bz2", ".gz", ".tar.xz", ".tar", ".tbz2", ".tgz", ".lzma2", + ".zip", ".7z", ".z", ".rar", ".iso", ".wim", ".cab", ".apm", ".arj", ".chm", + ".cpio", ".cramfs", ".deb", ".dmg", ".fat", ".hfs", ".lzh", ".lzma", ".mbr", + ".msi", ".mslz", ".nsis", ".ntfs", ".rpm", ".squashfs", ".udf", ".vhd", ".xar"] + def clean_download(path: str): if os.path.exists(path): LOGGER.info(f"Cleaning: {path}") @@ -47,80 +53,10 @@ def get_path_size(path: str): return total_size def get_base_name(orig_path: str): - if orig_path.endswith(".tar.bz2"): - return orig_path.rsplit(".tar.bz2", 1)[0] - elif orig_path.endswith(".tar.gz"): - return orig_path.rsplit(".tar.gz", 1)[0] - elif orig_path.endswith(".bz2"): - return orig_path.rsplit(".bz2", 1)[0] - elif orig_path.endswith(".gz"): - return orig_path.rsplit(".gz", 1)[0] - elif orig_path.endswith(".tar.xz"): - return orig_path.rsplit(".tar.xz", 1)[0] - elif orig_path.endswith(".tar"): - return orig_path.rsplit(".tar", 1)[0] - elif orig_path.endswith(".tbz2"): - return orig_path.rsplit("tbz2", 1)[0] - elif orig_path.endswith(".tgz"): - return orig_path.rsplit(".tgz", 1)[0] - elif orig_path.endswith(".zip"): - return orig_path.rsplit(".zip", 1)[0] - elif orig_path.endswith(".7z"): - return orig_path.rsplit(".7z", 1)[0] - elif orig_path.endswith(".Z"): - return orig_path.rsplit(".Z", 1)[0] - elif orig_path.endswith(".rar"): - return orig_path.rsplit(".rar", 1)[0] - elif orig_path.endswith(".iso"): - return orig_path.rsplit(".iso", 1)[0] - elif orig_path.endswith(".wim"): - return orig_path.rsplit(".wim", 1)[0] - elif orig_path.endswith(".cab"): - return orig_path.rsplit(".cab", 1)[0] - elif orig_path.endswith(".apm"): - return orig_path.rsplit(".apm", 1)[0] - elif orig_path.endswith(".arj"): - return orig_path.rsplit(".arj", 1)[0] - elif orig_path.endswith(".chm"): - return orig_path.rsplit(".chm", 1)[0] - elif orig_path.endswith(".cpio"): - return orig_path.rsplit(".cpio", 1)[0] - elif orig_path.endswith(".cramfs"): - return orig_path.rsplit(".cramfs", 1)[0] - elif orig_path.endswith(".deb"): - return orig_path.rsplit(".deb", 1)[0] - elif orig_path.endswith(".dmg"): - return orig_path.rsplit(".dmg", 1)[0] - elif orig_path.endswith(".fat"): - return orig_path.rsplit(".fat", 1)[0] - elif orig_path.endswith(".hfs"): - return orig_path.rsplit(".hfs", 1)[0] - elif orig_path.endswith(".lzh"): - return orig_path.rsplit(".lzh", 1)[0] - elif orig_path.endswith(".lzma"): - return orig_path.rsplit(".lzma", 1)[0] - elif orig_path.endswith(".lzma2"): - return orig_path.rsplit(".lzma2", 1)[0] - elif orig_path.endswith(".mbr"): - return orig_path.rsplit(".mbr", 1)[0] - elif orig_path.endswith(".msi"): - return orig_path.rsplit(".msi", 1)[0] - elif orig_path.endswith(".mslz"): - return orig_path.rsplit(".mslz", 1)[0] - elif orig_path.endswith(".nsis"): - return orig_path.rsplit(".nsis", 1)[0] - elif orig_path.endswith(".ntfs"): - return orig_path.rsplit(".ntfs", 1)[0] - elif orig_path.endswith(".rpm"): - return orig_path.rsplit(".rpm", 1)[0] - elif orig_path.endswith(".squashfs"): - return orig_path.rsplit(".squashfs", 1)[0] - elif orig_path.endswith(".udf"): - return orig_path.rsplit(".udf", 1)[0] - elif orig_path.endswith(".vhd"): - return orig_path.rsplit(".vhd", 1)[0] - elif orig_path.endswith(".xar"): - return orig_path.rsplit(".xar", 1)[0] + ext = [ext for ext in ARCH_EXT if orig_path.lower().endswith(ext)] + if len(ext) > 0: + ext = ext[0] + return re.split(ext + '$', orig_path, maxsplit=1, flags=re.I)[0] else: raise CompressExceptionHandler('Unsupported file format') diff --git a/bot/helper/status_utils/__init__.py b/bot/helper/status_utils/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/bot/helper/status_utils/__init__.py @@ -0,0 +1 @@ + diff --git a/bot/helper/status_utils/archive_status.py b/bot/helper/status_utils/archive_status.py index e29c02a2..2b58e997 100644 --- a/bot/helper/status_utils/archive_status.py +++ b/bot/helper/status_utils/archive_status.py @@ -1,34 +1,64 @@ -from bot.helper.ext_utils.bot_utils import TaskStatus, get_readable_file_size +import time + +from bot import LOGGER, DOWNLOAD_DIR +from bot.helper.ext_utils.bot_utils import TaskStatus, get_readable_file_size, get_readable_time +from bot.helper.ext_utils.fs_utils import get_path_size class ArchiveStatus: - def __init__(self, name, path, size): + def __init__(self, name, size, gid, listener): self.__name = name - self.__path = path self.__size = size + self.__gid = gid + self.__listener = listener + self.__uid = listener.uid + self.__start_time = time.time() + self.message = listener.message + + def gid(self) -> str: + return self.__gid - # The progress of the archive function cannot be tracked. So we just return - # dummy values. If this is possible in the future, we should implement it. + def progress_raw(self): + try: + return self.processed_bytes() / self.__size * 100 + except: + return 0 def progress(self): - return '0' + return f'{round(self.progress_raw(), 2)}%' + + def speed_raw(self): + return self.processed_bytes() / (time.time() - self.__start_time) def speed(self): - return '0' + return f'{get_readable_file_size(self.speed_raw())}/s' def name(self): return self.__name - def path(self): - return self.__path + def size_raw(self): + return self.__size def size(self): return get_readable_file_size(self.__size) def eta(self): - return '0s' + try: + seconds = (self.size_raw() - self.processed_bytes()) / self.speed_raw() + return f'{get_readable_time(seconds)}' + except: + return '-' def status(self): return TaskStatus.STATUS_ARCHIVING def processed_bytes(self): - return 0 + return get_path_size(f"{DOWNLOAD_DIR}{self.__uid}") - self.__size + + def download(self): + return self + + def cancel_task(self): + LOGGER.info(f"Cancelling archive: {self.__name}") + if self.__listener.suproc is not None: + self.__listener.suproc.kill() + self.__listener.onUploadError("The archive task has been cancelled") diff --git a/bot/helper/status_utils/clone_status.py b/bot/helper/status_utils/clone_status.py index bf4f0be4..bd5f4723 100644 --- a/bot/helper/status_utils/clone_status.py +++ b/bot/helper/status_utils/clone_status.py @@ -6,8 +6,8 @@ def __init__(self, obj, size, files, message, gid): self.__obj = obj self.__size = size self.__files = files - self.message = message self.__gid = gid + self.message = message def processed_bytes(self): return self.__obj.transferred_size @@ -36,7 +36,7 @@ def gid(self) -> str: def progress_raw(self): try: return self.__obj.transferred_size / self.__size * 100 - except ZeroDivisionError: + except: return 0 def progress(self): @@ -55,7 +55,7 @@ def eta(self): try: seconds = (self.__size - self.__obj.transferred_size) / self.speed_raw() return f'{get_readable_time(seconds)}' - except ZeroDivisionError: + except: return '-' def download(self): diff --git a/bot/helper/status_utils/download_status.py b/bot/helper/status_utils/download_status.py index 8201470b..9f3cfdc4 100644 --- a/bot/helper/status_utils/download_status.py +++ b/bot/helper/status_utils/download_status.py @@ -1,16 +1,11 @@ -from bot import DOWNLOAD_DIR from bot.helper.ext_utils.bot_utils import TaskStatus, get_readable_file_size, get_readable_time class DownloadStatus: def __init__(self, obj, size, listener, gid): self.__obj = obj self.__size = size - self.__uid = listener.uid - self.message = listener.message self.__gid = gid - - def path(self): - return f"{DOWNLOAD_DIR}{self.__uid}" + self.message = listener.message def processed_bytes(self): return self.__obj.downloaded_bytes @@ -33,7 +28,7 @@ def gid(self) -> str: def progress_raw(self): try: return self.__obj.downloaded_bytes / self.__size * 100 - except ZeroDivisionError: + except: return 0 def progress(self): @@ -52,7 +47,7 @@ def eta(self): try: seconds = (self.__size - self.__obj.downloaded_bytes) / self.speed_raw() return f'{get_readable_time(seconds)}' - except ZeroDivisionError: + except: return '-' def download(self): diff --git a/bot/helper/status_utils/extract_status.py b/bot/helper/status_utils/extract_status.py index d47beeb4..5d4744ae 100644 --- a/bot/helper/status_utils/extract_status.py +++ b/bot/helper/status_utils/extract_status.py @@ -1,34 +1,64 @@ -from bot.helper.ext_utils.bot_utils import TaskStatus, get_readable_file_size +import time + +from bot import LOGGER, DOWNLOAD_DIR +from bot.helper.ext_utils.bot_utils import TaskStatus, get_readable_file_size, get_readable_time +from bot.helper.ext_utils.fs_utils import get_path_size class ExtractStatus: - def __init__(self, name, path, size): + def __init__(self, name, size, gid, listener): self.__name = name - self.__path = path self.__size = size + self.__gid = gid + self.__listener = listener + self.__uid = listener.uid + self.__start_time = time.time() + self.message = listener.message + + def gid(self) -> str: + return self.__gid - # The progress of the extract function cannot be tracked. So we just return - # dummy values. If this is possible in the future, we should implement it. + def progress_raw(self): + try: + return self.processed_bytes() / self.__size * 100 + except: + return 0 def progress(self): - return '0' + return f'{round(self.progress_raw(), 2)}%' + + def speed_raw(self): + return self.processed_bytes() / (time.time() - self.__start_time) def speed(self): - return '0' + return f'{get_readable_file_size(self.speed_raw())}/s' def name(self): return self.__name - def path(self): - return self.__path + def size_raw(self): + return self.__size def size(self): return get_readable_file_size(self.__size) def eta(self): - return '0s' + try: + seconds = (self.size_raw() - self.processed_bytes()) / self.speed_raw() + return f'{get_readable_time(seconds)}' + except: + return '-' def status(self): return TaskStatus.STATUS_EXTRACTING def processed_bytes(self): - return 0 + return get_path_size(f"{DOWNLOAD_DIR}{self.__uid}") - self.__size + + def download(self): + return self + + def cancel_task(self): + LOGGER.info(f"Cancelling extract: {self.__name}") + if self.__listener.suproc is not None: + self.__listener.suproc.kill() + self.__listener.onUploadError("The extract task has been cancelled") diff --git a/bot/helper/status_utils/upload_status.py b/bot/helper/status_utils/upload_status.py index d9eb3b47..0df98e7c 100644 --- a/bot/helper/status_utils/upload_status.py +++ b/bot/helper/status_utils/upload_status.py @@ -1,17 +1,12 @@ -from bot import DOWNLOAD_DIR from bot.helper.ext_utils.bot_utils import TaskStatus, get_readable_file_size, get_readable_time class UploadStatus: def __init__(self, obj, size, gid, listener): self.__obj = obj self.__size = size - self.__uid = listener.uid self.__gid = gid self.message = listener.message - def path(self): - return f"{DOWNLOAD_DIR}{self.__uid}" - def processed_bytes(self): return self.__obj.uploaded_bytes diff --git a/bot/helper/telegram_helper/filters.py b/bot/helper/telegram_helper/filters.py index ea525a86..5e1a7fe5 100644 --- a/bot/helper/telegram_helper/filters.py +++ b/bot/helper/telegram_helper/filters.py @@ -6,19 +6,19 @@ class CustomFilters: class __OwnerFilter(MessageFilter): def filter(self, message: Message): - return bool(message.from_user.id == OWNER_ID) + return message.from_user.id == OWNER_ID owner_filter = __OwnerFilter() class __AuthorizedUserFilter(MessageFilter): def filter(self, message: Message): - id = message.from_user.id - return bool(id in AUTHORIZED_CHATS or id == OWNER_ID) + uid = message.from_user.id + return uid in AUTHORIZED_CHATS or uid == OWNER_ID authorized_user = __AuthorizedUserFilter() class __AuthorizedChat(MessageFilter): def filter(self, message: Message): - return bool(message.chat.id in AUTHORIZED_CHATS) + return message.chat.id in AUTHORIZED_CHATS authorized_chat = __AuthorizedChat() diff --git a/bot/helper/telegram_helper/message_utils.py b/bot/helper/telegram_helper/message_utils.py index 5109ae2b..2139b54c 100644 --- a/bot/helper/telegram_helper/message_utils.py +++ b/bot/helper/telegram_helper/message_utils.py @@ -1,4 +1,7 @@ +import time + from telegram import InlineKeyboardMarkup +from telegram.error import RetryAfter from telegram.message import Message from bot import bot, LOGGER, Interval, STATUS_UPDATE_INTERVAL, \ @@ -11,8 +14,13 @@ def sendMessage(text: str, bot, message: Message): reply_to_message_id=message.message_id, text=text, parse_mode='HTMl', disable_web_page_preview=True) + except RetryAfter as r: + LOGGER.warning(str(r)) + time.sleep(r.retry_after * 1.5) + return sendMessage(text, bot, message) except Exception as e: LOGGER.error(str(e)) + return def sendMarkup(text: str, bot, message: Message, reply_markup: InlineKeyboardMarkup): try: @@ -20,8 +28,13 @@ def sendMarkup(text: str, bot, message: Message, reply_markup: InlineKeyboardMar reply_to_message_id=message.message_id, text=text, reply_markup=reply_markup, parse_mode='HTMl', disable_web_page_preview=True) + except RetryAfter as r: + LOGGER.warning(str(r)) + time.sleep(r.retry_after * 1.5) + return sendMarkup(text, bot, message, reply_markup) except Exception as e: LOGGER.error(str(e)) + return def editMessage(text: str, message: Message, reply_markup=None): try: @@ -29,8 +42,13 @@ def editMessage(text: str, message: Message, reply_markup=None): chat_id=message.chat.id, reply_markup=reply_markup, parse_mode='HTMl', disable_web_page_preview=True) + except RetryAfter as r: + LOGGER.warning(str(r)) + time.sleep(r.retry_after * 1.5) + return editMessage(text, message, reply_markup) except Exception as e: LOGGER.error(str(e)) + return str(e) def deleteMessage(bot, message: Message): try: @@ -47,33 +65,42 @@ def sendLogFile(bot, message: Message): def delete_all_messages(): with status_reply_dict_lock: - for message in list(status_reply_dict.values()): + for data in list(status_reply_dict.values()): try: - deleteMessage(bot, message) - del status_reply_dict[message.chat.id] + deleteMessage(bot, data[0]) + del status_reply_dict[data[0].chat.id] except Exception as e: LOGGER.error(str(e)) -def update_all_messages(): +def update_all_messages(force=False): + with status_reply_dict_lock: + if not force and (not status_reply_dict or not Interval or time.time() - list(status_reply_dict.values())[0][1] < 3): + return + for chat_id in status_reply_dict: + status_reply_dict[chat_id][1] = time.time() msg = get_readable_message() + if msg is None: + return with status_reply_dict_lock: - for chat_id in list(status_reply_dict.keys()): - if status_reply_dict[chat_id] and msg != status_reply_dict[chat_id].text: - editMessage(msg, status_reply_dict[chat_id]) - status_reply_dict[chat_id].text = msg + for chat_id in status_reply_dict: + if status_reply_dict[chat_id] and msg != status_reply_dict[chat_id][0].text: + rmsg = editMessage(msg, status_reply_dict[chat_id][0]) + if rmsg == "Message to edit not found": + del status_reply_dict[chat_id] + return + status_reply_dict[chat_id][0].text = msg + status_reply_dict[chat_id][1] = time.time() def sendStatusMessage(msg, bot): - if len(Interval) == 0: - Interval.append(SetInterval(STATUS_UPDATE_INTERVAL, update_all_messages)) progress = get_readable_message() + if progress is None: + return with status_reply_dict_lock: - if msg.chat.id in list(status_reply_dict.keys()): - try: - message = status_reply_dict[msg.chat.id] - deleteMessage(bot, message) - del status_reply_dict[msg.chat.id] - except Exception as e: - LOGGER.error(str(e)) - del status_reply_dict[msg.chat.id] + if msg.chat.id in status_reply_dict: + message = status_reply_dict[msg.chat.id][0] + deleteMessage(bot, message) + del status_reply_dict[msg.chat.id] message = sendMessage(progress, bot, msg) - status_reply_dict[msg.chat.id] = message + status_reply_dict[msg.chat.id] = [message, time.time()] + if not Interval: + Interval.append(SetInterval(STATUS_UPDATE_INTERVAL, update_all_messages)) diff --git a/bot/modules/__init__.py b/bot/modules/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/bot/modules/__init__.py @@ -0,0 +1 @@ + diff --git a/bot/modules/auth.py b/bot/modules/auth.py index baadf99a..968d5bdf 100644 --- a/bot/modules/auth.py +++ b/bot/modules/auth.py @@ -7,13 +7,10 @@ from bot.helper.ext_utils.database import DatabaseHelper def authorize(update, context): - reply_message = None - message_ = None reply_message = update.message.reply_to_message - message_ = update.message.text.split(" ") - if len(message_) == 2: + if len(context.args) == 1: # Authorize an user in private - user_id = int(message_[1]) + user_id = int(context.args[0]) if user_id in AUTHORIZED_CHATS: msg = 'Already authorized' elif DATABASE_URL is not None: @@ -22,18 +19,7 @@ def authorize(update, context): else: AUTHORIZED_CHATS.add(user_id) msg = 'Authorization granted' - elif reply_message is None: - # Authorize a chat - chat_id = update.effective_chat.id - if chat_id in AUTHORIZED_CHATS: - msg = 'Already authorized' - elif DATABASE_URL is not None: - msg = DatabaseHelper().auth_user(chat_id) - AUTHORIZED_CHATS.add(chat_id) - else: - AUTHORIZED_CHATS.add(chat_id) - msg = 'Authorization granted' - else: + elif reply_message: # Authorize an user by replying user_id = reply_message.from_user.id if user_id in AUTHORIZED_CHATS: @@ -44,16 +30,24 @@ def authorize(update, context): else: AUTHORIZED_CHATS.add(user_id) msg = 'Authorization granted' + else: + # Authorize a chat + chat_id = update.effective_chat.id + if chat_id in AUTHORIZED_CHATS: + msg = 'Already authorized' + elif DATABASE_URL is not None: + msg = DatabaseHelper().auth_user(chat_id) + AUTHORIZED_CHATS.add(chat_id) + else: + AUTHORIZED_CHATS.add(chat_id) + msg = 'Authorization granted' sendMessage(msg, context.bot, update.message) def unauthorize(update, context): - reply_message = None - message_ = None reply_message = update.message.reply_to_message - message_ = update.message.text.split(" ") - if len(message_) == 2: + if len(context.args) == 1: # Unauthorize an user in private - user_id = int(message_[1]) + user_id = int(context.args[0]) if user_id in AUTHORIZED_CHATS: if DATABASE_URL is not None: msg = DatabaseHelper().unauth_user(user_id) @@ -62,26 +56,26 @@ def unauthorize(update, context): AUTHORIZED_CHATS.remove(user_id) else: msg = 'Already unauthorized' - elif reply_message is None: - # Unauthorize a chat - chat_id = update.effective_chat.id - if chat_id in AUTHORIZED_CHATS: + elif reply_message: + # Unauthorize an user by replying + user_id = reply_message.from_user.id + if user_id in AUTHORIZED_CHATS: if DATABASE_URL is not None: - msg = DatabaseHelper().unauth_user(chat_id) + msg = DatabaseHelper().unauth_user(user_id) else: msg = 'Authorization revoked' - AUTHORIZED_CHATS.remove(chat_id) + AUTHORIZED_CHATS.remove(user_id) else: msg = 'Already unauthorized' else: - # Unauthorize an user by replying - user_id = reply_message.from_user.id - if user_id in AUTHORIZED_CHATS: + # Unauthorize a chat + chat_id = update.effective_chat.id + if chat_id in AUTHORIZED_CHATS: if DATABASE_URL is not None: - msg = DatabaseHelper().unauth_user(user_id) + msg = DatabaseHelper().unauth_user(chat_id) else: msg = 'Authorization revoked' - AUTHORIZED_CHATS.remove(user_id) + AUTHORIZED_CHATS.remove(chat_id) else: msg = 'Already unauthorized' sendMessage(msg, context.bot, update.message) diff --git a/bot/modules/cancel.py b/bot/modules/cancel.py index 79d7dd9c..9ba4eb82 100644 --- a/bot/modules/cancel.py +++ b/bot/modules/cancel.py @@ -4,36 +4,29 @@ from bot.helper.telegram_helper.bot_commands import BotCommands from bot.helper.telegram_helper.filters import CustomFilters from bot.helper.telegram_helper.message_utils import sendMessage -from bot.helper.ext_utils.bot_utils import TaskStatus, getDownloadByGid +from bot.helper.ext_utils.bot_utils import getDownloadByGid def cancelNode(update, context): - args = update.message.text.split(" ", maxsplit=1) user_id = update.message.from_user.id - if len(args) > 1: - gid = args[1] + if len(context.args) == 1: + gid = context.args[0] dl = getDownloadByGid(gid) if not dl: return sendMessage(f"GID: {gid} not found", context.bot, update.message) elif update.message.reply_to_message: task_message = update.message.reply_to_message with download_dict_lock: - keys = list(download_dict.keys()) - if task_message.message_id in keys: + if task_message.message_id in download_dict: dl = download_dict[task_message.message_id] else: dl = None if not dl: return sendMessage("Not an active task", context.bot, update.message) - elif len(args) == 1: + elif len(context.args) == 0: return sendMessage("Send a GID along with command", context.bot, update.message) if OWNER_ID != user_id and dl.message.from_user.id != user_id: return sendMessage("Not your task", context.bot, update.message) - if dl.status() == TaskStatus.STATUS_ARCHIVING: - sendMessage("Archival in progress, the task cannot be cancelled", context.bot, update.message) - elif dl.status() == TaskStatus.STATUS_EXTRACTING: - sendMessage("Extraction in progress, the task cannot be cancelled", context.bot, update.message) - else: - dl.download().cancel_task() + dl.download().cancel_task() cancel_handler = CommandHandler(BotCommands.CancelCommand, cancelNode, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user, run_async=True) diff --git a/bot/modules/clone.py b/bot/modules/clone.py index ee190226..3e175e41 100644 --- a/bot/modules/clone.py +++ b/bot/modules/clone.py @@ -17,20 +17,20 @@ @new_thread def cloneNode(update, context): - args = update.message.text.split(" ", maxsplit=2) + args = update.message.text.split() reply_to = update.message.reply_to_message link = '' key = '' if len(args) > 1: - link = args[1] + link = args[1].strip() try: - key = args[2] + key = args[2].strip() except IndexError: pass - if reply_to is not None: - link = reply_to.text + if reply_to: + link = reply_to.text.split(maxsplit=1)[0].strip() try: - key = args[1] + key = args[1].strip() except IndexError: pass is_appdrive = is_appdrive_link(link) diff --git a/bot/modules/compress.py b/bot/modules/compress.py index 3aaf0589..7cd554f3 100644 --- a/bot/modules/compress.py +++ b/bot/modules/compress.py @@ -6,7 +6,6 @@ import threading from html import escape -from pathlib import PurePath from telegram.ext import CommandHandler from bot import LOGGER, dispatcher, DOWNLOAD_DIR, Interval, INDEX_URL, download_dict, download_dict_lock @@ -31,13 +30,14 @@ def __init__(self, bot, message, is_archive=False, is_extract=False, pswd=None): self.is_archive = is_archive self.is_extract = is_extract self.pswd = pswd + self.suproc = None def clean(self): try: Interval[0].cancel() - del Interval[0] + Interval.clear() delete_all_messages() - except IndexError: + except: pass def onDownloadComplete(self): @@ -45,33 +45,35 @@ def onDownloadComplete(self): download = download_dict[self.uid] name = str(download.name()).replace('/', '') gid = download.gid() - size = download.size_raw() m_path = f'{DOWNLOAD_DIR}{self.uid}/{name}' + size = get_path_size(m_path) if self.is_archive: - try: - with download_dict_lock: - download_dict[self.uid] = ArchiveStatus(name, m_path, size) - path = m_path + ".zip" - LOGGER.info(f"Archiving: {name}") - if self.pswd is not None: - subprocess.run(["7z", "a", "-mx=0", f"-p{self.pswd}", path, m_path]) - else: - subprocess.run(["7z", "a", "-mx=0", path, m_path]) - except FileNotFoundError: - LOGGER.info("File to archive not found") - self.onUploadError('Internal error') + path = m_path + ".zip" + with download_dict_lock: + download_dict[self.uid] = ArchiveStatus(name, size, gid, self) + LOGGER.info(f"Archiving: {name}") + if self.pswd is not None: + self.suproc = subprocess.Popen(["7z", "a", "-mx=0", f"-p{self.pswd}", path, m_path]) + else: + self.suproc = subprocess.Popen(["7z", "a", "-mx=0", path, m_path]) + self.suproc.wait() + if self.suproc.returncode == -9: return - try: - shutil.rmtree(m_path) - except: - os.remove(m_path) + elif self.suproc.returncode != 0: + LOGGER.error("Failed to archive the data") + path = f'{DOWNLOAD_DIR}{self.uid}/{name}' + if self.suproc.returncode == 0: + try: + shutil.rmtree(m_path) + except: + os.remove(m_path) elif self.is_extract: try: if os.path.isfile(m_path): path = get_base_name(m_path) LOGGER.info(f"Extracting: {name}") with download_dict_lock: - download_dict[self.uid] = ExtractStatus(name, m_path, size) + download_dict[self.uid] = ExtractStatus(name, size, gid, self) if os.path.isdir(m_path): for dirpath, subdir, files in os.walk(m_path, topdown=False): for file_ in files: @@ -79,22 +81,30 @@ def onDownloadComplete(self): or (file_.endswith(".rar") and not re.search(r'\.part\d+\.rar$', file_)): m_path = os.path.join(dirpath, file_) if self.pswd is not None: - result = subprocess.run(["7z", "x", f"-p{self.pswd}", m_path, f"-o{dirpath}", "-aot"]) + self.suproc = subprocess.Popen(["7z", "x", f"-p{self.pswd}", m_path, f"-o{dirpath}", "-aot"]) else: - result = subprocess.run(["7z", "x", m_path, f"-o{dirpath}", "-aot"]) - if result.returncode != 0: + self.suproc = subprocess.Popen(["7z", "x", m_path, f"-o{dirpath}", "-aot"]) + self.suproc.wait() + if self.suproc.returncode == -9: + return + elif self.suproc.returncode != 0: LOGGER.error("Failed to extract the archive") - for file_ in files: - if file_.endswith((".rar", ".zip", ".7z")) or re.search(r'\.r\d+$|\.7z\.\d+$|\.z\d+$|\.zip\.\d+$', file_): - del_path = os.path.join(dirpath, file_) - os.remove(del_path) + if self.suproc.returncode == 0: + for file_ in files: + if file_.endswith((".rar", ".zip", ".7z")) or \ + re.search(r'\.r\d+$|\.7z\.\d+$|\.z\d+$|\.zip\.\d+$', file_): + del_path = os.path.join(dirpath, file_) + os.remove(del_path) path = f'{DOWNLOAD_DIR}{self.uid}/{name}' else: if self.pswd is not None: - result = subprocess.run(["bash", "pextract", m_path, self.pswd]) + self.suproc = subprocess.Popen(["bash", "pextract", m_path, self.pswd]) else: - result = subprocess.run(["bash", "extract", m_path]) - if result.returncode == 0: + self.suproc = subprocess.Popen(["bash", "extract", m_path]) + self.suproc.wait() + if self.suproc.returncode == -9: + return + elif self.suproc.returncode == 0: os.remove(m_path) else: LOGGER.error("Failed to extract the archive") @@ -104,7 +114,7 @@ def onDownloadComplete(self): path = f'{DOWNLOAD_DIR}{self.uid}/{name}' else: path = f'{DOWNLOAD_DIR}{self.uid}/{name}' - up_name = PurePath(path).name + up_name = path.rsplit('/', 1)[-1] up_path = f'{DOWNLOAD_DIR}{self.uid}/{up_name}' size = get_path_size(up_path) LOGGER.info(f"Uploading: {up_name}") @@ -176,25 +186,24 @@ def onUploadError(self, error): def _compress(bot, message, is_archive=False, is_extract=False, pswd=None): mesg = message.text.split('\n') - message_args = mesg[0].split(" ", maxsplit=1) + message_args = mesg[0].split(maxsplit=1) reply_to = message.reply_to_message is_appdrive = False is_gdtot = False appdict = '' - try: - link = message_args[1] - if link.startswith("pswd: "): - raise IndexError - except: + if len(message_args) > 1: + link = message_args[1].strip() + if link.startswith("pswd:"): + link = '' + else: link = '' - link = re.split(r"pswd:| \|", link)[0] + link = re.split(r"pswd:|\|", link)[0] link = link.strip() - pswdMsg = mesg[0].split(' pswd: ') - if len(pswdMsg) > 1: - pswd = pswdMsg[1] - if reply_to is not None: - reply_text = reply_to.text - link = reply_text.strip() + pswd_arg = mesg[0].split(' pswd: ') + if len(pswd_arg) > 1: + pswd = pswd_arg[1] + if reply_to: + link = reply_to.text.split(maxsplit=1)[0].strip() is_appdrive = is_appdrive_link(link) is_gdtot = is_gdtot_link(link) if any([is_appdrive, is_gdtot]): diff --git a/bot/modules/count.py b/bot/modules/count.py index b4660bd2..c31a593f 100644 --- a/bot/modules/count.py +++ b/bot/modules/count.py @@ -9,13 +9,12 @@ @new_thread def countNode(update, context): - args = update.message.text.split(" ", maxsplit=1) reply_to = update.message.reply_to_message link = '' - if len(args) > 1: - link = args[1] - if reply_to is not None: - link = reply_to.text + if len(context.args) == 1: + link = context.args[0] + if reply_to: + link = reply_to.text.split(maxsplit=1)[0].strip() if is_gdrive_link(link): msg = sendMessage(f"Counting: {link}", context.bot, update.message) LOGGER.info(f"Counting: {link}") diff --git a/bot/modules/delete.py b/bot/modules/delete.py index 3cb6af3c..17f86374 100644 --- a/bot/modules/delete.py +++ b/bot/modules/delete.py @@ -9,13 +9,12 @@ @new_thread def deleteNode(update, context): - args = update.message.text.split(" ", maxsplit=1) reply_to = update.message.reply_to_message link = '' - if len(args) > 1: - link = args[1] - if reply_to is not None: - link = reply_to.text + if len(context.args) == 1: + link = context.args[0] + if reply_to: + link = reply_to.text.split(maxsplit=1)[0].strip() if is_gdrive_link(link): msg = sendMessage(f"Deleting: {link}", context.bot, update.message) LOGGER.info(f"Deleting: {link}") diff --git a/bot/modules/eval.py b/bot/modules/eval.py index 37274117..ccdb76d3 100644 --- a/bot/modules/eval.py +++ b/bot/modules/eval.py @@ -52,7 +52,7 @@ def cleanup_code(code): return code.strip('` \n') def do(func, bot, update): - content = update.message.text.split(' ', 1)[-1] + content = update.message.text.split(maxsplit=1)[-1] body = cleanup_code(content) env = namespace_of(update.message.chat_id, update, bot) diff --git a/bot/modules/list.py b/bot/modules/list.py index 6fc0e8b8..82d9c0d7 100644 --- a/bot/modules/list.py +++ b/bot/modules/list.py @@ -7,13 +7,13 @@ from bot.helper.telegram_helper.message_utils import sendMessage, editMessage def list_drive(update, context): - args = update.message.text.split(" ", maxsplit=1) + args = update.message.text.split(maxsplit=1) reply_to = update.message.reply_to_message query = '' if len(args) > 1: - query = args[1] + query = args[1].strip() if reply_to is not None: - query = reply_to.text + query = reply_to.text.strip() if query != '': reply = sendMessage(f"Search in progress...", context.bot, update.message) LOGGER.info(f"Finding: {query}") diff --git a/bot/modules/permission.py b/bot/modules/permission.py index 03aa2582..ef534562 100644 --- a/bot/modules/permission.py +++ b/bot/modules/permission.py @@ -9,27 +9,27 @@ @new_thread def permissionNode(update, context): - args = update.message.text.split(" ", maxsplit=2) + args = update.message.text.split() reply_to = update.message.reply_to_message link = '' access = '' if len(args) > 1: - link = args[1] + link = args[1].strip() try: - access = args[2] + access = args[2].strip() except IndexError: pass - if reply_to is not None: - link = reply_to.text + if reply_to: + link = reply_to.text.split(maxsplit=1)[0].strip() try: - access = args[1] + access = args[1].strip() except IndexError: pass if is_gdrive_link(link): msg = sendMessage(f"Setting permission: {link}", context.bot, update.message) LOGGER.info(f"Setting permission: {link}") gd = GoogleDriveHelper() - result = gd.setPerm(link, access) + result = gd.setPermission(link, access) deleteMessage(context.bot, msg) sendMessage(result, context.bot, update.message) else: diff --git a/bot/modules/shell.py b/bot/modules/shell.py index 4f1d4c8f..c6428ddf 100644 --- a/bot/modules/shell.py +++ b/bot/modules/shell.py @@ -5,18 +5,18 @@ from bot import dispatcher from bot.helper.telegram_helper.bot_commands import BotCommands from bot.helper.telegram_helper.filters import CustomFilters -from bot.helper.telegram_helper.message_utils import sendMessage def shell(update, context): message = update.effective_message - cmd = message.text.split(' ', 1) + cmd = message.text.split(maxsplit=1) if len(cmd) == 1: - return sendMessage('Send a command to execute', context.bot, update.message) + return message.reply_text('Send a command to execute', parse_mode='HTML') cmd = cmd[1] - process = subprocess.run(cmd, capture_output=True, shell=True) + process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) + stdout, stderr = process.communicate() reply = '' - stdout = process.stdout.decode('utf-8') - stderr = process.stderr.decode('utf-8') + stdout = stdout.decode() + stderr = stderr.decode() if len(stdout) != 0: reply += f"Stdout\n{stdout}\n" if len(stderr) != 0: @@ -31,9 +31,9 @@ def shell(update, context): reply_to_message_id=message.message_id, chat_id=message.chat_id) elif len(reply) != 0: - sendMessage(reply, context.bot, update.message) + message.reply_text(reply, parse_mode='HTML') else: - sendMessage('Executed', context.bot, update.message) + message.reply_text('Command executed', parse_mode='HTML') shell_handler = CommandHandler(BotCommands.ShellCommand, shell, filters=CustomFilters.owner_filter, run_async=True) diff --git a/bot/modules/status.py b/bot/modules/status.py index b8a1f0dd..db1bdbb4 100644 --- a/bot/modules/status.py +++ b/bot/modules/status.py @@ -2,22 +2,35 @@ from telegram.ext import CommandHandler -from bot import dispatcher, download_dict, download_dict_lock, status_reply_dict, status_reply_dict_lock -from bot.helper.telegram_helper.message_utils import sendMessage, deleteMessage, sendStatusMessage +from bot import dispatcher, Interval, STATUS_UPDATE_INTERVAL, download_dict, download_dict_lock, \ + status_reply_dict, status_reply_dict_lock +from bot.helper.ext_utils.bot_utils import SetInterval +from bot.helper.telegram_helper.message_utils import sendMessage, deleteMessage, \ + update_all_messages, sendStatusMessage from bot.helper.telegram_helper.bot_commands import BotCommands from bot.helper.telegram_helper.filters import CustomFilters def statusNode(update, context): with download_dict_lock: - if len(download_dict) == 0: - return sendMessage("No active task", context.bot, update.message) - index = update.effective_chat.id - with status_reply_dict_lock: - if index in status_reply_dict.keys(): - deleteMessage(context.bot, status_reply_dict[index]) - del status_reply_dict[index] - sendStatusMessage(update.message, context.bot) - deleteMessage(context.bot, update.message) + count = len(download_dict) + if count == 0: + return sendMessage("No active task", context.bot, update.message) + else: + index = update.effective_chat.id + with status_reply_dict_lock: + if index in status_reply_dict: + deleteMessage(context.bot, status_reply_dict[index][0]) + del status_reply_dict[index] + try: + if Interval: + Interval[0].cancel() + Interval.clear() + except: + pass + finally: + Interval.append(SetInterval(STATUS_UPDATE_INTERVAL, update_all_messages)) + sendStatusMessage(update.message, context.bot) + deleteMessage(context.bot, update.message) status_handler = CommandHandler(BotCommands.StatusCommand, statusNode, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user, run_async=True)