diff --git a/python/src/trezorlib/cli/settings.py b/python/src/trezorlib/cli/settings.py index fde3ca0bea..906c922573 100644 --- a/python/src/trezorlib/cli/settings.py +++ b/python/src/trezorlib/cli/settings.py @@ -18,7 +18,7 @@ from typing import TYPE_CHECKING, Optional, cast import click -from .. import device, messages, toif +from .. import device, firmware, 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 != toif.ToifMode.full_color: + if toif_image.mode != firmware.ToifMode.full_color: raise click.ClickException("Wrong image mode - should be full_color") return toif_image.to_bytes() diff --git a/python/src/trezorlib/firmware.py b/python/src/trezorlib/firmware.py index b2bbf0a094..eadad475c4 100644 --- a/python/src/trezorlib/firmware.py +++ b/python/src/trezorlib/firmware.py @@ -23,8 +23,7 @@ import construct as c import ecdsa from . import cosi, messages -from .toif import ToifStruct -from .tools import expect, session, EnumAdapter +from .tools import expect, session if TYPE_CHECKING: from .client import TrezorClient @@ -99,12 +98,43 @@ 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, @@ -129,7 +159,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" / ToifStruct, + "image" / Toif, "_end_offset" / c.Tell, "_min_header_len" / c.Check(c.this.header_len > (c.this._end_offset - c.this._start_offset) + 65), diff --git a/python/src/trezorlib/toif.py b/python/src/trezorlib/toif.py index 9382c5d3ab..db976c6a70 100644 --- a/python/src/trezorlib/toif.py +++ b/python/src/trezorlib/toif.py @@ -17,13 +17,11 @@ 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 .tools import EnumAdapter +from . import firmware try: # Explanation of having to use "Image.Image" in typing: @@ -38,22 +36,6 @@ 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() @@ -131,14 +113,17 @@ def _to_grayscale(data: bytes, right_hi: bool) -> bytes: @dataclass class Toif: - mode: ToifMode + mode: firmware.ToifMode size: Tuple[int, int] data: bytes def __post_init__(self) -> None: # checking the data size width, height = self.size - if self.mode is ToifMode.grayscale or self.mode is ToifMode.grayscale_eh: + if ( + 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 @@ -157,16 +142,16 @@ class Toif: uncompressed = _decompress(self.data) pil_mode: Literal["L", "RGB"] - if self.mode is ToifMode.grayscale: + if self.mode is firmware.ToifMode.grayscale: pil_mode = "L" raw_data = _to_grayscale(uncompressed, right_hi=False) - elif self.mode is ToifMode.grayscale_eh: + elif self.mode is firmware.ToifMode.grayscale_eh: pil_mode = "L" raw_data = _to_grayscale(uncompressed, right_hi=True) - elif self.mode is ToifMode.full_color: + elif self.mode is firmware.ToifMode.full_color: pil_mode = "RGB" raw_data = _to_rgb(uncompressed, little_endian=False) - else: # self.mode is ToifMode.full_color_le: + else: # self.mode is firmware.ToifMode.full_color_le: pil_mode = "RGB" raw_data = _to_rgb(uncompressed, little_endian=True) @@ -174,7 +159,7 @@ class Toif: def to_bytes(self) -> bytes: width, height = self.size - return ToifStruct.build( + return firmware.Toif.build( dict(format=self.mode, width=width, height=height, data=self.data) ) @@ -184,10 +169,7 @@ class Toif: def from_bytes(data: bytes) -> Toif: - return from_struct(ToifStruct.parse(data)) - - -def from_struct(parsed: c.Container) -> Toif: + parsed = firmware.Toif.parse(data) return Toif(parsed.format, (parsed.width, parsed.height), parsed.data) @@ -218,27 +200,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 = ToifMode.grayscale_eh + toif_mode = firmware.ToifMode.grayscale_eh toif_data = _from_pil_grayscale(image.getdata(), right_hi=True) else: - toif_mode = ToifMode.grayscale + toif_mode = firmware.ToifMode.grayscale toif_data = _from_pil_grayscale(image.getdata(), right_hi=False) elif image.mode == "LA": - toif_mode = ToifMode.grayscale + toif_mode = firmware.ToifMode.grayscale if image.size[0] % 2 != 0: raise ValueError("Only even-width grayscale images are supported") if not legacy_format: - toif_mode = ToifMode.grayscale_eh + toif_mode = firmware.ToifMode.grayscale_eh toif_data = _from_pil_grayscale_alpha(image.getdata(), right_hi=True) else: - toif_mode = ToifMode.grayscale + toif_mode = firmware.ToifMode.grayscale toif_data = _from_pil_grayscale_alpha(image.getdata(), right_hi=False) elif image.mode == "RGB": if not legacy_format: - toif_mode = ToifMode.full_color_le + toif_mode = firmware.ToifMode.full_color_le toif_data = _from_pil_rgb(image.getdata(), little_endian=True) else: - toif_mode = ToifMode.full_color + toif_mode = firmware.ToifMode.full_color toif_data = _from_pil_rgb(image.getdata(), little_endian=False) else: raise ValueError(f"Unsupported image mode: {image.mode}") diff --git a/python/src/trezorlib/tools.py b/python/src/trezorlib/tools.py index 5898e47fcd..f4d3df33c2 100644 --- a/python/src/trezorlib/tools.py +++ b/python/src/trezorlib/tools.py @@ -33,8 +33,6 @@ from typing import ( overload, ) -import construct - if TYPE_CHECKING: from .client import TrezorClient from .protobuf import MessageType @@ -374,18 +372,3 @@ 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