From 29c8df242dd563086bb648c1eaedac2a4a108693 Mon Sep 17 00:00:00 2001
From: Kun Lai <me@imlk.top>
Date: Sat, 24 Aug 2024 09:43:42 +0800
Subject: [PATCH] patch_itb.py: support both "data-position = " and "data = "

In some FIT file, the data is store in "data" field instead of "data-position".

```txt
		fdt-1 {
			description = "ARM64 OpenWrt cmcc_rax3000m-nand-ubootmod device tree blob";
			data = [d0 0d fe ed 00 00 57 f3 < ... very long ... > 00 72 65 64 5f 6c 65 64 00];
			type = "flat_dt";
			arch = "arm64";
			compression = "none";

			hash-1 {
				value = <0xa8ce7e2e>;
				algo = "crc32";
			};

			hash-2 {
				value = <0x41d92ad2 0xc528dfc 0x32c78d4d 0x68be20d8 0x939c2220>;
				algo = "sha1";
			};
		};
```

Signed-off-by: Kun Lai <me@imlk.top>
---
 patch_itb.py | 20 +++++---------------
 1 file changed, 5 insertions(+), 15 deletions(-)

diff --git a/patch_itb.py b/patch_itb.py
index 11a6c88..a15b617 100755
--- a/patch_itb.py
+++ b/patch_itb.py
@@ -130,7 +130,6 @@ def get_crc32_string(binary_file):
     orig_dtb_bytes = open(f"{build_dir}/orig.dtb", "rb").read()
     orig_sha1 = get_sha1_string(orig_dtb_bytes)
     orig_crc32 = get_crc32_string(orig_dtb_bytes)
-    dtb_offset = None
     with open(f"{build_dir}/orig_itb.its", "r") as f:
         orig_itb_its = f.read()
         if orig_itb_its.find(orig_sha1) == -1:
@@ -142,16 +141,6 @@ def get_crc32_string(binary_file):
         patched_sha1 = get_sha1_string(open(f"{build_dir}/patched.dtb", "rb").read())
         patched_crc32 = get_crc32_string(open(f"{build_dir}/patched.dtb", "rb").read())
         patched_itb_its = orig_itb_its.replace(orig_sha1, patched_sha1).replace(orig_crc32, patched_crc32)
-        # find dtb offset
-        itb_lines = patched_itb_its.splitlines()
-        data_position_idx = 0
-        for line in itb_lines:
-            if line.strip().startswith("data-position"):
-                if data_position_idx == FDT_OFFSET:
-                    dtb_offset = int(line.split()[2][1:-2], 16)
-                    break
-                data_position_idx += 1
-        assert dtb_offset is not None, "Failed to find dtb offset"
         with open(f"{build_dir}/patched_itb.its", "w") as f:
             f.write(patched_itb_its)
 
@@ -161,11 +150,12 @@ def get_crc32_string(binary_file):
     # patch itb file
     itb_file = bytearray(open(itb_path, "rb").read())
     new_itb_header = open(f"{build_dir}/patched_itb.itb", "rb").read()
+    itb_file[0:len(new_itb_header)] = new_itb_header[:]
     new_dtb = open(f"{build_dir}/patched.dtb", "rb").read()
-    for i in range(len(new_itb_header)):
-        itb_file[i] = new_itb_header[i]
-    for i in range(len(new_dtb)):
-        itb_file[dtb_offset + i] = new_dtb[i]
+    dtb_offset = itb_file.find(orig_dtb_bytes)
+    assert dtb_offset != -1, "Failed to find dtb offset"
+    itb_file[dtb_offset:dtb_offset+len(new_dtb)] = new_dtb[:]
+
     with open(patched_path, "wb") as f:
         f.write(itb_file)