From 6cd976fdee4e9b5774e447160797b09e6ae7e5f5 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Wed, 18 Dec 2019 13:14:44 +0100
Subject: [PATCH] python/firmware: support bootloader headers

---
 python/src/trezorlib/firmware.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/python/src/trezorlib/firmware.py b/python/src/trezorlib/firmware.py
index 7e00e59a28..b25e778947 100644
--- a/python/src/trezorlib/firmware.py
+++ b/python/src/trezorlib/firmware.py
@@ -86,6 +86,11 @@ class ToifMode(Enum):
     grayscale = b"g"
 
 
+class HeaderType(Enum):
+    FIRMWARE = b"TRZF"
+    BOOTLOADER = b"TRZB"
+
+
 class EnumAdapter(c.Adapter):
     def __init__(self, subcon, enum):
         self.enum = enum
@@ -160,7 +165,7 @@ VersionLong = c.Struct(
 
 FirmwareHeader = c.Struct(
     "_start_offset" / c.Tell,
-    "magic" / c.Const(b"TRZF"),
+    "magic" / EnumAdapter(c.Bytes(4), HeaderType),
     "header_len" / c.Int32ul,
     "expiry" / c.Int32ul,
     "code_length" / c.Rebuild(