bug fixes and improvements

- added group chat support
- fixed megatools not working without login
- cancellable upload tasks
- limit 1 process per user at a time
- fixed MeganzClient.full_cleanup throwing error when msg_id is None or not in glob_tmp
- prevent admins from getting banned
- prevent users from logging in public chats
- some text string updates

megadl/helpers/files.py

@@ -12,7 +12,7 @@
 from filetype import guess
 from filesplit.split import Split
 
-from megadl.lib.pyros import track_progress
+from megadl.helpers.pyros import track_progress
 from megadl.helpers.sysfncs import run_partial, run_on_shell
 
 
@@ -33,15 +33,21 @@ def fs_cleanup(cpath: str):
 
 
 # Guess the file type and send it accordingly
-async def send_as_guessed(client, file, chat_id, mid):
+async def send_as_guessed(client, file, chat_id, mid, **kwargs):
     """
     Upload the file to telegram as the way it should
     """
     ftype = await run_partial(guess, file)
     strtim = time()
     # Send file as a document if the file type could't be guessed
     if not ftype:
-        await client.send_document(chat_id, file)
+        await client.send_document(
+            chat_id,
+            file,
+            progress=track_progress,
+            progress_args=(client, [chat_id, mid], strtim),
+            **kwargs,
+        )
     else:
         fmime = ftype.mime
         # GIFs
@@ -51,6 +57,7 @@ async def send_as_guessed(client, file, chat_id, mid):
                 file,
                 progress=track_progress,
                 progress_args=(client, [chat_id, mid], strtim),
+                **kwargs,
             )
         # Images
         elif isit(r"\bimage\b", fmime):
@@ -59,6 +66,7 @@ async def send_as_guessed(client, file, chat_id, mid):
                 file,
                 progress=track_progress,
                 progress_args=(client, [chat_id, mid], strtim),
+                **kwargs,
             )
         # Audio
         elif isit(r"\baudio\b", fmime):
@@ -67,6 +75,7 @@ async def send_as_guessed(client, file, chat_id, mid):
                 file,
                 progress=track_progress,
                 progress_args=(client, [chat_id, mid], strtim),
+                **kwargs,
             )
         # Video
         elif isit(r"\bvideo\b", fmime):
@@ -92,6 +101,7 @@ async def send_as_guessed(client, file, chat_id, mid):
                 thumb=_thumb,
                 progress=track_progress,
                 progress_args=(client, [chat_id, mid], strtim),
+                **kwargs,
             )
         # Document
         else:
@@ -100,6 +110,7 @@ async def send_as_guessed(client, file, chat_id, mid):
                 file,
                 progress=track_progress,
                 progress_args=(client, [chat_id, mid], strtim),
+                **kwargs,
             )
 
 

megadl/helpers/mclient.py

@@ -88,11 +88,8 @@ def __init__(self):
             if len(_auths.split("|")) > 2
             else set(map(int, _auths.split("|")[1].split())).union({"*"})
         )
-
-        self.use_logs = (
-            {"dl_from", "up_to"} if os.getenv("USE_LOGS") in {"True", "true"} else set()
-        )
         self.log_chat = int(os.getenv("LOG_CHAT")) if os.getenv("LOG_CHAT") else None
+        self.use_logs = {"dl_from", "up_to"}
         self.is_public = True if self.database else False
 
         if self.is_public:
@@ -119,6 +116,7 @@ def __init__(self):
         print("> Setting up additional functions")
         self.listening = {}
         self.mega_running = {}
+        self.ddl_running = {}
         self.add_handler(MessageHandler(self.use_listner))
 
         print("--------------------")
@@ -132,6 +130,19 @@ async def cy_run(client: Client, msg: Message):
             can_use = False
             uid = msg.from_user.id
             try:
+                if func.__name__ in self.use_logs:
+                    # return if user has already started a process
+                    if uid in self.mega_running or uid in self.ddl_running:
+                        return await msg.reply(
+                            "`You've already started a process. Wait until it's finished before starting another one 🥱`"
+                        )
+                    # send logs to the log chat if available
+                    if self.log_chat:
+                        _frwded = await msg.forward(chat_id=self.log_chat)
+                        await _frwded.reply(
+                            f"**#UPLOAD_LOG** \n\n**From:** `{uid}` \n**Get history:** `/info {uid}`"
+                        )
+
                 # Check auth users
                 if self.database:
                     await self.database.add(uid)
@@ -143,18 +154,10 @@ async def cy_run(client: Client, msg: Message):
 
                 if not can_use:
                     await msg.reply(
-                        "You're not authorized to use this bot 🙅‍♂️ \n\n**Join @NexaBotsUpdates ❤️**"
+                        "`You're not authorized to use this bot 🙅‍♂️` \n\n**Join @NexaBotsUpdates ❤️**"
                     )
                     return msg.stop_propagation()
 
-                # send logs if needed
-                if func.__name__ in self.use_logs:
-                    if self.log_chat:
-                        _frwded = await msg.forward(chat_id=self.log_chat)
-                        await _frwded.reply(
-                            f"**#UPLOAD_LOG** \n\n**From:** `{uid}` \n**Get history:** `/info {uid}`"
-                        )
-
                 return await func(client, msg)
             # Floodwait handling
             except errors.FloodWait as e:
@@ -195,16 +198,24 @@ async def use_listner(self, _, msg: Message):
             lstn["task"].set_result(msg)
         return msg.continue_propagation()
 
-    async def full_cleanup(self, path: str, msg_id: int, chat_id: int = None):
+    async def full_cleanup(self, path: str, msg_id: int, user_id: int = None):
         """
         Delete the file/folder and the message
         """
-        fs_cleanup(path)
-        self.glob_tmp.pop(msg_id)
-        if chat_id and chat_id in self.mega_running:
-            self.mega_running.remove(chat_id)
-
-    async def send_files(self, files: list[str], chat_id: int, msg_id: int):
+        try:
+            fs_cleanup(path)
+            self.glob_tmp.pop(msg_id, None)
+
+            # idk why I didn't do self.<dict>.pop(<key>, None), whatever this works
+            if user_id:
+                if user_id in self.mega_running:
+                    self.mega_running.pop(user_id)
+                if user_id in self.ddl_running:
+                    self.ddl_running.pop(user_id)
+        except:
+            pass
+
+    async def send_files(self, files: list[str], chat_id: int, msg_id: int, **kwargs):
         """
         Send files with automatic floodwait handling and file splitting
         """
@@ -223,10 +234,10 @@ async def send_files(self, files: list[str], chat_id: int, msg_id: int):
                     splout = f"{self.tmp_loc}/splitted"
                     await splitit(file, splout)
                     for file in listfiles(splout):
-                        await send_as_guessed(self, file, chat_id, msg_id)
+                        await send_as_guessed(self, file, chat_id, msg_id, **kwargs)
                     fs_cleanup(splout)
                 else:
