From 7fc0874d68b8c48458ea92cde2682eebdd6bedda Mon Sep 17 00:00:00 2001 From: M1nd3r Date: Tue, 26 Nov 2024 10:11:06 +0100 Subject: [PATCH] test(python): fix upgrade_tests (partially) [no changelog] --- tests/upgrade_tests/test_firmware_upgrades.py | 68 +++++++++++-------- .../test_passphrase_consistency.py | 23 ++++--- 2 files changed, 51 insertions(+), 40 deletions(-) diff --git a/tests/upgrade_tests/test_firmware_upgrades.py b/tests/upgrade_tests/test_firmware_upgrades.py index e742e25ce5..4cf7e085fe 100644 --- a/tests/upgrade_tests/test_firmware_upgrades.py +++ b/tests/upgrade_tests/test_firmware_upgrades.py @@ -21,6 +21,7 @@ import pytest from shamir_mnemonic import shamir from trezorlib import btc, debuglink, device, exceptions, fido, models +from trezorlib import messages from trezorlib.debuglink import SessionDebugWrapper as Session from trezorlib.messages import ( ApplySettings, @@ -94,12 +95,14 @@ def test_upgrade_load_pin(gen: str, tag: str) -> None: assert not client.features.passphrase_protection assert client.features.initialized assert client.features.label == LABEL - client.use_pin_sequence([PIN]) - assert btc.get_address(client, "Bitcoin", PATH) == ADDRESS + session = Session(client.get_session()) + with client, session: + client.use_pin_sequence([PIN]) + assert btc.get_address(session, "Bitcoin", PATH) == ADDRESS with EmulatorWrapper(gen, tag) as emu: debuglink.load_device_by_mnemonic( - emu.client, + Session(emu.client.get_management_session()), mnemonic=MNEMONIC, pin=PIN, passphrase_protection=False, @@ -187,7 +190,6 @@ def test_upgrade_wipe_code(gen: str, tag: str): emu.client.use_pin_sequence([PIN, WIPE_CODE, WIPE_CODE]) session = Session(emu.client.get_management_session()) session.refresh_features() - raise Exception("client", str(emu.client.pin_callback)) device.change_wipe_code(session) device_id = emu.client.features.device_id @@ -220,7 +222,7 @@ def test_upgrade_reset(gen: str, tag: str): with EmulatorWrapper(gen, tag) as emu: device.reset( - emu.client, + emu.client.get_session(), strength=STRENGTH, passphrase_protection=False, pin_protection=False, @@ -228,13 +230,13 @@ def test_upgrade_reset(gen: str, tag: str): ) device_id = emu.client.features.device_id asserts(emu.client) - address = btc.get_address(emu.client, "Bitcoin", PATH) + address = btc.get_address(emu.client.get_session(), "Bitcoin", PATH) storage = emu.get_storage() with EmulatorWrapper(gen, storage=storage) as emu: assert device_id == emu.client.features.device_id asserts(emu.client) - assert btc.get_address(emu.client, "Bitcoin", PATH) == address + assert btc.get_address(emu.client.get_session(), "Bitcoin", PATH) == address @for_all() @@ -250,7 +252,7 @@ def test_upgrade_reset_skip_backup(gen: str, tag: str): with EmulatorWrapper(gen, tag) as emu: device.reset( - emu.client, + emu.client.get_session(), strength=STRENGTH, passphrase_protection=False, pin_protection=False, @@ -259,13 +261,13 @@ def test_upgrade_reset_skip_backup(gen: str, tag: str): ) device_id = emu.client.features.device_id asserts(emu.client) - address = btc.get_address(emu.client, "Bitcoin", PATH) + address = btc.get_address(emu.client.get_session(), "Bitcoin", PATH) storage = emu.get_storage() with EmulatorWrapper(gen, storage=storage) as emu: assert device_id == emu.client.features.device_id asserts(emu.client) - assert btc.get_address(emu.client, "Bitcoin", PATH) == address + assert btc.get_address(emu.client.get_session(), "Bitcoin", PATH) == address @for_all(legacy_minimum_version=(1, 7, 2)) @@ -281,7 +283,7 @@ def test_upgrade_reset_no_backup(gen: str, tag: str): with EmulatorWrapper(gen, tag) as emu: device.reset( - emu.client, + emu.client.get_session(), strength=STRENGTH, passphrase_protection=False, pin_protection=False, @@ -290,13 +292,13 @@ def test_upgrade_reset_no_backup(gen: str, tag: str): ) device_id = emu.client.features.device_id asserts(emu.client) - address = btc.get_address(emu.client, "Bitcoin", PATH) + address = btc.get_address(emu.client.get_session(), "Bitcoin", PATH) storage = emu.get_storage() with EmulatorWrapper(gen, storage=storage) as emu: assert device_id == emu.client.features.device_id asserts(emu.client) - assert btc.get_address(emu.client, "Bitcoin", PATH) == address + assert btc.get_address(emu.client.get_session(), "Bitcoin", PATH) == address # Although Shamir was introduced in 2.1.2 already, the debug instrumentation was not present until 2.1.9. @@ -309,7 +311,7 @@ def test_upgrade_shamir_recovery(gen: str, tag: Optional[str]): emu.client.watch_layout(True) debug = device_handler.debuglink() - device_handler.run(device.recover, pin_protection=False) + device_handler.run_with_session(device.recover, pin_protection=False) recovery_old.confirm_recovery(debug) recovery_old.select_number_of_words(debug) @@ -354,9 +356,10 @@ def test_upgrade_shamir_recovery(gen: str, tag: Optional[str]): @for_all("core", core_minimum_version=(2, 1, 9)) def test_upgrade_shamir_backup(gen: str, tag: Optional[str]): with EmulatorWrapper(gen, tag) as emu: + session = Session(emu.client.get_management_session()) # Generate a new encrypted master secret and record it. device.reset( - emu.client, + session, pin_protection=False, skip_backup=True, backup_type=BackupType.Slip39_Basic, @@ -366,14 +369,16 @@ def test_upgrade_shamir_backup(gen: str, tag: Optional[str]): mnemonic_secret = emu.client.debug.state().mnemonic_secret # Set passphrase_source = HOST. - resp = emu.client.call(ApplySettings(_passphrase_source=2, use_passphrase=True)) + session = Session(emu.client.get_session()) + resp = session.call(ApplySettings(_passphrase_source=2, use_passphrase=True)) assert isinstance(resp, Success) # Get a passphrase-less and a passphrased address. - address = btc.get_address(emu.client, "Bitcoin", PATH) - emu.client.init_device(new_session=True) - emu.client.use_passphrase("TREZOR") - address_passphrase = btc.get_address(emu.client, "Bitcoin", PATH) + address = btc.get_address(session, "Bitcoin", PATH) + if session.session_version == Session.CODEC_V1: + session.call(messages.Initialize(new_session=True)) + new_session = emu.client.get_session(passphrase="TREZOR") + address_passphrase = btc.get_address(new_session, "Bitcoin", PATH) assert emu.client.features.backup_availability == BackupAvailability.Required storage = emu.get_storage() @@ -386,7 +391,7 @@ def test_upgrade_shamir_backup(gen: str, tag: Optional[str]): with emu.client: IF = InputFlowSlip39BasicBackup(emu.client, False) emu.client.set_input_flow(IF.get()) - device.backup(emu.client) + device.backup(emu.client.get_session()) assert ( emu.client.features.backup_availability == BackupAvailability.NotAvailable ) @@ -407,10 +412,13 @@ def test_upgrade_shamir_backup(gen: str, tag: Optional[str]): assert ems.ciphertext == mnemonic_secret # Check that addresses are the same after firmware upgrade and backup. - assert btc.get_address(emu.client, "Bitcoin", PATH) == address - emu.client.init_device(new_session=True) - emu.client.use_passphrase("TREZOR") - assert btc.get_address(emu.client, "Bitcoin", PATH) == address_passphrase + assert btc.get_address(emu.client.get_session(), "Bitcoin", PATH) == address + assert ( + btc.get_address( + emu.client.get_session(passphrase="TREZOR"), "Bitcoin", PATH + ) + == address_passphrase + ) @for_all(legacy_minimum_version=(1, 8, 4), core_minimum_version=(2, 1, 9)) @@ -418,22 +426,22 @@ def test_upgrade_u2f(gen: str, tag: str): """Check U2F counter stayed the same after an upgrade.""" with EmulatorWrapper(gen, tag) as emu: debuglink.load_device_by_mnemonic( - emu.client, + emu.client.get_management_session(), mnemonic=MNEMONIC, pin="", passphrase_protection=False, label=LABEL, ) - - success = fido.set_counter(emu.client, 10) + session = emu.client.get_management_session() + success = fido.set_counter(session, 10) assert "U2F counter set" in success - counter = fido.get_next_counter(emu.client) + counter = fido.get_next_counter(session) assert counter == 11 storage = emu.get_storage() with EmulatorWrapper(gen, storage=storage) as emu: - counter = fido.get_next_counter(emu.client) + counter = fido.get_next_counter(session) assert counter == 12 diff --git a/tests/upgrade_tests/test_passphrase_consistency.py b/tests/upgrade_tests/test_passphrase_consistency.py index f70cde9673..be8c7dc74b 100644 --- a/tests/upgrade_tests/test_passphrase_consistency.py +++ b/tests/upgrade_tests/test_passphrase_consistency.py @@ -89,9 +89,8 @@ def test_passphrase_works(emulator: Emulator): messages.ButtonRequest, messages.Address, ] - emu_session = emulator.client.get_session(passphrase="") - with emulator.client, Session(emu_session) as session: - emulator.client.use_passphrase("TREZOR") + emu_session = emulator.client.get_session(passphrase="TREZOR") + with Session(emu_session) as session: session.set_expected_responses(expected_responses) btc.get_address(session, "Testnet", parse_path("44h/1h/0h/0/0")) @@ -133,14 +132,18 @@ def test_init_device(emulator: Emulator): messages.Address, ] - emu_session = emulator.client.get_session(passphrase="") - with emulator.client, Session(emu_session) as session: - emulator.client.use_passphrase("TREZOR") + emu_session = emulator.client.get_session(passphrase="TREZOR") + with Session(emu_session) as session: session.set_expected_responses(expected_responses) btc.get_address(session, "Testnet", parse_path("44h/1h/0h/0/0")) # in TT < 2.3.0 session_id will only be available after PassphraseStateRequest - session_id = emulator.client.features.session_id - # emulator.client.init_device() - btc.get_address(emulator.client, "Testnet", parse_path("44h/1h/0h/0/0")) - assert session_id == emulator.client.features.session_id + session_id = session.id + if session.session_version == Session.CODEC_V1: + session.call(messages.Initialize(session_id=session_id)) + btc.get_address( + session, + "Testnet", + parse_path("44h/1h/0h/0/0"), + ) + assert session_id == session.id