fixes: process cancel, cleanup and more

- moved GLOB_TMP to a class variable of MeganzClient
- renamed cleanup to fs_cleanup
- added full_cleanup method to MeganzClient
- fixed cancel button

megadl/__init__.py

@@ -13,7 +13,4 @@
 print("> Loading config")
 load_dotenv()
 
-from .helpers.mclient import MeganzClient
-
-# Dict to act as a temp key value database
-GLOB_TMP = {}
+from .helpers.mclient import MeganzClient

megadl/helpers/database.py

@@ -1,6 +1,5 @@
 from megadl.lib.aiomongo import AioMongo
 
-
 class Users:
     def __init__(self) -> None:
         self.mongoc = AioMongo()

megadl/helpers/files.py

@@ -25,7 +25,7 @@ def listfiles(rpath: str):
 
 
 # Clean up function
-def cleanup(cpath):
+def fs_cleanup(cpath: str):
     if path.isfile(cpath):
         remove(cpath)
     elif path.isdir(cpath):

megadl/helpers/mclient.py

@@ -17,7 +17,7 @@
 from pyrogram.handlers import MessageHandler
 
 from .database import Users
-from .files import send_as_guessed, cleanup, splitit, listfiles
+from .files import send_as_guessed, fs_cleanup, splitit, listfiles
 
 
 _emsg = """
@@ -66,6 +66,7 @@ def __init__(self):
 
         # Initializing mongodb
         print("> Initializing database")
+        self.glob_tmp = {}
         self.cipher = None
         self.is_public = True if self.database else False
 
@@ -100,7 +101,7 @@ def __init__(self):
         # other stuff
         print("> Setting up additional functions")
         self.add_handler(MessageHandler(self.use_listner))
-        self.listen_tasks = {}
+        self.listening = {}
         self.mega_running = {}
 
         print("--------------------")
@@ -122,6 +123,9 @@ async def fn_run(client: Client, msg: Message):
             except errors.FloodWait as e:
                 await xsleep(e.value)
                 return await func(client, msg)
+            # pyrogram message not modified error handling
+            except errors.MessageNotModified:
+                pass
             # Other exceptions
             except Exception as e:
                 logging.warning(_emsg.format(self.version, func.__module__, e))
@@ -134,9 +138,9 @@ async def ask(self, chat_id: int, text: str, *args, **kwargs):
         futr = woop.create_future()
 
         futr.add_done_callback(
-            functools.partial(lambda _, uid: self.listen_tasks.pop(uid, None), chat_id)
+            functools.partial(lambda _, uid: self.listening.pop(uid, None), chat_id)
         )
-        self.listen_tasks[chat_id] = {"task": futr}
+        self.listening[chat_id] = {"task": futr}
 
         # wait for 1 min
         try:
@@ -145,37 +149,53 @@ async def ask(self, chat_id: int, text: str, *args, **kwargs):
             await self.send_message(
                 chat_id, "Task was cancelled as you haven't answered for 1 minute"
             )
-            self.listen_tasks.pop(chat_id, None)
+            self.listening.pop(chat_id, None)
             return None
 
     async def use_listner(self, _, msg: Message):
-        lstn = self.listen_tasks.get(msg.chat.id)
+        lstn = self.listening.get(msg.chat.id)
         if lstn and not lstn["task"].done():
             lstn["task"].set_result(msg)
         return msg.continue_propagation()
 
+    async def full_cleanup(self, path: str, msg_id: int, chat_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):
         """
         Send files with automatic floodwait handling and file splitting
         """
-        for file in files:
-            # Split files larger than 2GB
-            if os.stat(file).st_size > int(os.getenv("TG_MAX_SIZE")):
-                await self.edit_message_text(
-                    chat_id,
-                    msg_id,
-                    """
-                    The file you're trying to upload exceeds telegram limits 😬.
-                    Trying to split the files 🔪...
-                    """,
-                )
-                splout = f"{self.tmp_loc}/splitted"
-                await splitit(file, splout)
-                for file in listfiles(splout):
+        if files:
+            for file in files:
+                # Split files larger than 2GB
+                if os.stat(file).st_size > int(os.getenv("TG_MAX_SIZE")):
+                    await self.edit_message_text(
+                        chat_id,
+                        msg_id,
+                        """
+                        The file you're trying to upload exceeds telegram limits 😬.
+                        Trying to split the files 🔪...
+                        """,
+                    )
+                    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)
+                    fs_cleanup(splout)
+                else:
                     await send_as_guessed(self, file, chat_id, msg_id)
-                cleanup(splout)
-            else:
-                await send_as_guessed(self, file, chat_id, msg_id)
+        else:
+            await self.edit_message_text(
+                chat_id,
+                msg_id,
+                "Couldn't find files to send. Maybe you've canceled the process?",
+            )
 
     async def send_document(self, *args, **kwargs):
         try:

megadl/helpers/sysfncs.py

@@ -4,8 +4,9 @@
 # Description: Contains functions related to shell, async execution
 
 
-from subprocess import Popen, PIPE
 from functools import partial
+from subprocess import Popen, PIPE
+from psutil import Process as PsuPrc
 from asyncio import get_running_loop
 
 
@@ -18,9 +19,7 @@ async def run_partial(func, *args, **kwargs):
         - func: function - Function to execute
     """
     loop = get_running_loop()
