From 0d7080102f28778c390b352b572cd96faea4f163 Mon Sep 17 00:00:00 2001 From: M1nd3r Date: Tue, 26 Nov 2024 11:47:13 +0100 Subject: [PATCH] test(python): fix persistence_tests [no changelog] --- python/src/trezorlib/debuglink.py | 18 +++++++-- tests/persistence_tests/test_safety_checks.py | 8 ++-- .../test_shamir_persistence.py | 18 +++++---- tests/persistence_tests/test_wipe_code.py | 40 ++++++++++++------- 4 files changed, 56 insertions(+), 28 deletions(-) diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py index df0587c351..2bffb86efe 100644 --- a/python/src/trezorlib/debuglink.py +++ b/python/src/trezorlib/debuglink.py @@ -1023,7 +1023,6 @@ message_filters = MessageFilterGenerator() class SessionDebugWrapper(Session): - session_version: int CODEC_V1: t.Final[int] = 1 THP_V2: t.Final[int] = 2 @@ -1031,14 +1030,18 @@ class SessionDebugWrapper(Session): self._session = session self.reset_debug_features() if isinstance(session, SessionV1): - self.session_version = 1 + self.client.session_version = 1 elif isinstance(session, SessionV2): - self.session_version = 2 + self.client.session_version = 2 elif isinstance(session, SessionDebugWrapper): raise Exception("Cannot wrap already wrapped session!") else: self.session_version = -1 # UNKNOWN + @property + def session_version(self) -> int: + return self.client.session_version + @property def client(self) -> TrezorClientDebugLink: assert isinstance(self._session.client, TrezorClientDebugLink) @@ -1279,6 +1282,7 @@ class TrezorClientDebugLink(TrezorClient): # by the device. def __init__(self, transport: "Transport", auto_interact: bool = True) -> None: + self._session_version: int = -1 try: debug_transport = transport.find_debug() self.debug = DebugLink(debug_transport, auto_interact) @@ -1305,6 +1309,14 @@ class TrezorClientDebugLink(TrezorClient): self.debug.version = self.version self.passphrase: str | None = None + @property + def session_version(self) -> int: + return self._session_version + + @session_version.setter + def session_version(self, value: int) -> None: + self._session_version = value + @property def layout_type(self) -> LayoutType: return self.debug.layout_type diff --git a/tests/persistence_tests/test_safety_checks.py b/tests/persistence_tests/test_safety_checks.py index 1cbf7d7551..c2552f04b3 100644 --- a/tests/persistence_tests/test_safety_checks.py +++ b/tests/persistence_tests/test_safety_checks.py @@ -20,16 +20,18 @@ from ..upgrade_tests import core_only def test_safety_checks_level_after_reboot( core_emulator: Emulator, set_level: SafetyCheckLevel, after_level: SafetyCheckLevel ): - device.wipe(core_emulator.client) + device.wipe(core_emulator.client.get_management_session()) + core_emulator.client = core_emulator.client.get_new_client() debuglink.load_device( - core_emulator.client, + core_emulator.client.get_management_session(), mnemonic=MNEMONIC12, pin="", passphrase_protection=False, label="SAFETYLEVEL", ) - device.apply_settings(core_emulator.client, safety_checks=set_level) + device.apply_settings(core_emulator.client.get_session(), safety_checks=set_level) + core_emulator.client.refresh_features() assert core_emulator.client.features.safety_checks == set_level core_emulator.restart() diff --git a/tests/persistence_tests/test_shamir_persistence.py b/tests/persistence_tests/test_shamir_persistence.py index e24f16eeb6..c8707aeb82 100644 --- a/tests/persistence_tests/test_shamir_persistence.py +++ b/tests/persistence_tests/test_shamir_persistence.py @@ -17,8 +17,11 @@ import pytest from trezorlib import device -from trezorlib.debuglink import DebugLink, LayoutType -from trezorlib.messages import RecoveryStatus +from trezorlib import messages +from trezorlib.debuglink import DebugLink, LayoutType, SessionDebugWrapper as Session +from trezorlib.messages import RecoveryStatus, WipeDevice + +from tests import buttons from ..click_tests import common, recovery from ..common import MNEMONIC_SLIP39_ADVANCED_20, MNEMONIC_SLIP39_BASIC_20_3of6 @@ -45,7 +48,7 @@ def test_abort(core_emulator: Emulator): assert features.recovery_status == RecoveryStatus.Nothing - device_handler.run(device.recover, pin_protection=False) + device_handler.run_with_session(device.recover, pin_protection=False) recovery.confirm_recovery(debug) layout = debug.read_layout() @@ -82,7 +85,7 @@ def test_recovery_single_reset(core_emulator: Emulator): assert features.initialized is False assert features.recovery_status == RecoveryStatus.Nothing - device_handler.run(device.recover, pin_protection=False) + device_handler.run_with_session(device.recover, pin_protection=False) recovery.confirm_recovery(debug) @@ -129,7 +132,7 @@ def test_recovery_on_old_wallet(core_emulator: Emulator): assert features.recovery_status == RecoveryStatus.Nothing # enter recovery mode - device_handler.run(device.recover, pin_protection=False) + device_handler.run_with_session(device.recover, pin_protection=False) recovery.confirm_recovery(debug) @@ -157,7 +160,8 @@ def test_recovery_on_old_wallet(core_emulator: Emulator): layout = debug.read_layout() # while keyboard is open, hit the device with Initialize/GetFeatures - device_handler.client.init_device() + if device_handler.client.session_version == Session.CODEC_V1: + device_handler.client.get_management_session().call(messages.Initialize()) device_handler.client.refresh_features() # try entering remaining 19 words @@ -207,7 +211,7 @@ def test_recovery_multiple_resets(core_emulator: Emulator): assert features.recovery_status == RecoveryStatus.Nothing # start device and recovery - device_handler.run(device.recover, pin_protection=False) + device_handler.run_with_session(device.recover, pin_protection=False) recovery.confirm_recovery(debug) diff --git a/tests/persistence_tests/test_wipe_code.py b/tests/persistence_tests/test_wipe_code.py index 2497a708f6..cb06eeb2cd 100644 --- a/tests/persistence_tests/test_wipe_code.py +++ b/tests/persistence_tests/test_wipe_code.py @@ -11,46 +11,55 @@ WIPE_CODE = "9876" def setup_device_legacy(client: Client, pin: str, wipe_code: str) -> None: - device.wipe(client) + device.wipe(client.get_management_session()) + client = client.get_new_client() debuglink.load_device( - client, MNEMONIC12, pin, passphrase_protection=False, label="WIPECODE" + client.get_management_session(), + MNEMONIC12, + pin, + passphrase_protection=False, + label="WIPECODE", ) with client: client.use_pin_sequence([PIN, WIPE_CODE, WIPE_CODE]) - device.change_wipe_code(client) + device.change_wipe_code(client.get_management_session()) def setup_device_core(client: Client, pin: str, wipe_code: str) -> None: - device.wipe(client) + device.wipe(client.get_management_session()) + client = client.get_new_client() debuglink.load_device( - client, MNEMONIC12, pin, passphrase_protection=False, label="WIPECODE" + client.get_management_session(), + MNEMONIC12, + pin, + passphrase_protection=False, + label="WIPECODE", ) with client: client.use_pin_sequence([pin, wipe_code, wipe_code]) - device.change_wipe_code(client) + device.change_wipe_code(client.get_management_session()) @core_only def test_wipe_code_activate_core(core_emulator: Emulator): # set up device setup_device_core(core_emulator.client, PIN, WIPE_CODE) - - core_emulator.client.init_device() + session = core_emulator.client.get_session() device_id = core_emulator.client.features.device_id # Initiate Change pin process - ret = core_emulator.client.call_raw(messages.ChangePin(remove=False)) + ret = session.call_raw(messages.ChangePin(remove=False)) assert isinstance(ret, messages.ButtonRequest) assert ret.name == "change_pin" core_emulator.client.debug.press_yes() - ret = core_emulator.client.call_raw(messages.ButtonAck()) + ret = session.call_raw(messages.ButtonAck()) # Enter the wipe code instead of the current PIN expected = message_filters.ButtonRequest(code=messages.ButtonRequestType.PinEntry) assert expected.match(ret) - core_emulator.client._raw_write(messages.ButtonAck()) + session._write(messages.ButtonAck()) core_emulator.client.debug.input(WIPE_CODE) # preserving screenshots even after it dies and starts again @@ -75,25 +84,26 @@ def test_wipe_code_activate_legacy(): # set up device setup_device_legacy(emu.client, PIN, WIPE_CODE) - emu.client.init_device() + session = emu.client.get_session() device_id = emu.client.features.device_id # Initiate Change pin process - ret = emu.client.call_raw(messages.ChangePin(remove=False)) + ret = session.call_raw(messages.ChangePin(remove=False)) assert isinstance(ret, messages.ButtonRequest) emu.client.debug.press_yes() - ret = emu.client.call_raw(messages.ButtonAck()) + ret = session.call_raw(messages.ButtonAck()) # Enter the wipe code instead of the current PIN assert isinstance(ret, messages.PinMatrixRequest) wipe_code_encoded = emu.client.debug.encode_pin(WIPE_CODE) - emu.client._raw_write(messages.PinMatrixAck(pin=wipe_code_encoded)) + session._write(messages.PinMatrixAck(pin=wipe_code_encoded)) # wait 30 seconds for emulator to shut down # this will raise a TimeoutError if the emulator doesn't die. emu.wait(30) emu.start() + emu.client.refresh_features() assert emu.client.features.initialized is False assert emu.client.features.pin_protection is False assert emu.client.features.wipe_code_protection is False