python/firmware: clarify firmware image types

python/src/trezorlib/cli/firmware.py

@@ -36,18 +36,18 @@ def validate_firmware(version, fw, expected_fingerprint=None):
     if version == firmware.FirmwareFormat.TREZOR_ONE:
         if fw.embedded_onev2:
             click.echo("Trezor One firmware with embedded v2 image (1.8.0 or later)")
-            _print_version(fw.embedded_onev2.firmware_header.version)
+            _print_version(fw.embedded_onev2.header.version)
         else:
             click.echo("Trezor One firmware image.")
     elif version == firmware.FirmwareFormat.TREZOR_ONE_V2:
         click.echo("Trezor One v2 firmware (1.8.0 or later)")
-        _print_version(fw.firmware_header.version)
+        _print_version(fw.header.version)
     elif version == firmware.FirmwareFormat.TREZOR_T:
         click.echo("Trezor T firmware image.")
         vendor = fw.vendor_header.vendor_string
         vendor_version = "{major}.{minor}".format(**fw.vendor_header.version)
         click.echo("Vendor header from {}, version {}".format(vendor, vendor_version))
-        _print_version(fw.firmware_header.version)
+        _print_version(fw.image.header.version)
 
     try:
         firmware.validate(version, fw, allow_unsigned=False)
