From a44d0da9f6da111d789155f88162abd0f1079a2d Mon Sep 17 00:00:00 2001
From: jaskaranSM <beerabassian32@gmail.com>
Date: Sat, 21 Mar 2020 14:28:38 +0530
Subject: [PATCH] Implement cancellation of downloads by gid

---
 bot/helper/ext_utils/bot_utils.py |  9 +++++++
 bot/modules/cancel_mirror.py      | 44 ++++++++++++++++++++-----------
 2 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py
index ea8f598ad..5e42b9af4 100644
--- a/bot/helper/ext_utils/bot_utils.py
+++ b/bot/helper/ext_utils/bot_utils.py
@@ -54,6 +54,14 @@ def get_readable_file_size(size_in_bytes) -> str:
     except IndexError:
         return 'File too large'
 
+def getDownloadByGid(gid):
+    with download_dict_lock:
+        for dl in download_dict.values():
+            if dl.status() == MirrorStatus.STATUS_DOWNLOADING:
+                if dl.download().gid == gid:
+                    return dl
+    return None
+
 
 def get_progress_bar_string(status):
     completed = status.processed_bytes() / 8
@@ -95,6 +103,7 @@ def get_readable_message():
                 if hasattr(download, 'is_torrent'):
                     msg += f"| P: {download.download().connections} " \
                            f"| S: {download.download().num_seeders}"
+                msg += f"\nGID: <code>{download.download().gid}</code>"
             msg += "\n\n"
         return msg
 
diff --git a/bot/modules/cancel_mirror.py b/bot/modules/cancel_mirror.py
index b69079d6c..fd8eae477 100644
--- a/bot/modules/cancel_mirror.py
+++ b/bot/modules/cancel_mirror.py
@@ -5,20 +5,36 @@
 from bot.helper.ext_utils.fs_utils import clean_download
 from bot.helper.telegram_helper.bot_commands import BotCommands
 from time import sleep
+from bot.helper.ext_utils.bot_utils import getDownloadByGid
 
 @run_async
 def cancel_mirror(bot,update):
-    mirror_message = update.message.reply_to_message
-    with download_dict_lock:
-        keys = download_dict.keys()
-        dl = download_dict[mirror_message.message_id]
-    if mirror_message is None or mirror_message.message_id not in keys:
-        if '/mirror' in mirror_message.text or '/tarmirror' in mirror_message.text:
-            msg = 'Message has already been cancelled'
-        else:
-            msg = 'Please reply to the /mirror message which was used to start the download to cancel it!'
-        sendMessage(msg, bot, update)
-        return
+    args = update.message.text.split(" ",maxsplit=1)
+    mirror_message = None
+    if len(args) > 1:
+        gid = args[1]
+        dl = getDownloadByGid(gid)
+        if not dl:
+            sendMessage(f"GID: <code>{gid}</code> not found.",bot,update)
+            return
+        with download_dict_lock:
+            keys = list(download_dict.keys())
+        mirror_message = dl._listener.message
+    elif update.message.reply_to_message:
+        mirror_message = update.message.reply_to_message
+        with download_dict_lock:
+            keys = list(download_dict.keys())
+            dl = download_dict[mirror_message.message_id]
+    if len(args) == 1:
+        if mirror_message is None or mirror_message.message_id not in keys:
+            if BotCommands.MirrorCommand in mirror_message.text or BotCommands.TarMirrorCommand in mirror_message.text:
+                msg = "Mirror already have been cancelled"
+                sendMessage(msg,bot,update)
+                return
+            else:
+                msg = "Please reply to the /mirror message which was used to start the download or /cancel gid to cancel it!"
+                sendMessage(msg,bot,update)
+                return
     if dl.status() == "Uploading":
         sendMessage("Upload in Progress, Don't Cancel it.", bot, update)
         return
@@ -31,13 +47,9 @@ def cancel_mirror(bot,update):
             downloads = aria2.get_downloads(download.followed_by_ids)
             aria2.pause(downloads)
         aria2.pause([download])
-
-    elif dl.status() == "Uploading":
-        sendMessage("Upload in Progress, Dont Cancel it.",bot,update)
-        return
     else:
         dl._listener.onDownloadError("Download stopped by user!")
-    sleep(1) #Wait a Second For Aria2 To free Resources. 
+    sleep(1)  # Wait a Second For Aria2 To free Resources.
     clean_download(f'{DOWNLOAD_DIR}{mirror_message.message_id}/')