-    return await loop.run_in_executor(
-        None, partial(func, *args, **kwargs)
-    )
+    return await loop.run_in_executor(None, partial(func, *args, **kwargs))
 
 
 def run_on_shell(cmd):
@@ -29,4 +28,14 @@ def run_on_shell(cmd):
     """
     run = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     shout = run.stdout.read()[:-1].decode("utf-8")
-    return shout
+    return shout
+
+
+async def kill_family(pid: int):
+    """
+    Murder a process and it's whole family using pid
+    """
+    running = PsuPrc(pid)
+    for child in running.children(recursive=True):
+        child.kill()
+    running.kill()

megadl/lib/megatools.py

@@ -161,16 +161,21 @@ def __shellExec(self, cmd: str, chat_id: int = None, msg_id: int = None, **kwarg
             shell=True,
             encoding="utf-8",
         )
-        self.client.mega_running[chat_id] = run
+        self.client.mega_running[chat_id] = run.pid
 
         try:
             # Live process info update
             while run.poll() is None:
                 out = run.stdout.readline()
                 if out != "":
-                    self.client.edit_message_text(
-                        chat_id, msg_id, f"**Process info:** \n`{out}`", **kwargs
-                    )
+                    try:
+                        self.client.edit_message_text(
+                            chat_id, msg_id, f"**Process info:** \n`{out}`", **kwargs
+                        )
+                        # run.terminate()
+                        # run.wait()
+                    except:
+                        pass
         except FileNotFoundError:
             pass
         sh_out = run.stdout.read()[:-1]

megadl/modules/callbacks.py

@@ -3,29 +3,30 @@
 # Project: https://github.com/Itz-fork/Mega.nz-Bot
 # Description: Contains global callbacks
 
-from shutil import rmtree
-
 from pyrogram import filters
 from pyrogram.types import CallbackQuery
 
-from megadl import MeganzClient, GLOB_TMP
+from megadl import MeganzClient
+from megadl.helpers.sysfncs import kill_family
 
 
-@MeganzClient.on_callback_query(filters.regex(r"cancelqcb"))
+@MeganzClient.on_callback_query(filters.regex(r"cancelqcb?.+"))
 @MeganzClient.handle_checks
 async def close_gb(client: MeganzClient, query: CallbackQuery):
+    _mid = int(query.data.split("-")[1])
     try:
         # Remove user from global temp db
-        dtmp = GLOB_TMP.pop(int(query.data.split("-")[1]))
+        dtmp = client.glob_tmp.get(_mid)
         # cancel if user has a download running
-        running = client.mega_running.get(query.message.chat.id)
-        if running:
-            running.kill()
+        prcid = client.mega_running.get(query.message.chat.id)
+        if prcid:
+            await kill_family(prcid)
         # Remove download folder of the user
-        rmtree(dtmp[1])
+        if dtmp:
+            await client.full_cleanup(dtmp[1], _mid)
     except:
         pass
     await query.edit_message_text(
-        "Process was canceled by the user",
+        "Process was canceled by the user ❌",
         reply_markup=None,
     )

megadl/modules/mega_dl.py

@@ -14,25 +14,25 @@
     InlineKeyboardMarkup,
 )
 
-from megadl import MeganzClient, GLOB_TMP
+from megadl import MeganzClient
 from megadl.lib.megatools import MegaTools
-from megadl.helpers.files import cleanup
 
 
 @MeganzClient.on_message(
     filters.regex(r"(https?:\/\/mega\.nz\/(file|folder|#)?.+)|(\/Root\/?.+)")
 )
 @MeganzClient.handle_checks
-async def dl_from(_: MeganzClient, msg: Message):
+async def dl_from(client: MeganzClient, msg: Message):
     # Push info to temp db
-    GLOB_TMP[msg.id] = [msg.text, f"{_.dl_loc}/{msg.id}"]
+    _mid = msg.id
+    client.glob_tmp[msg.id] = [msg.text, f"{client.dl_loc}/{_mid}"]
     await msg.reply(
         "Select what you want to do 🤗",
         reply_markup=InlineKeyboardMarkup(
             [
-                [InlineKeyboardButton("Download 💾", callback_data=f"dwn_mg-{msg.id}")],
-                [InlineKeyboardButton("Info ℹ️", callback_data=f"info_mg-{msg.id}")],
-                [InlineKeyboardButton("Cancel ❌", callback_data="cancelqcb")],
+                [InlineKeyboardButton("Download 💾", callback_data=f"dwn_mg-{_mid}")],
+                [InlineKeyboardButton("Info ℹ️", callback_data=f"info_mg-{_mid}")],
+                [InlineKeyboardButton("Cancel ❌", callback_data=f"cancelqcb-{_mid}")],
             ]
         ),
     )
@@ -42,7 +42,8 @@ async def dl_from(_: MeganzClient, msg: Message):
 @MeganzClient.handle_checks
 async def dl_from_cb(client: MeganzClient, query: CallbackQuery):
     # Access saved info
-    dtmp = GLOB_TMP.pop(int(query.data.split("-")[1]))
+    _mid = int(query.data.split("-")[1])
+    dtmp = client.glob_tmp.get(_mid)
     url = dtmp[0]
     dlid = dtmp[1]
     qcid = query.message.chat.id
@@ -64,18 +65,21 @@ async def dl_from_cb(client: MeganzClient, query: CallbackQuery):
             conf = f"--username {client.cipher.decrypt(udoc['email']).decode()} --password {client.cipher.decrypt(udoc['password']).decode()}"
     cli = MegaTools(client, conf)
 
-    f_list = await cli.download(
-        url,
-        qcid,
-        resp.id,
-        path=dlid,
-        reply_markup=InlineKeyboardMarkup(
-            [
-                [InlineKeyboardButton("Cancel ❌", callback_data="cancelqcb")],
-            ]
-        ),
-    )
+    f_list = None
     try:
+        f_list = await cli.download(
+            url,
+            qcid,
+            resp.id,
+            path=dlid,
+            reply_markup=InlineKeyboardMarkup(
+                [
+                    [InlineKeyboardButton("Cancel ❌", callback_data=f"cancelqcb-{_mid}")],
+                ]
+            ),
+        )
+        if not f_list:
+            return
         await query.edit_message_text("Successfully downloaded the content 🥳")
     except Exception as e:
         await query.edit_message_text(
@@ -88,14 +92,14 @@ async def dl_from_cb(client: MeganzClient, query: CallbackQuery):
     # Send file(s) to the user
     await resp.edit("Trying to upload now 📤...")
     await client.send_files(f_list, qcid, resp.id)
-    cleanup(dlid)
-    await resp.delete()
+    await client.full_cleanup(dlid, _mid)
+    # await resp.delete()
 
 
 @MeganzClient.on_callback_query(filters.regex(r"info_mg?.+"))
 @MeganzClient.handle_checks
-async def info_from_cb(_: MeganzClient, query: CallbackQuery):
-    url = GLOB_TMP.pop(int(query.data.split("-")[1]))[0]
+async def info_from_cb(client: MeganzClient, query: CallbackQuery):
+    url = client.glob_tmp.pop(int(query.data.split("-")[1]))[0]
     size, name = await MegaTools.file_info(url)
     await query.edit_message_text(
         f"""

