refactor(python): move out TOIF related structs to trezorlib.toif

python/src/trezorlib/cli/settings.py

@@ -18,7 +18,7 @@ from typing import TYPE_CHECKING, Optional, cast
 
 import click
 
-from .. import device, firmware, messages, toif
+from .. import device, messages, toif
 from . import AliasedGroup, ChoiceType, with_client
 
 if TYPE_CHECKING:
@@ -83,7 +83,7 @@ def image_to_tt(filename: str) -> bytes:
     if toif_image.size != (144, 144):
         raise click.ClickException("Wrong size of image - should be 144x144")
 
-    if toif_image.mode != firmware.ToifMode.full_color:
+    if toif_image.mode != toif.ToifMode.full_color:
         raise click.ClickException("Wrong image mode - should be full_color")
 
     return toif_image.to_bytes()

python/src/trezorlib/firmware.py

@@ -23,7 +23,8 @@ import construct as c
 import ecdsa
 
 from . import cosi, messages
-from .tools import expect, session
+from .toif import ToifStruct
+from .tools import expect, session, EnumAdapter
 
 if TYPE_CHECKING:
     from .client import TrezorClient
@@ -98,43 +99,12 @@ class Unsigned(FirmwareIntegrityError):
     pass
 
 
-class ToifMode(Enum):
-    full_color = b"f"  # big endian
-    grayscale = b"g"  # odd hi
-    full_color_le = b"F"  # little endian
-    grayscale_eh = b"G"  # even hi
-
-
 class HeaderType(Enum):
     FIRMWARE = b"TRZF"
     BOOTLOADER = b"TRZB"
 
 
-class EnumAdapter(c.Adapter):
-    def __init__(self, subcon: Any, enum: Any) -> None:
-        self.enum = enum
-        super().__init__(subcon)
-
-    def _encode(self, obj: Any, ctx: Any, path: Any):
-        return obj.value
-
-    def _decode(self, obj: Any, ctx: Any, path: Any):
-        try:
-            return self.enum(obj)
-        except ValueError:
-            return obj
-
-
 # fmt: off
