From 9f597a60720ce567d169b1b7f35f19baa29ea636 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Fri, 21 Apr 2023 12:03:11 +0200
Subject: [PATCH] feat(core/tools): make combine_firmware script a bit smarter
 and more opinionated

---
 core/Makefile               |  8 ++--
 core/tools/combine_firmware | 79 +++++++++++++++++++++++++++++--------
 2 files changed, 67 insertions(+), 20 deletions(-)

diff --git a/core/Makefile b/core/Makefile
index d1e3d51de8..39f1340058 100644
--- a/core/Makefile
+++ b/core/Makefile
@@ -321,10 +321,10 @@ sizecheck: ## check sizes of binary files
 
 combine: ## combine boardloader + bootloader + prodtest into one combined image
 	./tools/combine_firmware \
-		$(BOARDLOADER_START) $(BOARDLOADER_BUILD_DIR)/boardloader.bin \
-		$(BOOTLOADER_START) $(BOOTLOADER_BUILD_DIR)/bootloader.bin \
-		$(PRODTEST_START) $(PRODTEST_BUILD_DIR)/prodtest.bin \
-		> $(PRODTEST_BUILD_DIR)/combined.bin
+		$(BOARDLOADER_BUILD_DIR)/boardloader.bin \
+		$(BOOTLOADER_BUILD_DIR)/bootloader.bin \
+		$(PRODTEST_BUILD_DIR)/prodtest.bin \
+		$(PRODTEST_BUILD_DIR)/combined.bin
 
 upload: ## upload firmware using trezorctl
 	trezorctl firmware_update -f $(FIRMWARE_BUILD_DIR)/firmware.bin
diff --git a/core/tools/combine_firmware b/core/tools/combine_firmware
index 7f12286544..046d2f275d 100755
--- a/core/tools/combine_firmware
+++ b/core/tools/combine_firmware
@@ -1,26 +1,73 @@
 #!/usr/bin/env python3
+from __future__ import annotations
+
+import datetime
+import io
 import sys
+from pathlib import Path
+
+import click
 
 
-def pairwise(iterable):
-    a = iter(iterable)
-    return zip(a, a)
+BOARDLOADER_START = 0x08000000
+BOARDLOADER_END = 0x0800C000
+BOOTLOADER_START = 0x08020000
+FIRMWARE_START = 0x08040000
 
 
-files = sys.argv[1:]
-files = list(pairwise(files))
+@click.command()
+@click.argument(
+    "boardloader",
+    type=click.Path(exists=True, dir_okay=False, readable=True, path_type=Path),
+)
+@click.argument(
+    "bootloader",
+    type=click.Path(exists=True, dir_okay=False, readable=True, path_type=Path),
+)
+@click.argument(
+    "firmware",
+    type=click.Path(exists=True, dir_okay=False, readable=True, path_type=Path),
+)
+@click.argument(
+    "outfile",
+    type=click.Path(dir_okay=False, writable=True, path_type=Path),
+    required=False,
+)
+def main(
+    boardloader: Path, bootloader: Path, firmware: Path, outfile: Path | None
+) -> None:
+    if outfile is None:
+        today = datetime.date.today().strftime(r"%Y-%m-%d")
+        outfile = Path(f"combined-{today}.bin")
 
-offset = int(files[0][0], 16)
+    offset = BOARDLOADER_START
+    out_bytes = io.BytesIO()
 
-out = bytearray()
+    # write boardloader
+    offset += out_bytes.write(boardloader.read_bytes())
+    if offset > BOARDLOADER_END:
+        raise Exception("Boardloader too big")
 
-for addr, fn in files:
-    addr = int(addr, 16) - offset
-    data = open(fn, "rb").read()
-    if len(out) < addr:
-        out += b"\x00" * (addr - len(out))
-    if len(out) != addr:
-        raise Exception("Alignment failed")
-    out += data
+    # zero-pad until next section:
+    offset += out_bytes.write(b"\x00" * (BOOTLOADER_START - offset))
+    assert offset == BOOTLOADER_START
 
-sys.stdout.buffer.write(out)
+    # write bootlaoder
+    offset += out_bytes.write(bootloader.read_bytes())
+    if offset > FIRMWARE_START:
+        raise Exception("Bootloader too big")
+
+    # zero-pad until next section:
+    offset += out_bytes.write(b"\x00" * (FIRMWARE_START - offset))
+    assert offset == FIRMWARE_START
+
+    # write firmware
+    offset += out_bytes.write(firmware.read_bytes())
+
+    # write out contents
+    click.echo(f"Writing {outfile} ({offset - BOARDLOADER_START} bytes)")
+    outfile.write_bytes(out_bytes.getvalue())
+
+
+if __name__ == "__main__":
+    main()