287 lines
9.0 KiB
Python
287 lines
9.0 KiB
Python
from __future__ import annotations
|
|
import hashlib
|
|
import logging
|
|
import platform
|
|
import zipfile
|
|
from collections.abc import Callable
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
import requests
|
|
|
|
from .exception import HashMismatchException, DownloadException
|
|
|
|
DEFAULT_HASH_CHUNK_SIZE = 4096
|
|
DEFAULT_DOWNLOAD_CHUNK_SIZE = 4096
|
|
THREADED_DOWNLOAD_MAX_WORKERS = 10
|
|
|
|
logger = logging.getLogger("Launcher.CoreLib")
|
|
|
|
class FileObject:
|
|
def __init__(self, path: Path, url=None, sha1=None, sha256=None, md5=None, size=None):
|
|
self.__path = path
|
|
|
|
self.expected_sha1 = sha1
|
|
self.expected_sha256 = sha256
|
|
self.expected_md5 = md5
|
|
self.expected_size = size
|
|
|
|
self.__sha1 = None
|
|
self.__sha256 = None
|
|
self.__md5 = None
|
|
|
|
self.url = url
|
|
|
|
def __eq__(self, other):
|
|
"""
|
|
THIS OPERATION ONLY COMPARE AND FILE SIZE! Use
|
|
:param other:
|
|
:return:
|
|
"""
|
|
if self.path == other.path:
|
|
return True
|
|
|
|
if self.size == other.size:
|
|
return True
|
|
|
|
return False
|
|
|
|
def is_same(self, file: FileObject, hash_name="sha256") -> bool:
|
|
if hash_name == "sha1":
|
|
return self.sha1 == file.sha1
|
|
|
|
if hash_name == "md5":
|
|
return self.md5 == file.md5
|
|
|
|
if hash_name != "sha256":
|
|
logger.warning("Unsupported hash_name {}. Ignored and use sha256 instead.".format(hash_name))
|
|
|
|
return self.sha256 == file.sha256
|
|
|
|
@property
|
|
def sha1(self):
|
|
if self.__sha1 is None:
|
|
self.__sha1 = self._calculate_hash("sha1")
|
|
|
|
return self.__sha1
|
|
|
|
@property
|
|
def sha256(self):
|
|
if self.__sha256 is None:
|
|
self.__sha256 = self._calculate_hash("sha256")
|
|
|
|
return self.__sha256
|
|
|
|
@property
|
|
def md5(self):
|
|
if self.__md5 is None:
|
|
self.__md5 = self._calculate_hash("md5")
|
|
|
|
return self.__md5
|
|
|
|
def verify_sha1(self, expected_hash=None):
|
|
return self.sha1 == expected_hash if expected_hash else self.expected_sha1 == self.sha1
|
|
|
|
def verify_sha256(self, expected_hash=None):
|
|
return self.sha256 == expected_hash if expected_hash else self.expected_sha256 == self.sha256
|
|
|
|
def verify_md5(self, expected_hash=None):
|
|
return self.sha1 == expected_hash if expected_hash else self.expected_sha1 == self.sha1
|
|
|
|
def verify(self, algorithm, expected_hash=None):
|
|
if algorithm == "sha1":
|
|
return self.verify_sha1(expected_hash if expected_hash else self.expected_sha1)
|
|
|
|
if algorithm == "sha256":
|
|
return self.verify_sha256(expected_hash if expected_hash else self.expected_sha256)
|
|
|
|
if algorithm == "md5":
|
|
return self.verify_md5(expected_hash if expected_hash else self.expected_md5)
|
|
|
|
if algorithm == "size":
|
|
return self.verify_size(expected_hash if expected_hash else self.expected_size)
|
|
|
|
raise Exception("Unsupported algorithm: {}".format(algorithm))
|
|
|
|
def verify_size(self, expected_size):
|
|
return self.size == expected_size if expected_size else self.expected_size == self.size
|
|
|
|
def _calculate_hash(self, algorithm):
|
|
h = hashlib.new(algorithm)
|
|
|
|
with self.__path.open("rb") as f:
|
|
for chunk in iter(lambda: f.read(DEFAULT_HASH_CHUNK_SIZE), b""):
|
|
h.update(chunk)
|
|
|
|
return h.hexdigest()
|
|
|
|
@property
|
|
def path(self):
|
|
return self.__path
|
|
|
|
@property
|
|
def posix_path(self):
|
|
return self.__path.absolute().as_posix()
|
|
|
|
@property
|
|
def size(self):
|
|
return self.__path.stat().st_size
|
|
|
|
@property
|
|
def name(self):
|
|
return self.__path.name
|
|
|
|
@property
|
|
def extension(self):
|
|
return self.__path.suffix
|
|
|
|
def __hash__(self):
|
|
return hash(self.__path)
|
|
|
|
@property
|
|
def exists(self):
|
|
return self.__path.exists()
|
|
|
|
|
|
def download_file(file: FileObject, overwrite=True, progress_callback: Callable[int]=None) -> FileObject | Exception:
|
|
"""
|
|
Download file from url
|
|
:param file: File (Objected)
|
|
:param overwrite: Overwrite existing file
|
|
:param progress_callback: Callback function called with progress information
|
|
:return:
|
|
"""
|
|
|
|
file.path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
if file.exists and not overwrite:
|
|
raise FileExistsError("File {} already exists. Use overwrite=True to overwrite.".format(file))
|
|
|
|
if file.url is None:
|
|
raise ValueError("Target file {} URL cannot be None.".format(file))
|
|
|
|
logger.debug("Downloading file {} to {}".format(file.url, file))
|
|
|
|
try:
|
|
with requests.get(file.url, stream=True) as r:
|
|
current_chunk = 0
|
|
total_bytes = r.headers.get("content-length", None)
|
|
logger.debug("URL: {}".format(file.url))
|
|
logger.debug("Response code: {}".format(r.status_code))
|
|
|
|
with file.path.open("wb") as f:
|
|
for chunk in r.iter_content(chunk_size=DEFAULT_DOWNLOAD_CHUNK_SIZE):
|
|
if not chunk:
|
|
break
|
|
|
|
f.write(chunk)
|
|
current_chunk += 1
|
|
|
|
if callable(progress_callback) and total_bytes is not None:
|
|
current_progress = ((current_chunk * DEFAULT_DOWNLOAD_CHUNK_SIZE) / int(total_bytes)) * 100
|
|
progress_callback(file.name, current_progress)
|
|
except Exception as e:
|
|
raise DownloadException(
|
|
"An error occurred while downloading file {}. {}".format(file.url, e)
|
|
)
|
|
|
|
if file.expected_sha256 and not file.verify_sha256():
|
|
raise HashMismatchException("sha256", file.sha256, file.expected_sha256, file.posix_path)
|
|
elif file.expected_sha1 and not file.verify_sha1():
|
|
raise HashMismatchException("sha1", file.sha1, file.expected_sha1, file.posix_path)
|
|
elif file.expected_md5 and not file.verify_md5():
|
|
raise HashMismatchException("md5", file.md5, file.expected_md5, file.posix_path)
|
|
|
|
if any([file.expected_sha1, file.expected_md5, file.expected_sha256]):
|
|
logger.debug("File hash is match as the expected hash.")
|
|
else:
|
|
logger.warning("File hash is not checked.")
|
|
|
|
return file
|
|
|
|
def download_multiple_files(files: List[FileObject], download_callback: Callable[str]=None,
|
|
max_workers=THREADED_DOWNLOAD_MAX_WORKERS, progress_callback: Callable[str]=None) \
|
|
-> tuple[list[FileObject], list[FileObject]]:
|
|
"""
|
|
Download multiple files
|
|
:param download_callback:
|
|
:param files:
|
|
:param max_workers:
|
|
:param progress_callback:
|
|
:return:
|
|
fails: List of file objects that download failed.
|
|
successes: List of file objects that download succeeded.
|
|
"""
|
|
|
|
fails = []
|
|
successes = []
|
|
|
|
files_count = 0
|
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
future_to_file = {executor.submit(download_file, file): file for file in files}
|
|
|
|
for future in as_completed(future_to_file):
|
|
file = future_to_file[future]
|
|
|
|
try:
|
|
files_count += 1
|
|
future.result()
|
|
successes.append(file)
|
|
if callable(download_callback):
|
|
download_callback(file.posix_path)
|
|
|
|
if callable(progress_callback):
|
|
progress_callback(f"{len(files)} Files", files_count / len(files) * 100)
|
|
except Exception as exc:
|
|
logger.error("Unable to download file {}: {}".format(file.posix_path, exc))
|
|
fails.append(file)
|
|
|
|
|
|
return fails, successes
|
|
|
|
def unzip(filepath: Path, destination: Path, exclude: list[str] | None = None, flatten: bool=False):
|
|
exclude = exclude or []
|
|
|
|
with zipfile.ZipFile(filepath.as_posix(), "r") as archive:
|
|
for member in archive.namelist():
|
|
# Ignore exclude file.
|
|
if any(member.startswith(pattern) for pattern in exclude):
|
|
continue
|
|
|
|
if flatten:
|
|
member_path = Path(member)
|
|
|
|
if member.endswith("/") or not member_path.name:
|
|
continue
|
|
|
|
if any(member.startswith(pattern) for pattern in exclude):
|
|
continue
|
|
|
|
target = destination / member_path.name
|
|
|
|
with archive.open(member) as source:
|
|
with target.open("wb") as output:
|
|
output.write(source.read())
|
|
else:
|
|
archive.extract(member, path=destination)
|
|
|
|
def get_platform_info():
|
|
"""
|
|
Get information about the current platform.
|
|
IMPORTANT: All return values are lower-case
|
|
:return:
|
|
platform_name: str, architecture: str, os_version: str
|
|
"""
|
|
platform_name = platform.system().lower()
|
|
arch = platform.uname().machine.lower()
|
|
os_version = platform.version()
|
|
|
|
if platform_name == "darwin":
|
|
os_version = platform.mac_ver()[0]
|
|
elif platform_name == "linux":
|
|
os_version = platform.release()
|
|
|
|
return platform_name, arch, os_version
|
|
|