@@ -198,17 +198,18 @@ def firmware_update(
         click.echo("Please switch your device to bootloader mode.")
         sys.exit(1)
 
+    # bootloader for T1 does not export 'model', so we rely on major_version
     f = client.features
-    bootloader_onev2 = f.major_version == 1 and f.minor_version >= 8
+    bootloader_version = (f.major_version, f.minor_version, f.patch_version)
+    bootloader_onev2 = f.major_version == 1 and bootloader_version >= (1, 8, 0)
 
     if filename:
         data = open(filename, "rb").read()
     else:
         if not url:
-            bootloader_version = [f.major_version, f.minor_version, f.patch_version]
             version_list = [int(x) for x in version.split(".")] if version else None
             url, fp = find_best_firmware_version(
-                bootloader_version, version_list, beta, bitcoin_only
+                list(bootloader_version), version_list, beta, bitcoin_only
             )
             if not fingerprint:
                 fingerprint = fp

python/src/trezorlib/firmware.py

@@ -198,24 +198,37 @@ FirmwareHeader = c.Struct(
 )
 
 
-Firmware = c.Struct(
+"""Raw firmware image.
+
+Consists of firmware header and code block.
+This is the expected format of firmware binaries for Trezor One, or bootloader images
+for Trezor T."""
+FirmwareImage = c.Struct(
+    "header" / FirmwareHeader,
+    "_code_offset" / c.Tell,
+    "code" / c.Bytes(c.this.header.code_length),
+    c.Terminated,
+)
+
+
+"""Firmware image prefixed by a vendor header.
+
+This is the expected format of firmware binaries for Trezor T."""
+VendorFirmware = c.Struct(
     "vendor_header" / VendorHeader,
-    "firmware_header" / FirmwareHeader,
+    "image" / FirmwareImage,
-    "_code_offset" / c.Tell,
-    "code" / c.Bytes(c.this.firmware_header.code_length),
     c.Terminated,
 )
 
 
-FirmwareOneV2 = c.Struct(
+"""Legacy firmware image.
-    "firmware_header" / FirmwareHeader,
+Consists of a custom header and code block.
-    "_code_offset" / c.Tell,
+This is the expected format of firmware binaries for Trezor One pre-1.8.0.
-    "code" / c.Bytes(c.this.firmware_header.code_length),
-    c.Terminated,
-)
 
-
+The code block can optionally be interpreted as a new-style firmware image. That is the
-FirmwareOne = c.Struct(
+expected format of firmware binary for Trezor One version 1.8.0, which can be installed
+by both the older and the newer bootloader."""
+LegacyFirmware = c.Struct(
     "magic" / c.Const(b"TRZR"),
     "code_length" / c.Rebuild(c.Int32ul, c.len_(c.this.code)),
     "key_indexes" / c.Int8ul[V1_SIGNATURE_SLOTS],  # pylint: disable=E1136
@@ -228,7 +241,7 @@ FirmwareOne = c.Struct(
     "code" / c.Bytes(c.this.code_length),
     c.Terminated,
 
-    "embedded_onev2" / c.RestreamData(c.this.code, c.Optional(FirmwareOneV2)),
+    "embedded_onev2" / c.RestreamData(c.this.code, c.Optional(FirmwareImage)),
 )
 
 # fmt: on
@@ -240,20 +253,19 @@ class FirmwareFormat(Enum):
     TREZOR_ONE_V2 = 3
 
 
-FirmwareType = NewType("FirmwareType", c.Container)
+ParsedFirmware = Tuple[FirmwareFormat, c.Container]
-ParsedFirmware = Tuple[FirmwareFormat, FirmwareType]
 
 
 def parse(data: bytes) -> ParsedFirmware:
     if data[:4] == b"TRZR":
         version = FirmwareFormat.TREZOR_ONE
-        cls = FirmwareOne
+        cls = LegacyFirmware
     elif data[:4] == b"TRZV":
         version = FirmwareFormat.TREZOR_T
-        cls = Firmware
+        cls = VendorFirmware
     elif data[:4] == b"TRZF":
         version = FirmwareFormat.TREZOR_ONE_V2
-        cls = FirmwareOneV2
+        cls = FirmwareImage
     else:
         raise ValueError("Unrecognized firmware image type")
 
@@ -261,10 +273,10 @@ def parse(data: bytes) -> ParsedFirmware:
         fw = cls.parse(data)
     except Exception as e:
         raise FirmwareIntegrityError("Invalid firmware image") from e
-    return version, FirmwareType(fw)
+    return version, fw
 
 
-def digest_onev1(fw: FirmwareType) -> bytes:
+def digest_onev1(fw: c.Container) -> bytes:
     return hashlib.sha256(fw.code).digest()
 
 
@@ -286,7 +298,7 @@ def check_sig_v1(
         key_idx = key_indexes[i] - 1
         signature = signatures[i]
 
-        if key_idx not in V1_BOOTLOADER_KEYS:
+        if key_idx >= len(V1_BOOTLOADER_KEYS):
             # unknown pubkey
             raise InvalidSignatureError("Unknown key in slot {}".format(i))
 
@@ -310,60 +322,75 @@ def header_digest(
     return hash_function(header_bytes).digest()
 
 
-def digest_v2(fw: FirmwareType) -> bytes:
+def digest_v2(fw: c.Container) -> bytes:
     return header_digest(fw.image.header, FirmwareHeader, blake2s)
 
 
-def digest_onev2(fw: FirmwareType) -> bytes:
+def digest_onev2(fw: c.Container) -> bytes:
     return header_digest(fw.header, FirmwareHeader, hashlib.sha256)
 
 
-def validate_code_hashes(
+def calculate_code_hashes(
-    fw: FirmwareType,
+    code: bytes,
+    code_offset: int,
     hash_function: Callable = blake2s,
     chunk_size: int = V2_CHUNK_SIZE,
     padding_byte: bytes = None,
 ) -> None:
-    for i, expected_hash in enumerate(fw.firmware_header.hashes):
+    hashes = []
+    for i in range(16):
         if i == 0:
             # Because first chunk is sent along with headers, there is less code in it.
-            chunk = fw.code[: chunk_size - fw._code_offset]
+            chunk = code[: chunk_size - code_offset]
         else:
             # Subsequent chunks are shifted by the "missing header" size.
-            ptr = i * chunk_size - fw._code_offset
+            ptr = i * chunk_size - code_offset
-            chunk = fw.code[ptr : ptr + chunk_size]
+            chunk = code[ptr : ptr + chunk_size]
 
         # padding for last chunk
         if padding_byte is not None and i > 1 and chunk and len(chunk) < chunk_size:
             chunk += padding_byte[0:1] * (chunk_size - len(chunk))
 
-        if not chunk and expected_hash == b"\0" * 32:
+        if not chunk:
-            continue
+            hashes.append(b"\0" * 32)
-        chunk_hash = hash_function(chunk).digest()
+        else:
-        if chunk_hash != expected_hash:
+            hashes.append(hash_function(chunk).digest())
-            raise FirmwareIntegrityError("Invalid firmware data.")
+
+    return hashes
 
 
-def validate_onev2(fw: FirmwareType, allow_unsigned: bool = False) -> None:
+def validate_code_hashes(fw: c.Container, version: FirmwareFormat) -> None:
+    if version == FirmwareFormat.TREZOR_ONE_V2:
+        image = fw
+        hash_function = hashlib.sha256
+        chunk_size = ONEV2_CHUNK_SIZE
+        padding_byte = b"\xff"
+    else:
+        image = fw.image
+        hash_function = blake2s
+        chunk_size = V2_CHUNK_SIZE
+        padding_byte = None
+
+    expected_hashes = calculate_code_hashes(
+        image.code, image._code_offset, hash_function, chunk_size, padding_byte
+    )
+    if expected_hashes != image.header.hashes:
+        raise FirmwareIntegrityError("Invalid firmware data.")
+
+
+def validate_onev2(fw: c.Container, allow_unsigned: bool = False) -> None:
     try:
         check_sig_v1(
-            digest_onev2(fw),
+            digest_onev2(fw), fw.header.v1_key_indexes, fw.header.v1_signatures,
-            fw.firmware_header.v1_key_indexes,
-            fw.firmware_header.v1_signatures,
         )
     except Unsigned:
         if not allow_unsigned:
             raise
 
-    validate_code_hashes(
+    validate_code_hashes(fw, FirmwareFormat.TREZOR_ONE_V2)
-        fw,
-        hash_function=hashlib.sha256,
-        chunk_size=ONEV2_CHUNK_SIZE,
-        padding_byte=b"\xFF",
-    )
 
 
-def validate_onev1(fw: FirmwareType, allow_unsigned: bool = False) -> None:
+def validate_onev1(fw: c.Container, allow_unsigned: bool = False) -> None:
     try:
         check_sig_v1(digest_onev1(fw), fw.key_indexes, fw.signatures)
     except Unsigned:
@@ -373,7 +400,7 @@ def validate_onev1(fw: FirmwareType, allow_unsigned: bool = False) -> None:
         validate_onev2(fw.embedded_onev2, allow_unsigned)
 
 
-def validate_v2(fw: FirmwareType, skip_vendor_header: bool = False) -> None:
+def validate_v2(fw: c.Container, skip_vendor_header: bool = False) -> None:
     vendor_fingerprint = header_digest(fw.vendor_header, VendorHeader)
     fingerprint = digest_v2(fw)
 
@@ -399,23 +426,23 @@ def validate_v2(fw: FirmwareType, skip_vendor_header: bool = False) -> None:
 
     try:
         cosi.verify_m_of_n(
-            fw.firmware_header.signature,
+            fw.image.header.signature,
             fingerprint,
             fw.vendor_header.vendor_sigs_required,
             fw.vendor_header.vendor_sigs_n,
-            fw.firmware_header.sigmask,
+            fw.image.header.sigmask,
             fw.vendor_header.pubkeys,
         )
     except Exception:
         raise InvalidSignatureError("Invalid firmware signature.")
 
     # XXX expiry is not used now
-    # if time.gmtime(fw.firmware_header.expiry) < now:
+    # if time.gmtime(fw.image.header.expiry) < now:
     #     raise ValueError("Firmware header expired.")
-    validate_code_hashes(fw)
+    validate_code_hashes(fw, FirmwareFormat.TREZOR_T)
 
 
-def digest(version: FirmwareFormat, fw: FirmwareType) -> bytes:
+def digest(version: FirmwareFormat, fw: c.Container) -> bytes:
     if version == FirmwareFormat.TREZOR_ONE:
         return digest_onev1(fw)
     elif version == FirmwareFormat.TREZOR_ONE_V2:
@@ -427,7 +454,7 @@ def digest(version: FirmwareFormat, fw: FirmwareType) -> bytes:
 
 
 def validate(
-    version: FirmwareFormat, fw: FirmwareType, allow_unsigned: bool = False
+    version: FirmwareFormat, fw: c.Container, allow_unsigned: bool = False
 ) -> None:
     if version == FirmwareFormat.TREZOR_ONE:
         return validate_onev1(fw, allow_unsigned)