mirror of
https://github.com/spalencsar/bearhub.git
synced 2026-07-07 00:04:15 +02:00
324 lines
13 KiB
Python
324 lines
13 KiB
Python
import os
|
|
import re
|
|
import shutil
|
|
import time
|
|
import traceback
|
|
from io import StringIO, BytesIO
|
|
from logging import Logger
|
|
from math import floor
|
|
from pathlib import Path
|
|
from threading import Thread
|
|
from typing import Optional, Tuple
|
|
|
|
from bauh.api.abstract.download import FileDownloader
|
|
from bauh.api.abstract.handler import ProcessWatcher
|
|
from bauh.api.http import HttpClient
|
|
from bauh.commons.html import bold
|
|
from bauh.commons.system import ProcessHandler, SimpleProcess
|
|
from bauh.commons.view_utils import get_human_size_str
|
|
from bauh.view.util.translation import I18n
|
|
|
|
RE_HAS_EXTENSION = re.compile(r'.+\.\w+$')
|
|
|
|
|
|
class SelfFileDownloader(FileDownloader):
|
|
|
|
def __init__(self, logger: Logger, i18n: I18n, http_client: HttpClient,
|
|
check_ssl: bool):
|
|
self._logger = logger
|
|
self._i18n = i18n
|
|
self._client = http_client
|
|
self._ssl = check_ssl
|
|
|
|
def is_multithreaded(self) -> bool:
|
|
return False
|
|
|
|
def can_work(self) -> bool:
|
|
return True
|
|
|
|
def get_supported_multithreaded_clients(self) -> Tuple[str, ...]:
|
|
return tuple()
|
|
|
|
def is_multithreaded_client_available(self, name: str) -> bool:
|
|
return False
|
|
|
|
def list_available_multithreaded_clients(self) -> Tuple[str, ...]:
|
|
return tuple()
|
|
|
|
def get_supported_clients(self) -> Tuple[str, ...]:
|
|
return tuple()
|
|
|
|
def download(self, file_url: str, watcher: Optional[ProcessWatcher], output_path: str, cwd: str,
|
|
root_password: Optional[str] = None, substatus_prefix: str = None, display_file_size: bool = True,
|
|
max_threads: int = None, known_size: int = None) -> bool:
|
|
try:
|
|
res = self._client.get(url=file_url, ignore_ssl=not self._ssl, stream=True)
|
|
except Exception:
|
|
return False
|
|
|
|
try:
|
|
content_length = int(res.headers.get("content-length", 0))
|
|
except Exception:
|
|
content_length = 0
|
|
self._logger.warning(f"Could not retrieve the content-length for file '{file_url}'")
|
|
|
|
file_name = file_url.split("/")[-1]
|
|
msg = StringIO()
|
|
msg.write(f"{substatus_prefix} " if substatus_prefix else "")
|
|
msg.write(f"{self._i18n['downloading']} {bold(file_name)}")
|
|
base_msg = msg.getvalue()
|
|
|
|
byte_stream = BytesIO()
|
|
total_downloaded = 0
|
|
known_size = content_length and content_length > 0
|
|
total_size_str = get_human_size_str(content_length) if known_size > 0 else "?"
|
|
|
|
for data in res.iter_content(chunk_size=1024):
|
|
byte_stream.write(data)
|
|
total_downloaded += len(data)
|
|
perc = f"({(total_downloaded / content_length) * 100:.2f}%) " if known_size > 0 else ""
|
|
watcher.change_substatus(f"{perc}{base_msg} ({get_human_size_str(total_downloaded)} / {total_size_str})")
|
|
|
|
self._logger.info(f"Writing downloaded file content to disk: {output_path}")
|
|
|
|
try:
|
|
with open(output_path, "wb+") as f:
|
|
f.write(byte_stream.getvalue())
|
|
except Exception:
|
|
self._logger.error(f"Unexpected exception when saving downloaded content to disk: {output_path}")
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
class AdaptableFileDownloader(FileDownloader):
|
|
|
|
def __init__(self, logger: Logger, multithread_enabled: bool, i18n: I18n, http_client: HttpClient,
|
|
multithread_client: str, check_ssl: bool):
|
|
self.logger = logger
|
|
self.multithread_enabled = multithread_enabled
|
|
self.i18n = i18n
|
|
self.http_client = http_client
|
|
self.supported_multithread_clients = ("aria2", "axel")
|
|
self.multithread_client = multithread_client
|
|
self.check_ssl = check_ssl
|
|
self._self_downloader = SelfFileDownloader(logger=logger,
|
|
i18n=i18n,
|
|
http_client=http_client,
|
|
check_ssl=check_ssl)
|
|
|
|
@staticmethod
|
|
def is_aria2c_available() -> bool:
|
|
return bool(shutil.which('aria2c'))
|
|
|
|
@staticmethod
|
|
def is_axel_available() -> bool:
|
|
return bool(shutil.which('axel'))
|
|
|
|
def _get_aria2c_process(self, url: str, output_path: str, cwd: str, root_password: Optional[str], threads: int) -> SimpleProcess:
|
|
cmd = ['aria2c', url,
|
|
'--no-conf',
|
|
'-x', '16',
|
|
'--enable-color=false',
|
|
'--stderr=true',
|
|
'--summary-interval=0',
|
|
'--disable-ipv6',
|
|
'-k', '1M',
|
|
'--allow-overwrite=true',
|
|
'-c',
|
|
'-t', '5',
|
|
'--max-file-not-found=3',
|
|
'--file-allocation=none',
|
|
'--console-log-level=error']
|
|
|
|
if threads > 1:
|
|
cmd.append('-s')
|
|
cmd.append(str(threads))
|
|
|
|
if output_path:
|
|
output_split = output_path.split('/')
|
|
cmd.append('-d')
|
|
cmd.append('/'.join(output_split[:-1]))
|
|
cmd.append('-o')
|
|
cmd.append(output_split[-1])
|
|
|
|
return SimpleProcess(cmd=cmd, root_password=root_password, cwd=cwd)
|
|
|
|
def _get_axel_process(self, url: str, output_path: str, cwd: str, root_password: Optional[str], threads: int) -> SimpleProcess:
|
|
cmd = ['axel', url, '-n', str(threads), '-4', '-c', '-T', '5']
|
|
|
|
if not self.check_ssl:
|
|
cmd.append('-k')
|
|
|
|
if output_path:
|
|
cmd.append(f'--output={output_path}')
|
|
|
|
return SimpleProcess(cmd=cmd, cwd=cwd, root_password=root_password)
|
|
|
|
def _rm_bad_file(self, file_name: str, output_path: str, cwd, handler: ProcessHandler, root_password: Optional[str]):
|
|
to_delete = output_path if output_path else f'{cwd}/{file_name}'
|
|
|
|
if to_delete and os.path.exists(to_delete):
|
|
self.logger.info(f'Removing downloaded file {to_delete}')
|
|
success, _ = handler.handle_simple(SimpleProcess(['rm', '-rf', to_delete], root_password=root_password))
|
|
return success
|
|
|
|
def _concat_file_size(self, file_url: str, base_substatus: StringIO, watcher: ProcessWatcher):
|
|
watcher.change_substatus(f'{base_substatus.getvalue()} ( ? Mb )')
|
|
|
|
try:
|
|
size = self.http_client.get_content_length(file_url)
|
|
|
|
if size:
|
|
base_substatus.write(f' ( {size} )')
|
|
watcher.change_substatus(base_substatus.getvalue())
|
|
except Exception:
|
|
pass
|
|
|
|
def _get_appropriate_threads_number(self, max_threads: int, known_size: int) -> int:
|
|
if max_threads and max_threads > 0:
|
|
threads = max_threads
|
|
elif known_size:
|
|
threads = 16 if known_size >= 16000000 else floor(known_size / 1000000)
|
|
|
|
if threads <= 0:
|
|
threads = 1
|
|
else:
|
|
threads = 16
|
|
|
|
return threads
|
|
|
|
def _download_with_threads(self, client: str, file_url: str, output_path: str, cwd: str,
|
|
max_threads: int, known_size: int, display_file_size: bool, handler: ProcessHandler,
|
|
root_password: Optional[str] = None, substatus_prefix: Optional[str] = None) \
|
|
-> Tuple[float, bool]:
|
|
|
|
threads = self._get_appropriate_threads_number(max_threads, known_size)
|
|
|
|
if client == 'aria2':
|
|
start_time = time.time()
|
|
process = self._get_aria2c_process(file_url, output_path, cwd, root_password, threads)
|
|
downloader = 'aria2'
|
|
else:
|
|
start_time = time.time()
|
|
process = self._get_axel_process(file_url, output_path, cwd, root_password, threads)
|
|
downloader = 'axel'
|
|
|
|
name = file_url.split('/')[-1]
|
|
|
|
if output_path and not RE_HAS_EXTENSION.match(name) and RE_HAS_EXTENSION.match(output_path):
|
|
name = output_path.split('/')[-1]
|
|
|
|
if handler.watcher:
|
|
msg = StringIO()
|
|
msg.write(f'{substatus_prefix} ' if substatus_prefix else '')
|
|
msg.write(f"{bold('[{}]'.format(downloader))} {self.i18n['downloading']} {bold(name)}")
|
|
|
|
if display_file_size:
|
|
if known_size:
|
|
msg.write(f' ( {get_human_size_str(known_size)} )')
|
|
handler.watcher.change_substatus(msg.getvalue())
|
|
else:
|
|
Thread(target=self._concat_file_size, args=(file_url, msg, handler.watcher), daemon=True).start()
|
|
else:
|
|
msg.write(' ( ? Mb )')
|
|
handler.watcher.change_substatus(msg.getvalue())
|
|
|
|
success, _ = handler.handle_simple(process)
|
|
return start_time, success
|
|
|
|
def download(self, file_url: str, watcher: ProcessWatcher, output_path: str = None, cwd: str = None, root_password: Optional[str] = None, substatus_prefix: str = None, display_file_size: bool = True, max_threads: int = None, known_size: int = None) -> bool:
|
|
self.logger.info(f'Downloading {file_url}')
|
|
handler = ProcessHandler(watcher)
|
|
file_name = file_url.split('/')[-1]
|
|
|
|
final_cwd = cwd if cwd else '.'
|
|
|
|
success = False
|
|
start_time = time.time()
|
|
try:
|
|
if output_path:
|
|
if os.path.exists(output_path):
|
|
self.logger.info(f'Removing old file found before downloading: {output_path}')
|
|
os.remove(output_path)
|
|
self.logger.info(f'Old file {output_path} removed')
|
|
else:
|
|
output_dir = os.path.dirname(output_path)
|
|
|
|
try:
|
|
Path(output_dir).mkdir(exist_ok=True, parents=True)
|
|
except OSError:
|
|
self.logger.error(f"Could not make download directory '{output_dir}'")
|
|
watcher.print(self.i18n['error.mkdir'].format(dir=output_dir))
|
|
return False
|
|
|
|
threaded_client = self.get_available_multithreaded_tool()
|
|
if threaded_client:
|
|
start_time, success = self._download_with_threads(client=threaded_client, file_url=file_url,
|
|
output_path=output_path,
|
|
cwd=final_cwd, max_threads=max_threads,
|
|
known_size=known_size, handler=handler,
|
|
display_file_size=display_file_size,
|
|
root_password=root_password)
|
|
else:
|
|
start_time = time.time()
|
|
success = self._self_downloader.download(file_url=file_url, watcher=watcher, output_path=output_path,
|
|
cwd=cwd, root_password=root_password,
|
|
substatus_prefix=substatus_prefix,
|
|
display_file_size=display_file_size, max_threads=max_threads,
|
|
known_size=known_size)
|
|
except Exception:
|
|
traceback.print_exc()
|
|
self._rm_bad_file(file_name, output_path, final_cwd, handler, root_password)
|
|
|
|
final_time = time.time()
|
|
self.logger.info(f'{file_name} download took {(final_time - start_time) / 60:.2f} minutes')
|
|
|
|
if not success:
|
|
self.logger.error(f"Could not download '{file_name}'")
|
|
self._rm_bad_file(file_name, output_path, final_cwd, handler, root_password)
|
|
|
|
return success
|
|
|
|
def is_multithreaded(self) -> bool:
|
|
return bool(self.get_available_multithreaded_tool())
|
|
|
|
def get_available_multithreaded_tool(self) -> str:
|
|
if self.multithread_enabled:
|
|
if self.multithread_client is None or self.multithread_client not in self.supported_multithread_clients:
|
|
for client in self.supported_multithread_clients:
|
|
if self.is_multithreaded_client_available(client):
|
|
return client
|
|
else:
|
|
possible_clients = {*self.supported_multithread_clients}
|
|
|
|
if self.is_multithreaded_client_available(self.multithread_client):
|
|
return self.multithread_client
|
|
else:
|
|
possible_clients.remove(self.multithread_client)
|
|
|
|
for client in possible_clients:
|
|
if self.is_multithreaded_client_available(client):
|
|
return client
|
|
|
|
def can_work(self) -> bool:
|
|
return True
|
|
|
|
def get_supported_multithreaded_clients(self) -> Tuple[str, ...]:
|
|
return self.supported_multithread_clients
|
|
|
|
def is_multithreaded_client_available(self, name: str) -> bool:
|
|
if name == 'aria2':
|
|
return self.is_aria2c_available()
|
|
elif name == 'axel':
|
|
return self.is_axel_available()
|
|
else:
|
|
return False
|
|
|
|
def list_available_multithreaded_clients(self) -> Tuple[str, ...]:
|
|
return tuple(c for c in self.supported_multithread_clients if self.is_multithreaded_client_available(c))
|
|
|
|
def get_supported_clients(self) -> Tuple[str, ...]:
|
|
return "self", "aria2", "axel"
|