Introduce message_utils module managing telegram messages

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

bot/__init__.py

@@ -7,7 +7,7 @@ logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s
                     level=logging.INFO)
 
 config = configparser.ConfigParser()
-config.read('bot/conf.ini')
+config.read('bot/config.ini')
 
 
 def getConfig(name: str):

bot/helper/download_status.py

@@ -1,4 +1,4 @@
-from bot import aria2
+from bot import aria2, download_list, DOWNLOAD_DIR
 
 
 def get_download(gid):
@@ -11,9 +11,10 @@ class DownloadStatus:
     STATUS_WAITING = "Queued"
     STATUS_FAILED = "Failed. Cleaning download"
 
-    def __init__(self, gid):
+    def __init__(self, gid, update_id):
         self.__gid = gid
         self.__download = get_download(gid)
+        self.__uid = update_id
 
     def __update(self):
         self.__download = get_download(self.__gid)
@@ -29,6 +30,9 @@ class DownloadStatus:
     def name(self):
         return self.__download.name
 
+    def path(self):
+        return "{}{}/{}".format(DOWNLOAD_DIR, self.__uid, self.name())
+
     def size(self):
         return self.__download.total_length_string()
 
@@ -45,6 +49,8 @@ class DownloadStatus:
             # If download exists and is complete the it must be uploading
             # otherwise the gid would have been removed from the download_list
             status = DownloadStatus.STATUS_UPLOADING
+        elif self.__download.has_failed:
+            status = DownloadStatus.STATUS_FAILED
         elif self.__download.is_active:
             status = DownloadStatus.STATUS_DOWNLOADING
         return status

bot/helper/download_tools.py

@@ -1,13 +1,12 @@
 from time import sleep
-from bot import LOGGER, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, download_list, aria2
+from bot import DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, aria2
-from bot.helper.listeners import *
 from .download_status import DownloadStatus
 from .bot_utils import *
 
 
 class DownloadHelper:
 
-    def __init__(self, listener: MirrorListeners):
+    def __init__(self, listener=None):
         self.__listener = listener
         self.__is_torrent = False
 
@@ -33,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)
+        download_list[self.__listener.update.update_id] = DownloadStatus(download.gid, self.__listener.update.update_id)
         self.__listener.onDownloadStarted(link)
         self.__update_download_status()
 
@@ -57,17 +56,18 @@ 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)
+            download_list[self.__listener.update.update_id] = DownloadStatus(new_gid, self.__listener.update.update_id)
 
         # Start tracking the actual download
         previous = None
         download = self.__get_download()
         while not download.is_complete:
             status_list = get_download_status_list()
+            index = get_download_index(status_list, self.__get_download().gid)
             if self.__get_download().has_failed:
-                self.__listener.onDownloadError(self.__get_download().error_message)
+                self.__listener.onDownloadError(self.__get_download().error_message, status_list[index])
                 break
-            # TODO: Find a better way
+            # TODO: Find a better way to differentiate between 2 list of objects
             progress_str_list = get_download_str()
             if progress_str_list != previous:
                 self.__listener.onDownloadProgress(status_list, index)

bot/helper/fs_utils.py

@@ -5,6 +5,7 @@ import os
 import pathlib
 import mimetypes
 
+
 def clean_download(path: str):
     if os.path.exists(path):
         shutil.rmtree(path)
@@ -13,8 +14,8 @@ def clean_download(path: str):
 def exit_clean_up(signal, frame):
     try:
         LOGGER.info("Please wait, while we clean up the downloads and stop running downloads")
-        aria2.autopurge()
+        for download in aria2.get_downloads():
-        aria2.remove_all(True)
+            download.remove()
         shutil.rmtree(DOWNLOAD_DIR)
         sys.exit(0)
     except KeyboardInterrupt:

bot/helper/gdriveTools.py

@@ -4,15 +4,14 @@ from google.auth.transport.requests import Request
 from googleapiclient.http import MediaFileUpload
 import pickle
 import os
-from bot import LOGGER, parent_id, DOWNLOAD_DIR, download_list
+from bot import LOGGER, parent_id, DOWNLOAD_DIR
-from .listeners import MirrorListeners
+from .fs_utils import get_mime_type
-from .fs_utils import clean_download, get_mime_type
 from .bot_utils import *
 
 
 class GoogleDriveHelper:
 
-    def __init__(self, listener: MirrorListeners):
+    def __init__(self, listener=None):
         self.__G_DRIVE_TOKEN_FILE = "token.pickle"
         # Check https://developers.google.com/drive/scopes for all available scopes
         self.__OAUTH_SCOPE = "https://www.googleapis.com/auth/drive.file"
@@ -68,7 +67,7 @@ class GoogleDriveHelper:
                 link = g_drive_link
             except Exception as e:
                 LOGGER.error(str(e))
-                pass
+                self.__listener.onUploadError(str(e), _list, index)
         else:
             try:
                 dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_id)
@@ -77,14 +76,10 @@ class GoogleDriveHelper:
                 link = "https://drive.google.com/folderview?id={}".format(dir_id)
             except Exception as e:
                 LOGGER.error(str(e))
-                self.__listener.onUploadError(str(e))
+                self.__listener.onUploadError(str(e), _list, index)
-                clean_download(file_dir)
-                raise Exception('Error: {}'.format(str(e)))
-        del download_list[self.__listener.update.update_id]
         LOGGER.info(download_list)
         self.__listener.onUploadComplete(link, _list, index)
         LOGGER.info("Deleting downloaded file/folder..")
-        clean_download(file_dir)
         return link
 
     def create_directory(self, directory_name,parent_id):

