From 42533ac92c1f3fe37e34c16854997ab4b4223df5 Mon Sep 17 00:00:00 2001 From: Thorsten Date: Fri, 16 Feb 2024 16:51:24 +0100 Subject: [PATCH] Make SquashFs a contextmanager fix small stuff --- lib/inputstreamhelper/unsquash.py | 26 +++++++++--------- lib/inputstreamhelper/widevine/arm_lacros.py | 28 ++++++++++---------- lib/inputstreamhelper/widevine/repo.py | 5 ++++ 3 files changed, 32 insertions(+), 27 deletions(-) diff --git a/lib/inputstreamhelper/unsquash.py b/lib/inputstreamhelper/unsquash.py index b6838985..5b25e59b 100644 --- a/lib/inputstreamhelper/unsquash.py +++ b/lib/inputstreamhelper/unsquash.py @@ -26,7 +26,7 @@ from .kodiutils import log -class ZstdDecompressor: +class ZstdDecompressor: # pylint: disable=too-few-public-methods """ zstdandard decompressor class @@ -53,7 +53,7 @@ def decompress(self, comp_data, comp_size, outsize=8*2**10): @dataclass(frozen=True) -class SBlk: +class SBlk: # pylint: disable=too-many-instance-attributes """superblock as dataclass, does some checks after initialization""" s_magic: int inodes: int @@ -129,7 +129,7 @@ class BasicFileInode: fragment: int offset: int file_size: int - block_list: tuple[int] + block_list: tuple # once we remove support for python below 3.9 this can be: tuple[int] @dataclass(frozen=True) @@ -174,11 +174,17 @@ class SquashFs: """ def __init__(self, fpath): self.zdecomp = ZstdDecompressor() - self.imfile = open(fpath, "rb") + self.imfile = open(fpath, "rb") # pylint: disable=consider-using-with # we have our own context manager self.sblk = self._get_sblk() self.frag_entries = self._get_fragment_table() log(0, "squashfs image initialized") + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.imfile.close() + def _get_sblk(self): """ Read and check the superblock. @@ -404,7 +410,7 @@ def read_file_blocks(self, filename): for bsize in inode.block_list: compressed, size = self._get_size(bsize) - if not curr_pos == self.imfile.tell(): + if curr_pos != self.imfile.tell(): log(3, "Pointer not at correct position. Moving.") self.imfile.seek(curr_pos) @@ -431,11 +437,5 @@ def extract_file(self, filename, target_dir): Extracts file to """ with open(os.path.join(target_dir, filename), "wb") as outfile: - try: - for block in self.read_file_blocks(filename): - outfile.write(block) - - return True - except FileNotFoundError as err: - log(4, err) - return False + for block in self.read_file_blocks(filename): + outfile.write(block) diff --git a/lib/inputstreamhelper/widevine/arm_lacros.py b/lib/inputstreamhelper/widevine/arm_lacros.py index 1ddd71e8..46267e8f 100644 --- a/lib/inputstreamhelper/widevine/arm_lacros.py +++ b/lib/inputstreamhelper/widevine/arm_lacros.py @@ -30,29 +30,29 @@ def extract_widevine_lacros(dl_path, backup_path, img_version): progress = progress_dialog() progress.create(heading=localize(30043), message=localize(30044)) # Extracting Widevine CDM, prepping image + fnames = (config.WIDEVINE_CDM_FILENAME[system_os()], config.WIDEVINE_MANIFEST_FILE, "LICENSE") # Here it's not LICENSE.txt, as defined in the config.py + bpath = os.path.join(backup_path, img_version) + if not exists(bpath): + mkdirs(bpath) + try: - sfs = SquashFs(dl_path) - except IOError as err: - log(4, "SquashFs raised IOError") + with SquashFs(dl_path) as sfs: + for num, fname in enumerate(fnames): + sfs.extract_file(fname, bpath) + progress.update(int(90 / len(fnames) * (num + 1)), localize(30048)) # Extracting from image + + except (IOError, FileNotFoundError) as err: + log(4, "SquashFs raised an error") log(4, err) return False - fnames = (config.WIDEVINE_CDM_FILENAME[system_os()], config.WIDEVINE_MANIFEST_FILE, "LICENSE") # Here it's not LICENSE.txt, as defined in the config.py - mkdirs(os.path.join(backup_path, img_version)) - for num, fname in enumerate(fnames): - extracted = sfs.extract_file(fname, os.path.join(backup_path, img_version)) - if not extracted: - log(4, f"{fname} not found in {os.path.basename(dl_path)}") - return False - - progress.update(int(90 / len(fnames) * (num + 1)), localize(30048)) # Extracting from image - with open_file(os.path.join(backup_path, img_version, config.WIDEVINE_MANIFEST_FILE), "r") as manifest_file: + with open_file(os.path.join(bpath, config.WIDEVINE_MANIFEST_FILE), "r") as manifest_file: manifest_json = json.load(manifest_file) manifest_json.update({"img_version": img_version}) - with open_file(os.path.join(backup_path, img_version, config.WIDEVINE_MANIFEST_FILE), "w") as manifest_file: + with open_file(os.path.join(bpath, config.WIDEVINE_MANIFEST_FILE), "w") as manifest_file: json.dump(manifest_json, manifest_file, indent=2) log(0, f"Successfully extracted all files from lacros image {os.path.basename(dl_path)}") diff --git a/lib/inputstreamhelper/widevine/repo.py b/lib/inputstreamhelper/widevine/repo.py index 8838fb1f..7ed5afc7 100644 --- a/lib/inputstreamhelper/widevine/repo.py +++ b/lib/inputstreamhelper/widevine/repo.py @@ -29,6 +29,9 @@ def widevines_available_from_repo(): if http_status == 200: available_cdms.append({'version': cdm_version, 'url': cdm_url}) + if not available_cdms: + log(4, "could not find any available cdm in repo") + return available_cdms @@ -36,7 +39,9 @@ def latest_widevine_available_from_repo(available_cdms=None): """Returns the latest available Widevine CDM version and url from Google's library CDM repository""" if not available_cdms: available_cdms = widevines_available_from_repo() + latest = available_cdms[-1] # That's probably correct, but the following for loop makes sure + for cdm in available_cdms: if parse_version(cdm['version']) > parse_version(latest['version']): latest = cdm