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

Implement Centralized Aria Websocket Listener

This commit is contained in:
jaskaranSM 2020-05-15 13:34:09 +05:30
parent f4290210b1
commit cdefe12348
5 changed files with 89 additions and 64 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
bot/helper/ext_utils/bot_utils.py
View File

@ -61,7 +61,7 @@ def get_readable_file_size(size_in_bytes) -> str:
def getDownloadByGid(gid):
with download_dict_lock:
for dl in download_dict.values():
if dl.status() == MirrorStatus.STATUS_DOWNLOADING or dl.status() == MirrorStatus.STATUS_WAITING:
if dl.status() != MirrorStatus.STATUS_UPLOADING and dl.status() != MirrorStatus.STATUS_ARCHIVING:
if dl.gid() == gid:
return dl
return None
@ -135,3 +135,13 @@ def is_magnet(url: str):
if magnet:
return True
return False
def new_thread(fn):
"""To use as decorator to make a function call threaded.
Needs import
from threading import Thread"""
def wrapper(*args, **kwargs):
thread = threading.Thread(target=fn, args=args, kwargs=kwargs)
thread.start()
return thread
return wrapper

97
bot/helper/mirror_utils/download_utils/aria2_download.py
View File

@ -1,85 +1,78 @@
from bot import aria2
from bot import aria2, download_dict_lock
from bot.helper.ext_utils.bot_utils import *
from .download_helper import DownloadHelper
from bot.helper.mirror_utils.status_utils.aria_download_status import AriaDownloadStatus
from bot.helper.telegram_helper.message_utils import *
import threading
from aria2p import API
from time import sleep
class AriaDownloadHelper(DownloadHelper):
def __init__(self, listener):
def __init__(self):
super().__init__()
self.gid = None
self.__listener = listener
self._resource_lock = threading.RLock()
@new_thread
def __onDownloadStarted(self, api, gid):
with self._resource_lock:
LOGGER.info(f"onDownloadStart: {gid}")
if self.gid == gid:
download = api.get_download(gid)
self.name = download.name
update_all_messages()
LOGGER.info(f"onDownloadStart: {gid}")
update_all_messages()
def __onDownloadComplete(self, api: API, gid):
with self._resource_lock:
LOGGER.info(f"onDownloadComplete: {gid}")
if self.gid == gid:
download = api.get_download(gid)
if download.followed_by_ids:
self.gid = download.followed_by_ids[0]
with download_dict_lock:
download_dict[self.__listener.uid] = AriaDownloadStatus(self, self.__listener)
if download.is_torrent:
download_dict[self.__listener.uid].is_torrent = True
update_all_messages()
LOGGER.info(f'Changed gid from {gid} to {self.gid}')
else:
self.__listener.onDownloadComplete()
LOGGER.info(f"onDownloadComplete: {gid}")
dl = getDownloadByGid(gid)
download = api.get_download(gid)
if download.followed_by_ids:
new_gid = download.followed_by_ids[0]
with download_dict_lock:
download_dict[dl.uid()] = AriaDownloadStatus(new_gid,dl.getListener())
if download.is_torrent:
download_dict[dl.uid()].is_torrent = True
update_all_messages()
LOGGER.info(f'Changed gid from {gid} to {new_gid}')
else:
if dl: threading.Thread(target=dl.getListener().onDownloadComplete).start()
@new_thread
def __onDownloadPause(self, api, gid):
if self.gid == gid:
LOGGER.info("Called onDownloadPause")
self.__listener.onDownloadError('Download stopped by user!')
LOGGER.info(f"onDownloadPause: {gid}")
dl = getDownloadByGid(gid)
dl.getListener().onDownloadError('Download stopped by user!')
@new_thread
def __onDownloadStopped(self, api, gid):
if self.gid == gid:
LOGGER.info("Called on_download_stop")
self.__listener.onDownloadError('Download stopped by user!')
LOGGER.info(f"onDownloadStop: {gid}")
dl = getDownloadByGid(gid)
if dl: dl.getListener().onDownloadError('Download stopped by user!')
@new_thread
def __onDownloadError(self, api, gid):
with self._resource_lock:
if self.gid == gid:
download = api.get_download(gid)
error = download.error_message
LOGGER.info(f"Download Error: {error}")
self.__listener.onDownloadError(error)
sleep(0.5) #sleep for split second to ensure proper dl gid update from onDownloadComplete
LOGGER.info(f"onDownloadError: {gid}")
dl = getDownloadByGid(gid)
download = api.get_download(gid)
error = download.error_message
LOGGER.info(f"Download Error: {error}")
if dl: dl.getListener().onDownloadError(error)
def add_download(self, link: str, path):
def start_listener(self):
aria2.listen_to_notifications(threaded=True, on_download_start=self.__onDownloadStarted,
on_download_error=self.__onDownloadError,
on_download_pause=self.__onDownloadPause,
on_download_stop=self.__onDownloadStopped,
on_download_complete=self.__onDownloadComplete)
def add_download(self, link: str, path,listener):
if is_magnet(link):
download = aria2.add_magnet(link, {'dir': path})
else:
download = aria2.add_uris([link], {'dir': path})
with self._resource_lock:
self.gid = download.gid
if download.error_message: #no need to proceed further at this point
listener.onDownloadError(download.error_message)
return
with download_dict_lock:
download_dict[self.__listener.uid] = AriaDownloadStatus(self, self.__listener)
LOGGER.info(f"Started: {self.gid} DIR:{download.dir} ")
download_dict[listener.uid] = AriaDownloadStatus(download.gid,listener)
LOGGER.info(f"Started: {download.gid} DIR:{download.dir} ")
def cancel_download(self):
download = aria2.get_download(self.gid)
if download.is_waiting:
aria2.remove([download])
self.__listener.onDownloadError("Cancelled by user")
return
if len(download.followed_by_ids) != 0:
downloads = aria2.get_downloads(download.followed_by_ids)
aria2.pause(downloads)
aria2.pause([download])

