diff --git a/fact_extractor/plugins/unpacking/uboot/code/fit.py b/fact_extractor/plugins/unpacking/uboot/code/fit.py index bf08ef96..f03c0462 100644 --- a/fact_extractor/plugins/unpacking/uboot/code/fit.py +++ b/fact_extractor/plugins/unpacking/uboot/code/fit.py @@ -2,49 +2,83 @@ This plugin unpacks Flattened Image Trees. """ -from contextlib import suppress +from __future__ import annotations + +import bz2 +import gzip +import lzma from pathlib import Path import libfdt as fdt -from common_helper_files import write_binary_to_file NAME = 'FIT' MIME_PATTERNS = ['linux/device-tree'] -VERSION = '0.2.0' +VERSION = '0.2.1' TRAILING_DATA_MIN_SIZE = 100 +DECOMPRESSORS = { + 'gzip': gzip.decompress, + 'lzma': lzma.decompress, + 'bzip2': bz2.decompress, + # FixMe: add lzo and lz4 decompressors +} + +try: + from compression import zstd + + DECOMPRESSORS['zstd'] = zstd.decompress +except ImportError: + pass # zstd decompression (zstd was added in Python 3.14) -def unpack_function(file_path, tmp_dir): + +def unpack_function(file_path: str, tmp_dir: str) -> dict: file = Path(file_path) + dtb = fdt.Fdt(file.read_bytes()) + root_offset = dtb.path_offset('/') + output = extract_nodes(dtb, root_offset, tmp_dir, Path()) + output.append('successfully unpacked FIT image') + return {'output': '\n'.join(output), 'size': dtb.size_dt_struct()} + + +def extract_nodes(dtb: fdt.Fdt, offset: int, tmp_dir: str, path: Path) -> list[str]: try: - with file.open('rb') as f: - fit_data = f.read() + child_offset = dtb.first_subnode(offset) + except fdt.FdtException: + return [] # no child nodes in this node + + output = [] + while True: + try: + name = dtb.get_name(child_offset) + current_path = path / name - dtb = fdt.Fdt(fit_data) - root_offset = dtb.path_offset('/') - subnode_offset = dtb.first_subnode(root_offset) - while True: try: - component_offset = dtb.first_subnode(subnode_offset) - while True: - try: - outfile = Path(tmp_dir) / dtb.get_name(component_offset) - with suppress(TypeError): - data = dtb.getprop(component_offset, 'data') - if data: - write_binary_to_file(bytes(data), outfile) - component_offset = dtb.next_subnode(component_offset) - except fdt.FdtException: - break - subnode_offset = dtb.next_subnode(subnode_offset) - except fdt.FdtException: - break - except OSError as io_error: - return {'output': f'failed to read file: {io_error!s}'} - message = 'successfully unpacked FIT image' - - return {'output': message, 'size': dtb.size_dt_struct()} + data = _read_from_dtb_at(dtb, child_offset) + output_path = (Path(tmp_dir) / current_path).with_suffix('.bin') + output_path.parent.mkdir(exist_ok=True) + output_path.write_bytes(data) + output.append(f'unpacked data from node {current_path} ({len(data)} bytes)') + except (TypeError, fdt.FdtException): + pass # no "data" entry + + # recurse through child nodes + output.extend(extract_nodes(dtb, child_offset, tmp_dir, current_path)) + + child_offset = dtb.next_subnode(child_offset) + except fdt.FdtException: + break + return output + + +def _read_from_dtb_at(dtb: fdt.Fdt, offset: int) -> bytes: + data = bytes(dtb.getprop(offset, 'data')) + try: + compression = dtb.getprop(offset, 'compression').as_str() + decompressor = DECOMPRESSORS[compression] + return decompressor(data) + except (fdt.FdtException, KeyError): + return data # ----> Do not edit below this line <---- diff --git a/fact_extractor/plugins/unpacking/uboot/test/data/test.dtb b/fact_extractor/plugins/unpacking/uboot/test/data/test.dtb new file mode 100644 index 00000000..961ca61e Binary files /dev/null and b/fact_extractor/plugins/unpacking/uboot/test/data/test.dtb differ diff --git a/fact_extractor/plugins/unpacking/uboot/test/test_plugin_fit.py b/fact_extractor/plugins/unpacking/uboot/test/test_plugin_fit.py index d987d64a..85bd1243 100644 --- a/fact_extractor/plugins/unpacking/uboot/test/test_plugin_fit.py +++ b/fact_extractor/plugins/unpacking/uboot/test/test_plugin_fit.py @@ -11,13 +11,31 @@ class TestFITUnpacker(TestUnpackerBase): def test_unpacker_selection_generic(self): self.check_unpacker_selection('linux/device-tree', 'FIT') - def test_extraction(self): + def test_extraction_itb(self): test_file_path = Path(TEST_DATA_DIR) / 'fit.itb' extracted_files, meta_data = self.unpacker.extract_files_from_file(str(test_file_path), self.tmp_dir.name) assert meta_data['plugin_used'] == 'FIT', 'wrong plugin applied' + assert meta_data.get('error') is None + assert meta_data.get('output') is not None assert len(extracted_files) == EXPECTED_FILE_COUNT, 'not all files extracted' - assert all( - Path(element).name in ['kernel', 'fdt', 'rootfs', 'trailing.bin'] for element in extracted_files - ), 'not all files extracted' + for element in extracted_files: + assert Path(element).name in ['kernel.bin', 'fdt.bin', 'rootfs.bin', 'trailing.bin'] + + def test_extraction_dtb(self): + test_file_path = Path(TEST_DATA_DIR) / 'test.dtb' + extracted_files, meta_data = self.unpacker.extract_files_from_file(str(test_file_path), self.tmp_dir.name) + + assert meta_data['plugin_used'] == 'FIT', 'wrong plugin applied' + assert meta_data.get('error') is None + assert meta_data.get('output') is not None + + files = {p.name: p for f in extracted_files if (p := Path(f))} + assert 'kernel.bin' in files + assert files['kernel.bin'].read_bytes() == b'Hello' + assert 'ramdisk.bin' in files + assert files['ramdisk.bin'].read_bytes() == b'World', 'decompression failed' + assert 'nested.bin' in files + assert files['nested.bin'].read_bytes() == b'Nested', 'nested file not extracted' + assert 'unpacked data from node images/ramdisk/nested (6 bytes)' in meta_data['output']