Implement mirror_status: invoked by /status

Signed-off-by: lzzy12 <jhashivam2020@gmail.com>

bot/__init__.py

@@ -33,6 +33,9 @@ aria2 = aria2p.API(
 
 DOWNLOAD_DIR = None
 BOT_TOKEN = None
+
+status_reply_dict = {}
+download_dict = {}
 try:
     BOT_TOKEN = getConfig('BOT_TOKEN')
     parent_id = getConfig('GDRIVE_FOLDER_ID')
@@ -40,7 +43,6 @@ try:
     if DOWNLOAD_DIR[-1] != '/' or DOWNLOAD_DIR[-1] != '\\':
         DOWNLOAD_DIR = DOWNLOAD_DIR + '/'
     DOWNLOAD_STATUS_UPDATE_INTERVAL = int(getConfig('DOWNLOAD_STATUS_UPDATE_INTERVAL'))
-    download_list = {}
 
 except KeyError as e:
     LOGGER.error("One or more env variables missing! Exiting now")

bot/helper/bot_utils.py

@@ -1,12 +1,12 @@
-from bot import download_list
+from bot import download_dict
 
 
 def get_download(update_id):
-    return download_list[update_id].download()
+    return download_dict[update_id].download()
 
 
 def get_download_status_list():
-    return list(download_list.values())
+    return list(download_dict.values())
 
 
 def get_download_index(_list, gid):
@@ -19,6 +19,19 @@ def get_download_index(_list, gid):
 
 def get_download_str():
     result = ""
-    for status in list(download_list.values()):
+    for status in list(download_dict.values()):
         result += (status.progress() + status.speed() + status.status())
     return result
+
+
+def get_readable_message(progress_list: list = download_dict.values()):
+    msg = ""
+    for status in progress_list:
+        msg += "<b>Name:</b> {}\n" \
+               "<b>status:</b> {}\n" \
+               "<b>Downloaded:</b> {} of {}\n" \
+               "<b>Speed:</b> {}\n" \
+               "<b>ETA:</b> {}\n\n".format(status.name(), status.status(),
+                                           status.progress(), status.size(),
+                                           status.speed(), status.eta())
+    return msg

bot/helper/download_status.py

@@ -1,4 +1,4 @@
-from bot import aria2, download_list, DOWNLOAD_DIR
+from bot import aria2, download_dict, DOWNLOAD_DIR
 
 
 def get_download(gid):

bot/helper/download_tools.py

@@ -32,7 +32,7 @@ class DownloadHelper:
         else:
             self.__listener.onDownloadError("No download URL or URL malformed")
             return
-        download_list[self.__listener.update.update_id] = DownloadStatus(download.gid, self.__listener.update.update_id)
+        download_dict[self.__listener.update.update_id] = DownloadStatus(download.gid, self.__listener.update.update_id)
         self.__listener.onDownloadStarted(link)
         self.__update_download_status()
 
@@ -56,7 +56,7 @@ class DownloadHelper:
                 sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL)
                 new_gid = self.__get_followed_download_gid()
                 self.__listener.onDownloadProgress(get_download_status_list(), index)
-            download_list[self.__listener.update.update_id] = DownloadStatus(new_gid, self.__listener.update.update_id)
+            download_dict[self.__listener.update.update_id] = DownloadStatus(new_gid, self.__listener.update.update_id)
 
         # Start tracking the actual download
         previous = None

bot/helper/fs_utils.py

@@ -8,7 +8,10 @@ import mimetypes
 
 def clean_download(path: str):
     if os.path.exists(path):
-        shutil.rmtree(path)
+        if os.path.isfile(path):
+            os.remove(path)
+        else:
+            shutil.rmtree(path)
 
 
 def exit_clean_up(signal, frame):

bot/helper/gdriveTools.py

@@ -77,7 +77,7 @@ class GoogleDriveHelper:
             except Exception as e:
                 LOGGER.error(str(e))
                 self.__listener.onUploadError(str(e), _list, index)