33
bot/helper/mirror_utils/status_utils/aria_download_status.py
View File

@ -1,4 +1,4 @@
from bot import aria2, DOWNLOAD_DIR
from bot import aria2, DOWNLOAD_DIR, LOGGER
from bot.helper.ext_utils.bot_utils import MirrorStatus
from .status import Status
@ -9,13 +9,12 @@ def get_download(gid):
class AriaDownloadStatus(Status):
def __init__(self, obj, listener):
def __init__(self, gid, listener):
super().__init__()
self.upload_name = None
self.is_archiving = False
self.obj = obj
self.__gid = obj.gid
self.__download = get_download(obj.gid)
self.__gid = gid
self.__download = get_download(self.__gid)
self.__uid = listener.uid
self.__listener = listener
self.message = listener.message
@ -75,7 +74,16 @@ class AriaDownloadStatus(Status):
return self.__download
def download(self):
return self.obj
return self
def updateName(self,name):
self.__name = name
def updateGid(self,gid):
self.__gid = gid
def getListener(self):
return self.__listener
def uid(self):
return self.__uid
@ -83,3 +91,16 @@ class AriaDownloadStatus(Status):
def gid(self):
self.__update()
return self.__gid
def cancel_download(self):
LOGGER.info(f"Cancelling Download: {self.name()}")
download = self.aria_download()
if download.is_waiting:
aria2.remove([download])
self.__listener.onDownloadError("Cancelled by user")
return
if len(download.followed_by_ids) != 0:
downloads = aria2.get_downloads(download.followed_by_ids)
aria2.pause(downloads)
aria2.pause([download])

2
bot/helper/telegram_helper/message_utils.py
View File

@ -66,7 +66,7 @@ def update_all_messages():
msg = get_readable_message()
with status_reply_dict_lock:
for chat_id in list(status_reply_dict.keys()):
if msg != status_reply_dict[chat_id].text:
if status_reply_dict[chat_id] and msg != status_reply_dict[chat_id].text:
try:
editMessage(msg, status_reply_dict[chat_id])
except Exception as e:

9
bot/modules/mirror.py
View File

@ -6,7 +6,7 @@ from bot import dispatcher, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, downl
from bot.helper.ext_utils import fs_utils, bot_utils
from bot.helper.ext_utils.bot_utils import setInterval
from bot.helper.ext_utils.exceptions import DirectDownloadLinkException
from bot.helper.mirror_utils.download_utils import aria2_download
from bot.helper.mirror_utils.download_utils.aria2_download import AriaDownloadHelper
from bot.helper.mirror_utils.download_utils.direct_link_generator import direct_link_generator
from bot.helper.mirror_utils.download_utils.telegram_downloader import TelegramDownloadHelper
from bot.helper.mirror_utils.status_utils import listeners
@ -19,6 +19,9 @@ from bot.helper.telegram_helper.message_utils import *
import pathlib
import os
ariaDlManager = AriaDownloadHelper()
ariaDlManager.start_listener()
class MirrorListener(listeners.MirrorListeners):
@ -82,7 +85,6 @@ class MirrorListener(listeners.MirrorListeners):
del download_dict[self.uid]
LOGGER.info(f"Deleting folder: {download.path()}")
fs_utils.clean_download(download.path())
LOGGER.info(f"Deleting {download.name()} from download_dict.")
LOGGER.info(str(download_dict))
except Exception as e:
LOGGER.error(str(e))
@ -185,8 +187,7 @@ def _mirror(bot, update, isTar=False):
except DirectDownloadLinkException as e:
LOGGER.info(f'{link}: {e}')
listener = MirrorListener(bot, update, isTar, tag)
aria = aria2_download.AriaDownloadHelper(listener)
aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/')
ariaDlManager.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/',listener)
sendStatusMessage(update, bot)
if len(Interval) == 0:
Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages))