From c99fd824b3051cca290d26db2b78a62d092df572 Mon Sep 17 00:00:00 2001 From: grdddj Date: Fri, 21 Jul 2023 11:38:28 +0200 Subject: [PATCH] chore(tests): refactor recovery input flows [no changelog] --- tests/common.py | 134 +-- .../test_recovery_bip39_dryrun.py | 12 +- .../reset_recovery/test_recovery_bip39_t2.py | 6 +- .../test_recovery_slip39_advanced.py | 7 +- .../test_recovery_slip39_advanced_dryrun.py | 2 +- .../test_recovery_slip39_basic.py | 26 +- .../test_recovery_slip39_basic_dryrun.py | 8 +- .../test_reset_recovery_bip39.py | 4 +- tests/input_flows.py | 884 +++++------------- tests/input_flows_helpers.py | 249 +++++ 10 files changed, 543 insertions(+), 789 deletions(-) create mode 100644 tests/input_flows_helpers.py diff --git a/tests/common.py b/tests/common.py index c6c83e13d..4d6cd9d09 100644 --- a/tests/common.py +++ b/tests/common.py @@ -23,7 +23,7 @@ from unittest import mock import pytest -from trezorlib import btc, tools +from trezorlib import btc, messages, tools from trezorlib.messages import ButtonRequestType if TYPE_CHECKING: @@ -32,6 +32,9 @@ if TYPE_CHECKING: from _pytest.mark.structures import MarkDecorator +BRGeneratorType = Generator[None, messages.ButtonRequest, None] + + # fmt: off # 1 2 3 4 5 6 7 8 9 10 11 12 MNEMONIC12 = "alcohol woman abuse must during monitor noble actual mixed trade anger aisle" @@ -129,135 +132,6 @@ def generate_entropy( return entropy_stripped -def recovery_enter_shares( - debug: "DebugLink", - shares: list[str], - groups: bool = False, - click_info: bool = False, -) -> Generator[None, "ButtonRequest", None]: - if debug.model == "T": - yield from recovery_enter_shares_tt( - debug, shares, groups=groups, click_info=click_info - ) - elif debug.model == "R": - yield from recovery_enter_shares_tr(debug, shares, groups=groups) - else: - raise ValueError(f"Unknown model: {debug.model}") - - -def recovery_enter_shares_tt( - debug: "DebugLink", - shares: list[str], - groups: bool = False, - click_info: bool = False, -) -> Generator[None, "ButtonRequest", None]: - """Perform the recovery flow for a set of Shamir shares. - - For use in an input flow function. - Example: - - def input_flow(): - yield # start recovery - client.debug.press_yes() - yield from recovery_enter_shares(client.debug, SOME_SHARES) - """ - word_count = len(shares[0].split(" ")) - - # Input word number - br = yield - assert br.code == ButtonRequestType.MnemonicWordCount - assert "number of words" in debug.wait_layout().text_content() - debug.input(str(word_count)) - # Homescreen - proceed to share entry - yield - assert "Enter any share" in debug.wait_layout().text_content() - debug.press_yes() - # Enter shares - for share in shares: - br = yield - assert br.code == ButtonRequestType.MnemonicInput - # Enter mnemonic words - for word in share.split(" "): - debug.input(word) - - if groups: - # Confirm share entered - yield - debug.press_yes() - - # Homescreen - continue - # or Homescreen - confirm success - yield - - if click_info: - # Moving through the INFO button - debug.press_info() - yield - debug.swipe_up() - debug.press_yes() - - # Finishing with current share - debug.press_yes() - - -def recovery_enter_shares_tr( - debug: "DebugLink", - shares: list[str], - groups: bool = False, -) -> Generator[None, "ButtonRequest", None]: - """Perform the recovery flow for a set of Shamir shares. - - For use in an input flow function. - Example: - - def input_flow(): - yield # start recovery - client.debug.press_yes() - yield from recovery_enter_shares(client.debug, SOME_SHARES) - """ - word_count = len(shares[0].split(" ")) - - # Homescreen - proceed to word number selection - yield - assert "number of words" in debug.wait_layout().text_content() - debug.press_yes() - # Input word number - br = yield - assert "NUMBER OF WORDS" in debug.wait_layout().title() - assert br.code == ButtonRequestType.MnemonicWordCount - debug.input(str(word_count)) - # Homescreen - proceed to share entry - yield - assert "Enter any share" in debug.wait_layout().text_content() - debug.press_right() - debug.press_right() - debug.press_yes() - - # Enter shares - for index, share in enumerate(shares): - br = yield - assert br.code == ButtonRequestType.MnemonicInput - assert "MnemonicKeyboard" in debug.wait_layout().all_components() - - # Enter mnemonic words - for word in share.split(" "): - debug.input(word) - - if groups: - # Confirm share entered - yield - debug.press_yes() - - # Homescreen - continue - # or Homescreen - confirm success - yield - - # Finishing with current share - debug.press_yes() - - yield - - def click_through( debug: "DebugLink", screens: int, code: Optional[ButtonRequestType] = None ) -> Generator[None, "ButtonRequest", None]: diff --git a/tests/device_tests/reset_recovery/test_recovery_bip39_dryrun.py b/tests/device_tests/reset_recovery/test_recovery_bip39_dryrun.py index a647769ad..bae94538b 100644 --- a/tests/device_tests/reset_recovery/test_recovery_bip39_dryrun.py +++ b/tests/device_tests/reset_recovery/test_recovery_bip39_dryrun.py @@ -51,19 +51,19 @@ def do_recover_legacy(client: Client, mnemonic: list[str], **kwargs: Any): return ret -def do_recover_core(client: Client, mnemonic: list[str], **kwargs: Any): +def do_recover_core(client: Client, mnemonic: list[str], mismatch: bool = False): with client: client.watch_layout() - IF = InputFlowBip39RecoveryDryRun(client, mnemonic) + IF = InputFlowBip39RecoveryDryRun(client, mnemonic, mismatch=mismatch) client.set_input_flow(IF.get()) - return device.recover(client, dry_run=True, **kwargs) + return device.recover(client, dry_run=True) -def do_recover(client: Client, mnemonic: list[str]): +def do_recover(client: Client, mnemonic: list[str], mismatch: bool = False): if client.features.model == "1": return do_recover_legacy(client, mnemonic) else: - return do_recover_core(client, mnemonic) + return do_recover_core(client, mnemonic, mismatch) @pytest.mark.setup_client(mnemonic=MNEMONIC12) @@ -77,7 +77,7 @@ def test_seed_mismatch(client: Client): with pytest.raises( exceptions.TrezorFailure, match="does not match the one in the device" ): - do_recover(client, ["all"] * 12) + do_recover(client, ["all"] * 12, mismatch=True) @pytest.mark.skip_t2 diff --git a/tests/device_tests/reset_recovery/test_recovery_bip39_t2.py b/tests/device_tests/reset_recovery/test_recovery_bip39_t2.py index 582389318..21c4922c2 100644 --- a/tests/device_tests/reset_recovery/test_recovery_bip39_t2.py +++ b/tests/device_tests/reset_recovery/test_recovery_bip39_t2.py @@ -20,7 +20,7 @@ from trezorlib import device, exceptions, messages from trezorlib.debuglink import TrezorClientDebugLink as Client from ...common import MNEMONIC12 -from ...input_flows import InputFlowBip39RecoveryNoPIN, InputFlowBip39RecoveryPIN +from ...input_flows import InputFlowBip39Recovery pytestmark = pytest.mark.skip_t1 @@ -28,7 +28,7 @@ pytestmark = pytest.mark.skip_t1 @pytest.mark.setup_client(uninitialized=True) def test_tt_pin_passphrase(client: Client): with client: - IF = InputFlowBip39RecoveryPIN(client, MNEMONIC12.split(" ")) + IF = InputFlowBip39Recovery(client, MNEMONIC12.split(" "), pin="654") client.set_input_flow(IF.get()) device.recover( client, @@ -48,7 +48,7 @@ def test_tt_pin_passphrase(client: Client): @pytest.mark.setup_client(uninitialized=True) def test_tt_nopin_nopassphrase(client: Client): with client: - IF = InputFlowBip39RecoveryNoPIN(client, MNEMONIC12.split(" ")) + IF = InputFlowBip39Recovery(client, MNEMONIC12.split(" ")) client.set_input_flow(IF.get()) device.recover( client, diff --git a/tests/device_tests/reset_recovery/test_recovery_slip39_advanced.py b/tests/device_tests/reset_recovery/test_recovery_slip39_advanced.py index ac0eb1d5f..5b6a4dbb0 100644 --- a/tests/device_tests/reset_recovery/test_recovery_slip39_advanced.py +++ b/tests/device_tests/reset_recovery/test_recovery_slip39_advanced.py @@ -24,7 +24,8 @@ from ...input_flows import ( InputFlowSlip39AdvancedRecovery, InputFlowSlip39AdvancedRecoveryAbort, InputFlowSlip39AdvancedRecoveryNoAbort, - InputFlowSlip39AdvancedRecoveryTwoSharesWarning, + InputFlowSlip39AdvancedRecoveryShareAlreadyEntered, + InputFlowSlip39AdvancedRecoveryThresholdReached, ) pytestmark = pytest.mark.skip_t1 @@ -119,7 +120,7 @@ def test_same_share(client: Client): second_share = MNEMONIC_SLIP39_ADVANCED_20[1].split(" ")[:4] with client: - IF = InputFlowSlip39AdvancedRecoveryTwoSharesWarning( + IF = InputFlowSlip39AdvancedRecoveryShareAlreadyEntered( client, first_share, second_share ) client.set_input_flow(IF.get()) @@ -135,7 +136,7 @@ def test_group_threshold_reached(client: Client): second_share = MNEMONIC_SLIP39_ADVANCED_20[0].split(" ")[:3] with client: - IF = InputFlowSlip39AdvancedRecoveryTwoSharesWarning( + IF = InputFlowSlip39AdvancedRecoveryThresholdReached( client, first_share, second_share ) client.set_input_flow(IF.get()) diff --git a/tests/device_tests/reset_recovery/test_recovery_slip39_advanced_dryrun.py b/tests/device_tests/reset_recovery/test_recovery_slip39_advanced_dryrun.py index 1e733b163..659eb2a53 100644 --- a/tests/device_tests/reset_recovery/test_recovery_slip39_advanced_dryrun.py +++ b/tests/device_tests/reset_recovery/test_recovery_slip39_advanced_dryrun.py @@ -67,7 +67,7 @@ def test_2of3_invalid_seed_dryrun(client: Client): TrezorFailure, match=r"The seed does not match the one in the device" ): IF = InputFlowSlip39AdvancedRecoveryDryRun( - client, INVALID_SHARES_SLIP39_ADVANCED_20 + client, INVALID_SHARES_SLIP39_ADVANCED_20, mismatch=True ) client.set_input_flow(IF.get()) device.recover( diff --git a/tests/device_tests/reset_recovery/test_recovery_slip39_basic.py b/tests/device_tests/reset_recovery/test_recovery_slip39_basic.py index 0bbf6fb2b..b80569c71 100644 --- a/tests/device_tests/reset_recovery/test_recovery_slip39_basic.py +++ b/tests/device_tests/reset_recovery/test_recovery_slip39_basic.py @@ -26,10 +26,9 @@ from ...common import ( from ...input_flows import ( InputFlowSlip39BasicRecovery, InputFlowSlip39BasicRecoveryAbort, + InputFlowSlip39BasicRecoveryInvalidFirstShare, + InputFlowSlip39BasicRecoveryInvalidSecondShare, InputFlowSlip39BasicRecoveryNoAbort, - InputFlowSlip39BasicRecoveryPIN, - InputFlowSlip39BasicRecoveryRetryFirst, - InputFlowSlip39BasicRecoveryRetrySecond, InputFlowSlip39BasicRecoverySameShare, InputFlowSlip39BasicRecoveryWrongNthWord, ) @@ -63,7 +62,7 @@ def test_secret(client: Client, shares: list[str], secret: str): client.set_input_flow(IF.get()) ret = device.recover(client, pin_protection=False, label="label") - # Workflow succesfully ended + # Workflow successfully ended assert ret == messages.Success(message="Device recovered") assert client.features.pin_protection is False assert client.features.passphrase_protection is False @@ -76,8 +75,8 @@ def test_secret(client: Client, shares: list[str], secret: str): @pytest.mark.setup_client(uninitialized=True) def test_recover_with_pin_passphrase(client: Client): with client: - IF = InputFlowSlip39BasicRecoveryPIN( - client, MNEMONIC_SLIP39_BASIC_20_3of6, "654" + IF = InputFlowSlip39BasicRecovery( + client, MNEMONIC_SLIP39_BASIC_20_3of6, pin="654" ) client.set_input_flow(IF.get()) ret = device.recover( @@ -116,17 +115,20 @@ def test_noabort(client: Client): @pytest.mark.setup_client(uninitialized=True) -def test_ask_word_number(client: Client): +def test_invalid_mnemonic_first_share(client: Client): with client: - IF = InputFlowSlip39BasicRecoveryRetryFirst(client) + IF = InputFlowSlip39BasicRecoveryInvalidFirstShare(client) client.set_input_flow(IF.get()) with pytest.raises(exceptions.Cancelled): device.recover(client, pin_protection=False, label="label") client.init_device() assert client.features.initialized is False + +@pytest.mark.setup_client(uninitialized=True) +def test_invalid_mnemonic_second_share(client: Client): with client: - IF = InputFlowSlip39BasicRecoveryRetrySecond( + IF = InputFlowSlip39BasicRecoveryInvalidSecondShare( client, MNEMONIC_SLIP39_BASIC_20_3of6 ) client.set_input_flow(IF.get()) @@ -149,11 +151,9 @@ def test_wrong_nth_word(client: Client, nth_word: int): @pytest.mark.setup_client(uninitialized=True) def test_same_share(client: Client): - first_share = MNEMONIC_SLIP39_BASIC_20_3of6[0].split(" ") - # second share is first 4 words of first - second_share = MNEMONIC_SLIP39_BASIC_20_3of6[0].split(" ")[:4] + share = MNEMONIC_SLIP39_BASIC_20_3of6[0].split(" ") with client: - IF = InputFlowSlip39BasicRecoverySameShare(client, first_share, second_share) + IF = InputFlowSlip39BasicRecoverySameShare(client, share) client.set_input_flow(IF.get()) with pytest.raises(exceptions.Cancelled): device.recover(client, pin_protection=False, label="label") diff --git a/tests/device_tests/reset_recovery/test_recovery_slip39_basic_dryrun.py b/tests/device_tests/reset_recovery/test_recovery_slip39_basic_dryrun.py index f5be6efc1..aee56a73f 100644 --- a/tests/device_tests/reset_recovery/test_recovery_slip39_basic_dryrun.py +++ b/tests/device_tests/reset_recovery/test_recovery_slip39_basic_dryrun.py @@ -20,7 +20,7 @@ from trezorlib import device, messages from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.exceptions import TrezorFailure -from ...input_flows import InputFlowSlip39BasicRecovery +from ...input_flows import InputFlowSlip39BasicRecoveryDryRun pytestmark = pytest.mark.skip_t1 @@ -39,7 +39,7 @@ INVALID_SHARES_20_2of3 = [ @pytest.mark.setup_client(mnemonic=SHARES_20_2of3[0:2]) def test_2of3_dryrun(client: Client): with client: - IF = InputFlowSlip39BasicRecovery(client, SHARES_20_2of3[1:3], dry_run=True) + IF = InputFlowSlip39BasicRecoveryDryRun(client, SHARES_20_2of3[1:3]) client.set_input_flow(IF.get()) ret = device.recover( client, @@ -62,7 +62,9 @@ def test_2of3_invalid_seed_dryrun(client: Client): with client, pytest.raises( TrezorFailure, match=r"The seed does not match the one in the device" ): - IF = InputFlowSlip39BasicRecovery(client, INVALID_SHARES_20_2of3, dry_run=True) + IF = InputFlowSlip39BasicRecoveryDryRun( + client, INVALID_SHARES_20_2of3, mismatch=True + ) client.set_input_flow(IF.get()) device.recover( client, diff --git a/tests/device_tests/reset_recovery/test_reset_recovery_bip39.py b/tests/device_tests/reset_recovery/test_reset_recovery_bip39.py index e9b3a99b4..fa6957ccd 100644 --- a/tests/device_tests/reset_recovery/test_reset_recovery_bip39.py +++ b/tests/device_tests/reset_recovery/test_reset_recovery_bip39.py @@ -22,7 +22,7 @@ from trezorlib.messages import BackupType from trezorlib.tools import parse_path from ...common import WITH_MOCK_URANDOM -from ...input_flows import InputFlowBip39RecoveryNoPIN, InputFlowBip39ResetBackup +from ...input_flows import InputFlowBip39Recovery, InputFlowBip39ResetBackup @pytest.mark.skip_t1 @@ -67,7 +67,7 @@ def reset(client: Client, strength: int = 128, skip_backup: bool = False) -> str def recover(client: Client, mnemonic: str): words = mnemonic.split(" ") with client: - IF = InputFlowBip39RecoveryNoPIN(client, words) + IF = InputFlowBip39Recovery(client, words) client.set_input_flow(IF.get()) client.watch_layout() ret = device.recover(client, pin_protection=False, label="label") diff --git a/tests/input_flows.py b/tests/input_flows.py index b359823b1..7e23541e3 100644 --- a/tests/input_flows.py +++ b/tests/input_flows.py @@ -9,8 +9,10 @@ only on the actual tests and data-assertions, not on the lower-level input flow details. """ +from __future__ import annotations + import time -from typing import Callable, Generator, Optional +from typing import Callable, Generator from trezorlib import messages from trezorlib.debuglink import ( @@ -22,21 +24,18 @@ from trezorlib.debuglink import ( from . import buttons from .common import ( + BRGeneratorType, check_pin_backoff_time, click_info_button_tt, click_through, read_and_confirm_mnemonic, - recovery_enter_shares, ) - -GeneratorType = Generator[None, messages.ButtonRequest, None] +from .input_flows_helpers import BackupFlow, PinFlow, RecoveryFlow B = messages.ButtonRequestType -def swipe_if_necessary( - debug: DebugLink, br_code: Optional[messages.ButtonRequestType] = None -) -> GeneratorType: +def swipe_if_necessary(debug: DebugLink, br_code: B | None = None) -> BRGeneratorType: br = yield if br_code is not None: assert br.code == br_code @@ -49,11 +48,14 @@ class InputFlowBase: def __init__(self, client: Client): self.client = client self.debug: DebugLink = client.debug + self.PIN = PinFlow(self.client) + self.REC = RecoveryFlow(self.client) + self.BAK = BackupFlow(self.client) - def model(self) -> Optional[str]: + def model(self) -> str | None: return self.client.features.model - def get(self) -> Callable[[], GeneratorType]: + def get(self) -> Callable[[], BRGeneratorType]: self.client.watch_layout(True) # There could be one common input flow for all models @@ -66,11 +68,11 @@ class InputFlowBase: else: raise ValueError("Unknown model") - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: """Special for TT""" raise NotImplementedError - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: """Special for TR""" raise NotImplementedError @@ -83,7 +85,7 @@ class InputFlowBase: def all_components(self) -> list[str]: return self.debug.wait_layout().all_components() - def title(self) -> list[str]: + def title(self) -> str: return self.debug.wait_layout().title() @@ -93,7 +95,7 @@ class InputFlowSetupDevicePINWIpeCode(InputFlowBase): self.pin = pin self.wipe_code = wipe_code - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: yield # do you want to set/change the wipe code? self.debug.press_yes() @@ -122,7 +124,7 @@ class InputFlowNewCodeMismatch(InputFlowBase): self.first_code = first_code self.second_code = second_code - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: yield # do you want to set/change the pin/wipe code? self.debug.press_yes() @@ -130,17 +132,8 @@ class InputFlowNewCodeMismatch(InputFlowBase): yield from swipe_if_necessary(self.debug) # code info self.debug.press_yes() - def input_two_different_pins(): - yield # enter new PIN/wipe_code - assert "PinKeyboard" in self.all_components() - self.debug.input(self.first_code) - if self.debug.model == "R": - yield # Please re-enter PIN to confirm - assert "re-enter PIN" in self.text_content() - self.debug.press_yes() - yield # enter new PIN/wipe_code again (but different) - assert "PinKeyboard" in self.all_components() - self.debug.input(self.second_code) + def input_two_different_pins() -> BRGeneratorType: + yield from self.PIN.setup_new_pin(self.first_code, self.second_code) yield from input_two_different_pins() @@ -166,24 +159,13 @@ class InputFlowCodeChangeFail(InputFlowBase): self.new_pin_1 = new_pin_1 self.new_pin_2 = new_pin_2 - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: yield # do you want to change pin? self.debug.press_yes() yield # enter current pin self.debug.input(self.current_pin) - yield # enter new pin - assert "PinKeyboard" in self.all_components() - self.debug.input(self.new_pin_1) - - if self.debug.model == "R": - yield # Please re-enter PIN to confirm - assert "re-enter PIN" in self.text_content() - self.debug.press_yes() - - yield # enter new pin again (but different) - assert "PinKeyboard" in self.all_components() - self.debug.input(self.new_pin_2) + yield from self.PIN.setup_new_pin(self.new_pin_1, self.new_pin_2) yield # PIN mismatch self.debug.press_yes() # try again @@ -198,7 +180,7 @@ class InputFlowWrongPIN(InputFlowBase): super().__init__(client) self.wrong_pin = wrong_pin - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: yield # do you want to change pin? self.debug.press_yes() yield # enter wrong current pin @@ -213,7 +195,7 @@ class InputFlowPINBackoff(InputFlowBase): self.wrong_pin = wrong_pin self.good_pin = good_pin - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: """Inputting some bad PINs and finally the correct one""" yield # PIN entry for attempt in range(3): @@ -229,7 +211,7 @@ class InputFlowSignMessagePagination(InputFlowBase): super().__init__(client) self.message_read = "" - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: # collect screen contents into `message_read`. # Using a helper debuglink function to assemble the final text. layouts: list[LayoutContent] = [] @@ -251,7 +233,7 @@ class InputFlowSignMessagePagination(InputFlowBase): self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: # confirm address yield self.debug.press_yes() @@ -269,7 +251,7 @@ class InputFlowShowAddressQRCode(InputFlowBase): def __init__(self, client: Client): super().__init__(client) - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: yield self.debug.click(buttons.CORNER_BUTTON, wait=True) # synchronize; TODO get rid of this once we have single-global-layout @@ -283,7 +265,7 @@ class InputFlowShowAddressQRCode(InputFlowBase): self.debug.press_no(wait=True) self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: yield # Go into details self.debug.press_right() @@ -299,7 +281,7 @@ class InputFlowShowAddressQRCodeCancel(InputFlowBase): def __init__(self, client: Client): super().__init__(client) - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: yield self.debug.click(buttons.CORNER_BUTTON, wait=True) # synchronize; TODO get rid of this once we have single-global-layout @@ -310,7 +292,7 @@ class InputFlowShowAddressQRCodeCancel(InputFlowBase): self.debug.press_no(wait=True) self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: yield # Go into details self.debug.press_right() @@ -331,7 +313,7 @@ class InputFlowShowMultisigXPUBs(InputFlowBase): self.xpubs = xpubs self.index = index - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: yield # show address layout = self.debug.wait_layout() assert "RECEIVE ADDRESS\n(MULTISIG)" == layout.title() @@ -363,7 +345,7 @@ class InputFlowShowMultisigXPUBs(InputFlowBase): # show address self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: yield # show address layout = self.debug.wait_layout() assert "RECEIVE ADDRESS (MULTISIG)" in layout.title() @@ -407,7 +389,7 @@ class InputFlowPaymentRequestDetails(InputFlowBase): super().__init__(client) self.outputs = outputs - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: yield # request to see details self.debug.wait_layout() self.debug.press_info() @@ -437,7 +419,7 @@ class InputFlowSignTxHighFee(InputFlowBase): super().__init__(client) self.finished = False - def go_through_all_screens(self, screens: list[B]) -> GeneratorType: + def go_through_all_screens(self, screens: list[B]) -> BRGeneratorType: for expected in screens: br = yield assert br.code == expected @@ -445,7 +427,7 @@ class InputFlowSignTxHighFee(InputFlowBase): self.finished = True - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: screens = [ B.ConfirmOutput, B.ConfirmOutput, @@ -454,7 +436,7 @@ class InputFlowSignTxHighFee(InputFlowBase): ] yield from self.go_through_all_screens(screens) - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: screens = [ B.ConfirmOutput, B.FeeOverThreshold, @@ -517,12 +499,12 @@ class InputFlowSignTxInformation(InputFlowBase): assert "fee rate" in content assert "71.56 sat" in content - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: content = yield from sign_tx_go_to_info(self.client) self.assert_content(content) self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: content = yield from sign_tx_go_to_info_tr(self.client) self.assert_content(content.lower()) self.debug.press_yes() @@ -538,12 +520,12 @@ class InputFlowSignTxInformationMixed(InputFlowBase): assert "fee rate" in content assert "18.33 sat" in content - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: content = yield from sign_tx_go_to_info(self.client) self.assert_content(content) self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: content = yield from sign_tx_go_to_info_tr(self.client) self.assert_content(content.lower()) self.debug.press_yes() @@ -553,11 +535,11 @@ class InputFlowSignTxInformationCancel(InputFlowBase): def __init__(self, client: Client): super().__init__(client) - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: yield from sign_tx_go_to_info(self.client) self.debug.press_no() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: yield from sign_tx_go_to_info_tr(self.client) self.debug.press_left() @@ -566,7 +548,7 @@ class InputFlowSignTxInformationReplacement(InputFlowBase): def __init__(self, client: Client): super().__init__(client) - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: yield # confirm txid self.debug.press_yes() yield # confirm address @@ -583,7 +565,7 @@ class InputFlowSignTxInformationReplacement(InputFlowBase): self.debug.click(buttons.CORNER_BUTTON, wait=True) self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: yield # confirm txid self.debug.press_right() self.debug.press_right() @@ -601,7 +583,7 @@ def lock_time_input_flow_tt( debug: DebugLink, layout_assert_func: Callable[[DebugLink], None], double_confirm: bool = False, -) -> GeneratorType: +) -> BRGeneratorType: yield # confirm output debug.wait_layout() debug.press_yes() @@ -622,7 +604,7 @@ def lock_time_input_flow_tt( def lock_time_input_flow_tr( debug: DebugLink, layout_assert_func: Callable[[DebugLink], None] -) -> GeneratorType: +) -> BRGeneratorType: yield # confirm output debug.wait_layout() debug.swipe_up() @@ -642,7 +624,7 @@ class InputFlowLockTimeBlockHeight(InputFlowBase): super().__init__(client) self.block_height = block_height - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: def assert_func(debug: DebugLink) -> None: layout_text = debug.wait_layout().text_content() assert "blockheight" in layout_text @@ -650,7 +632,7 @@ class InputFlowLockTimeBlockHeight(InputFlowBase): yield from lock_time_input_flow_tt(self.debug, assert_func, double_confirm=True) - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: def assert_func(debug: DebugLink) -> None: assert "blockheight" in debug.wait_layout().text_content() debug.press_right() @@ -664,7 +646,7 @@ class InputFlowLockTimeDatetime(InputFlowBase): super().__init__(client) self.lock_time_str = lock_time_str - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: def assert_func(debug: DebugLink): layout_text = debug.wait_layout().text_content() assert "Locktime" in layout_text @@ -672,7 +654,7 @@ class InputFlowLockTimeDatetime(InputFlowBase): yield from lock_time_input_flow_tt(self.debug, assert_func) - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: def assert_func(debug: DebugLink): assert "Locktime" in debug.wait_layout().text_content() debug.press_right() @@ -695,7 +677,7 @@ class InputFlowEIP712ShowMore(InputFlowBase): elif self.model() == "R": self.debug.press_right() - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: """Triggers show more wherever possible""" yield # confirm address self.debug.press_yes() @@ -742,7 +724,7 @@ class InputFlowEIP712Cancel(InputFlowBase): def __init__(self, client: Client): super().__init__(client) - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: """Clicks cancelling button""" yield # confirm address self.debug.press_yes() @@ -756,7 +738,7 @@ class InputFlowEthereumSignTxSkip(InputFlowBase): super().__init__(client) self.cancel = cancel - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: yield # confirm address self.debug.press_yes() yield # confirm amount @@ -782,7 +764,7 @@ class InputFlowEthereumSignTxScrollDown(InputFlowBase): super().__init__(client) self.cancel = cancel - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: yield # confirm address self.debug.wait_layout() self.debug.press_yes() @@ -813,7 +795,7 @@ class InputFlowEthereumSignTxScrollDown(InputFlowBase): yield # hold to confirm self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: yield # confirm address self.debug.wait_layout() self.debug.press_yes() @@ -843,7 +825,7 @@ class InputFlowEthereumSignTxGoBack(InputFlowBase): super().__init__(client) self.cancel = cancel - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: br = yield # confirm address self.debug.wait_layout() self.debug.press_yes() @@ -904,7 +886,7 @@ class InputFlowBip39Backup(InputFlowBase): super().__init__(client) self.mnemonic = None - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: # 1. Confirm Reset yield from click_through(self.debug, screens=1, code=B.ResetDevice) @@ -918,7 +900,7 @@ class InputFlowBip39ResetBackup(InputFlowBase): self.mnemonic = None # NOTE: same as above, just two more YES - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: # 1. Confirm Reset # 2. Backup your seed # 3. Confirm warning @@ -933,23 +915,12 @@ class InputFlowBip39ResetPIN(InputFlowBase): super().__init__(client) self.mnemonic = None - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: br = yield # Confirm Reset assert br.code == B.ResetDevice self.debug.press_yes() - yield # Enter new PIN - assert "PinKeyboard" in self.all_components() - self.debug.input("654") - - if self.debug.model == "R": - yield # Re-enter PIN - assert "re-enter PIN" in self.text_content() - self.debug.press_yes() - - yield # Confirm PIN - assert "PinKeyboard" in self.all_components() - self.debug.input("654") + yield from self.PIN.setup_new_pin("654") br = yield # Confirm entropy assert br.code == B.ResetDevice @@ -980,7 +951,7 @@ class InputFlowBip39ResetFailedCheck(InputFlowBase): super().__init__(client) self.mnemonic = None - def input_flow_common(self) -> GeneratorType: + def input_flow_common(self) -> BRGeneratorType: # 1. Confirm Reset # 2. Backup your seed # 3. Confirm warning @@ -1031,7 +1002,7 @@ class InputFlowSlip39BasicBackup(InputFlowBase): self.mnemonics: list[str] = [] self.click_info = click_info - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: yield # 1. Checklist self.debug.press_yes() if self.click_info: @@ -1056,7 +1027,7 @@ class InputFlowSlip39BasicBackup(InputFlowBase): assert br.code == B.Success self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: yield # 1. Checklist self.debug.press_yes() yield # 1.5 Number of shares info @@ -1087,7 +1058,7 @@ class InputFlowSlip39BasicResetRecovery(InputFlowBase): super().__init__(client) self.mnemonics: list[str] = [] - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: # 1. Confirm Reset # 2. Backup your seed # 3. Confirm warning @@ -1105,7 +1076,7 @@ class InputFlowSlip39BasicResetRecovery(InputFlowBase): assert br.code == B.Success self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: yield # Confirm Reset self.debug.press_yes() yield # Backup your seed @@ -1159,7 +1130,7 @@ class InputFlowSlip39AdvancedBackup(InputFlowBase): self.mnemonics: list[str] = [] self.click_info = click_info - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: yield # 1. Checklist self.debug.press_yes() if self.click_info: @@ -1193,7 +1164,7 @@ class InputFlowSlip39AdvancedBackup(InputFlowBase): assert br.code == B.Success self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: yield # 1. Checklist self.debug.press_yes() yield # 2. Set and confirm group count @@ -1230,7 +1201,7 @@ class InputFlowSlip39AdvancedResetRecovery(InputFlowBase): self.mnemonics: list[str] = [] self.click_info = click_info - def input_flow_tt(self) -> GeneratorType: + def input_flow_tt(self) -> BRGeneratorType: # 1. Confirm Reset # 2. Backup your seed # 3. Confirm warning @@ -1251,7 +1222,7 @@ class InputFlowSlip39AdvancedResetRecovery(InputFlowBase): assert br.code == B.Success self.debug.press_yes() - def input_flow_tr(self) -> GeneratorType: + def input_flow_tr(self) -> BRGeneratorType: yield # Wallet backup self.debug.press_yes() yield # Wallet creation @@ -1286,225 +1257,67 @@ class InputFlowSlip39AdvancedResetRecovery(InputFlowBase): self.debug.press_yes() -def enter_recovery_seed_dry_run_tt( - debug: DebugLink, mnemonic: list[str] -) -> GeneratorType: - yield - assert "check the recovery seed" in debug.wait_layout().text_content() - debug.click(buttons.OK) - - yield - assert "SelectWordCount" in debug.wait_layout().all_components() - # click the correct number - word_option_offset = 6 - word_options = (12, 18, 20, 24, 33) - index = word_option_offset + word_options.index(len(mnemonic)) - debug.click(buttons.grid34(index % 3, index // 3)) - - yield - assert "Enter your backup" in debug.wait_layout().text_content() - debug.click(buttons.OK) - - yield - for word in mnemonic: - assert debug.wait_layout().main_component() == "MnemonicKeyboard" - debug.input(word) - - class InputFlowBip39RecoveryDryRun(InputFlowBase): - def __init__(self, client: Client, mnemonic: list[str]): + def __init__(self, client: Client, mnemonic: list[str], mismatch: bool = False): super().__init__(client) self.mnemonic = mnemonic + self.mismatch = mismatch - def input_flow_tt(self) -> GeneratorType: - yield from enter_recovery_seed_dry_run_tt(self.debug, self.mnemonic) - - yield - self.debug.wait_layout() - self.debug.click(buttons.OK) - - def input_flow_tr(self) -> GeneratorType: - yield - assert "check the recovery seed" in self.text_content() - self.debug.press_yes() - - yield from enter_recovery_seed_tr(self.debug, self.mnemonic) - - yield - assert "is valid" in self.text_content() - self.debug.press_yes() + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_dry_run() + yield from self.REC.setup_bip39_recovery(len(self.mnemonic)) + yield from self.REC.input_mnemonic(self.mnemonic) + if self.mismatch: + yield from self.REC.warning_bip39_dryrun_mismatch() + else: + yield from self.REC.success_bip39_dry_run_valid() class InputFlowBip39RecoveryDryRunInvalid(InputFlowBase): def __init__(self, client: Client): super().__init__(client) + self.invalid_mnemonic = ["stick"] * 12 - def input_flow_tt(self) -> GeneratorType: - mnemonic = ["stick"] * 12 - yield from enter_recovery_seed_dry_run_tt(self.debug, mnemonic) - - br = yield - assert br.code == messages.ButtonRequestType.Warning - assert "Invalid recovery seed" in self.text_content() - self.debug.click(buttons.OK) + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_dry_run() + yield from self.REC.setup_bip39_recovery(len(self.invalid_mnemonic)) + yield from self.REC.input_mnemonic(self.invalid_mnemonic) + yield from self.REC.warning_invalid_recovery_seed() yield - assert "SelectWordCount" in self.all_components() self.client.cancel() - def input_flow_tr(self) -> GeneratorType: - yield - assert "check the recovery seed" in self.text_content() - self.debug.press_right() - mnemonic = ["stick"] * 12 - yield from enter_recovery_seed_tr(self.debug, mnemonic) - - br = yield - assert br.code == messages.ButtonRequestType.Warning - assert "Invalid recovery seed" in self.text_content() - self.debug.press_middle() - - yield # retry screen - assert "number of words" in self.text_content() - self.debug.press_left() - - yield - assert "abort the backup check" in self.text_content() - self.debug.press_right() - - -def bip39_recovery_possible_pin_tt( - debug: DebugLink, mnemonic: list[str], pin: Optional[str] -) -> GeneratorType: - yield - assert "By continuing you agree to" in debug.wait_layout().text_content() - debug.press_yes() - - # PIN when requested - if pin is not None: - yield - assert "PinKeyboard" in debug.wait_layout().all_components() - debug.input(pin) - - yield - assert "PinKeyboard" in debug.wait_layout().all_components() - debug.input(pin) - - yield - assert "SelectWordCount" in debug.wait_layout().all_components() - debug.input(str(len(mnemonic))) - - yield - assert "Enter your backup" in debug.wait_layout().text_content() - debug.press_yes() - - yield - for word in mnemonic: - assert debug.wait_layout().main_component() == "MnemonicKeyboard" - debug.input(word) - - yield - assert "Wallet recovered successfully" in debug.wait_layout().text_content() - debug.press_yes() - - -def bip39_recovery_possible_pin_tr( - debug: DebugLink, mnemonic: list[str], pin: Optional[str] -) -> GeneratorType: - yield - assert "By continuing you agree" in debug.wait_layout().text_content() - debug.press_right() - assert "trezor.io/tos" in debug.wait_layout().text_content() - debug.press_yes() - - # PIN when requested - if pin is not None: - yield - assert "PinKeyboard" in debug.wait_layout().all_components() - debug.input("654") - - yield - assert "re-enter PIN" in debug.wait_layout().text_content() - debug.press_right() - - yield - assert "PinKeyboard" in debug.wait_layout().all_components() - debug.input("654") - - yield from enter_recovery_seed_tr(debug, mnemonic) - - yield - assert "Wallet recovered successfully" in debug.wait_layout().text_content() - debug.press_yes() - - -def enter_recovery_seed_tr(debug: DebugLink, mnemonic: list[str]) -> GeneratorType: - yield - assert "number of words" in debug.wait_layout().text_content() - debug.press_yes() - - yield - assert "NUMBER OF WORDS" in debug.wait_layout().title() - debug.input(str(len(mnemonic))) - - yield - assert "Enter your backup" in debug.wait_layout().text_content() - # Paginate to see info - debug.press_right() - debug.press_right() - debug.press_yes() - - yield - assert "MnemonicKeyboard" in debug.wait_layout().all_components() - for index, word in enumerate(mnemonic): - assert f"WORD {index + 1}" in debug.wait_layout().title() - debug.input(word) - - -class InputFlowBip39RecoveryPIN(InputFlowBase): - def __init__(self, client: Client, mnemonic: list[str]): +class InputFlowBip39Recovery(InputFlowBase): + def __init__(self, client: Client, mnemonic: list[str], pin: str | None = None): super().__init__(client) self.mnemonic = mnemonic + self.pin = pin - def input_flow_tt(self) -> GeneratorType: - yield from bip39_recovery_possible_pin_tt(self.debug, self.mnemonic, pin="654") - - def input_flow_tr(self) -> GeneratorType: - yield from bip39_recovery_possible_pin_tr(self.debug, self.mnemonic, pin="654") - - -class InputFlowBip39RecoveryNoPIN(InputFlowBase): - def __init__(self, client: Client, mnemonic: list[str]): - super().__init__(client) - self.mnemonic = mnemonic - - def input_flow_tt(self) -> GeneratorType: - yield from bip39_recovery_possible_pin_tt(self.debug, self.mnemonic, pin=None) - - def input_flow_tr(self) -> GeneratorType: - yield from bip39_recovery_possible_pin_tr(self.debug, self.mnemonic, pin=None) + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() + if self.pin is not None: + yield from self.PIN.setup_new_pin(self.pin) + yield from self.REC.setup_bip39_recovery(len(self.mnemonic)) + yield from self.REC.input_mnemonic(self.mnemonic) + yield from self.REC.success_wallet_recovered() class InputFlowSlip39AdvancedRecoveryDryRun(InputFlowBase): - def __init__(self, client: Client, shares: list[str]): + def __init__(self, client: Client, shares: list[str], mismatch: bool = False): super().__init__(client) self.shares = shares + self.mismatch = mismatch + self.word_count = len(shares[0].split(" ")) - def input_flow_common(self) -> GeneratorType: - yield # Confirm Dryrun - assert "check the recovery seed" in self.text_content() - self.debug.press_yes() - # run recovery flow - yield from recovery_enter_shares(self.debug, self.shares, groups=True) - - -def confirm_recovery(debug: DebugLink) -> GeneratorType: - yield # Confirm Recovery - assert "By continuing you agree" in debug.wait_layout().text_content() - if debug.model == "R": - debug.press_right() - debug.press_yes() + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_dry_run() + yield from self.REC.setup_slip39_recovery(self.word_count) + yield from self.REC.input_all_slip39_shares(self.shares, has_groups=True) + if self.mismatch: + yield from self.REC.warning_slip39_dryrun_mismatch() + else: + yield from self.REC.success_slip39_dryrun_valid() class InputFlowSlip39AdvancedRecovery(InputFlowBase): @@ -1512,358 +1325,200 @@ class InputFlowSlip39AdvancedRecovery(InputFlowBase): super().__init__(client) self.shares = shares self.click_info = click_info + self.word_count = len(shares[0].split(" ")) - def input_flow_common(self) -> GeneratorType: - yield from confirm_recovery(self.debug) - yield from recovery_enter_shares( - self.debug, self.shares, groups=True, click_info=self.click_info + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() + yield from self.REC.setup_slip39_recovery(self.word_count) + yield from self.REC.input_all_slip39_shares( + self.shares, has_groups=True, click_info=self.click_info ) + yield from self.REC.success_wallet_recovered() class InputFlowSlip39AdvancedRecoveryAbort(InputFlowBase): def __init__(self, client: Client): super().__init__(client) - def input_flow_common(self) -> GeneratorType: - yield from confirm_recovery(self.debug) + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() if self.debug.model == "T": - # Need to choose the word amount for TT - yield - self.debug.input("12") - yield - assert "Enter your backup" in self.text_content() - else: - yield - assert "number of words" in self.text_content() - self.debug.press_no() - yield # Homescreen - confirm abort - assert "abort the recovery" in self.text_content() - self.debug.press_yes() + yield from self.REC.input_number_of_words(20) + yield from self.REC.abort_recovery(True) class InputFlowSlip39AdvancedRecoveryNoAbort(InputFlowBase): def __init__(self, client: Client, shares: list[str]): super().__init__(client) self.shares = shares + self.word_count = len(shares[0].split(" ")) - def input_flow_tt(self) -> GeneratorType: - yield from confirm_recovery(self.debug) - yield - self.debug.input("20") - yield # Homescreen - abort process - assert "Enter any share" in self.text_content() - self.debug.press_no() - yield # Homescreen - go back to process - assert "abort the recovery" in self.text_content() - self.debug.press_no() - yield - assert "Enter any share" in self.text_content() - self.debug.press_yes() - - # TODO: make it reusable - for index, share in enumerate(self.shares): - yield - assert "MnemonicKeyboard" in self.debug.wait_layout().all_components() - for word in share.split(" "): - self.debug.input(word) - - yield - if index == len(self.shares) - 1: - assert "Wallet recovered" in self.text_content() - else: - assert "You have entered" in self.text_content() - self.debug.press_yes() - yield - assert "More shares needed" in self.text_content() - self.debug.press_yes() - - yield - - def input_flow_tr(self) -> GeneratorType: - yield from confirm_recovery(self.debug) - yield # Homescreen - abort process - assert "number of words" in self.text_content() - self.debug.press_no() - yield # Homescreen - go back to process - assert "abort the recovery" in self.text_content() - self.debug.press_right() - self.debug.press_no() - yield from recovery_enter_shares(self.debug, self.shares, groups=True) + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() + if self.debug.model == "T": + yield from self.REC.input_number_of_words(self.word_count) + yield from self.REC.abort_recovery(False) + else: + yield from self.REC.abort_recovery(False) + yield from self.REC.tr_recovery_homescreen() + yield from self.REC.input_number_of_words(self.word_count) + yield from self.REC.enter_any_share() + yield from self.REC.input_all_slip39_shares(self.shares, has_groups=True) + yield from self.REC.success_wallet_recovered() -class InputFlowSlip39AdvancedRecoveryTwoSharesWarning(InputFlowBase): - def __init__(self, client: Client, first_share: list[str], second_share: list[str]): +class InputFlowSlip39AdvancedRecoveryThresholdReached(InputFlowBase): + def __init__( + self, + client: Client, + first_share: list[str], + second_share: list[str], + ): super().__init__(client) self.first_share = first_share self.second_share = second_share - def input_flow_common(self) -> GeneratorType: - yield from confirm_recovery(self.debug) - yield from slip39_recovery_setup_and_first_share(self.debug, self.first_share) + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() + yield from self.REC.setup_slip39_recovery(len(self.first_share)) + yield from self.REC.input_mnemonic(self.first_share) + yield from self.REC.success_share_group_entered() + yield from self.REC.success_more_shares_needed() + yield from self.REC.input_mnemonic(self.second_share) + yield from self.REC.warning_group_threshold_reached() - yield # Continue to next share - assert "You have entered" in self.text_content() - self.debug.press_yes() - yield # Homescreen - next share - assert "More shares needed" in self.text_content() - self.debug.press_yes() - yield # Enter next share - assert "MnemonicKeyboard" in self.debug.wait_layout().all_components() - for word in self.second_share: - self.debug.input(word) - - br = yield - assert br.code == messages.ButtonRequestType.Warning - assert ( - "Share already entered" in self.text_content() - or "Group threshold reached" in self.text_content() - ) - self.debug.press_yes() yield - self.client.cancel() -def slip39_recovery_possible_pin( - debug: DebugLink, shares: list[str], pin: Optional[str], dry_run: bool = False -) -> GeneratorType: - yield # Confirm Recovery/Dryrun - if dry_run: - assert "check the recovery seed" in debug.wait_layout().text_content() - else: - assert "By continuing you agree" in debug.wait_layout().text_content() - if debug.model == "R": - debug.press_right() - debug.press_yes() +class InputFlowSlip39AdvancedRecoveryShareAlreadyEntered(InputFlowBase): + def __init__( + self, + client: Client, + first_share: list[str], + second_share: list[str], + ): + super().__init__(client) + self.first_share = first_share + self.second_share = second_share - if pin is not None: - yield # Enter PIN - assert "PinKeyboard" in debug.wait_layout().all_components() - debug.input(pin) - if debug.model == "R": - yield # Reenter PIN - assert "re-enter PIN" in debug.wait_layout().text_content() - debug.press_yes() - yield # Enter PIN again - assert "PinKeyboard" in debug.wait_layout().all_components() - debug.input(pin) + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() + yield from self.REC.setup_slip39_recovery(len(self.first_share)) + yield from self.REC.input_mnemonic(self.first_share) + yield from self.REC.success_share_group_entered() + yield from self.REC.success_more_shares_needed() + yield from self.REC.input_mnemonic(self.second_share) + yield from self.REC.warning_share_already_entered() - # Proceed with recovery - yield from recovery_enter_shares(debug, shares) + yield + self.client.cancel() + + +class InputFlowSlip39BasicRecoveryDryRun(InputFlowBase): + def __init__(self, client: Client, shares: list[str], mismatch: bool = False): + super().__init__(client) + self.shares = shares + self.mismatch = mismatch + self.word_count = len(shares[0].split(" ")) + + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_dry_run() + yield from self.REC.setup_slip39_recovery(self.word_count) + yield from self.REC.input_all_slip39_shares(self.shares) + if self.mismatch: + yield from self.REC.warning_slip39_dryrun_mismatch() + else: + yield from self.REC.success_slip39_dryrun_valid() class InputFlowSlip39BasicRecovery(InputFlowBase): - def __init__(self, client: Client, shares: list[str], dry_run: bool = False): - super().__init__(client) - self.shares = shares - self.dry_run = dry_run - - def input_flow_common(self) -> GeneratorType: - yield from slip39_recovery_possible_pin( - self.debug, self.shares, pin=None, dry_run=self.dry_run - ) - - -class InputFlowSlip39BasicRecoveryPIN(InputFlowBase): - def __init__(self, client: Client, shares: list[str], pin: str): + def __init__(self, client: Client, shares: list[str], pin: str | None = None): super().__init__(client) self.shares = shares self.pin = pin + self.word_count = len(shares[0].split(" ")) - def input_flow_common(self) -> GeneratorType: - yield from slip39_recovery_possible_pin(self.debug, self.shares, pin=self.pin) + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() + if self.pin is not None: + yield from self.PIN.setup_new_pin(self.pin) + yield from self.REC.setup_slip39_recovery(self.word_count) + yield from self.REC.input_all_slip39_shares(self.shares) + yield from self.REC.success_wallet_recovered() class InputFlowSlip39BasicRecoveryAbort(InputFlowBase): def __init__(self, client: Client): super().__init__(client) - def input_flow_tt(self) -> GeneratorType: - yield from confirm_recovery(self.debug) - yield - assert "number of words" in self.text_content() - self.debug.input("20") - yield - assert "Enter any share" in self.text_content() - self.debug.press_no() - yield - assert "abort the recovery" in self.text_content() - self.debug.press_yes() - yield - - def input_flow_tr(self) -> GeneratorType: - yield from confirm_recovery(self.debug) - yield # Homescreen - abort process - assert "number of words" in self.text_content() - self.debug.press_no() - yield # Homescreen - confirm abort - assert "abort the recovery" in self.text_content() - self.debug.press_yes() + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() + if self.debug.model == "T": + yield from self.REC.input_number_of_words(20) + yield from self.REC.abort_recovery(True) class InputFlowSlip39BasicRecoveryNoAbort(InputFlowBase): def __init__(self, client: Client, shares: list[str]): super().__init__(client) self.shares = shares + self.word_count = len(shares[0].split(" ")) - def input_flow_tt(self) -> GeneratorType: - yield from confirm_recovery(self.debug) - yield - assert "number of words" in self.text_content() - self.debug.input("20") - yield - assert "Enter any share" in self.text_content() - self.debug.press_no() - yield - assert "abort the recovery" in self.text_content() - self.debug.press_no() - yield - assert "Enter any share" in self.text_content() - self.debug.press_yes() - # run recovery flow - # TODO: make this a reusable function - for index, share in enumerate(self.shares): - yield - assert "MnemonicKeyboard" in self.debug.wait_layout().all_components() - for word in share.split(" "): - self.debug.input(word) + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() - yield - if index == len(self.shares) - 1: - assert "Wallet recovered" in self.text_content() - else: - assert "shares entered successfully" in self.text_content() - self.debug.press_yes() + if self.debug.model == "T": + yield from self.REC.input_number_of_words(self.word_count) + yield from self.REC.abort_recovery(False) + else: + yield from self.REC.abort_recovery(False) + yield from self.REC.tr_recovery_homescreen() + yield from self.REC.input_number_of_words(self.word_count) - yield - - def input_flow_tr(self) -> GeneratorType: - yield from confirm_recovery(self.debug) - yield # Homescreen - abort process - assert "number of words" in self.text_content() - self.debug.press_no() - yield # Homescreen - go back to process - assert "abort the recovery" in self.text_content() - self.debug.press_right() - self.debug.press_no() - # run recovery flow - yield from recovery_enter_shares(self.debug, self.shares) + yield from self.REC.enter_any_share() + yield from self.REC.input_all_slip39_shares(self.shares) + yield from self.REC.success_wallet_recovered() -def slip39_recovery_setup_and_first_share( - debug: DebugLink, first_share: list[str] -) -> GeneratorType: - if debug.model == "R": - yield - assert "number of words" in debug.wait_layout().text_content() - debug.press_yes() - yield - assert "NUMBER OF WORDS" in debug.wait_layout().title() - debug.input(str(len(first_share))) - yield # Homescreen - proceed to share entry - assert "Enter any share" in debug.wait_layout().text_content() - debug.press_right(wait=True) - debug.press_right(wait=True) - debug.press_yes() - yield # Enter first share - assert "MnemonicKeyboard" in debug.wait_layout().all_components() - for index, word in enumerate(first_share): - assert f"WORD {index + 1}" in debug.wait_layout().title() - debug.input(word) - else: - yield # Enter number of words - assert "number of words" in debug.wait_layout().text_content() - debug.input(str(len(first_share))) - yield # Homescreen - proceed to share entry - assert "Enter any share" in debug.wait_layout().text_content() - debug.press_yes() - yield # Enter first share - assert debug.wait_layout().main_component() == "MnemonicKeyboard" - for index, word in enumerate(first_share): - assert f"Type word {index + 1}" in debug.wait_layout().text_content() - debug.input(word) - - -class InputFlowSlip39BasicRecoveryRetryFirst(InputFlowBase): +class InputFlowSlip39BasicRecoveryInvalidFirstShare(InputFlowBase): def __init__(self, client: Client): super().__init__(client) + self.first_invalid = ["slush"] * 20 + self.second_invalid = ["slush"] * 33 - def input_flow_common(self) -> GeneratorType: - yield from confirm_recovery(self.debug) + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() + yield from self.REC.setup_slip39_recovery(len(self.first_invalid)) + yield from self.REC.input_mnemonic(self.first_invalid) + yield from self.REC.warning_invalid_recovery_share() + yield from self.REC.setup_slip39_recovery(len(self.second_invalid)) + yield from self.REC.input_mnemonic(self.second_invalid) + yield from self.REC.warning_invalid_recovery_share() - first_share = ["slush"] * 20 - yield from slip39_recovery_setup_and_first_share(self.debug, first_share) - - br = yield # Invalid share - assert br.code == messages.ButtonRequestType.Warning - assert "Invalid recovery share" in self.text_content() - self.debug.press_yes() - - first_share = ["slush"] * 33 - yield from slip39_recovery_setup_and_first_share(self.debug, first_share) - - br = yield # Invalid share - assert br.code == messages.ButtonRequestType.Warning - assert "Invalid recovery share" in self.text_content() - self.debug.press_yes() - - yield # Homescreen - assert "number of words" in self.text_content() - if self.debug.model == "R": - self.debug.press_yes() - yield - assert "NUMBER OF WORDS" in self.title() - - # Cancelling the recovery process - # (needs to be manual to be compatible with the next input-flow) - - self.debug.input("20") - yield # Homescreen - proceed to share entry - assert "Enter any share" in self.text_content() - - self.debug.press_no() - yield # Confirm abort - assert "abort the recovery process" in self.text_content() - if self.debug.model == "R": - self.debug.press_right(wait=True) - self.debug.press_yes() + yield + self.client.cancel() -class InputFlowSlip39BasicRecoveryRetrySecond(InputFlowBase): +class InputFlowSlip39BasicRecoveryInvalidSecondShare(InputFlowBase): def __init__(self, client: Client, shares: list[str]): super().__init__(client) self.shares = shares + self.first_share = shares[0].split(" ") + self.invalid_share = self.first_share[:3] + ["slush"] * 17 + self.second_share = shares[1].split(" ") - def input_flow_common(self) -> GeneratorType: - yield from confirm_recovery(self.debug) - - # First valid share - first_share = self.shares[0].split(" ") - yield from slip39_recovery_setup_and_first_share(self.debug, first_share) - - yield # More shares needed - assert "more shares needed" in self.text_content() - self.debug.press_yes() - - yield # Enter another share - assert "MnemonicKeyboard" in self.debug.wait_layout().all_components() - invalid_share = first_share[:3] + ["slush"] * 17 - for word in invalid_share: - self.debug.input(word) - - yield # Invalid share - assert "Invalid recovery share" in self.text_content() - self.debug.press_yes() - - yield # Proceed to next share - assert "MnemonicKeyboard" in self.debug.wait_layout().all_components() - second_share = self.shares[1].split(" ") - for word in second_share: - self.debug.input(word) - - yield # More shares needed - assert "1 more share needed" in self.text_content() + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() + yield from self.REC.setup_slip39_recovery(len(self.first_share)) + yield from self.REC.input_mnemonic(self.first_share) + yield from self.REC.success_more_shares_needed(2) + yield from self.REC.input_mnemonic(self.invalid_share) + yield from self.REC.warning_invalid_recovery_share() + yield from self.REC.input_mnemonic(self.second_share) + yield from self.REC.success_more_shares_needed(1) + yield self.client.cancel() @@ -1872,58 +1527,35 @@ class InputFlowSlip39BasicRecoveryWrongNthWord(InputFlowBase): super().__init__(client) self.share = share self.nth_word = nth_word + # Invalid share - just enough words to trigger the warning + self.modified_share = share[:nth_word] + [self.share[-1]] - def input_flow_common(self) -> GeneratorType: - yield from confirm_recovery(self.debug) - - # First complete share - yield from slip39_recovery_setup_and_first_share(self.debug, self.share) - - yield # Continue to next share - assert "more shares needed" in self.text_content() - self.debug.press_yes() - yield # Enter next share - assert "MnemonicKeyboard" in self.debug.wait_layout().all_components() - for i, word in enumerate(self.share): - if i < self.nth_word: - self.debug.input(word) - else: - self.debug.input(self.share[-1]) - break - - br = yield - assert br.code == messages.ButtonRequestType.Warning - assert "entered a share from another" in self.text_content() - self.debug.press_yes() + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() + yield from self.REC.setup_slip39_recovery(len(self.share)) + yield from self.REC.input_mnemonic(self.share) + yield from self.REC.success_more_shares_needed() + yield from self.REC.input_mnemonic(self.modified_share) + yield from self.REC.warning_share_from_another_shamir() yield self.client.cancel() class InputFlowSlip39BasicRecoverySameShare(InputFlowBase): - def __init__(self, client: Client, first_share: list[str], second_share: list[str]): + def __init__(self, client: Client, share: list[str]): super().__init__(client) - self.first_share = first_share - self.second_share = second_share + self.share = share + # Second duplicate share - only 4 words are needed to verify it + self.duplicate_share = self.share[:4] - def input_flow_common(self) -> GeneratorType: - yield from confirm_recovery(self.debug) - - # First complete share - yield from slip39_recovery_setup_and_first_share(self.debug, self.first_share) - - yield # Continue to next share - assert "more shares needed" in self.text_content() - self.debug.press_yes() - yield # Enter next share - assert "MnemonicKeyboard" in self.debug.wait_layout().all_components() - for word in self.second_share: - self.debug.input(word) - - br = yield - assert br.code == messages.ButtonRequestType.Warning - assert "Share already entered" in self.text_content() - self.debug.press_yes() + def input_flow_common(self) -> BRGeneratorType: + yield from self.REC.confirm_recovery() + yield from self.REC.setup_slip39_recovery(len(self.share)) + yield from self.REC.input_mnemonic(self.share) + yield from self.REC.success_more_shares_needed() + yield from self.REC.input_mnemonic(self.duplicate_share) + yield from self.REC.warning_share_already_entered() yield self.client.cancel() @@ -1933,12 +1565,8 @@ class InputFlowResetSkipBackup(InputFlowBase): def __init__(self, client: Client): super().__init__(client) - def input_flow_common(self) -> GeneratorType: - yield # Confirm Recovery - assert "By continuing you agree to" in self.text_content() - if self.debug.model == "R": - self.debug.press_right() - self.debug.press_yes() + def input_flow_common(self) -> BRGeneratorType: + yield from self.BAK.confirm_new_wallet() yield # Skip Backup assert "New wallet created" in self.text_content() if self.debug.model == "R": diff --git a/tests/input_flows_helpers.py b/tests/input_flows_helpers.py new file mode 100644 index 000000000..d196f1c11 --- /dev/null +++ b/tests/input_flows_helpers.py @@ -0,0 +1,249 @@ +from trezorlib import messages +from trezorlib.debuglink import TrezorClientDebugLink as Client + +from .common import BRGeneratorType + +B = messages.ButtonRequestType + + +class PinFlow: + def __init__(self, client: Client): + self.client = client + self.debug = self.client.debug + + def setup_new_pin( + self, pin: str, second_different_pin: str | None = None + ) -> BRGeneratorType: + yield # Enter PIN + assert "PinKeyboard" in self.debug.wait_layout().all_components() + self.debug.input(pin) + if self.debug.model == "R": + yield # Reenter PIN + assert "re-enter PIN" in self.debug.wait_layout().text_content() + self.debug.press_yes() + yield # Enter PIN again + assert "PinKeyboard" in self.debug.wait_layout().all_components() + if second_different_pin is not None: + self.debug.input(second_different_pin) + else: + self.debug.input(pin) + + +class BackupFlow: + def __init__(self, client: Client): + self.client = client + self.debug = self.client.debug + + def confirm_new_wallet(self) -> BRGeneratorType: + yield + assert "By continuing you agree" in self.debug.wait_layout().text_content() + if self.debug.model == "R": + self.debug.press_right() + self.debug.press_yes() + + +class RecoveryFlow: + def __init__(self, client: Client): + self.client = client + self.debug = self.client.debug + + def confirm_recovery(self) -> BRGeneratorType: + yield + assert "By continuing you agree" in self.debug.wait_layout().text_content() + if self.debug.model == "R": + self.debug.press_right() + self.debug.press_yes() + + def confirm_dry_run(self) -> BRGeneratorType: + yield + assert "check the recovery seed" in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def setup_slip39_recovery(self, num_words: int) -> BRGeneratorType: + if self.debug.model == "R": + yield from self.tr_recovery_homescreen() + yield from self.input_number_of_words(num_words) + yield from self.enter_any_share() + + def setup_bip39_recovery(self, num_words: int) -> BRGeneratorType: + if self.debug.model == "R": + yield from self.tr_recovery_homescreen() + yield from self.input_number_of_words(num_words) + yield from self.enter_your_backup() + + def tr_recovery_homescreen(self) -> BRGeneratorType: + yield + assert "number of words" in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def enter_your_backup(self) -> BRGeneratorType: + yield + assert "Enter your backup" in self.debug.wait_layout().text_content() + if ( + self.debug.model == "R" + and "BACKUP CHECK" not in self.debug.wait_layout().title() + ): + # Normal recovery has extra info (not dry run) + self.debug.press_right(wait=True) + self.debug.press_right(wait=True) + self.debug.press_yes() + + def enter_any_share(self) -> BRGeneratorType: + yield + assert "Enter any share" in self.debug.wait_layout().text_content() + if ( + self.debug.model == "R" + and "BACKUP CHECK" not in self.debug.wait_layout().title() + ): + # Normal recovery has extra info (not dry run) + self.debug.press_right(wait=True) + self.debug.press_right(wait=True) + self.debug.press_yes() + + def abort_recovery(self, confirm: bool) -> BRGeneratorType: + yield + if self.debug.model == "R": + assert "number of words" in self.debug.wait_layout().text_content() + else: + assert "Enter any share" in self.debug.wait_layout().text_content() + self.debug.press_no() + + yield + assert "abort the recovery" in self.debug.wait_layout().text_content() + if self.debug.model == "R": + self.debug.press_right() + if confirm: + self.debug.press_yes() + else: + self.debug.press_no() + + def input_number_of_words(self, num_words: int) -> BRGeneratorType: + br = yield + assert br.code == B.MnemonicWordCount + if self.debug.model == "R": + assert "NUMBER OF WORDS" in self.debug.wait_layout().title() + else: + assert "number of words" in self.debug.wait_layout().text_content() + self.debug.input(str(num_words)) + + def warning_invalid_recovery_seed(self) -> BRGeneratorType: + br = yield + assert br.code == B.Warning + assert "Invalid recovery seed" in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def warning_invalid_recovery_share(self) -> BRGeneratorType: + br = yield + assert br.code == B.Warning + assert "Invalid recovery share" in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def warning_group_threshold_reached(self) -> BRGeneratorType: + br = yield + assert br.code == B.Warning + assert "Group threshold reached" in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def warning_share_already_entered(self) -> BRGeneratorType: + br = yield + assert br.code == B.Warning + assert "Share already entered" in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def warning_share_from_another_shamir(self) -> BRGeneratorType: + br = yield + assert br.code == B.Warning + assert ( + "You have entered a share from another Shamir Backup" + in self.debug.wait_layout().text_content() + ) + self.debug.press_yes() + + def success_share_group_entered(self) -> BRGeneratorType: + yield + assert "You have entered" in self.debug.wait_layout().text_content() + assert "Group" in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def success_wallet_recovered(self) -> BRGeneratorType: + br = yield + assert br.code == B.Success + assert ( + "Wallet recovered successfully" in self.debug.wait_layout().text_content() + ) + self.debug.press_yes() + + def success_bip39_dry_run_valid(self) -> BRGeneratorType: + br = yield + assert br.code == B.Success + assert "recovery seed is valid" in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def success_slip39_dryrun_valid(self) -> BRGeneratorType: + br = yield + assert br.code == B.Success + assert "recovery shares are valid" in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def warning_slip39_dryrun_mismatch(self) -> BRGeneratorType: + br = yield + assert br.code == B.Warning + assert "do not match" in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def warning_bip39_dryrun_mismatch(self) -> BRGeneratorType: + br = yield + assert br.code == B.Warning + assert "does not match" in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def success_more_shares_needed( + self, count_needed: int | None = None + ) -> BRGeneratorType: + yield + assert ( + "1 more share needed" in self.debug.wait_layout().text_content().lower() + or "more shares needed" in self.debug.wait_layout().text_content().lower() + ) + if count_needed is not None: + assert str(count_needed) in self.debug.wait_layout().text_content() + self.debug.press_yes() + + def input_mnemonic(self, mnemonic: list[str]) -> BRGeneratorType: + br = yield + assert br.code == B.MnemonicInput + assert "MnemonicKeyboard" in self.debug.wait_layout().all_components() + for index, word in enumerate(mnemonic): + if self.debug.model == "R": + assert f"WORD {index + 1}" in self.debug.wait_layout().title() + else: + assert ( + f"Type word {index + 1}" in self.debug.wait_layout().text_content() + ) + self.debug.input(word) + + def input_all_slip39_shares( + self, + shares: list[str], + has_groups: bool = False, + click_info: bool = False, + ) -> BRGeneratorType: + for index, share in enumerate(shares): + mnemonic = share.split(" ") + yield from self.input_mnemonic(mnemonic) + + if index < len(shares) - 1: + if has_groups: + yield from self.success_share_group_entered() + if self.debug.model == "T" and click_info: + yield from self.tt_click_info() + yield from self.success_more_shares_needed() + + def tt_click_info( + self, + ) -> BRGeneratorType: + # Moving through the INFO button + self.debug.press_info() + yield + self.debug.swipe_up() + self.debug.press_yes()