-        LOGGER.info(download_list)
+        LOGGER.info(download_dict)
         self.__listener.onUploadComplete(link, _list, index)
         LOGGER.info("Deleting downloaded file/folder..")
         return link

bot/helper/message_utils.py

@@ -1,4 +1,3 @@
-from telegram.ext import CommandHandler, run_async
 from telegram.error import BadRequest
 from telegram.message import Message
 from telegram.update import Update

bot/mirror.py

@@ -1,27 +1,14 @@
-
+from telegram.ext import CommandHandler, run_async
 from bot.helper import download_tools, gdriveTools, listeners
 from bot import LOGGER, dispatcher
 from bot.helper import fs_utils
-from bot.helper.download_status import DownloadStatus
+from bot import download_dict, status_reply_dict, DOWNLOAD_STATUS_UPDATE_INTERVAL
-from bot import download_list
 from bot.helper.message_utils import *
+from time import sleep
+from bot.helper.bot_utils import get_readable_message
 LOGGER.info('mirror.py')
 
 
-def get_readable_message(progress_list: list):
-    msg = ""
-    LOGGER.info(progress_list)
-    for status in progress_list:
-        msg += "<b>Name:</b> {}\n" \
-               "<b>status:</b> {}\n" \
-               "<b>Downloaded:</b> {} of {}\n" \
-               "<b>Speed:</b> {}\n" \
-               "<b>ETA:</b> {}\n\n".format(status.name(), status.status(),
-                                           status.progress(), status.size(),
-                                           status.speed(), status.eta())
-    return msg
-
-
 class MirrorListener(listeners.MirrorListeners):
     def __init__(self, context, update, reply_message):
         super().__init__(context, update, reply_message)
@@ -44,6 +31,7 @@ class MirrorListener(listeners.MirrorListeners):
     def onDownloadError(self, error, progress_status_list: list, index: int):
         LOGGER.error(error)
         editMessage(error, self.context, self.reply_message)
+        del download_dict[self.update.update_id]
         fs_utils.clean_download(progress_status_list[index].path())
 
     def onUploadStarted(self, progress_status_list: list, index: int):
@@ -54,11 +42,13 @@ class MirrorListener(listeners.MirrorListeners):
         msg = '<a href="{}">{}</a>'.format(link, progress_status_list[index].name())
         deleteMessage(self.context, self.reply_message)
         sendMessage(msg, self.context, self.update)
+        del download_dict[self.update.update_id]
         fs_utils.clean_download(progress_status_list[index].path())
 
     def onUploadError(self, error: str, progress_status: list, index: int):
         LOGGER.error(error)
         editMessage(error, self.context, self.reply_message)
+        del download_dict[self.update.update_id]
         fs_utils.clean_download(progress_status[index].path())
 
 
@@ -73,9 +63,31 @@ def mirror(update, context):
 
 
 @run_async
-def cancel_mirror(update, context):
+def mirror_status(update, context):
-    pass
+    try:
+        deleteMessage(context, status_reply_dict[update.effective_chat])
+        del status_reply_dict[update.effective_chat]
+    except KeyError:
+        pass
+
+    while True:
+        message = get_readable_message()
+        if len(message) == 0:
+            message = "No active downloads"
+            try:
+                deleteMessage(context, status_reply_dict[update.effective_chat])
+                del status_reply_dict[update.effective_chat]
+            except KeyError:
+                pass
+            break
+        try:
+            editMessage(message, context, status_reply_dict[update.effective_chat])
+        except KeyError:
+            status_reply_dict[update.effective_chat] = sendMessage(message, context, update)
+        sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL)
 
 
 mirror_handler = CommandHandler('mirror', mirror)
+mirror_status_handler = CommandHandler('status', mirror_status)
 dispatcher.add_handler(mirror_handler)
+dispatcher.add_handler(mirror_status_handler)