From 72d14a370c85d6d126aca0c544f58d11c50e35f0 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Mon, 20 Jan 2025 11:34:12 +0100
Subject: [PATCH] feat(python): allow `trezorctl firmware` commands to work
 with unknown models

[no changelog]
---
 python/src/trezorlib/cli/firmware.py      | 16 +++++++++++-----
 python/src/trezorlib/firmware/__init__.py |  3 +++
 python/src/trezorlib/firmware/core.py     |  8 ++++++++
 python/src/trezorlib/firmware/legacy.py   |  3 +++
 4 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/python/src/trezorlib/cli/firmware.py b/python/src/trezorlib/cli/firmware.py
index 4376a4f283..923c20b3d6 100644
--- a/python/src/trezorlib/cli/firmware.py
+++ b/python/src/trezorlib/cli/firmware.py
@@ -95,9 +95,9 @@ def _print_firmware_model(hw_model: Union[bytes, fw_models.Model]) -> None:
         pass
 
     assert isinstance(hw_model, bytes)
-    if all(0x20 <= b < 0x80 for b in hw_model):  # isascii
+    if hw_model.isascii():
         model_name = hw_model.decode("ascii")
-        click.echo(f"Unknown hardware model: {model_name}")
+        click.echo(f"Unrecognized hardware model: {model_name}")
         return
 
     click.echo(f"Suspicious hardware model code: {hw_model.hex()} ({hw_model!r})")
@@ -404,7 +404,7 @@ def validate_firmware(
     fingerprint: Optional[str] = None,
     model: Optional[TrezorModel] = None,
     bootloader_onev2: Optional[bool] = None,
-    prompt_unsigned: bool = True,
+    verify_only: bool = False,
 ) -> None:
     """Validate the firmware through multiple tests.
 
@@ -419,8 +419,14 @@ def validate_firmware(
         sys.exit(2)
 
     print_firmware_version(fw)
+    if not fw.model():
+        click.echo("Cannot validate firmware for unrecognized model.")
+        if not verify_only:
+            click.echo("(Hint: use --skip-check to skip validation.)")
+        sys.exit(3)
+
     validate_fingerprint(fw, fingerprint)
-    validate_signatures(fw, prompt_unsigned=prompt_unsigned)
+    validate_signatures(fw, prompt_unsigned=not verify_only)
 
     if model is not None and bootloader_onev2 is not None:
         check_device_match(fw, model, bootloader_onev2)
@@ -548,7 +554,7 @@ def verify(
         fingerprint=fingerprint,
         bootloader_onev2=bootloader_onev2,
         model=model,
-        prompt_unsigned=False,
+        verify_only=True,
     )
 
 
diff --git a/python/src/trezorlib/firmware/__init__.py b/python/src/trezorlib/firmware/__init__.py
index 4cfc11dd40..8d44536272 100644
--- a/python/src/trezorlib/firmware/__init__.py
+++ b/python/src/trezorlib/firmware/__init__.py
@@ -21,6 +21,7 @@ from typing_extensions import Protocol, TypeGuard
 
 from .. import messages
 from ..tools import session
+from .models import Model
 from .core import VendorFirmware
 from .legacy import LegacyFirmware, LegacyV2Firmware
 
@@ -50,6 +51,8 @@ if t.TYPE_CHECKING:
 
         def digest(self) -> bytes: ...
 
+        def model(self) -> Model | None: ...
+
 
 def parse(data: bytes) -> "FirmwareType":
     try:
diff --git a/python/src/trezorlib/firmware/core.py b/python/src/trezorlib/firmware/core.py
index ecae74bc6c..ee342f3607 100644
--- a/python/src/trezorlib/firmware/core.py
+++ b/python/src/trezorlib/firmware/core.py
@@ -165,6 +165,11 @@ class FirmwareImage(Struct):
         header.v1_signatures = [b"\x00" * 64] * consts.V1_SIGNATURE_SLOTS
         return hash_params.hash_function(header.build()).digest()
 
+    def model(self) -> Model | None:
+        if isinstance(self.header.hw_model, Model):
+            return self.header.hw_model
+        return None
+
 
 class VendorFirmware(Struct):
     """Firmware image prefixed by a vendor header.
@@ -203,3 +208,6 @@ class VendorFirmware(Struct):
         # now = time.gmtime()
         # if time.gmtime(fw.vendor_header.expiry) < now:
         #     raise ValueError("Vendor header expired.")
+
+    def model(self) -> Model | None:
+        return self.firmware.model()
diff --git a/python/src/trezorlib/firmware/legacy.py b/python/src/trezorlib/firmware/legacy.py
index a4557489e4..8635a7539a 100644
--- a/python/src/trezorlib/firmware/legacy.py
+++ b/python/src/trezorlib/firmware/legacy.py
@@ -202,3 +202,6 @@ class LegacyFirmware(Struct):
 
         if self.embedded_v2:
             self.embedded_v2.verify()
+
+    def model(self) -> Model | None:
+        return Model.T1B1