-Toif = c.Struct(
-    "magic" / c.Const(b"TOI"),
-    "format" / EnumAdapter(c.Bytes(1), ToifMode),
-    "width" / c.Int16ul,
-    "height" / c.Int16ul,
-    "data" / c.Prefixed(c.Int32ul, c.GreedyBytes),
-)
-
-
 VendorTrust = c.Transformed(c.BitStruct(
     "_reserved" / c.Default(c.BitsInteger(9), 0),
     "show_vendor_string" / c.Flag,
@@ -159,7 +129,7 @@ VendorHeader = c.Struct(
     "_reserved" / c.Padding(14),
     "pubkeys" / c.Bytes(32)[c.this.sig_n],
     "text" / c.Aligned(4, c.PascalString(c.Int8ul, "utf-8")),
-    "image" / Toif,
+    "image" / ToifStruct,
     "_end_offset" / c.Tell,
 
     "_min_header_len" / c.Check(c.this.header_len > (c.this._end_offset - c.this._start_offset) + 65),

python/src/trezorlib/toif.py

@@ -17,11 +17,13 @@
 import struct
 import zlib
 from dataclasses import dataclass
+from enum import Enum
 from typing import Sequence, Tuple
 
+import construct as c
 from typing_extensions import Literal
 
-from . import firmware
+from .tools import EnumAdapter
 
 try:
     # Explanation of having to use "Image.Image" in typing:
@@ -36,6 +38,22 @@ except ImportError:
 RGBPixel = Tuple[int, int, int]
 
 
+class ToifMode(Enum):
+    full_color = b"f"  # big endian
+    grayscale = b"g"  # odd hi
+    full_color_le = b"F"  # little endian
+    grayscale_eh = b"G"  # even hi
+
+
+ToifStruct = c.Struct(
+    "magic" / c.Const(b"TOI"),
+    "format" / EnumAdapter(c.Bytes(1), ToifMode),
+    "width" / c.Int16ul,
+    "height" / c.Int16ul,
+    "data" / c.Prefixed(c.Int32ul, c.GreedyBytes),
+)
+
+
 def _compress(data: bytes) -> bytes:
     z = zlib.compressobj(level=9, wbits=-10)
     return z.compress(data) + z.flush()
@@ -113,17 +131,14 @@ def _to_grayscale(data: bytes, right_hi: bool) -> bytes:
 
 @dataclass
 class Toif:
-    mode: firmware.ToifMode
+    mode: ToifMode
     size: Tuple[int, int]
     data: bytes
 
     def __post_init__(self) -> None:
         # checking the data size
         width, height = self.size
-        if (
+        if self.mode is ToifMode.grayscale or self.mode is ToifMode.grayscale_eh:
-            self.mode is firmware.ToifMode.grayscale
-            or self.mode is firmware.ToifMode.grayscale_eh
-        ):
             expected_size = width * height // 2
         else:
             expected_size = width * height * 2
@@ -142,16 +157,16 @@ class Toif:
         uncompressed = _decompress(self.data)
 
         pil_mode: Literal["L", "RGB"]
-        if self.mode is firmware.ToifMode.grayscale:
+        if self.mode is ToifMode.grayscale:
             pil_mode = "L"
             raw_data = _to_grayscale(uncompressed, right_hi=False)
-        elif self.mode is firmware.ToifMode.grayscale_eh:
+        elif self.mode is ToifMode.grayscale_eh:
             pil_mode = "L"
             raw_data = _to_grayscale(uncompressed, right_hi=True)
-        elif self.mode is firmware.ToifMode.full_color:
+        elif self.mode is ToifMode.full_color:
             pil_mode = "RGB"
             raw_data = _to_rgb(uncompressed, little_endian=False)
-        else:  # self.mode is firmware.ToifMode.full_color_le:
+        else:  # self.mode is ToifMode.full_color_le:
             pil_mode = "RGB"
             raw_data = _to_rgb(uncompressed, little_endian=True)
 
@@ -159,7 +174,7 @@ class Toif:
 
     def to_bytes(self) -> bytes:
         width, height = self.size
-        return firmware.Toif.build(
+        return ToifStruct.build(
             dict(format=self.mode, width=width, height=height, data=self.data)
         )
 
@@ -169,7 +184,10 @@ class Toif:
 
 
 def from_bytes(data: bytes) -> Toif:
-    parsed = firmware.Toif.parse(data)
+    return from_struct(ToifStruct.parse(data))
+
+
+def from_struct(parsed: c.Container) -> Toif:
     return Toif(parsed.format, (parsed.width, parsed.height), parsed.data)
 
 
@@ -200,27 +218,27 @@ def from_image(
         if image.size[0] % 2 != 0:
             raise ValueError("Only even-width grayscale images are supported")
         if not legacy_format:
-            toif_mode = firmware.ToifMode.grayscale_eh
+            toif_mode = ToifMode.grayscale_eh
             toif_data = _from_pil_grayscale(image.getdata(), right_hi=True)
         else:
-            toif_mode = firmware.ToifMode.grayscale
+            toif_mode = ToifMode.grayscale
             toif_data = _from_pil_grayscale(image.getdata(), right_hi=False)
     elif image.mode == "LA":
-        toif_mode = firmware.ToifMode.grayscale
+        toif_mode = ToifMode.grayscale
         if image.size[0] % 2 != 0:
             raise ValueError("Only even-width grayscale images are supported")
         if not legacy_format:
-            toif_mode = firmware.ToifMode.grayscale_eh
+            toif_mode = ToifMode.grayscale_eh
             toif_data = _from_pil_grayscale_alpha(image.getdata(), right_hi=True)
         else:
-            toif_mode = firmware.ToifMode.grayscale
+            toif_mode = ToifMode.grayscale
             toif_data = _from_pil_grayscale_alpha(image.getdata(), right_hi=False)
     elif image.mode == "RGB":
         if not legacy_format:
-            toif_mode = firmware.ToifMode.full_color_le
+            toif_mode = ToifMode.full_color_le
             toif_data = _from_pil_rgb(image.getdata(), little_endian=True)
         else:
-            toif_mode = firmware.ToifMode.full_color
+            toif_mode = ToifMode.full_color
             toif_data = _from_pil_rgb(image.getdata(), little_endian=False)
     else:
         raise ValueError(f"Unsupported image mode: {image.mode}")

python/src/trezorlib/tools.py

@@ -33,6 +33,8 @@ from typing import (
     overload,
 )
 
+import construct
+
 if TYPE_CHECKING:
     from .client import TrezorClient
     from .protobuf import MessageType
@@ -372,3 +374,18 @@ def descriptor_checksum(desc: str) -> str:
     for j in range(0, 8):
         ret[j] = CHECKSUM_CHARSET[(c >> (5 * (7 - j))) & 31]
     return "".join(ret)
+
+
+class EnumAdapter(construct.Adapter):
+    def __init__(self, subcon: Any, enum: Any) -> None:
+        self.enum = enum
+        super().__init__(subcon)
+
+    def _encode(self, obj: Any, ctx: Any, path: Any):
+        return obj.value
+
+    def _decode(self, obj: Any, ctx: Any, path: Any):
+        try:
+            return self.enum(obj)
+        except ValueError:
+            return obj