diff --git a/bot/__init__.py b/bot/__init__.py
index 38fc273..57ea6d8 100644
--- a/bot/__init__.py
+++ b/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")
diff --git a/bot/helper/bot_utils.py b/bot/helper/bot_utils.py
index b6cf8cd..2573057 100644
--- a/bot/helper/bot_utils.py
+++ b/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 += "Name: {}\n" \
+ "status: {}\n" \
+ "Downloaded: {} of {}\n" \
+ "Speed: {}\n" \
+ "ETA: {}\n\n".format(status.name(), status.status(),
+ status.progress(), status.size(),
+ status.speed(), status.eta())
+ return msg
diff --git a/bot/helper/download_status.py b/bot/helper/download_status.py
index 171f772..d1f4f37 100644
--- a/bot/helper/download_status.py
+++ b/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):
diff --git a/bot/helper/download_tools.py b/bot/helper/download_tools.py
index a779175..ac74c45 100644
--- a/bot/helper/download_tools.py
+++ b/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
diff --git a/bot/helper/fs_utils.py b/bot/helper/fs_utils.py
index bd840a7..50c34dc 100644
--- a/bot/helper/fs_utils.py
+++ b/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):
diff --git a/bot/helper/gdriveTools.py b/bot/helper/gdriveTools.py
index 59bc607..3c0b884 100644
--- a/bot/helper/gdriveTools.py
+++ b/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
diff --git a/bot/helper/message_utils.py b/bot/helper/message_utils.py
index 605463d..65ba56f 100644
--- a/bot/helper/message_utils.py
+++ b/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
diff --git a/bot/mirror.py b/bot/mirror.py
index 4d2d673..2ac19b9 100644
--- a/bot/mirror.py
+++ b/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_list
+from bot import download_dict, status_reply_dict, DOWNLOAD_STATUS_UPDATE_INTERVAL
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 += "Name: {}\n" \
- "status: {}\n" \
- "Downloaded: {} of {}\n" \
- "Speed: {}\n" \
- "ETA: {}\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 = '{}'.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):
- pass
+def mirror_status(update, context):
+ 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)