From 5d9bb0ebfcd194b31adcaffa17bcec0698a56d1b Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Fri, 12 Apr 2024 14:20:32 +0200
Subject: [PATCH] tools(core): nicer build_vendorheader for manual use

also will show you required header size
---
 core/embed/vendorheader/generate.sh |  3 +--
 core/tools/build_vendorheader       | 39 +++++++++++++++++++++++------
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/core/embed/vendorheader/generate.sh b/core/embed/vendorheader/generate.sh
index 14a9b1fe2..2a63dc8ef 100755
--- a/core/embed/vendorheader/generate.sh
+++ b/core/embed/vendorheader/generate.sh
@@ -23,8 +23,7 @@ for MODEL in ${MODELS[@]}; do
     echo "Generating vendor headers for $MODEL"
     # construct all vendor headers
     for fn in *.json; do
-        name=$(echo $fn | sed 's/vendor_\(.*\)\.json/\1/')
-        $BUILDVH $QUIET $CHECK vendor_${name}.json vendor_${name}.toif vendorheader_${name}_unsigned.bin
+        $BUILDVH $QUIET $CHECK $fn
     done
 
     TMPDIR=$(mktemp -d)
diff --git a/core/tools/build_vendorheader b/core/tools/build_vendorheader
index 14fda0c67..871777561 100755
--- a/core/tools/build_vendorheader
+++ b/core/tools/build_vendorheader
@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+from __future__ import annotations
+
 import json
 from pathlib import Path
 
@@ -7,7 +9,7 @@ import click
 from trezorlib import firmware, toif
 
 
-def minimum_header_len(spec):
+def minimum_header_len(spec, quiet):
     spec = spec.copy()
     spec["header_len"] = 512000
     reparsed = firmware.VendorHeader.SUBCON.parse(
@@ -15,24 +17,45 @@ def minimum_header_len(spec):
     )
     data_length = reparsed._end_offset - reparsed._start_offset
     # data length + 65 for signatures, rounded up to nearest multiple of 512
+    if not quiet:
+        click.echo(f"Minimum header length: {data_length + 65} bytes.")
     return (data_length + 65 + 511) // 512 * 512
 
 
+FilePath = click.Path(dir_okay=False, path_type=Path)
+
+
 @click.command()
-@click.argument("specfile", type=click.File("r"))
-@click.argument("image", type=click.File("rb"))
-@click.argument("outfile", type=click.Path(dir_okay=False, writable=True, path_type=Path))
+@click.argument("specfile", type=FilePath)
+@click.option("-i", "--image", type=FilePath)
+@click.option("-o", "--outfile", type=FilePath)
 @click.option("-c", "--check", is_flag=True, help="Check but do not write header.")
 @click.option("-q", "--quiet", is_flag=True, help="Do not print anything.")
-def build_vendorheader(specfile, image, outfile, check: bool, quiet: bool):
+def build_vendorheader(
+    specfile: Path, image: Path | None, outfile: Path | None, check: bool, quiet: bool
+):
     if quiet:
         echo = lambda *args, **kwargs: None
     else:
         echo = click.echo
 
-    spec = json.load(specfile)
+    if not specfile.exists():
+        raise click.ClickException(f"Spec file {specfile.name} does not exist.")
+
+    if image is None:
+        image = specfile.with_suffix(".toif")
+    if not image.exists():
+        raise click.ClickException(f"Image file {image.name} does not exist.")
+
+    if outfile is None:
+        vh_stem = specfile.stem
+        if vh_stem.startswith("vendor_"):
+            vh_stem = vh_stem[len("vendor_") :]
+        outfile = specfile.parent / f"vendorheader_{vh_stem}_unsigned.bin"
+
+    spec = json.loads(specfile.read_text())
     spec["pubkeys"] = [bytes.fromhex(k) for k in spec["pubkeys"]]
-    spec["image"] = toif.ToifStruct.parse(image.read())
+    spec["image"] = toif.ToifStruct.parse(image.read_bytes())
     spec["sigmask"] = 0
     spec["signature"] = b"\x00" * 64
     if spec["hw_model"] is None:
@@ -40,7 +63,7 @@ def build_vendorheader(specfile, image, outfile, check: bool, quiet: bool):
     else:
         spec["hw_model"] = spec["hw_model"].encode("ascii")
 
-    min_length = minimum_header_len(spec)
+    min_length = minimum_header_len(spec, quiet)
     if "header_len" not in spec:
         spec["header_len"] = min_length
     elif spec["header_len"] < min_length: