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

Random fixes all around the source 

Signed-off-by: lzzy12 <jhashivam2020@gmail.com>
This commit is contained in:
lzzy12 2019-10-11 00:07:46 +05:30
parent c683c7732f
commit 9d83cd31c7
6 changed files with 52 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
bot/__init__.py
View File

@ -3,6 +3,7 @@ import configparser
import aria2p
import threading
from telegram.ext import Updater
import os
logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
level=logging.INFO)
@ -25,12 +26,12 @@ except KeyError:
pass
aria2 = aria2p.API(
aria2p.Client(
host="http://localhost",
port=6800,
secret="",
)
)
aria2p.Client(
host="http://localhost",
port=6800,
secret="",
)
)
DOWNLOAD_DIR = None
BOT_TOKEN = None
@ -41,11 +42,12 @@ status_reply_dict = {}
download_dict = {}
# Stores list of users and chats the bot is authorized to use in
AUTHORIZED_CHATS = []
with open('authorized_chats.txt', 'r+') as f:
lines = f.readlines()
for line in lines:
LOGGER.info(line.split())
AUTHORIZED_CHATS.append(int(line.split()[0]))
if os.path.exists('authorized_chats.txt'):
with open('authorized_chats.txt', 'r+') as f:
lines = f.readlines()
for line in lines:
# LOGGER.info(line.split())
AUTHORIZED_CHATS.append(int(line.split()[0]))
try:
BOT_TOKEN = getConfig('BOT_TOKEN')
parent_id = getConfig('GDRIVE_FOLDER_ID')

8
bot/helper/ext_utils/exceptions.py
View File

@ -4,6 +4,14 @@ class DriveAuthError(Exception):
class KillThreadException(Exception):
""" Custom Exception class for killing thread as soon as they aren't needed"""
def __init__(self, message, error=None):
super().__init__(message)
self.error = error
class DownloadCancelled(Exception):
def __init__(self, message, error=None):
super().__init__(message)
self.error = error

15
bot/helper/mirror_utils/download_tools.py
View File

@ -24,13 +24,12 @@ class DownloadHelper:
return False
def add_download(self, link: str):
download = None
if self.is_url(link):
if link.endswith('.torrent'):
self.__is_torrent = True
download = aria2.add_uris([link], {'dir': DOWNLOAD_DIR + str(self.__listener.message.message_id)})
download = aria2.add_uris([link], {'dir': DOWNLOAD_DIR + str(self.__listener.uid)})
elif self.is_magnet(link):
download = aria2.add_magnet(link, {'dir': DOWNLOAD_DIR + str(self.__listener.message.message_id)})
download = aria2.add_magnet(link, {'dir': DOWNLOAD_DIR + str(self.__listener.uid)})
self.__is_torrent = True
else:
_list = get_download_status_list()
@ -38,12 +37,12 @@ class DownloadHelper:
return
with download_dict_lock:
download_dict[self.__listener.message.message_id] = DownloadStatus(download.gid,
self.__listener.message.message_id)
self.__listener.uid)
self.__listener.onDownloadStarted(link)
self.__update_download_status()
def __get_download(self):
return get_download(self.__listener.message.message_id)
return get_download(self.__listener.uid)
def __get_followed_download_gid(self):
download = self.__get_download()
@ -59,8 +58,8 @@ class DownloadHelper:
if self.__is_torrent:
# Waiting for the actual gid
new_gid = None
download = self.__get_download()
while new_gid is None:
download = self.__get_download()
if download.has_failed:
self.__listener.onDownloadError(download.error_message, status_list, index)
return
@ -69,12 +68,12 @@ class DownloadHelper:
return
sleep(DOWNLOAD_STATUS_UPDATE_INTERVAL)
if should_update:
# Check every few seconds
new_gid = self.__get_followed_download_gid()
try:
self.__listener.onDownloadProgress(get_download_status_list(), index)
except KillThreadException:
should_update = False
# Check every few seconds
new_gid = self.__get_followed_download_gid()
with download_dict_lock:
download_dict[self.__listener.message.message_id] = DownloadStatus(new_gid,
self.__listener.message.message_id)

10
bot/modules/cancel_mirror.py
View File

@ -1,7 +1,9 @@
from telegram.ext import CommandHandler, run_async
from bot.helper.telegram_helper.message_utils import *
from bot import download_dict, aria2, dispatcher, download_dict_lock
from bot import download_dict, aria2, dispatcher, download_dict_lock, DOWNLOAD_DIR
from bot.helper.telegram_helper.filters import CustomFilters
from bot.helper.ext_utils.fs_utils import clean_download
@run_async
def cancel_mirror(update: Update, context):
@ -10,13 +12,15 @@ def cancel_mirror(update: Update, context):
keys = download_dict.keys()
download = download_dict[mirror_message.message_id].download()
if mirror_message is None or mirror_message.message_id not in keys:
if '/mirror' in mirror_message.text:
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, context, update)
return
if len(download.followed_by_ids) != 0:
downloads = aria2.get_downloads(download.followed_by_ids)
aria2.pause(downloads)
aria2.pause([download])
sendMessage("Download canceled", context, update)

21
bot/modules/mirror.py
View File

@ -9,6 +9,8 @@ from bot.helper.ext_utils.bot_utils import get_readable_message, MirrorStatus
from bot.helper.ext_utils.exceptions import KillThreadException
from bot.helper.telegram_helper.filters import CustomFilters
import pathlib
import threading
class MirrorListener(listeners.MirrorListeners):
def __init__(self, context, update, reply_message, isTar=False):
@ -19,9 +21,10 @@ class MirrorListener(listeners.MirrorListeners):
LOGGER.info("Adding link: " + link)
def onDownloadProgress(self, progress_status_list: list, index: int):
if progress_status_list[index].status() == MirrorStatus.STATUS_CANCELLED:
raise KillThreadException('Mirror cancelled by user')
msg = get_readable_message(progress_status_list)
if progress_status_list[index].status() == MirrorStatus.STATUS_CANCELLED:
editMessage(msg, self.context, self.reply_message)
raise KillThreadException('Mirror cancelled by user')
# LOGGER.info("Editing message")
try:
editMessage(msg, self.context, self.reply_message)
@ -50,9 +53,9 @@ class MirrorListener(listeners.MirrorListeners):
deleteMessage(self.context, status_reply_dict[self.update.effective_chat.id])
del status_reply_dict[self.update.effective_chat.id]
if index is not None:
fs_utils.clean_download(progress_status_list[index].path())
with download_dict_lock:
del download_dict[self.message.message_id]
fs_utils.clean_download(progress_status_list[index].path())
msg = f"@{self.message.from_user.username} your download has been stopped due to: {error}"
sendMessage(msg, self.context, self.update)
@ -92,13 +95,20 @@ class MirrorListener(listeners.MirrorListeners):
def _mirror(update, context, isTar=False):
message_args = update.message.text.split(' ')
link = message_args[1]
try:
link = message_args[1]
except KeyError:
link = ''
LOGGER.info(link)
link = link.strip()
if len(link) == 0 and update.message.reply_to_message is not None:
document = update.message.reply_to_message.document
if document is not None and document.mime_type == "application/x-bittorrent":
link = document.get_file().file_path
else:
sendMessage('No download source provided', context, update)
return
reply_msg = sendMessage('Starting Download', context, update)
index = update.effective_chat.id
with status_reply_dict_lock:
@ -107,7 +117,8 @@ def _mirror(update, context, isTar=False):
status_reply_dict[index] = reply_msg
listener = MirrorListener(context, update, reply_msg, isTar)
aria = download_tools.DownloadHelper(listener)
aria.add_download(link)
t = threading.Thread(target=aria.add_download, args=(link,))
t.start()
@run_async

1
bot/modules/mirror_status.py
View File

@ -6,6 +6,7 @@ from bot.helper.ext_utils.bot_utils import get_readable_message
from telegram.error import BadRequest
from bot.helper.telegram_helper.filters import CustomFilters
@run_async
def mirror_status(update: Update, context):
message = get_readable_message()