From c6e18263cd983a9a85e133bd89150d0df85b1f75 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B6rg=20Stucke?= <joerg.stucke@fkie.fraunhofer.de>
Date: Thu, 9 Apr 2026 13:17:43 +0200
Subject: [PATCH] fix: FIT unpacker fixes

* fixed the problem that the unpacker did not extract data from nested nodes (depth >2)

* fixed the issue that the unpacker stops after the first node that does not contain a 'data' entry (even though there could be more subnodes on the same level with such entries)
---
 .../plugins/unpacking/uboot/code/fit.py       |  92 ++++++++++++------
 .../unpacking/uboot/test/data/test.dtb        | Bin 0 -> 685 bytes
 .../unpacking/uboot/test/test_plugin_fit.py   |  26 ++++-
 3 files changed, 85 insertions(+), 33 deletions(-)
 create mode 100644 fact_extractor/plugins/unpacking/uboot/test/data/test.dtb

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 0000000000000000000000000000000000000000..961ca61e96b2910a212be4634f83989526690d35
GIT binary patch
literal 685
zcmZuuOG*Pl5Uq)*5fd~n0^-8KjZp{^+$e|}K@EZn$pwZ?CuuV29%ed-Ucs}t^$>1!
z=P5japw_F-Xc8T$sn_q-r>pw4`uR>&dL*JF_Ktmnt;V*_hVPj96Ax$p!Q`~Sehs*z
zau=SrJE0y$eI;rjF!$KlW*oVwX>9COBIIsB1sC3GH?U(DB1%%I#4Z@B%&0^N#8-f%
z*VT^zei7QGN)n3zUwfH-Z7v^ibz@C&2kDVTqgdy|BCiS{J#KylrqE}Tt+M@4rrp~&
zUHV#mIe2adbaDFhfy(|vT$!ylS66!X?{zwf3k>f}#d1PU9msQ`tfUdE)D6PJJqeNi
zOAIU5vw8o`gt=t4t~EW~A7_!%*7&zTFYIOqTC>H7_dJE6Z-G3U&i(s9MV<Y*9b@*^
o)a=Dw;8$XmcQc*hHY88&?ovfvJ4!Q^=fbL9G)^2%suqR50j_~|bN~PV

literal 0
HcmV?d00001

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']