feat: add cancel button for megatools

megadl/helpers/mclient.py

@@ -100,7 +100,8 @@ def __init__(self):
         # other stuff
         print("> Setting up additional functions")
         self.add_handler(MessageHandler(self.use_listner))
-        self.tasks = {}
+        self.listen_tasks = {}
+        self.mega_running = {}
 
         print("--------------------")
 
@@ -133,9 +134,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.tasks.pop(uid, None), chat_id)
+            functools.partial(lambda _, uid: self.listen_tasks.pop(uid, None), chat_id)
         )
-        self.tasks[chat_id] = {"task": futr}
+        self.listen_tasks[chat_id] = {"task": futr}
 
         # wait for 1 min
         try:
@@ -144,11 +145,11 @@ 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.tasks.pop(chat_id, None)
+            self.listen_tasks.pop(chat_id, None)
             return None
 
     async def use_listner(self, _, msg: Message):
-        lstn = self.tasks.get(msg.chat.id)
+        lstn = self.listen_tasks.get(msg.chat.id)
         if lstn and not lstn["task"].done():
             lstn["task"].set_result(msg)
         return msg.continue_propagation()

megadl/lib/megatools.py

@@ -42,6 +42,7 @@ async def download(
         chat_id: int,
         message_id: int,
         path: str = "MegaDownloads",
+        **kwargs,
     ) -> list:
         """
         Download file/folder from given link
@@ -62,11 +63,18 @@ async def download(
 
         else:
             cmd = f'megacopy --no-ask-password {self.config} -l "{path}" -r "{url}" --download'
-        await run_partial(self.__shellExec, cmd, chat_id=chat_id, msg_id=message_id)
+        await run_partial(
+            self.__shellExec, cmd, chat_id=chat_id, msg_id=message_id, **kwargs
+        )
         return listfiles(path)
 
     async def upload(
-        self, path: str, chat_id: int, message_id: int, to_path: str = "MegaBot"
+        self,
+        path: str,
+        chat_id: int,
+        message_id: int,
+        to_path: str = "MegaBot",
+        **kwargs,
     ) -> str:
         """
         Upload files/folders
@@ -97,7 +105,9 @@ async def upload(
                 run_on_shell, f'megamkdir {self.config} "/Root/{to_path}"'
             )
         # Upload
-        await run_partial(self.__shellExec, cmd, chat_id=chat_id, msg_id=message_id)
+        await run_partial(
+            self.__shellExec, cmd, chat_id=chat_id, msg_id=message_id, **kwargs
+        )
         # Generate link
         ulink = await run_partial(
             run_on_shell,
@@ -143,22 +153,23 @@ async def file_info(url: str) -> list[str]:
         fname = decrypt_attr(base64_url_decode(data["at"]), tk)["n"]
         return [fsize, fname]
 
-    def __shellExec(self, cmd: str, chat_id: int = None, msg_id: int = None):
-        print(cmd)
+    def __shellExec(self, cmd: str, chat_id: int = None, msg_id: int = None, **kwargs):
         run = subprocess.Popen(
             cmd,
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             shell=True,
             encoding="utf-8",
         )
+        self.client.mega_running[chat_id] = run
+
         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}`"
+                        chat_id, msg_id, f"**Process info:** \n`{out}`", **kwargs
                     )
         except FileNotFoundError:
             pass

megadl/modules/callbacks.py

@@ -11,12 +11,16 @@
 from megadl import MeganzClient, GLOB_TMP
 
 
-@MeganzClient.on_callback_query(filters.regex(r"closeqcb"))
+@MeganzClient.on_callback_query(filters.regex(r"cancelqcb"))
 @MeganzClient.handle_checks
-async def close_gb(_: MeganzClient, query: CallbackQuery):
+async def close_gb(client: MeganzClient, query: CallbackQuery):
     try:
         # Remove user from global temp db
         dtmp = GLOB_TMP.pop(int(query.data.split("-")[1]))
+        # cancel if user has a download running
+        running = client.mega_running.get(query.message.chat.id)
+        if running:
+            running.kill()
         # Remove download folder of the user
         rmtree(dtmp[1])
     except:

megadl/modules/mega_dl.py

@@ -32,7 +32,7 @@ async def dl_from(_: MeganzClient, msg: Message):
             [
                 [InlineKeyboardButton("Download 💾", callback_data=f"dwn_mg-{msg.id}")],
                 [InlineKeyboardButton("Info ℹ️", callback_data=f"info_mg-{msg.id}")],
-                [InlineKeyboardButton("Close ❌", callback_data="closeqcb")],
+                [InlineKeyboardButton("Cancel ❌", callback_data="cancelqcb")],
             ]
         ),
     )
@@ -64,7 +64,17 @@ 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)
+    f_list = await cli.download(
+        url,
+        qcid,
+        resp.id,
+        path=dlid,
+        reply_markup=InlineKeyboardMarkup(
+            [
+                [InlineKeyboardButton("Cancel ❌", callback_data="cancelqcb")],
+            ]
+        ),
+    )
     try:
         await query.edit_message_text("Successfully downloaded the content 🥳")
     except Exception as e:

megadl/modules/mega_up.py

@@ -38,7 +38,7 @@ async def to_up(_: MeganzClient, msg: Message):
         reply_markup=InlineKeyboardMarkup(
             [
                 [InlineKeyboardButton("Download 💾", callback_data=f"up_tgdl-{msg.id}")],
-                [InlineKeyboardButton("Close ❌", callback_data="closeqcb")],
+                [InlineKeyboardButton("Cancel ❌", callback_data="cancelqcb")],
             ]
         ),
     )