From b5a162f15d1715faa5529fa998cfcdb29f0cccbc Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Sat, 27 Jun 2020 23:40:15 +0530 Subject: [PATCH] Use subprocess.run to call extract shell script Signed-off-by: lzzy12 --- bot/helper/ext_utils/bot_utils.py | 2 ++ bot/helper/ext_utils/exceptions.py | 6 +++++ bot/helper/ext_utils/fs_utils.py | 21 +++++++++-------- bot/modules/mirror.py | 37 +++++++++++++++--------------- extract | 6 +++-- 5 files changed, 42 insertions(+), 30 deletions(-) diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index a015fb7..15a7feb 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/bot/helper/ext_utils/bot_utils.py @@ -137,9 +137,11 @@ def is_magnet(url: str): return True return False + def is_mega_link(url: str): return "mega.nz" in url + def new_thread(fn): """To use as decorator to make a function call threaded. Needs import diff --git a/bot/helper/ext_utils/exceptions.py b/bot/helper/ext_utils/exceptions.py index 25ff87f..a2f600c 100644 --- a/bot/helper/ext_utils/exceptions.py +++ b/bot/helper/ext_utils/exceptions.py @@ -1,2 +1,8 @@ class DirectDownloadLinkException(Exception): + """Not method found for extracting direct download link from the http link""" + pass + + +class NotSupportedExtractionArchive(Exception): + """The archive format use is trying to extract is not supported""" pass diff --git a/bot/helper/ext_utils/fs_utils.py b/bot/helper/ext_utils/fs_utils.py index 90951d6..8044c76 100644 --- a/bot/helper/ext_utils/fs_utils.py +++ b/bot/helper/ext_utils/fs_utils.py @@ -5,6 +5,7 @@ import os import pathlib import magic import tarfile +from .exceptions import NotSupportedExtractionArchive def clean_download(path: str): @@ -60,25 +61,25 @@ def get_base_name(orig_path: str): if orig_path.endswith(".tar.bz2"): return orig_path.replace(".tar.bz2", "") elif orig_path.endswith(".tar.gz"): - return orig_path.replace(".tar.gz","") + return orig_path.replace(".tar.gz", "") elif orig_path.endswith(".bz2"): - return orig_path.replace(".bz2","") + return orig_path.replace(".bz2", "") elif orig_path.endswith(".gz"): - return orig_path.replace(".gz","") + return orig_path.replace(".gz", "") elif orig_path.endswith(".tar"): - return orig_path.replace(".tar","") + return orig_path.replace(".tar", "") elif orig_path.endswith(".tbz2"): - return orig_path.replace("tbz2","") + return orig_path.replace("tbz2", "") elif orig_path.endswith(".tgz"): - return orig_path.replace(".tgz","") + return orig_path.replace(".tgz", "") elif orig_path.endswith(".zip"): - return orig_path.replace(".zip","") + return orig_path.replace(".zip", "") elif orig_path.endswith(".Z"): - return orig_path.replace(".Z","") + return orig_path.replace(".Z", "") elif orig_path.endswith(".rar"): - return orig_path.replace(".rar","") + return orig_path.replace(".rar", "") else: - return "unsupported" + raise NotSupportedExtractionArchive('File format not supported for extraction') def get_mime_type(file_path): diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index dc1c3fa..e86d0b3 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -5,7 +5,7 @@ from bot import Interval, INDEX_URL from bot import dispatcher, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, download_dict, download_dict_lock 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.ext_utils.exceptions import DirectDownloadLinkException, NotSupportedExtractionArchive from bot.helper.mirror_utils.download_utils.aria2_download import AriaDownloadHelper from bot.helper.mirror_utils.download_utils.mega_downloader import MegaDownloadHelper from bot.helper.mirror_utils.download_utils.direct_link_generator import direct_link_generator @@ -20,6 +20,8 @@ from bot.helper.telegram_helper.filters import CustomFilters from bot.helper.telegram_helper.message_utils import * import pathlib import os +import subprocess +import threading ariaDlManager = AriaDownloadHelper() ariaDlManager.start_listener() @@ -66,30 +68,29 @@ class MirrorListener(listeners.MirrorListeners): return elif self.extract: download.is_extracting = True - - path = fs_utils.get_base_name(m_path) - if path != "unsupported": + try: + path = fs_utils.get_base_name(m_path) LOGGER.info( - f"Extracting : {download_dict[self.uid].name()} " + f"Extracting : {name} " ) - download_dict[self.uid] = ExtractStatus(name, m_path, size) - os.system(f"extract '{m_path}'") - if not os.path.exists(path): - self.onUploadError("Cannot extract file, check integrity of the file") - return + with download_dict_lock: + download_dict[self.uid] = ExtractStatus(name, m_path, size) + archive_result = subprocess.run(["extract", m_path]) + if archive_result.returncode == 0: + threading.Thread(target=os.remove, args=(m_path,)).start() + LOGGER.info(f"Deleting archive : {m_path}") + else: + LOGGER.warning('Unable to extract archive! Uploading anyway') + path = f'{DOWNLOAD_DIR}{self.uid}/{name}' LOGGER.info( f'got path : {path}' ) - try: - os.remove(m_path) - LOGGER.info(f"Deleting archive : {m_path}") - except Exception as e: - LOGGER.error(str(e)) - else: + + except NotSupportedExtractionArchive: LOGGER.info("Not any valid archive, uploading file as it is.") - path = f'{DOWNLOAD_DIR}{self.uid}/{download_dict[self.uid].name()}' + path = f'{DOWNLOAD_DIR}{self.uid}/{name}' else: - path = f'{DOWNLOAD_DIR}{self.uid}/{download_dict[self.uid].name()}' + path = f'{DOWNLOAD_DIR}{self.uid}/{name}' up_name = pathlib.PurePath(path).name LOGGER.info(f"Upload Name : {up_name}") drive = gdriveTools.GoogleDriveHelper(up_name, self) diff --git a/extract b/extract index f124d7c..2a5aa6c 100755 --- a/extract +++ b/extract @@ -25,8 +25,10 @@ extract () { *.rar) a_dir=`expr "$arg" : '\(.*\).rar'` ; mkdir "$a_dir" ; - 7z x "$arg" -o"$a_dir" ;; - *) echo "'$arg' cannot be extracted via extract()" ;; + 7z x "$arg" -o"$a_dir" ;; + *) echo "'$arg' cannot be extracted via extract()" 1>&2 + exit 1 ;; + esac cd - || exit }