diff --git a/python/src/trezorlib/cli/settings.py b/python/src/trezorlib/cli/settings.py index 906c92257..fde3ca0be 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, 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() diff --git a/python/src/trezorlib/firmware.py b/python/src/trezorlib/firmware.py index eadad475c..b2bbf0a09 100644 --- a/python/src/trezorlib/firmware.py +++ b/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), diff --git a/python/src/trezorlib/toif.py b/python/src/trezorlib/toif.py index db976c6a7..9382c5d3a 100644 --- a/python/src/trezorlib/toif.py +++ b/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 ( - self.mode is firmware.ToifMode.grayscale - or self.mode is firmware.ToifMode.grayscale_eh - ): + if self.mode is ToifMode.grayscale or self.mode is 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}") diff --git a/python/src/trezorlib/tools.py b/python/src/trezorlib/tools.py index f4d3df33c..5898e47fc 100644 --- a/python/src/trezorlib/tools.py +++ b/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