Handle sender_chat for some functions

Signed-off-by: anasty17 <e.anastayyar@gmail.com>
rwcats · Nov 10, 2023 · 7c84d13 · 7c84d13
1 parent 4a5cb67
commit 7c84d13
Show file tree

Hide file tree

Showing 10 changed files with 51 additions and 21 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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: <code>{link}</code>")
         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
@@ -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
@@ -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
@@ -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