bot/helper/listeners.py

@@ -13,7 +13,7 @@ class MirrorListeners:
     def onDownloadComplete(self, progress_status_list: list, index: int):
         raise NotImplementedError
 
-    def onDownloadError(self, error: str):
+    def onDownloadError(self, error: str, progress_status_list: list, index: int):
         raise NotImplementedError
 
     def onUploadStarted(self, progress_status_list: list, index: int):
@@ -22,5 +22,5 @@ class MirrorListeners:
     def onUploadComplete(self, link: str, progress_status_list: list, index: int):
         raise NotImplementedError
 
-    def onUploadError(self, error: str):
+    def onUploadError(self, error: str, progress_status: list, index: int):
         raise NotImplementedError

bot/helper/message_utils.py

@@ -0,0 +1,24 @@
+from telegram.ext import CommandHandler, run_async
+from telegram.error import BadRequest
+from telegram.message import Message
+from telegram.update import Update
+
+
+def sendMessage(text: str, context, update: Update):
+    return context.bot.send_message(update.message.chat_id,
+                                    reply_to_message_id=update.message.message_id,
+                                    text=text, parse_mode='HTMl')
+
+
+def editMessage(text: str, context, message: Message):
+    try:
+        context.bot.edit_message_text(text=text, message_id=message.message_id,
+                                      chat_id=message.chat.id,
+                                      parse_mode='HTMl')
+    except BadRequest:
+        pass
+
+
+def deleteMessage(context, message: Message):
+    context.bot.delete_message(chat_id=message.chat.id,
+                               message_id=message.message_id)

bot/mirror.py

@@ -1,8 +1,10 @@
-from telegram.ext import CommandHandler, run_async
+
-from telegram.error import BadRequest
 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.helper.message_utils import *
 LOGGER.info('mirror.py')
 
 
@@ -29,64 +31,51 @@ class MirrorListener(listeners.MirrorListeners):
 
     def onDownloadProgress(self, progress_status_list: list, index: int):
         msg = get_readable_message(progress_status_list)
-        if self.reply_message.text != msg:
         LOGGER.info("Editing message")
-            self.context.bot.edit_message_text(text=msg, message_id=self.reply_message.message_id,
+        editMessage(msg, self.context, self.reply_message)
-                                               chat_id=self.reply_message.chat.id,
-                                               parse_mode='HTMl')
 
     def onDownloadComplete(self, progress_status_list, index: int):
         msg = get_readable_message(progress_status_list)
         LOGGER.info("Download completed")
-        try:
+        editMessage(msg, self.context, self.reply_message)
-            LOGGER.info("Editing message")
-            self.context.bot.edit_message_text(text=msg, message_id=self.reply_message.message_id,
-                                               chat_id=self.reply_message.chat.id,
-                                               parse_mode='HTMl')
-        except BadRequest:
-            pass
         gdrive = gdriveTools.GoogleDriveHelper(self)
         gdrive.upload(progress_status_list[index].name())
 
-    def onDownloadError(self, error):
+    def onDownloadError(self, error, progress_status_list: list, index: int):
         LOGGER.error(error)
-        self.context.bot.edit_message_text(text=error, message_id=self.reply_message.message_id,
+        editMessage(error, self.context, self.reply_message)
-                                           chat_id=self.reply_message.chat.id,
+        fs_utils.clean_download(progress_status_list[index].path())
-                                           parse_mode='HTMl')
 
     def onUploadStarted(self, progress_status_list: list, index: int):
-        try:
         msg = get_readable_message(progress_status_list)
-            self.context.bot.edit_message_text(text=msg, message_id=self.reply_message.message_id,
+        editMessage(msg, self.context, self.reply_message)
-                                               chat_id=self.reply_message.chat.id,
-                                               parse_mode='HTMl')
-        except BadRequest:
-            pass
 
     def onUploadComplete(self, link: str, progress_status_list: list, index: int):
         msg = '<a href="{}">{}</a>'.format(link, progress_status_list[index].name())
-        self.context.bot.delete_message(chat_id=self.reply_message.chat.id, message_id=self.reply_message.message_id)
+        deleteMessage(self.context, self.reply_message)
-        self.context.bot.send_message(chat_id=self.update.message.chat_id,
+        sendMessage(msg, self.context, self.update)
-                                      reply_to_message_id=self.update.message.message_id,
+        fs_utils.clean_download(progress_status_list[index].path())
-                                      text=msg, parse_mode='HTMl')
 
-    def onUploadError(self, error: str):
+    def onUploadError(self, error: str, progress_status: list, index: int):
         LOGGER.error(error)
-        self.context.bot.edit_message_text(text=error, message_id=self.reply_message.message_id,
+        editMessage(error, self.context, self.reply_message)
-                                           chat_id=self.reply_message.chat.id,
+        fs_utils.clean_download(progress_status[index].path())
-                                           parse_mode='HTMl')
 
 
 @run_async
 def mirror(update, context):
     message = update.message.text
     link = message.replace('/mirror', '')[1:]
-    reply_msg = context.bot.send_message(chat_id=update.message.chat_id, reply_to_message_id=update.message.message_id,
+    reply_msg = sendMessage('Starting Download', context, update)
-                                         text="Starting Download")
     listener = MirrorListener(context, update, reply_msg)
     aria = download_tools.DownloadHelper(listener)
     aria.add_download(link)
 
 
+@run_async
+def cancel_mirror(update, context):
+    pass
+
+
 mirror_handler = CommandHandler('mirror', mirror)
 dispatcher.add_handler(mirror_handler)