From b84964c069a07fab75ba8336ef52c33e36636337 Mon Sep 17 00:00:00 2001 From: yshalsager Date: Wed, 30 Oct 2024 18:02:57 +0300 Subject: [PATCH] chore: format using ruff Signed-off-by: yshalsager --- .flake8 | 9 - xiaomi_flashable_firmware_creator/__init__.py | 1 + xiaomi_flashable_firmware_creator/__main__.py | 3 +- .../extractors/handlers/android_one_zip.py | 20 +- .../extractors/zip_extractor.py | 22 +- .../firmware_creator.py | 262 ++++++++---------- .../helpers/misc.py | 49 ++-- .../helpers/mock_zip.py | 25 +- .../tests/test_creator.py | 36 +-- .../tests/test_misc.py | 6 +- .../xiaomi_flashable_firmware_creator.py | 19 +- 11 files changed, 202 insertions(+), 250 deletions(-) delete mode 100644 .flake8 diff --git a/.flake8 b/.flake8 deleted file mode 100644 index 25a529a..0000000 --- a/.flake8 +++ /dev/null @@ -1,9 +0,0 @@ -[flake8] -max-line-length = 88 -max-complexity = 18 -extend-ignore = - E203 - E266 - E501 - W503 -select = ["B", "C", "E", "F", "W", "T4"] \ No newline at end of file diff --git a/xiaomi_flashable_firmware_creator/__init__.py b/xiaomi_flashable_firmware_creator/__init__.py index 5c6461e..12a23d6 100644 --- a/xiaomi_flashable_firmware_creator/__init__.py +++ b/xiaomi_flashable_firmware_creator/__init__.py @@ -1,4 +1,5 @@ """Xiaomi Flashable Firmware Creator initialization.""" + from pathlib import Path work_dir = Path(__package__).parent.absolute() diff --git a/xiaomi_flashable_firmware_creator/__main__.py b/xiaomi_flashable_firmware_creator/__main__.py index 55a3b2c..04e5b18 100644 --- a/xiaomi_flashable_firmware_creator/__main__.py +++ b/xiaomi_flashable_firmware_creator/__main__.py @@ -1,5 +1,6 @@ """Xiaomi Flashable Firmware Creator main entry point.""" + from xiaomi_flashable_firmware_creator.xiaomi_flashable_firmware_creator import main -if __name__ == "__main__": +if __name__ == '__main__': main() diff --git a/xiaomi_flashable_firmware_creator/extractors/handlers/android_one_zip.py b/xiaomi_flashable_firmware_creator/extractors/handlers/android_one_zip.py index 896719c..25d96e6 100644 --- a/xiaomi_flashable_firmware_creator/extractors/handlers/android_one_zip.py +++ b/xiaomi_flashable_firmware_creator/extractors/handlers/android_one_zip.py @@ -28,16 +28,12 @@ def prepare(self) -> List[str]: Extract payload from zip in order to get partition names :return: a list of partition names strings """ - self.payload_file = open( - self.extractor.extract("payload.bin", self._tmp_dir), "rb" - ) + self.payload_file = open(self.extractor.extract('payload.bin', self._tmp_dir), 'rb') self.payload = Payload(self.payload_file) self.payload.Init() - self.partitions = { - i.partition_name: i for i in self.payload.manifest.partitions - } + self.partitions = {i.partition_name: i for i in self.payload.manifest.partitions} self.all_partitions = set(self.partitions.keys()) - self.files = sorted([f"firmware-update/{i}.img" for i in self.all_partitions]) + self.files = sorted([f'firmware-update/{i}.img' for i in self.all_partitions]) return self.files def extract(self, files_to_extract: List[str]): @@ -46,15 +42,15 @@ def extract(self, files_to_extract: List[str]): :param files_to_extract: a list of files to extract :return: """ - Path(self._tmp_dir / "firmware-update").mkdir(parents=True, exist_ok=True) + Path(self._tmp_dir / 'firmware-update').mkdir(parents=True, exist_ok=True) files_to_extract: set = set(self.files).intersection(set(files_to_extract)) for file in files_to_extract: # partition: RepeatedCompositeContainer = self.partitions.get(file.split('/')[-1].split('.')[0]) - partition = self.partitions.get(file.split("/")[-1].split(".")[0]) + partition = self.partitions.get(file.split('/')[-1].split('.')[0]) with open( - Path(self._tmp_dir / f"firmware-update/{partition.partition_name}.img"), - "wb", + Path(self._tmp_dir / f'firmware-update/{partition.partition_name}.img'), + 'wb', ) as out_f: parse_payload(self.payload, partition, out_f) self.payload_file.close() - Path(self._tmp_dir / "payload.bin").unlink() + Path(self._tmp_dir / 'payload.bin').unlink() diff --git a/xiaomi_flashable_firmware_creator/extractors/zip_extractor.py b/xiaomi_flashable_firmware_creator/extractors/zip_extractor.py index e0a0de7..db39091 100644 --- a/xiaomi_flashable_firmware_creator/extractors/zip_extractor.py +++ b/xiaomi_flashable_firmware_creator/extractors/zip_extractor.py @@ -31,19 +31,15 @@ def __init__(self, zip_file, tmp_dir): :param zip_file: a path object or a string to a zip that contains a full recovery ROM. :param tmp_dir: output directory to place the extracted zip in. """ - self.zip_url = ( - zip_file if "http" in zip_file or "ota.d.miui.com" in zip_file else "" - ) + self.zip_url = zip_file if 'http' in zip_file or 'ota.d.miui.com' in zip_file else '' self.zip_file_path = ( - Path(zip_file) if not self.zip_url and isinstance(zip_file, str) else "" + Path(zip_file) if not self.zip_url and isinstance(zip_file, str) else '' ) self.files = [] - self._extractor = ( - RemoteZip(self.zip_url) if self.zip_url else ZipFile(self.zip_file_path) - ) + self._extractor = RemoteZip(self.zip_url) if self.zip_url else ZipFile(self.zip_file_path) self.handler = ( AndroidOneZip(self.zip_file_path, tmp_dir, self._extractor) - if "payload.bin" in str(self._extractor.namelist()) + if 'payload.bin' in str(self._extractor.namelist()) else StandardZip(self.zip_file_path, tmp_dir, self._extractor) ) @@ -53,9 +49,7 @@ def exists(self) -> bool: :return: True if zip file exists, False otherwise. """ - return ( - self.zip_file_path.exists() if self.zip_file_path else head(self.zip_url).ok - ) + return self.zip_file_path.exists() if self.zip_file_path else head(self.zip_url).ok def get_files_list(self): """ @@ -71,11 +65,7 @@ def get_file_name(self) -> str: :return: a string of the input zip file name. """ - return ( - self.zip_file_path.name - if self.zip_file_path - else self.zip_url.split("/")[-1] - ) + return self.zip_file_path.name if self.zip_file_path else self.zip_url.split('/')[-1] def prepare(self): if isinstance(self.handler, AndroidOneZip): diff --git a/xiaomi_flashable_firmware_creator/firmware_creator.py b/xiaomi_flashable_firmware_creator/firmware_creator.py index 5087d7d..6f1ce83 100644 --- a/xiaomi_flashable_firmware_creator/firmware_creator.py +++ b/xiaomi_flashable_firmware_creator/firmware_creator.py @@ -46,7 +46,7 @@ class FlashableFirmwareCreator: update_script: str is_android_one: bool - def __init__(self, input_file, _extract_mode, out_dir=""): + def __init__(self, input_file, _extract_mode, out_dir=''): """ Initialize FlashableFirmwareCreator. @@ -57,43 +57,41 @@ def __init__(self, input_file, _extract_mode, out_dir=""): :param out_dir: The output directory to store the extracted file in. """ self.input_file = input_file - _tmp_sub_dir = f"{Path(input_file).stem.split('_')[-2]}_" # set tmp subdirectory to ROM's hash + _tmp_sub_dir = ( + f"{Path(input_file).stem.split('_')[-2]}_" # set tmp subdirectory to ROM's hash + ) self._tmp_dir = ( - Path(out_dir) / "tmp" / _tmp_sub_dir - if out_dir - else work_dir / "tmp" / _tmp_sub_dir + Path(out_dir) / 'tmp' / _tmp_sub_dir if out_dir else work_dir / 'tmp' / _tmp_sub_dir ) self._out_dir = self._tmp_dir.parent.parent.absolute() - self._flashing_script_dir = ( - self._tmp_dir.absolute() / "META-INF/com/google/android" - ) + self._flashing_script_dir = self._tmp_dir.absolute() / 'META-INF/com/google/android' self.host = gethostname() self.datetime = datetime.now() self.extract_mode = self.get_extract_mode(_extract_mode) - self.update_script = "" + self.update_script = '' self.is_android_one = False self.firmware_excluded_files = [ - "dtbo", - "logo", - "splash", - "vbmeta", - "boot.img", - "recovery.img", - "system", - "vendor", - "product", - "odm", - "exaid", - "mi_ext", - "dynamic_partitions_op_list", - "metadata.pb", + 'dtbo', + 'logo', + 'splash', + 'vbmeta', + 'boot.img', + 'recovery.img', + 'system', + 'vendor', + 'product', + 'odm', + 'exaid', + 'mi_ext', + 'dynamic_partitions_op_list', + 'metadata.pb', ] self.vendor_excluded_files = [ - "vbmeta", - "system", - "product.", - "odm.", - "metadata.pb", + 'vbmeta', + 'system', + 'product.', + 'odm.', + 'metadata.pb', ] self.extractor = ZipExtractor(self.input_file, self._tmp_dir) self.init() @@ -116,11 +114,11 @@ def init(self): self._tmp_dir.mkdir(parents=True, exist_ok=True) self._flashing_script_dir.mkdir(parents=True, exist_ok=True) if not self.extractor.exists(): - raise FileNotFoundError(f"input file {self.input_file} does not exist!") + raise FileNotFoundError(f'input file {self.input_file} does not exist!') self.extractor.get_files_list() if not self.is_valid_rom(): rmtree(self._tmp_dir) - raise RuntimeError(f"{self.input_file} is not a valid ROM file. Exiting..") + raise RuntimeError(f'{self.input_file} is not a valid ROM file. Exiting..') @staticmethod def get_extract_mode(extract_mode) -> ProcessTypes: @@ -131,15 +129,15 @@ def get_extract_mode(extract_mode) -> ProcessTypes: :return: ProcessTypes enum """ modes = { - "firmware": ProcessTypes.firmware, - "nonarb": ProcessTypes.non_arb_firmware, - "firmwareless": ProcessTypes.firmware_less, - "vendor": ProcessTypes.vendor, + 'firmware': ProcessTypes.firmware, + 'nonarb': ProcessTypes.non_arb_firmware, + 'firmwareless': ProcessTypes.firmware_less, + 'vendor': ProcessTypes.vendor, } try: return modes[extract_mode] except KeyError as err: - print("Unknown process!") + print('Unknown process!') raise err def is_valid_rom(self) -> bool: @@ -150,9 +148,9 @@ def is_valid_rom(self) -> bool: in contents list, False otherwise. """ return ( - "META-INF/com/google/android/update-binary" in self.extractor.files - or "META-INF/com/google/android/updater-script" in self.extractor.files - or "payload.bin" in self.extractor.files + 'META-INF/com/google/android/update-binary' in self.extractor.files + or 'META-INF/com/google/android/updater-script' in self.extractor.files + or 'payload.bin' in self.extractor.files ) def get_rom_type(self): @@ -162,14 +160,9 @@ def get_rom_type(self): :return: An enum of ROM type. Either qcom or mtk. """ files = str(self.extractor.files) - if "lk.img" in files or "preloader.img" in files: + if 'lk.img' in files or 'preloader.img' in files: self.type = ZipTypes.mtk - elif ( - "firmware-update" in files - or "rpm" in files - or "tz" in files - or "keymaster" in files - ): + elif 'firmware-update' in files or 'rpm' in files or 'tz' in files or 'keymaster' in files: self.type = ZipTypes.qcom else: raise RuntimeError("Can't detect rom type. It's not qcom or mtk!'") @@ -188,10 +181,8 @@ def get_files_list(self) -> List[str]: i for i in self.extractor.files if ( - i.startswith("META-INF/") - and ( - i.endswith("updater-script") or i.endswith("update-binary") - ) + i.startswith('META-INF/') + and (i.endswith('updater-script') or i.endswith('update-binary')) ) or all(n not in i for n in self.firmware_excluded_files) ] @@ -200,8 +191,7 @@ def get_files_list(self) -> List[str]: n for n in self.extractor.files if all( - file not in n - for file in self.firmware_excluded_files + ["file_contexts"] + file not in n for file in self.firmware_excluded_files + ['file_contexts'] ) ] ) @@ -209,19 +199,15 @@ def get_files_list(self) -> List[str]: return [ n for n in self.extractor.files - if "dspso.bin" in n - or n.startswith("firmware-update/BTFM.bin") - or n.startswith("firmware-update/NON-HLOS.bin") - or n.startswith("META-INF/") - and (n.endswith("updater-script") or n.endswith("update-binary")) + if 'dspso.bin' in n + or n.startswith('firmware-update/BTFM.bin') + or n.startswith('firmware-update/NON-HLOS.bin') + or n.startswith('META-INF/') + and (n.endswith('updater-script') or n.endswith('update-binary')) ] if self.extract_mode is ProcessTypes.firmware_less: return ( - [ - n - for n in self.extractor.files - if not n.startswith("firmware-update/") - ] + [n for n in self.extractor.files if not n.startswith('firmware-update/')] if self.type is ZipTypes.qcom else [] ) @@ -239,9 +225,9 @@ def get_updater_script_lines(self, invalid_files: set[str]) -> str: :return: a string of the updater-script lines """ - original_updater_script = Path(self._flashing_script_dir / "updater-script") + original_updater_script = Path(self._flashing_script_dir / 'updater-script') if not original_updater_script.exists(): - raise FileNotFoundError("updater-script not found!") + raise FileNotFoundError('updater-script not found!') original_updater_script = original_updater_script.read_text().splitlines() lines = [] if self.extract_mode is ProcessTypes.firmware: @@ -249,24 +235,22 @@ def get_updater_script_lines(self, invalid_files: set[str]) -> str: [ line for line in original_updater_script - if ("getprop" in line and "ro.build.date.utc" not in line) - or "Target" in line + if ('getprop' in line and 'ro.build.date.utc' not in line) + or 'Target' in line or ( - "firmware-update" in line - and line.split("/")[1].split('"')[0] not in invalid_files + 'firmware-update' in line + and line.split('/')[1].split('"')[0] not in invalid_files ) and ( - "ro.product" in line - or all( - file not in line for file in self.firmware_excluded_files - ) + 'ro.product' in line + or all(file not in line for file in self.firmware_excluded_files) ) ] if self.type is ZipTypes.qcom else [ line for line in original_updater_script - if "ro.product" in line + if 'ro.product' in line or all(file not in line for file in self.firmware_excluded_files) ] ) @@ -274,23 +258,23 @@ def get_updater_script_lines(self, invalid_files: set[str]) -> str: lines = [ line for line in original_updater_script - if ("getprop" in line and "ro.build.date.utc" not in line) - or "Target" in line - or "modem" in line - or "bluetooth" in line - or "dsp" in line + if ('getprop' in line and 'ro.build.date.utc' not in line) + or 'Target' in line + or 'modem' in line + or 'bluetooth' in line + or 'dsp' in line ] elif self.extract_mode is ProcessTypes.firmware_less: lines = ( [ line for line in original_updater_script - if ("getprop" in line and "ro.build.date.utc" not in line) - or "Target" in line - or "boot.img" in line - or "system" in line - or "vendor" in line - or "product" in line + if ('getprop' in line and 'ro.build.date.utc' not in line) + or 'Target' in line + or 'boot.img' in line + or 'system' in line + or 'vendor' in line + or 'product' in line ] if self.type is ZipTypes.qcom else [] @@ -300,24 +284,24 @@ def get_updater_script_lines(self, invalid_files: set[str]) -> str: [ line for line in original_updater_script - if ("getprop" in line and "ro.build.date.utc" not in line) - or "Target" in line - or "dynamic_partitions_op_list" in line + if ('getprop' in line and 'ro.build.date.utc' not in line) + or 'Target' in line + or 'dynamic_partitions_op_list' in line or ( - "firmware-update" in line - and line.split("/")[1].split('"')[0] not in invalid_files + 'firmware-update' in line + and line.split('/')[1].split('"')[0] not in invalid_files ) - and "vbmeta" not in line - or "vendor" in line + and 'vbmeta' not in line + or 'vendor' in line ] if self.type is ZipTypes.qcom else [ line for line in original_updater_script - if "system" not in line and "boot.img" not in line + if 'system' not in line and 'boot.img' not in line ] ) - return "\n".join(lines) + return '\n'.join(lines) def generate_updater_script(self, invalid_files: set[str]): """ @@ -326,54 +310,50 @@ def generate_updater_script(self, invalid_files: set[str]): :return: """ template = ScriptTemplate( - Path( - Path(__file__).parent / "templates/recovery_updater_script" - ).read_text() + Path(Path(__file__).parent / 'templates/recovery_updater_script').read_text() ) lines = self.get_updater_script_lines(invalid_files) if not lines: - raise RuntimeError("Could not extract lines from updater-script!") + raise RuntimeError('Could not extract lines from updater-script!') if self.extract_mode is ProcessTypes.firmware: - process = "Normal Firmware" + process = 'Normal Firmware' elif self.extract_mode is ProcessTypes.non_arb_firmware: - process = "non-ARB firmware" + process = 'non-ARB firmware' elif self.extract_mode is ProcessTypes.firmware_less: - process = "firmware-less ROM" + process = 'firmware-less ROM' elif self.extract_mode is ProcessTypes.vendor: - process = "firmware+vendor" + process = 'firmware+vendor' else: - process = "Unknown" # This should never happen + process = 'Unknown' # This should never happen updater_script = template.substitute( datetime=self.datetime, host=self.host, process=process, lines=lines ) # correct some updater-script lines that exist in old devices' file - if "/firmware/image/sec.dat" in updater_script: + if '/firmware/image/sec.dat' in updater_script: updater_script = updater_script.replace( - "/firmware/image/sec.dat", "/dev/block/bootdevice/by-name/sec" + '/firmware/image/sec.dat', '/dev/block/bootdevice/by-name/sec' ) - if "/firmware/image/splash.img" in updater_script: + if '/firmware/image/splash.img' in updater_script: updater_script = updater_script.replace( - "/firmware/image/splash.img", "/dev/block/bootdevice/by-name/splash" + '/firmware/image/splash.img', '/dev/block/bootdevice/by-name/splash' ) self.update_script = updater_script - write_text_to_file( - f"{str(self._flashing_script_dir)}/updater-script", updater_script - ) + write_text_to_file(f'{str(self._flashing_script_dir)}/updater-script', updater_script) # Use modified dynamic_partitions_op_list with resize vendor line only if ( self.extract_mode is ProcessTypes.vendor - and "dynamic_partitions_op_list" in updater_script + and 'dynamic_partitions_op_list' in updater_script ): original_dynamic_partitions_list = Path( - self._tmp_dir / "dynamic_partitions_op_list" + self._tmp_dir / 'dynamic_partitions_op_list' ).read_text() vendor_resize = re.search( - r"(resize vendor .*$)", original_dynamic_partitions_list, re.M + r'(resize vendor .*$)', original_dynamic_partitions_list, re.M ) if vendor_resize: write_text_to_file( - f"{str(self._tmp_dir)}/dynamic_partitions_op_list", + f'{str(self._tmp_dir)}/dynamic_partitions_op_list', vendor_resize.group(1), ) @@ -392,27 +372,23 @@ def generate_updater_script(self, invalid_files: set[str]): def generate_ab_updater_script(self, invalid_files: set[str]): script_template = ScriptTemplate( - Path( - Path(__file__).parent / "templates/recovery_ab_updater_script" - ).read_text() + Path(Path(__file__).parent / 'templates/recovery_ab_updater_script').read_text() ) flashing_template = ScriptTemplate( - Path(Path(__file__).parent / "templates/partition_flashing").read_text() + Path(Path(__file__).parent / 'templates/partition_flashing').read_text() ) lines = [ - flashing_template.substitute(partition=file.split("/")[-1].split(".")[0]) + flashing_template.substitute(partition=file.split('/')[-1].split('.')[0]) for file in self.get_files_list() - if file.startswith("firmware-update/") and file not in invalid_files + if file.startswith('firmware-update/') and file not in invalid_files ] updater_script = script_template.substitute( datetime=self.datetime, host=self.host, zip_name=self.extractor.get_file_name(), - lines="\n".join(lines), - ) - write_text_to_file( - f"{str(self._flashing_script_dir)}/updater-script", updater_script + lines='\n'.join(lines), ) + write_text_to_file(f'{str(self._flashing_script_dir)}/updater-script', updater_script) def generate_flashing_script(self, invalid_files: set[str]): """ @@ -422,8 +398,8 @@ def generate_flashing_script(self, invalid_files: set[str]): if self.is_android_one is True: self.generate_ab_updater_script(invalid_files) copy2( - Path(Path(__file__).parent / "binaries/update-binary"), - f"{str(self._flashing_script_dir)}/update-binary", + Path(Path(__file__).parent / 'binaries/update-binary'), + f'{str(self._flashing_script_dir)}/update-binary', ) else: self.generate_updater_script(invalid_files) @@ -435,31 +411,29 @@ def make_zip(self) -> str: Also, name it according to process parameters. :return: """ - out = Path(f"{self._out_dir}/result.zip") - make_archive(str(out.with_suffix("").absolute()), "zip", self._tmp_dir) + out = Path(f'{self._out_dir}/result.zip') + make_archive(str(out.with_suffix('').absolute()), 'zip', self._tmp_dir) if not out.exists(): - raise RuntimeError("Could not create result zip file!") + raise RuntimeError('Could not create result zip file!') if not self.is_android_one: codename = extract_codename(self.update_script) else: file_name = self.extractor.get_file_name() codename = cleanup_codename( - file_name.split("-")[0] if "OS2." in file_name else file_name.split("_")[1] + file_name.split('-')[0] if 'OS2.' in file_name else file_name.split('_')[1] ).lower() - zip_prefix = "" + zip_prefix = '' if self.extract_mode is ProcessTypes.firmware: - zip_prefix = "fw" + zip_prefix = 'fw' elif self.extract_mode is ProcessTypes.non_arb_firmware: - zip_prefix = "fw-non-arb" + zip_prefix = 'fw-non-arb' elif self.extract_mode is ProcessTypes.firmware_less: - zip_prefix = "fw-less" + zip_prefix = 'fw-less' elif self.extract_mode is ProcessTypes.vendor: - zip_prefix = "fw-vendor" + zip_prefix = 'fw-vendor' else: pass # This should never happen - zip_name = ( - f"{self._out_dir}/{zip_prefix}_{codename}_{self.extractor.get_file_name()}" - ) + zip_name = f'{self._out_dir}/{zip_prefix}_{codename}_{self.extractor.get_file_name()}' out.rename(zip_name) return zip_name @@ -469,25 +443,23 @@ def extract(self) -> tuple[set[str], set[str]]: :return: a tuple of a set of valid files to extract and a set of zero length invalid files. """ - if hasattr(self.extractor, "handler") and isinstance( - self.extractor.handler, AndroidOneZip - ): + if hasattr(self.extractor, 'handler') and isinstance(self.extractor.handler, AndroidOneZip): self.is_android_one = True self.extractor.prepare() self.get_rom_type() files_to_extract = self.get_files_list() if not files_to_extract or ( len(files_to_extract) == 2 - and "updater-script" in str(files_to_extract) - and "update-binary" in str(files_to_extract) + and 'updater-script' in str(files_to_extract) + and 'update-binary' in str(files_to_extract) ): self.cleanup() self.close() - raise RuntimeError("Nothing found to extract!") + raise RuntimeError('Nothing found to extract!') self.extractor.extract(files_to_extract) # Filter out zero length invalid files invalid_files = set() - firmware_update_dir = Path(self._tmp_dir / "firmware-update") + firmware_update_dir = Path(self._tmp_dir / 'firmware-update') if not firmware_update_dir.exists(): return set(files_to_extract), invalid_files for file in firmware_update_dir.iterdir(): @@ -521,9 +493,9 @@ def auto(self) -> str: :return: output zip file name as a string. """ _, invalid_files = self.extract() - print("Generating flashing script...") + print('Generating flashing script...') self.generate_flashing_script(invalid_files) - print("Creating new zip file..") + print('Creating new zip file..') new_zip = self.make_zip() self.cleanup() self.close() diff --git a/xiaomi_flashable_firmware_creator/helpers/misc.py b/xiaomi_flashable_firmware_creator/helpers/misc.py index 9f7a5e1..bae54c7 100644 --- a/xiaomi_flashable_firmware_creator/helpers/misc.py +++ b/xiaomi_flashable_firmware_creator/helpers/misc.py @@ -1,4 +1,5 @@ """Miscellaneous functions used by the tool.""" + import re from pathlib import Path from string import Template @@ -18,15 +19,15 @@ def extract_codename(updater_script) -> str: :return: extracted codename if found or 'codename' if not found. """ pattern = re.compile( - r"(?:/[\w\d_-]+/([\w\d]+):\d)|" - r"(?:\(\"ro\.product\.device\"\) == \"([\w\d]+)\")|" - r"(?:get_device_compatible\(\"([\w\d]+)\"\))" + r'(?:/[\w_-]+/(\w+):\d)|' + r'(?:\(\"ro\.product\.device\"\) == \"(\w+)\")|' + r'(?:get_device_compatible\(\"(\w+)\"\))' ) match = pattern.search(updater_script) if match: codename = [i for i in match.groups() if i is not None] return codename[0] - return "codename" + return 'codename' def cleanup_codename(codename: str) -> str: @@ -35,30 +36,30 @@ def cleanup_codename(codename: str) -> str: :param codename: codename from miui zip :return: clean codename """ - if "SPROUT" in codename: - codename = codename.replace("SPROUT", "") - if codename.endswith("PRE"): - codename = codename.replace("PRE", "") - if "EEAGlobal" in codename: - return codename.replace("EEAGlobal", "") - if "IDGlobal" in codename and codename not in ["CUPIDGlobal"]: - return codename.replace("IDGlobal", "") - if "INGlobal" in codename and codename not in ["CHOPINGlobal"]: - return codename.replace("INGlobal", "") - if "RUGlobal" in codename: - return codename.replace("RUGlobal", "") - if "TRGlobal" in codename: - return codename.replace("TRGlobal", "") - if "TWGlobal" in codename: - return codename.replace("TWGlobal", "") - if "Global" in codename: - return codename.replace("Global", "") + if 'SPROUT' in codename: + codename = codename.replace('SPROUT', '') + if codename.endswith('PRE'): + codename = codename.replace('PRE', '') + if 'EEAGlobal' in codename: + return codename.replace('EEAGlobal', '') + if 'IDGlobal' in codename and codename not in ['CUPIDGlobal']: + return codename.replace('IDGlobal', '') + if 'INGlobal' in codename and codename not in ['CHOPINGlobal']: + return codename.replace('INGlobal', '') + if 'RUGlobal' in codename: + return codename.replace('RUGlobal', '') + if 'TRGlobal' in codename: + return codename.replace('TRGlobal', '') + if 'TWGlobal' in codename: + return codename.replace('TWGlobal', '') + if 'Global' in codename: + return codename.replace('Global', '') return codename class ScriptTemplate(Template): - delimiter = "[-]" + delimiter = '[-]' def write_text_to_file(file: Union[str, Path], text: str): - Path(file).write_bytes(text.encode("utf-8")) + Path(file).write_bytes(text.encode('utf-8')) diff --git a/xiaomi_flashable_firmware_creator/helpers/mock_zip.py b/xiaomi_flashable_firmware_creator/helpers/mock_zip.py index 116fc2d..79470c1 100644 --- a/xiaomi_flashable_firmware_creator/helpers/mock_zip.py +++ b/xiaomi_flashable_firmware_creator/helpers/mock_zip.py @@ -1,4 +1,5 @@ """Mock a zip file and create a placeholder files of its contents.""" + from pathlib import Path from shutil import make_archive, rmtree from sys import argv @@ -17,36 +18,36 @@ def mock_zip(zip_file: str, out_dir: Union[str, Path]): :type out_dir: str or Path """ if isinstance(out_dir, str): - out_dir = Path(out_dir) / "mock_tmp" + out_dir = Path(out_dir) / 'mock_tmp' else: - out_dir = out_dir / "mock_tmp" + out_dir = out_dir / 'mock_tmp' out_dir = out_dir.absolute() out_dir.mkdir(parents=True, exist_ok=True) zip_file = Path(zip_file) - with ZipFile(zip_file, "r") as zipfile: + with ZipFile(zip_file, 'r') as zipfile: files = zipfile.namelist() for item in files: item_path = Path(out_dir / item) - if item.endswith("/"): + if item.endswith('/'): if not item_path.exists(): item_path.mkdir(parents=True, exist_ok=True) else: if not item_path.exists(): item_path.parent.mkdir(parents=True, exist_ok=True) - if item.endswith("updater-script"): + if item.endswith('updater-script'): zipfile.extract(item, out_dir) else: - with open(f"{out_dir}/{item}", "wb") as out: - out.write(b"") + with open(f'{out_dir}/{item}', 'wb') as out: + out.write(b'') - mocked_zip = f"{out_dir.parent}/mocked_{zip_file.stem}" - make_archive(mocked_zip, zip_file.suffix.split(".")[-1], out_dir) - out = Path(f"{mocked_zip}.zip") + mocked_zip = f'{out_dir.parent}/mocked_{zip_file.stem}' + make_archive(mocked_zip, zip_file.suffix.split('.')[-1], out_dir) + out = Path(f'{mocked_zip}.zip') if not out.exists(): - raise RuntimeError("Could not create mocked zip file!") + raise RuntimeError('Could not create mocked zip file!') rmtree(out_dir) -if __name__ == "__main__": +if __name__ == '__main__': mock_zip(argv[1], argv[2]) diff --git a/xiaomi_flashable_firmware_creator/tests/test_creator.py b/xiaomi_flashable_firmware_creator/tests/test_creator.py index 0120ea4..aeaf2cf 100644 --- a/xiaomi_flashable_firmware_creator/tests/test_creator.py +++ b/xiaomi_flashable_firmware_creator/tests/test_creator.py @@ -19,8 +19,8 @@ def setUp(self): :return: """ self.work_dir = Path(__file__).parent - self.out_dir = self.work_dir / "out" - self.files = self.work_dir.glob("files/*/*.zip") + self.out_dir = self.work_dir / 'out' + self.files = self.work_dir.glob('files/*/*.zip') @staticmethod def run_extractor(firmware_creator): @@ -30,7 +30,7 @@ def run_extractor(firmware_creator): :param firmware_creator: FlashableFirmwareCreator object :return: """ - print("Unzipping ROM...") + print('Unzipping ROM...') firmware_creator.auto() def test_firmware(self): @@ -41,9 +41,9 @@ def test_firmware(self): """ for file in self.files: firmware_creator = FlashableFirmwareCreator( - str(file.absolute()), "firmware", self.out_dir + str(file.absolute()), 'firmware', self.out_dir ) - print(f"Testing {file.name}") + print(f'Testing {file.name}') self.run_extractor(firmware_creator) def test_firmwareless(self): @@ -54,13 +54,13 @@ def test_firmwareless(self): """ for file in self.files: firmware_creator = FlashableFirmwareCreator( - str(file.absolute()), "firmwareless", self.out_dir + str(file.absolute()), 'firmwareless', self.out_dir ) - print(f"Testing {file.name}") + print(f'Testing {file.name}') try: self.run_extractor(firmware_creator) except RuntimeError as err: - if str(err) != "Nothing found to extract!": + if str(err) != 'Nothing found to extract!': raise err def test_nonarb(self): @@ -71,13 +71,13 @@ def test_nonarb(self): """ for file in self.files: firmware_creator = FlashableFirmwareCreator( - str(file.absolute()), "nonarb", self.out_dir + str(file.absolute()), 'nonarb', self.out_dir ) - print(f"Testing {file.name}") + print(f'Testing {file.name}') try: self.run_extractor(firmware_creator) except RuntimeError as err: - if str(err) != "Nothing found to extract!": + if str(err) != 'Nothing found to extract!': raise err def test_vendor(self): @@ -88,9 +88,9 @@ def test_vendor(self): """ for file in self.files: firmware_creator = FlashableFirmwareCreator( - str(file.absolute()), "vendor", self.out_dir + str(file.absolute()), 'vendor', self.out_dir ) - print(f"Testing {file.name}") + print(f'Testing {file.name}') self.run_extractor(firmware_creator) def test_date_assertion(self): @@ -101,15 +101,15 @@ def test_date_assertion(self): """ for file in self.files: firmware_creator = FlashableFirmwareCreator( - str(file.absolute()), "firmware", self.out_dir + str(file.absolute()), 'firmware', self.out_dir ) - print(f"Testing {file.name}") + print(f'Testing {file.name}') firmware_creator.extract() firmware_creator.generate_flashing_script([]) update_script = Path( - firmware_creator._flashing_script_dir / "updater-script" + firmware_creator._flashing_script_dir / 'updater-script' ).read_text() - self.assertNotIn("ro.build.date.utc", update_script) + self.assertNotIn('ro.build.date.utc', update_script) firmware_creator.cleanup() firmware_creator.close() @@ -122,5 +122,5 @@ def tearDown(self): rmtree(self.out_dir) -if __name__ == "__main__": +if __name__ == '__main__': unittest.main() diff --git a/xiaomi_flashable_firmware_creator/tests/test_misc.py b/xiaomi_flashable_firmware_creator/tests/test_misc.py index d95bb0b..7c58768 100644 --- a/xiaomi_flashable_firmware_creator/tests/test_misc.py +++ b/xiaomi_flashable_firmware_creator/tests/test_misc.py @@ -22,7 +22,7 @@ def setUp(self): :return: """ self.work_dir = Path(__file__).parent - self.files = self.work_dir.glob("files/updater-scripts/*") + self.files = self.work_dir.glob('files/updater-scripts/*') def test_extract_codename(self): """ @@ -33,8 +33,8 @@ def test_extract_codename(self): for file in self.files: updater_script = file.read_text() codename = extract_codename(updater_script) - self.assertNotEqual("codename", codename) + self.assertNotEqual('codename', codename) -if __name__ == "__main__": +if __name__ == '__main__': unittest.main() diff --git a/xiaomi_flashable_firmware_creator/xiaomi_flashable_firmware_creator.py b/xiaomi_flashable_firmware_creator/xiaomi_flashable_firmware_creator.py index 9bb4568..69cfc80 100644 --- a/xiaomi_flashable_firmware_creator/xiaomi_flashable_firmware_creator.py +++ b/xiaomi_flashable_firmware_creator/xiaomi_flashable_firmware_creator.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 """Xiaomi Flashable Firmware Creator cli module.""" from argparse import ArgumentParser @@ -17,15 +16,15 @@ def arg_parse() -> tuple[str, str, str]: """ output = None parser = ArgumentParser( - prog="python3 -m xiaomi_flashable_firmware_creator", - description="Xiaomi Flashable Firmware Creator", + prog='python3 -m xiaomi_flashable_firmware_creator', + description='Xiaomi Flashable Firmware Creator', ) group = parser.add_mutually_exclusive_group(required=True) - group.add_argument("-F", "--firmware", help="Create normal Firmware zip") - group.add_argument("-N", "--nonarb", help="Create non-ARB Firmware zip") - group.add_argument("-L", "--firmwareless", help="Create Firmware-less zip") - group.add_argument("-V", "--vendor", help="Create Firmware+Vendor zip") - parser.add_argument("-o", "--output", help="Output directory") + group.add_argument('-F', '--firmware', help='Create normal Firmware zip') + group.add_argument('-N', '--nonarb', help='Create non-ARB Firmware zip') + group.add_argument('-L', '--firmwareless', help='Create Firmware-less zip') + group.add_argument('-V', '--vendor', help='Create Firmware+Vendor zip') + parser.add_argument('-o', '--output', help='Output directory') args = parser.parse_args() if args.output: output = args.output @@ -39,6 +38,6 @@ def main(): """Xiaomi Flashable Firmware Creator main module.""" zip_, process, output = arg_parse() firmware_creator = FlashableFirmwareCreator(zip_, process, output) - print("Unzipping MIUI ROM...") + print('Unzipping MIUI ROM...') new_zip = firmware_creator.auto() - print(f"All done! Output file is {new_zip}") + print(f'All done! Output file is {new_zip}')