-                    await send_as_guessed(self, file, chat_id, msg_id)
+                    await send_as_guessed(self, file, chat_id, msg_id, **kwargs)
         else:
             await self.edit_message_text(
                 chat_id,

megadl/lib/pyros.py → megadl/helpers/pyros.py

@@ -9,7 +9,7 @@
 
 # Porogress bar for pyrogram
 # Credits: SpEcHiDe's AnyDL-Bot
-async def track_progress(current, total, client, ides, start):
+async def track_progress(current, total, client, ides, start, **kwargs):
     now = time()
     diff = now - start
     if round(diff % 10.00) == 0 or current == total:
@@ -31,7 +31,7 @@ async def track_progress(current, total, client, ides, start):
         tmp = f"{progress}{humanbytes(current)} of {humanbytes(total)}\n**Speed:** {humanbytes(speed)}/s\n**ETA:** {estimated_total_time if estimated_total_time != '' else '0 s'}\n"
         try:
             await client.edit_message_text(
-                ides[0], ides[1], f"{tmp}\n\n**Powered by @NexaBotsUpdates**"
+                ides[0], ides[1], f"{tmp}\n\n**Powered by @NexaBotsUpdates**", **kwargs
             )
         except:
             pass

megadl/lib/README.md

@@ -5,5 +5,4 @@ Helper libraries for Mega.nz-Bot
 # Extensions
 - [`megatools`](megatools.py) - Wrapper for megatools cli with extended features
 - [`aiomongo`](aiomongo.py) - Async wrapper for pymongo's insert, find, update and delete operations
-- [`pyros`](pyros.py) - Tools and helper functions related to pyrogram
 - [`ddl`](ddl.py) - Downloader for direct download links and gdrive

megadl/lib/ddl.py

@@ -4,13 +4,14 @@
 # Description: Downloader for direct download links and gdrive
 
 import os
+import asyncio
 
 from time import time
 from re import match, sub
 from aiohttp import ClientSession
 from aiofiles import open as async_open
 
-from .pyros import track_progress
+from ..helpers.pyros import track_progress
 
 
 class Downloader:
@@ -20,27 +21,35 @@ class Downloader:
     Supports
         - Direct download links
         - Google Drive links (shared files)
+
+    Arguments:
+        - `client` - Pyrogram client object
     """
 
-    def __init__(self) -> None:
+    def __init__(self, client) -> None:
         self.gdrive_regex = r"https://drive\.google\.com/file/d/(.*?)/.*?\?usp=sharing"
+        self.tg_client = client
 
-    async def download(self, url: str, path: str, client, ides: tuple[int, int]) -> str:
+    async def download(
+        self, url: str, path: str, ides: tuple[int, int, int], **kwargs
+    ) -> str:
         """
         Download a file from direct / gdrive link
 
-        Parameters:
+        Arguments:
 
             - `url` - Url to the file
             - `path` - Output path
-            - `client` - Pyrogram client object
-            - `ides` - Tuple of ids (chat_id, msg_id)
+            - `ides` - Tuple of ids (chat_id, msg_id, user_id)
         """
         if match(self.gdrive_regex, url):
-            gurl = await self._parse_gdrive(url)
-            return await self.from_ddl(url=gurl, path=path, client=client, ides=ides)
-        else:
-            return await self.from_ddl(url=url, path=path, client=client, ides=ides)
+            url = await self._parse_gdrive(url)
+
+        dl_task = asyncio.create_task(
+            self.from_ddl(url=url, path=path, ides=(ides[0], ides[1]), **kwargs)
+        )
+        self.tg_client.ddl_running[ides[2]] = dl_task
+        return await dl_task
 
     async def _parse_gdrive(self, url: str):
         return sub(
@@ -49,14 +58,15 @@ async def _parse_gdrive(self, url: str):
             url,
         )
 
-    async def from_ddl(self, url: str, path: str, client, ides: tuple[int, int]) -> str:
+    async def from_ddl(
+        self, url: str, path: str, ides: tuple[int, int], **kwargs
+    ) -> str:
         """
         Download files from a direct download link
 
         Arguments:
             - `url` - Url of the file
             - `path` - Output path
-            - `client` - Pyrogram client object
             - `ides` - chat id and message id where the action takes place (chat_id, msg_id)
         """
         # Create folder if it doesn't exist
@@ -79,8 +89,12 @@ async def from_ddl(self, url: str, path: str, client, ides: tuple[int, int]) ->
                         await file.write(chunk)
                         curr += len(chunk)
                         # Make sure everything is present before calling track_progress
-                        if None not in {ides, client, total}:
-                            await track_progress(curr, total, client, ides, st)
+                        if None not in {ides, self.tg_client, total}:
+                            await track_progress(
+                                curr, total, self.tg_client, ides, st, **kwargs
+                            )
+
+                        await asyncio.sleep(0)
             return wpath