evanferrao
/
dump-mirrorbot
mirror of https://github.com/Box-boi/dump-mirrorbot.git
1
0
Fork 0
Code Issues Actions Packages Projects Releases Wiki Activity

Introduce message_utils module managing telegram messages 

Signed-off-by: lzzy12 <jhashivam2020@gmail.com>
This commit is contained in:
lzzy12 2019-09-27 19:12:25 +05:30
parent 86d624214d
commit c2f1fbadff
8 changed files with 77 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
bot/__init__.py
View File

@ -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):

10
bot/helper/download_status.py
View File

@ -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

14
bot/helper/download_tools.py
View File

@ -1,13 +1,12 @@
from time import sleep
from bot import LOGGER, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, download_list, aria2
from bot.helper.listeners import *
from bot import DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, aria2
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)

5
bot/helper/fs_utils.py
View File

@ -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()
aria2.remove_all(True)
for download in aria2.get_downloads():
download.remove()
shutil.rmtree(DOWNLOAD_DIR)
sys.exit(0)
except KeyboardInterrupt:

19
bot/helper/gdriveTools.py
View File

@ -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 .listeners import MirrorListeners
from .fs_utils import clean_download, get_mime_type
from bot import LOGGER, parent_id, DOWNLOAD_DIR
from .fs_utils import 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"
@ -63,28 +62,24 @@ class GoogleDriveHelper:
if os.path.isfile(file_path):
mime_type = get_mime_type(file_path)
try:
g_drive_link = self.upload_file(file_path, file_name, mime_type,parent_id)
g_drive_link = self.upload_file(file_path, file_name, mime_type, parent_id)
LOGGER.info("Uploaded To G-Drive: " + file_path)
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)
dir_id = self.create_directory(os.path.basename(os.path.abspath(file_name)), parent_id)
self.upload_dir(file_path,dir_id)
LOGGER.info("Uploaded To G-Drive: " + file_name)
link = "https://drive.google.com/folderview?id={}".format(dir_id)
except Exception as e:
LOGGER.error(str(e))
self.__listener.onUploadError(str(e))
clean_download(file_dir)
raise Exception('Error: {}'.format(str(e)))
del download_list[self.__listener.update.update_id]
self.__listener.onUploadError(str(e), _list, index)
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):

4
bot/helper/listeners.py
View File

@ -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

24
bot/helper/message_utils.py Normal file
View File

@ -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)

61
bot/mirror.py
View File

@ -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,
chat_id=self.reply_message.chat.id,
parse_mode='HTMl')
LOGGER.info("Editing message")
editMessage(msg, self.context, self.reply_message)
def onDownloadComplete(self, progress_status_list, index: int):
msg = get_readable_message(progress_status_list)
LOGGER.info("Download completed")
try:
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
editMessage(msg, self.context, self.reply_message)
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,
chat_id=self.reply_message.chat.id,
parse_mode='HTMl')
editMessage(error, self.context, self.reply_message)
fs_utils.clean_download(progress_status_list[index].path())
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,
chat_id=self.reply_message.chat.id,
parse_mode='HTMl')
except BadRequest:
pass
msg = get_readable_message(progress_status_list)
editMessage(msg, self.context, self.reply_message)
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)
self.context.bot.send_message(chat_id=self.update.message.chat_id,
reply_to_message_id=self.update.message.message_id,
text=msg, parse_mode='HTMl')
deleteMessage(self.context, self.reply_message)
sendMessage(msg, self.context, self.update)
fs_utils.clean_download(progress_status_list[index].path())
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,
chat_id=self.reply_message.chat.id,
parse_mode='HTMl')
editMessage(error, self.context, self.reply_message)
fs_utils.clean_download(progress_status[index].path())
@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,
text="Starting Download")
reply_msg = sendMessage('Starting Download', context, update)
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)