diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py index e5326a2df..843bfca6d 100644 --- a/python/src/trezorlib/debuglink.py +++ b/python/src/trezorlib/debuglink.py @@ -21,7 +21,7 @@ import textwrap import time from copy import deepcopy from datetime import datetime -from enum import IntEnum +from enum import Enum, IntEnum, auto from itertools import zip_longest from pathlib import Path from typing import ( @@ -65,6 +65,25 @@ EXPECTED_RESPONSES_CONTEXT_LINES = 3 LOG = logging.getLogger(__name__) +class LayoutType(Enum): + T1 = auto() + TT = auto() + TR = auto() + Mercury = auto() + + @classmethod + def from_model(cls, model: models.TrezorModel) -> "LayoutType": + if model in (models.T2T1,): + return cls.TT + if model in (models.T2B1, models.T3B1): + return cls.TR + if model in (models.T3T1,): + return cls.Mercury + if model in (models.T1B1,): + return cls.T1 + raise ValueError(f"Unknown model: {model}") + + class UnstructuredJSONReader: """Contains data-parsing helpers for JSON data that have unknown structure.""" @@ -404,6 +423,11 @@ class DebugLink: """Differences in handling debug events and LayoutContent.""" return self.version < (2, 6, 1) + @property + def layout_type(self) -> LayoutType: + assert self.model is not None + return LayoutType.from_model(self.model) + def set_screen_text_file(self, file_path: Optional[Path]) -> None: if file_path is not None: file_path.write_bytes(b"") @@ -995,6 +1019,10 @@ class TrezorClientDebugLink(TrezorClient): self.debug.model = self.model self.debug.version = self.version + @property + def layout_type(self) -> LayoutType: + return self.debug.layout_type + def reset_debug_features(self) -> None: """Prepare the debugging client for a new testcase.