From 7c84d1333a2dfedf3d12b9491b96bb97b695c816 Mon Sep 17 00:00:00 2001 From: anasty17 Date: Fri, 10 Nov 2023 13:14:52 +0200 Subject: [PATCH] Handle sender_chat for some functions Signed-off-by: anasty17 --- bot/helper/common.py | 17 ++++++++++++----- .../download_utils/mega_download.py | 16 ++++++++++++---- bot/helper/mirror_utils/gdrive_utils/clone.py | 9 ++++++++- bot/modules/authorize.py | 8 ++++---- bot/modules/cancel_task.py | 2 +- bot/modules/eval.py | 2 +- bot/modules/gd_count.py | 5 +++-- bot/modules/gd_delete.py | 3 ++- bot/modules/mirror_leech.py | 5 ++++- bot/modules/status.py | 5 ++++- 10 files changed, 51 insertions(+), 21 deletions(-) diff --git a/bot/helper/common.py b/bot/helper/common.py index 6c0803e37dd..5374ad3a4fa 100644 --- a/bot/helper/common.py +++ b/bot/helper/common.py @@ -58,7 +58,8 @@ class TaskConfig: def __init__(self, message): self.message = message self.mid = self.message.id - self.user_id = self.message.from_user.id + self.user = self.message.from_user or self.message.sender_chat + self.user_id = self.user.id self.user_dict = user_data.get(self.user_id, {}) self.sameDir = {} self.bulk = [] @@ -282,12 +283,12 @@ async def beforeStart(self): async def getTag(self, text: list): if len(text) > 1 and text[1].startswith("Tag: "): self.tag, id_ = text[1].split("Tag: ")[1].split() - self.message.from_user = await self.client.get_users(id_) + self.user = self.message.from_user = await self.client.get_users(id_) try: await self.message.unpin() except: pass - if username := self.message.from_user.username: + if username := self.user.username: self.tag = f"@{username}" else: self.tag = self.message.from_user.mention @@ -331,7 +332,10 @@ async def run_multi(self, input_list, folder_name, obj): ) if folder_name: self.sameDir["tasks"].add(nextmsg.id) - nextmsg.from_user = self.message.from_user + if self.message.from_user: + nextmsg.from_user = self.user + else: + nextmsg.sender_chat = self.user obj( self.client, nextmsg, @@ -360,7 +364,10 @@ async def initBulk(self, input_list, bulk_start, bulk_end, obj): nextmsg = await self.client.get_messages( chat_id=self.message.chat.id, message_ids=nextmsg.id ) - nextmsg.from_user = self.message.from_user + if self.message.from_user: + nextmsg.from_user = self.user + else: + nextmsg.sender_chat = self.user obj( self.client, nextmsg, diff --git a/bot/helper/mirror_utils/download_utils/mega_download.py b/bot/helper/mirror_utils/download_utils/mega_download.py index 98d0eacf1cf..737456a7b7e 100644 --- a/bot/helper/mirror_utils/download_utils/mega_download.py +++ b/bot/helper/mirror_utils/download_utils/mega_download.py @@ -1,7 +1,13 @@ from secrets import token_urlsafe from aiofiles.os import makedirs from threading import Event -from mega import MegaApi, MegaListener, MegaRequest, MegaTransfer, MegaError +from mega import ( + MegaApi, + MegaListener, + MegaRequest, + MegaTransfer, + MegaError, +) from bot import ( LOGGER, @@ -45,7 +51,9 @@ def speed(self): def downloaded_bytes(self): return self._bytes_transferred - def onRequestFinish(self, api, request, error): + def onRequestFinish(self, api: MegaApi, request: MegaRequest, error): + if self.is_cancelled: + return if str(error).lower() != "no error": self.error = error.copy() LOGGER.error(f"Mega onRequestFinishError: {self.error}") @@ -84,7 +92,7 @@ def onTransferUpdate(self, api: MegaApi, transfer: MegaTransfer): self._speed = transfer.getSpeed() self._bytes_transferred = transfer.getTransferredBytes() - def onTransferFinish(self, api: MegaApi, transfer: MegaTransfer, error): + def onTransferFinish(self, api, transfer: MegaTransfer, error): try: if self.is_cancelled: self.continue_event.set() @@ -94,7 +102,7 @@ def onTransferFinish(self, api: MegaApi, transfer: MegaTransfer, error): except Exception as e: LOGGER.error(e) - def onTransferTemporaryError(self, api, transfer, error): + def onTransferTemporaryError(self, api, transfer: MegaTransfer, error: MegaError): filen = transfer.getFileName() state = transfer.getState() errStr = error.toString() diff --git a/bot/helper/mirror_utils/gdrive_utils/clone.py b/bot/helper/mirror_utils/gdrive_utils/clone.py index f1021f4209a..868e0d5241c 100644 --- a/bot/helper/mirror_utils/gdrive_utils/clone.py +++ b/bot/helper/mirror_utils/gdrive_utils/clone.py @@ -44,7 +44,14 @@ def clone(self): try: file_id = self.getIdFromUrl(self.listener.link) except (KeyError, IndexError): - return "Google Drive ID could not be found in the provided link", None, None, None, None, None + return ( + "Google Drive ID could not be found in the provided link", + None, + None, + None, + None, + None, + ) self.service = self.authorize() msg = "" LOGGER.info(f"File ID: {file_id}") diff --git a/bot/modules/authorize.py b/bot/modules/authorize.py index 9097e6e8694..6d40971b952 100644 --- a/bot/modules/authorize.py +++ b/bot/modules/authorize.py @@ -14,7 +14,7 @@ async def authorize(client, message): if len(msg) > 1: id_ = int(msg[1].strip()) elif reply_to := message.reply_to_message: - id_ = reply_to.from_user.id + id_ = reply_to.from_user.id if reply_to.from_user else reply_to.sender_chat.id else: id_ = message.chat.id if id_ in user_data and user_data[id_].get("is_auth"): @@ -32,7 +32,7 @@ async def unauthorize(client, message): if len(msg) > 1: id_ = int(msg[1].strip()) elif reply_to := message.reply_to_message: - id_ = reply_to.from_user.id + id_ = reply_to.from_user.id if reply_to.from_user else reply_to.sender_chat.id else: id_ = message.chat.id if id_ not in user_data or user_data[id_].get("is_auth"): @@ -51,7 +51,7 @@ async def addSudo(client, message): if len(msg) > 1: id_ = int(msg[1].strip()) elif reply_to := message.reply_to_message: - id_ = reply_to.from_user.id + id_ = reply_to.from_user.id if reply_to.from_user else reply_to.sender_chat.id if id_: if id_ in user_data and user_data[id_].get("is_sudo"): msg = "Already Sudo!" @@ -71,7 +71,7 @@ async def removeSudo(client, message): if len(msg) > 1: id_ = int(msg[1].strip()) elif reply_to := message.reply_to_message: - id_ = reply_to.from_user.id + id_ = reply_to.from_user.id if reply_to.from_user else reply_to.sender_chat.id if id_ and id_ not in user_data or user_data[id_].get("is_sudo"): update_user_ldata(id_, "is_sudo", False) if DATABASE_URL: diff --git a/bot/modules/cancel_task.py b/bot/modules/cancel_task.py index c76b74292a6..3ac6fb0c4b0 100644 --- a/bot/modules/cancel_task.py +++ b/bot/modules/cancel_task.py @@ -16,7 +16,7 @@ async def cancel_task(_, message): - user_id = message.from_user.id + user_id = message.from_user.id if message.from_user else message.sender_chat.id msg = message.text.split() if len(msg) > 1: gid = msg[1] diff --git a/bot/modules/eval.py b/bot/modules/eval.py index 43479ff5348..d7ee7492d3e 100644 --- a/bot/modules/eval.py +++ b/bot/modules/eval.py @@ -30,7 +30,7 @@ def namespace_of(message): def log_input(message): LOGGER.info( - f"IN: {message.text} (user={message.from_user.id}, chat={message.chat.id})" + f"IN: {message.text} (user={message.from_user.id if message.from_user else message.sender_chat.id}, chat={message.chat.id})" ) diff --git a/bot/modules/gd_count.py b/bot/modules/gd_count.py index eb20b09652c..a243e4dd703 100644 --- a/bot/modules/gd_count.py +++ b/bot/modules/gd_count.py @@ -14,7 +14,8 @@ @new_task async def countNode(_, message): args = message.text.split() - if username := message.from_user.username: + user = message.from_user or message.sender_chat + if username := user.username: tag = f"@{username}" else: tag = message.from_user.mention @@ -26,7 +27,7 @@ async def countNode(_, message): if is_gdrive_link(link): msg = await sendMessage(message, f"Counting: {link}") name, mime_type, size, files, folders = await sync_to_async( - gdCount().count, link, message.from_user.id + gdCount().count, link, user.id ) if mime_type is None: await sendMessage(message, name) diff --git a/bot/modules/gd_delete.py b/bot/modules/gd_delete.py index 6db1c5370f3..fe800533b25 100644 --- a/bot/modules/gd_delete.py +++ b/bot/modules/gd_delete.py @@ -13,6 +13,7 @@ @new_task async def deletefile(_, message): args = message.text.split() + user = message.from_user or message.sender_chat if len(args) > 1: link = args[1] elif reply_to := message.reply_to_message: @@ -21,7 +22,7 @@ async def deletefile(_, message): link = "" if is_gdrive_link(link): LOGGER.info(link) - msg = await sync_to_async(gdDelete().deletefile, link, message.from_user.id) + msg = await sync_to_async(gdDelete().deletefile, link, user.id) else: msg = ( "Send Gdrive link along with command or by replying to the link by command" diff --git a/bot/modules/mirror_leech.py b/bot/modules/mirror_leech.py index 95a6b1c2d54..496c610a41b 100644 --- a/bot/modules/mirror_leech.py +++ b/bot/modules/mirror_leech.py @@ -186,7 +186,10 @@ async def newEvent(self): nextmsg = await self.client.get_messages( chat_id=self.message.chat.id, message_ids=nextmsg.id ) - nextmsg.from_user = self.message.from_user + if self.message.from_user: + nextmsg.from_user = self.user + else: + nextmsg.sender_chat = self.user Mirror( self.client, nextmsg, diff --git a/bot/modules/status.py b/bot/modules/status.py index 191ebc99e20..4a4ffe97f03 100644 --- a/bot/modules/status.py +++ b/bot/modules/status.py @@ -94,6 +94,7 @@ async def status_pages(_, query): "Clone": 0, "CheckUp": 0, "Pause": 0, + "SamVid": 0, } dl_speed = 0 up_speed = 0 @@ -126,10 +127,12 @@ async def status_pages(_, query): tasks["CheckUp"] += 1 elif tstatus == MirrorStatus.STATUS_PAUSED: tasks["Pause"] += 1 + elif tstatus == MirrorStatus.STATUS_SAMVID: + tasks["SamVid"] += 1 msg = f"""DL: {tasks['Download']} | UP: {tasks['Upload']} | SD: {tasks['Seed']} | AR: {tasks['Archive']} EX: {tasks['Extract']} | SP: {tasks['Split']} | QD: {tasks['QueueDl']} | QU: {tasks['QueueUp']} -CL: {tasks['Clone']} | CH: {tasks['CheckUp']} | PA:{tasks['Pause']} +CL: {tasks['Clone']} | CH: {tasks['CheckUp']} | PA:{tasks['Pause']} | SV:{tasks['SamVid']} ODLS: {get_readable_file_size(dl_speed)}/s OULS: {get_readable_file_size(up_speed)}/s