megadl/modules/mega_up.py

@@ -17,8 +17,6 @@
 from megadl.lib.megatools import MegaTools
 from megadl.lib.pyros import track_progress
 
-from megadl.helpers.files import cleanup
-
 
 # Respond only to Documents, Photos, Videos, GIFs, Audio and to urls other than mega
 @MeganzClient.on_message(
@@ -33,12 +31,13 @@
 )
 @MeganzClient.handle_checks
 async def to_up(_: MeganzClient, msg: Message):
+    _mid = msg.id
     await msg.reply(
         "Select what you want to do 🤗",
         reply_markup=InlineKeyboardMarkup(
             [
-                [InlineKeyboardButton("Download 💾", callback_data=f"up_tgdl-{msg.id}")],
-                [InlineKeyboardButton("Cancel ❌", callback_data="cancelqcb")],
+                [InlineKeyboardButton("Download 💾", callback_data=f"up_tgdl-{_mid}")],
+                [InlineKeyboardButton("Cancel ❌", callback_data=f"cancelqcb-{_mid}")],
             ]
         ),
     )
@@ -48,10 +47,11 @@ async def to_up(_: MeganzClient, msg: Message):
 @MeganzClient.handle_checks
 async def to_up_cb(client: MeganzClient, query: CallbackQuery):
     # Get message content
+    _mid = int(query.data.split("-")[1])
     qcid = query.message.chat.id
     qmid = query.message.id
     strtim = time()
-    msg = await client.get_messages(qcid, int(query.data.split("-")[1]))
+    msg = await client.get_messages(qcid, _mid)
     # Status msg
     await client.edit_message_text(
         qcid, qmid, "Trying to download the file 📬", reply_markup=None
@@ -86,4 +86,4 @@ async def to_up_cb(client: MeganzClient, query: CallbackQuery):
             [[InlineKeyboardButton("Visit 🔗", url=limk)]]
         ),
     )
-    cleanup(dl_path)
+    await client.full_cleanup(dl_path, _mid)