From e5b6d43a34c523d8c97d49d52b973d27671a572c Mon Sep 17 00:00:00 2001 From: grdddj Date: Mon, 16 Jan 2023 15:09:36 +0100 Subject: [PATCH] chore(tests): centralize all model-dependent input flows [no changelog] --- tests/common.py | 108 +- .../bitcoin/test_getaddress_show.py | 54 +- .../device_tests/bitcoin/test_signmessage.py | 68 +- tests/device_tests/bitcoin/test_signtx.py | 114 +- .../bitcoin/test_signtx_payreq.py | 31 +- .../ethereum/test_sign_typed_data.py | 63 +- tests/device_tests/ethereum/test_signtx.py | 98 +- .../test_recovery_bip39_dryrun.py | 182 +- .../reset_recovery/test_recovery_bip39_t2.py | 157 +- .../test_recovery_slip39_advanced.py | 218 +-- .../test_recovery_slip39_advanced_dryrun.py | 59 +- .../test_recovery_slip39_basic.py | 440 +---- .../test_recovery_slip39_basic_dryrun.py | 42 +- .../reset_recovery/test_reset_backup.py | 201 +- .../reset_recovery/test_reset_bip39_t2.py | 149 +- .../test_reset_recovery_bip39.py | 110 +- .../test_reset_recovery_slip39_advanced.py | 117 +- .../test_reset_recovery_slip39_basic.py | 105 +- .../test_reset_slip39_advanced.py | 93 +- .../reset_recovery/test_reset_slip39_basic.py | 84 +- tests/device_tests/test_msg_backup_device.py | 191 +- .../test_msg_change_wipe_code_t2.py | 18 +- tests/device_tests/test_msg_changepin_t2.py | 54 +- tests/device_tests/test_pin.py | 25 +- tests/input_flows.py | 1738 +++++++++++++++++ tests/persistence_tests/test_wipe_code.py | 17 +- 26 files changed, 2070 insertions(+), 2466 deletions(-) create mode 100644 tests/input_flows.py diff --git a/tests/common.py b/tests/common.py index 5281b13f39..223aaaed3a 100644 --- a/tests/common.py +++ b/tests/common.py @@ -15,8 +15,9 @@ # If not, see . import json +import time from pathlib import Path -from typing import TYPE_CHECKING, Generator, List, Optional +from typing import TYPE_CHECKING, Generator, Optional import pytest @@ -24,7 +25,7 @@ from trezorlib import btc, tools from trezorlib.messages import ButtonRequestType if TYPE_CHECKING: - from trezorlib.debuglink import LayoutLines + from trezorlib.debuglink import LayoutContent from trezorlib.debuglink import DebugLink, TrezorClientDebugLink as Client from trezorlib.messages import ButtonRequest from _pytest.mark.structures import MarkDecorator @@ -125,7 +126,23 @@ def generate_entropy( def recovery_enter_shares( debug: "DebugLink", - shares: List[str], + 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]: @@ -181,7 +198,7 @@ def recovery_enter_shares( def recovery_enter_shares_tr( debug: "DebugLink", - shares: List[str], + shares: list[str], groups: bool = False, ) -> Generator[None, "ButtonRequest", None]: """Perform the recovery flow for a set of Shamir shares. @@ -236,7 +253,7 @@ def recovery_enter_shares_tr( def click_through( - debug: "DebugLink", screens: int, code: ButtonRequestType = None + debug: "DebugLink", screens: int, code: Optional[ButtonRequestType] = None ) -> Generator[None, "ButtonRequest", None]: """Click through N dialog screens. @@ -259,6 +276,19 @@ def click_through( def read_and_confirm_mnemonic( debug: "DebugLink", choose_wrong: bool = False +) -> Generator[None, "ButtonRequest", Optional[str]]: + if debug.model == "T": + mnemonic = yield from read_and_confirm_mnemonic_tt(debug, choose_wrong) + elif debug.model == "R": + mnemonic = yield from read_and_confirm_mnemonic_tr(debug, choose_wrong) + else: + raise ValueError(f"Unknown model: {debug.model}") + + return mnemonic + + +def read_and_confirm_mnemonic_tt( + debug: "DebugLink", choose_wrong: bool = False ) -> Generator[None, "ButtonRequest", Optional[str]]: """Read a given number of mnemonic words from Trezor T screen and correctly answer confirmation questions. Return the full mnemonic. @@ -271,7 +301,7 @@ def read_and_confirm_mnemonic( mnemonic = yield from read_and_confirm_mnemonic(client.debug) """ - mnemonic: List[str] = [] + mnemonic: list[str] = [] br = yield assert br.pages is not None for _ in range(br.pages - 1): @@ -294,19 +324,46 @@ def read_and_confirm_mnemonic( return " ".join(mnemonic) +def read_and_confirm_mnemonic_tr( + debug: "DebugLink", choose_wrong: bool = False +) -> Generator[None, "ButtonRequest", Optional[str]]: + mnemonic: list[str] = [] + br = yield + assert br.pages is not None + for _ in range(br.pages): + layout = debug.wait_layout() + + words = ModelRLayout(layout).get_mnemonic_words() + mnemonic.extend(words) + debug.press_right() + + yield # Select correct words... + debug.press_right() + + for _ in range(3): + index = debug.read_reset_word_pos() + if choose_wrong: + debug.input(mnemonic[(index + 1) % len(mnemonic)]) + return None + else: + debug.input(mnemonic[index]) + + return " ".join(mnemonic) + + class ModelRLayout: """Layout shortcuts for Model R.""" - def __init__(self, layout: "LayoutLines") -> None: + def __init__(self, layout: "LayoutContent") -> None: self.layout = layout - def get_mnemonic_words(self) -> List[str]: + def get_mnemonic_words(self) -> list[str]: """Extract mnemonic words from the layout lines. Example input: [..., '4 must', '5 during', '6 monitor', ...] Example output: ['must', 'during', 'monitor'] """ - words: List[str] = [] + words: list[str] = [] for line in self.layout.lines: if " " in line: number, word = line.split(" ", 1) @@ -335,31 +392,18 @@ class ModelRLayout: return buttons.split("[Select(")[1].split(")]")[0] -def read_and_confirm_mnemonic_tr( - debug: "DebugLink", choose_wrong: bool = False -) -> Generator[None, "ButtonRequest", Optional[str]]: - mnemonic: List[str] = [] - br = yield - assert br.pages is not None - for _ in range(br.pages): - layout = debug.wait_layout() +def click_info_button(debug: "DebugLink"): + """Click Shamir backup info button and return back.""" + debug.press_info() + yield # Info screen with text + debug.press_yes() - words = ModelRLayout(layout).get_mnemonic_words() - mnemonic.extend(words) - debug.press_right() - yield # Select correct words... - debug.press_right() - - for _ in range(3): - index = debug.read_reset_word_pos() - if choose_wrong: - debug.input(mnemonic[(index + 1) % len(mnemonic)]) - return None - else: - debug.input(mnemonic[index]) - - return " ".join(mnemonic) +def check_PIN_backoff_time(attempts: int, start: float) -> None: + """Helper to assert the exponentially growing delay after incorrect PIN attempts""" + expected = (2**attempts) - 1 + got = round(time.time() - start, 2) + assert got >= expected def get_test_address(client: "Client") -> str: diff --git a/tests/device_tests/bitcoin/test_getaddress_show.py b/tests/device_tests/bitcoin/test_getaddress_show.py index 5bec93560e..f8e435a6f1 100644 --- a/tests/device_tests/bitcoin/test_getaddress_show.py +++ b/tests/device_tests/bitcoin/test_getaddress_show.py @@ -20,6 +20,8 @@ from trezorlib import btc, messages, tools from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.exceptions import TrezorFailure +from ...input_flows import InputFlowShowAddressQRCode, InputFlowShowMultisigXPUBs + VECTORS = ( # path, script_type, address ( "m/44h/0h/12h/0/0", @@ -48,22 +50,19 @@ VECTORS = ( # path, script_type, address def test_show( client: Client, path: str, script_type: messages.InputScriptType, address: str ): - def input_flow_tt(): + def input_flow_t1(): yield client.debug.press_no() yield client.debug.press_yes() - def input_flow_tr(): - yield - client.debug.press_right() - client.debug.press_yes() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + # This is the only place where even T1 is using input flow + if client.features.model == "1": + client.set_input_flow(input_flow_t1) + else: + IF = InputFlowShowAddressQRCode(client) + client.set_input_flow(IF.get()) assert ( btc.get_address( client, @@ -220,41 +219,10 @@ def test_show_multisig_xpubs( ) for i in range(3): - - def input_flow(): - yield # show address - layout = client.debug.wait_layout() # TODO: do not need to *wait* now? - assert layout.get_title() == "MULTISIG 2 OF 3" - assert layout.get_content().replace(" ", "") == address - - client.debug.press_no() - yield # show QR code - assert "Painter" in client.debug.wait_layout().text - - # Three xpub pages with the same testing logic - for xpub_num in range(3): - expected_title = f"XPUB #{xpub_num + 1} " + ( - "(yours)" if i == xpub_num else "(cosigner)" - ) - - client.debug.press_no() - yield # show XPUB#{xpub_num} - layout1 = client.debug.wait_layout() - assert layout1.get_title() == expected_title - client.debug.swipe_up() - - layout2 = client.debug.wait_layout() - assert layout2.get_title() == expected_title - content = (layout1.get_content() + layout2.get_content()).replace( - " ", "" - ) - assert content == xpubs[xpub_num] - - client.debug.press_yes() - with client: client.watch_layout() - client.set_input_flow(input_flow) + IF = InputFlowShowMultisigXPUBs(client, address, xpubs, i) + client.set_input_flow(IF.get()) btc.get_address( client, "Bitcoin", diff --git a/tests/device_tests/bitcoin/test_signmessage.py b/tests/device_tests/bitcoin/test_signmessage.py index 6aa811e603..4fcc12622d 100644 --- a/tests/device_tests/bitcoin/test_signmessage.py +++ b/tests/device_tests/bitcoin/test_signmessage.py @@ -14,22 +14,20 @@ # You should have received a copy of the License along with this library. # If not, see . +from typing import Any import pytest from trezorlib import btc, messages -from trezorlib.debuglink import ( - LayoutContent, - TrezorClientDebugLink as Client, - message_filters, - multipage_content, -) +from trezorlib.debuglink import TrezorClientDebugLink as Client, message_filters from trezorlib.tools import parse_path +from ...input_flows import InputFlowSignMessagePagination + S = messages.InputScriptType -def case(id, *args, altcoin=False, skip_t1=False): +def case(id: str, *args: Any, altcoin: bool = False, skip_t1: bool = False): marks = [] if altcoin: marks.append(pytest.mark.altcoin) @@ -273,7 +271,14 @@ VECTORS = ( # case name, coin_name, path, script_type, address, message, signat "coin_name, path, script_type, no_script_type, address, message, signature", VECTORS ) def test_signmessage( - client, coin_name, path, script_type, no_script_type, address, message, signature + client: Client, + coin_name: str, + path: str, + script_type: messages.InputScriptType, + no_script_type: bool, + address: str, + message: str, + signature: str, ): sig = btc.sign_message( client, @@ -301,50 +306,9 @@ MESSAGE_LENGTHS = ( @pytest.mark.skip_t1 @pytest.mark.parametrize("message", MESSAGE_LENGTHS) def test_signmessage_pagination(client: Client, message: str): - message_read = "" - - def input_flow_tt(): - # collect screen contents into `message_read`. - # Using a helper debuglink function to assemble the final text. - nonlocal message_read - layouts: list[LayoutContent] = [] - - # confirm address - br = yield - client.debug.wait_layout() - client.debug.press_yes() - - br = yield - for i in range(br.pages): - layout = client.debug.wait_layout() - layouts.append(layout) - - if i < br.pages - 1: - client.debug.swipe_up() - - message_read = multipage_content(layouts) - - client.debug.press_yes() - - def input_flow_tr(): - # confirm address - yield - client.debug.press_yes() - - br = yield - # TODO: try load the message_read the same way as in model T - if br.pages is not None: - for i in range(br.pages): - if i < br.pages - 1: - client.debug.swipe_up() - client.debug.press_yes() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) - client.debug.watch_layout(True) + IF = InputFlowSignMessagePagination(client) + client.set_input_flow(IF.get()) btc.sign_message( client, coin_name="Bitcoin", @@ -358,7 +322,7 @@ def test_signmessage_pagination(client: Client, message: str): expected_message = ( ("Confirm message: " + message).replace("\n", "").replace(" ", "") ) - message_read = message_read.replace(" ", "") + message_read = IF.message_read.replace(" ", "") assert expected_message == message_read diff --git a/tests/device_tests/bitcoin/test_signtx.py b/tests/device_tests/bitcoin/test_signtx.py index 8ecf8f90c6..e3faaaec8b 100644 --- a/tests/device_tests/bitcoin/test_signtx.py +++ b/tests/device_tests/bitcoin/test_signtx.py @@ -23,6 +23,11 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.exceptions import TrezorFailure from trezorlib.tools import H_, parse_path +from ...input_flows import ( + InputFlowLockTimeBlockHeight, + InputFlowLockTimeDatetime, + InputFlowSignTxHighFee, +) from ...tx_cache import TxCache from .signtx import ( assert_tx_matches, @@ -657,33 +662,13 @@ def test_fee_high_hardfail(client: Client): client, safety_checks=messages.SafetyCheckLevel.PromptTemporarily ) with client: - tt = client.features.model == "T" - finished = False - - def input_flow(): - nonlocal finished - for expected in ( - B.ConfirmOutput, - (tt, B.ConfirmOutput), - B.FeeOverThreshold, - B.SignTx, - (tt, B.SignTx), - ): - if isinstance(expected, tuple): - is_valid, expected = expected - if not is_valid: - continue - br = yield - assert br.code == expected - client.debug.press_yes() - finished = True - - client.set_input_flow(input_flow) + IF = InputFlowSignTxHighFee(client) + client.set_input_flow(IF.get()) _, serialized_tx = btc.sign_tx( client, "Testnet", [inp1], [out1], prev_txes=TX_CACHE_TESTNET ) - assert finished + assert IF.finished # Transaction does not exist on the blockchain, not using assert_tx_matches() assert ( @@ -1486,47 +1471,9 @@ def test_lock_time_blockheight(client: Client): script_type=messages.OutputScriptType.PAYTOADDRESS, ) - def input_flow_tt(): - yield # confirm output - client.debug.wait_layout() - client.debug.press_yes() - yield # confirm output - client.debug.wait_layout() - client.debug.press_yes() - - yield # confirm locktime - layout = client.debug.wait_layout() - assert "blockheight" in layout.text - assert "499999999" in layout.text - client.debug.press_yes() - - yield # confirm transaction - client.debug.press_yes() - yield # confirm transaction - client.debug.press_yes() - - def input_flow_tr(): - yield # confirm output - client.debug.wait_layout() - client.debug.swipe_up() - client.debug.wait_layout() - client.debug.press_yes() - - yield # confirm locktime - layout = client.debug.wait_layout() - assert "blockheight" in layout.text - assert "499999999" in layout.text - client.debug.press_yes() - - yield # confirm transaction - client.debug.press_yes() - with client: - if client.debug.model == "T": - client.set_input_flow(input_flow_tt) - elif client.debug.model == "R": - client.set_input_flow(input_flow_tr) - client.watch_layout(True) + IF = InputFlowLockTimeBlockHeight(client, "499999999") + client.set_input_flow(IF.get()) btc.sign_tx( client, @@ -1559,50 +1506,13 @@ def test_lock_time_datetime(client: Client, lock_time_str: str): script_type=messages.OutputScriptType.PAYTOADDRESS, ) - def input_flow_tt(): - yield # confirm output - client.debug.wait_layout() - client.debug.press_yes() - yield # confirm output - client.debug.wait_layout() - client.debug.press_yes() - - yield # confirm locktime - layout = client.debug.wait_layout() - assert lock_time_str in layout.text - - client.debug.press_yes() - - yield # confirm transaction - client.debug.press_yes() - yield # confirm transaction - client.debug.press_yes() - - def input_flow_tr(): - yield # confirm output - client.debug.wait_layout() - client.debug.swipe_up() - client.debug.wait_layout() - client.debug.press_yes() - - yield # confirm locktime - layout = client.debug.wait_layout() - assert lock_time_str in layout.text - client.debug.press_yes() - - yield # confirm transaction - client.debug.press_yes() - lock_time_naive = datetime.strptime(lock_time_str, "%Y-%m-%d %H:%M:%S") lock_time_utc = lock_time_naive.replace(tzinfo=timezone.utc) lock_time_timestamp = int(lock_time_utc.timestamp()) with client: - if client.debug.model == "T": - client.set_input_flow(input_flow_tt) - elif client.debug.model == "R": - client.set_input_flow(input_flow_tr) - client.watch_layout(True) + IF = InputFlowLockTimeDatetime(client, lock_time_str) + client.set_input_flow(IF.get()) btc.sign_tx( client, diff --git a/tests/device_tests/bitcoin/test_signtx_payreq.py b/tests/device_tests/bitcoin/test_signtx_payreq.py index df958e0934..d199c7adee 100644 --- a/tests/device_tests/bitcoin/test_signtx_payreq.py +++ b/tests/device_tests/bitcoin/test_signtx_payreq.py @@ -23,6 +23,7 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.exceptions import TrezorFailure from trezorlib.tools import parse_path +from ...input_flows import InputFlowPaymentRequestDetails from .payment_req import CoinPurchaseMemo, RefundMemo, TextMemo, make_payment_request from .signtx import forge_prevtx @@ -194,35 +195,9 @@ def test_payment_request_details(client: Client): ) ] - def input_flow(): - yield # request to see details - client.debug.wait_layout() - client.debug.press_info() - - yield # confirm first output - layout = client.debug.wait_layout() - assert outputs[0].address[:16] in layout.text - client.debug.press_yes() - yield # confirm first output - client.debug.wait_layout() - client.debug.press_yes() - - yield # confirm second output - layout = client.debug.wait_layout() - assert outputs[1].address[:16] in layout.text - client.debug.press_yes() - yield # confirm second output - client.debug.wait_layout() - client.debug.press_yes() - - yield # confirm transaction - client.debug.press_yes() - yield # confirm transaction - client.debug.press_yes() - with client: - client.set_input_flow(input_flow) - client.watch_layout(True) + IF = InputFlowPaymentRequestDetails(client, outputs) + client.set_input_flow(IF.get()) _, serialized_tx = btc.sign_tx( client, diff --git a/tests/device_tests/ethereum/test_sign_typed_data.py b/tests/device_tests/ethereum/test_sign_typed_data.py index 5679a3c7c5..3dfa60e9f6 100644 --- a/tests/device_tests/ethereum/test_sign_typed_data.py +++ b/tests/device_tests/ethereum/test_sign_typed_data.py @@ -21,8 +21,7 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.tools import parse_path from ...common import parametrize_using_common_fixtures - -SHOW_MORE = (143, 167) +from ...input_flows import InputFlowEIP712Cancel, InputFlowEIP712ShowMore pytestmark = [pytest.mark.altcoin, pytest.mark.ethereum] @@ -95,65 +94,12 @@ DATA = { } -def _confirm_show_more(client: Client) -> None: - """Model-specific, either clicks a screen or presses a button.""" - if client.features.model == "T": - client.debug.click(SHOW_MORE) - elif client.features.model == "R": - client.debug.press_yes() - - -def input_flow_show_more(client: Client): - """Clicks show_more button wherever possible""" - yield # confirm domain - client.debug.wait_layout() - _confirm_show_more(client) - - # confirm domain properties - for _ in range(4): - yield - client.debug.press_yes() - - yield # confirm message - client.debug.wait_layout() - _confirm_show_more(client) - - yield # confirm message.from - client.debug.wait_layout() - _confirm_show_more(client) - - # confirm message.from properties - for _ in range(2): - yield - client.debug.press_yes() - - yield # confirm message.to - client.debug.wait_layout() - _confirm_show_more(client) - - # confirm message.to properties - for _ in range(2): - yield - client.debug.press_yes() - - yield # confirm message.contents - client.debug.press_yes() - - yield # confirm final hash - client.debug.press_yes() - - -def input_flow_cancel(client: Client): - """Clicks cancelling button""" - yield # confirm domain - client.debug.press_no() - - @pytest.mark.skip_t1 def test_ethereum_sign_typed_data_show_more_button(client: Client): with client: client.watch_layout() - client.set_input_flow(input_flow_show_more(client)) + IF = InputFlowEIP712ShowMore(client) + client.set_input_flow(IF.get()) ethereum.sign_typed_data( client, parse_path("m/44h/60h/0h/0/0"), @@ -169,7 +115,8 @@ def test_ethereum_sign_typed_data_cancel(client: Client): with client, pytest.raises(exceptions.Cancelled): client.watch_layout() - client.set_input_flow(input_flow_cancel(client)) + IF = InputFlowEIP712Cancel(client) + client.set_input_flow(IF.get()) ethereum.sign_typed_data( client, parse_path("m/44h/60h/0h/0/0"), diff --git a/tests/device_tests/ethereum/test_signtx.py b/tests/device_tests/ethereum/test_signtx.py index ec2e520ed7..c023a7a2e9 100644 --- a/tests/device_tests/ethereum/test_signtx.py +++ b/tests/device_tests/ethereum/test_signtx.py @@ -22,10 +22,14 @@ from trezorlib.exceptions import TrezorFailure from trezorlib.tools import parse_path from ...common import parametrize_using_common_fixtures +from ...input_flows import ( + InputFlowEthereumSignTxGoBack, + InputFlowEthereumSignTxScrollDown, + InputFlowEthereumSignTxSkip, +) TO_ADDR = "0x1d1c328764a41bda0492b66baa30c4a339ff85ef" -SHOW_ALL = (143, 167) -GO_BACK = (16, 220) + pytestmark = [pytest.mark.altcoin, pytest.mark.ethereum] @@ -329,103 +333,19 @@ def test_sanity_checks_eip1559(client: Client): def input_flow_skip(client: Client, cancel: bool = False): - yield # confirm address - client.debug.press_yes() - yield # confirm amount - client.debug.wait_layout() - client.debug.press_yes() - yield # confirm data - if cancel: - client.debug.press_no() - else: - client.debug.press_yes() - yield # gas price - client.debug.press_yes() - yield # maximum fee - client.debug.press_yes() - yield # hold to confirm - client.debug.press_yes() + return InputFlowEthereumSignTxSkip(client, cancel).get() def input_flow_scroll_down(client: Client, cancel: bool = False): if client.features.model == "R": pytest.skip("Freezes") - - yield # confirm address - client.debug.wait_layout() - client.debug.press_yes() - yield # confirm amount - client.debug.wait_layout() - if client.features.model == "R": - client.debug.swipe_up() # extra confirmation screen for model R - client.debug.press_yes() - - # TODO: For model R, we do not have the "Show all" button, so users - # do not have the possibility to skip scrolling. - if client.features.model == "T": - yield # confirm data - client.debug.wait_layout() - client.debug.click(SHOW_ALL) - - br = yield # paginated data - for i in range(br.pages): - client.debug.wait_layout() - if i < br.pages - 1: - client.debug.swipe_up() - - client.debug.press_yes() - yield # confirm data - if cancel: - client.debug.press_no() - else: - client.debug.press_yes() - yield # gas price - client.debug.press_yes() - yield # maximum fee - client.debug.press_yes() - yield # hold to confirm - client.debug.press_yes() + return InputFlowEthereumSignTxScrollDown(client, cancel).get() def input_flow_go_back(client: Client, cancel: bool = False): if client.features.model == "R": pytest.skip("Go back not supported for model R") - - br = yield # confirm address - client.debug.wait_layout() - client.debug.press_yes() - br = yield # confirm amount - - client.debug.wait_layout() - client.debug.press_yes() - br = yield # confirm data - client.debug.wait_layout() - client.debug.click(SHOW_ALL) - - br = yield # paginated data - for i in range(br.pages): - client.debug.wait_layout() - if i == 2: - client.debug.click(GO_BACK) - yield # confirm data - client.debug.wait_layout() - if cancel: - client.debug.press_no() - else: - client.debug.press_yes() - yield # confirm address - client.debug.wait_layout() - client.debug.press_yes() - yield # confirm amount - client.debug.wait_layout() - client.debug.press_yes() - yield # hold to confirm - client.debug.wait_layout() - client.debug.press_yes() - return - - elif i < br.pages - 1: - client.debug.swipe_up() + return InputFlowEthereumSignTxGoBack(client, cancel).get() HEXDATA = "0123456789abcd000023456789abcd010003456789abcd020000456789abcd030000056789abcd040000006789abcd050000000789abcd060000000089abcd070000000009abcd080000000000abcd090000000001abcd0a0000000011abcd0b0000000111abcd0c0000001111abcd0d0000011111abcd0e0000111111abcd0f0000000002abcd100000000022abcd110000000222abcd120000002222abcd130000022222abcd140000222222abcd15" 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 a57c6d2fad..5006dda51a 100644 --- a/tests/device_tests/reset_recovery/test_recovery_bip39_dryrun.py +++ b/tests/device_tests/reset_recovery/test_recovery_bip39_dryrun.py @@ -14,18 +14,21 @@ # You should have received a copy of the License along with this library. # If not, see . -from typing import Any, List +from typing import Any import pytest from trezorlib import device, exceptions, messages from trezorlib.debuglink import TrezorClientDebugLink as Client -from ... import buttons from ...common import MNEMONIC12 +from ...input_flows import ( + InputFlowBip39RecoveryDryRun, + InputFlowBip39RecoveryDryRunInvalid, +) -def do_recover_legacy(client: Client, mnemonic: List[str], **kwargs: Any): +def do_recover_legacy(client: Client, mnemonic: list[str], **kwargs: Any): def input_callback(_): word, pos = client.debug.read_recovery_word() if pos != 0 and pos is not None: @@ -49,95 +52,17 @@ def do_recover_legacy(client: Client, mnemonic: List[str], **kwargs: Any): return ret -def do_recover_core(client: Client, mnemonic: List[str], **kwargs: Any): - layout = client.debug.wait_layout - - def input_flow(): - yield - assert "check the recovery seed" in layout().get_content() - client.debug.click(buttons.OK) - - yield - assert "select the number of words" in layout().get_content() - client.debug.click(buttons.OK) - - yield - assert "SelectWordCount" in layout().text - # click the number - word_option_offset = 6 - word_options = (12, 18, 20, 24, 33) - index = word_option_offset + word_options.index(len(mnemonic)) - client.debug.click(buttons.grid34(index % 3, index // 3)) - - yield - assert "enter your recovery seed" in layout().get_content() - client.debug.click(buttons.OK) - - yield - for word in mnemonic: - client.debug.wait_layout() - client.debug.input(word) - - yield - client.debug.wait_layout() - client.debug.click(buttons.OK) - +def do_recover_core(client: Client, mnemonic: list[str], **kwargs: Any): with client: client.watch_layout() - client.set_input_flow(input_flow) + IF = InputFlowBip39RecoveryDryRun(client, mnemonic) + client.set_input_flow(IF.get()) return device.recover(client, dry_run=True, show_tutorial=False, **kwargs) -def do_recover_r(client: Client, mnemonic: List[str], **kwargs: Any): - layout = client.debug.wait_layout - - def input_flow(): - yield - assert "check the recovery seed" in layout().text - client.debug.press_right() - - yield - assert "select the number of words" in layout().text - client.debug.press_yes() - - yield - yield - assert "NUMBER OF WORDS" in layout().text - word_options = (12, 18, 20, 24, 33) - index = word_options.index(len(mnemonic)) - for _ in range(index): - client.debug.press_right() - client.debug.input(str(len(mnemonic))) - - yield - assert "enter your recovery seed" in layout().text - client.debug.press_yes() - - yield - client.debug.press_right() - yield - for word in mnemonic: - assert "WORD" in layout().text - client.debug.input(word) - - client.debug.wait_layout() - client.debug.press_right() - yield - client.debug.press_yes() - yield - - with client: - client.watch_layout() - client.set_input_flow(input_flow) - return device.recover(client, dry_run=True, show_tutorial=False, **kwargs) - - -def do_recover(client: Client, mnemonic: List[str]): - +def do_recover(client: Client, mnemonic: list[str]): if client.features.model == "1": return do_recover_legacy(client, mnemonic) - elif client.features.model == "R": - return do_recover_r(client, mnemonic) else: return do_recover_core(client, mnemonic) @@ -165,93 +90,10 @@ def test_invalid_seed_t1(client: Client): @pytest.mark.skip_t1 def test_invalid_seed_core(client: Client): - layout = client.debug.wait_layout - - def input_flow_tt(): - yield - assert "check the recovery seed" in layout().get_content() - client.debug.click(buttons.OK) - - yield - assert "select the number of words" in layout().get_content() - client.debug.click(buttons.OK) - - yield - assert "SelectWordCount" in layout().text - # select 12 words - client.debug.click(buttons.grid34(0, 2)) - - yield - assert "enter your recovery seed" in layout().get_content() - client.debug.click(buttons.OK) - - yield - for _ in range(12): - assert layout().text == "< MnemonicKeyboard >" - client.debug.input("stick") - - br = yield - assert br.code == messages.ButtonRequestType.Warning - assert "invalid recovery seed" in layout().get_content() - client.debug.click(buttons.OK) - - yield - # retry screen - assert "select the number of words" in layout().get_content() - client.debug.click(buttons.CANCEL) - - yield - assert "ABORT SEED CHECK" == layout().get_title() - client.debug.click(buttons.OK) - - def input_flow_tr(): - yield - assert "check the recovery seed" in layout().text - client.debug.press_right() - - yield - assert "select the number of words" in layout().text - client.debug.press_yes() - - yield - yield - assert "NUMBER OF WORDS" in layout().text - # select 12 words - client.debug.press_middle() - - yield - assert "enter your recovery seed" in layout().text - client.debug.press_yes() - - yield - assert "WORD ENTERING" in layout().text - client.debug.press_yes() - - yield - for _ in range(12): - assert "WORD" in layout().text - client.debug.input("stick") - - br = yield - assert br.code == messages.ButtonRequestType.Warning - assert "invalid recovery seed" in layout().text - client.debug.press_right() - - yield - # retry screen - assert "select the number of words" in layout().text - client.debug.press_left() - - yield - assert "abort" in layout().text - client.debug.press_right() - with client: client.watch_layout() - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowBip39RecoveryDryRunInvalid(client) + client.set_input_flow(IF.get()) with pytest.raises(exceptions.Cancelled): return device.recover(client, dry_run=True, show_tutorial=False) 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 00f8c63aa3..b05b162bdb 100644 --- a/tests/device_tests/reset_recovery/test_recovery_bip39_t2.py +++ b/tests/device_tests/reset_recovery/test_recovery_bip39_t2.py @@ -20,99 +20,16 @@ from trezorlib import device, exceptions, messages from trezorlib.debuglink import TrezorClientDebugLink as Client from ...common import MNEMONIC12 +from ...input_flows import InputFlowBip39RecoveryNoPIN, InputFlowBip39RecoveryPIN pytestmark = pytest.mark.skip_t1 @pytest.mark.setup_client(uninitialized=True) def test_tt_pin_passphrase(client: Client): - layout = client.debug.wait_layout - mnemonic = MNEMONIC12.split(" ") - - def input_flow_tt(): - yield - assert "Do you really want to recover a wallet?" in layout().get_content() - client.debug.press_yes() - - yield - assert layout().text == "< PinKeyboard >" - client.debug.input("654") - - yield - assert layout().text == "< PinKeyboard >" - client.debug.input("654") - - yield - assert "select the number of words" in layout().get_content() - client.debug.press_yes() - - yield - assert "SelectWordCount" in layout().text - client.debug.input(str(len(mnemonic))) - - yield - assert "enter your recovery seed" in layout().get_content() - client.debug.press_yes() - - yield - for word in mnemonic: - assert layout().text == "< MnemonicKeyboard >" - client.debug.input(word) - - yield - assert "You have successfully recovered your wallet." in layout().get_content() - client.debug.press_yes() - - def input_flow_tr(): - yield - assert "By continuing you agree" in layout().text - client.debug.press_right() - assert "trezor.io/tos" in layout().text - client.debug.press_yes() - - yield - assert "ENTER" in layout().text - client.debug.input("654") - - yield - assert "re-enter PIN to confirm" in layout().text - client.debug.press_right() - - yield - assert "ENTER" in layout().text - client.debug.input("654") - - yield - assert "select the number of words" in layout().text - client.debug.press_yes() - - yield - yield - assert "NUMBER OF WORDS" in layout().text - client.debug.input(str(len(mnemonic))) - - yield - assert "enter your recovery seed" in layout().text - client.debug.press_yes() - - yield - assert "WORD ENTERING" in layout().text - client.debug.press_right() - - yield - for word in mnemonic: - assert "WORD" in layout().text - client.debug.input(word) - - yield - assert "You have finished recovering your wallet." in layout().text - client.debug.press_yes() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowBip39RecoveryPIN(client, MNEMONIC12.split(" ")) + client.set_input_flow(IF.get()) client.watch_layout() device.recover( client, @@ -132,73 +49,9 @@ def test_tt_pin_passphrase(client: Client): @pytest.mark.setup_client(uninitialized=True) def test_tt_nopin_nopassphrase(client: Client): - layout = client.debug.wait_layout - mnemonic = MNEMONIC12.split(" ") - - def input_flow_tt(): - yield - assert "Do you really want to recover a wallet?" in layout().get_content() - client.debug.press_yes() - - yield - assert "select the number of words" in layout().get_content() - client.debug.press_yes() - - yield - assert "SelectWordCount" in layout().text - client.debug.input(str(len(mnemonic))) - - yield - assert "enter your recovery seed" in layout().get_content() - client.debug.press_yes() - - yield - for word in mnemonic: - assert layout().text == "< MnemonicKeyboard >" - client.debug.input(word) - - yield - assert "You have successfully recovered your wallet." in layout().get_content() - client.debug.press_yes() - - def input_flow_tr(): - yield - assert "By continuing you agree" in layout().text - client.debug.press_right() - assert "trezor.io/tos" in layout().text - client.debug.press_yes() - - yield - assert "select the number of words" in layout().text - client.debug.press_yes() - - yield - yield - assert "NUMBER OF WORDS" in layout().text - client.debug.input(str(len(mnemonic))) - - yield - assert "enter your recovery seed" in layout().text - client.debug.press_yes() - - yield - assert "WORD ENTERING" in layout().text - client.debug.press_right() - - yield - for word in mnemonic: - assert "WORD" in layout().text - client.debug.input(word) - - yield - assert "You have finished recovering your wallet." in layout().text - client.debug.press_yes() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowBip39RecoveryNoPIN(client, MNEMONIC12.split(" ")) + client.set_input_flow(IF.get()) client.watch_layout() 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 655cfca8f3..1a391faa86 100644 --- a/tests/device_tests/reset_recovery/test_recovery_slip39_advanced.py +++ b/tests/device_tests/reset_recovery/test_recovery_slip39_advanced.py @@ -19,11 +19,12 @@ import pytest from trezorlib import device, exceptions, messages from trezorlib.debuglink import TrezorClientDebugLink as Client -from ...common import ( - MNEMONIC_SLIP39_ADVANCED_20, - MNEMONIC_SLIP39_ADVANCED_33, - recovery_enter_shares, - recovery_enter_shares_tr, +from ...common import MNEMONIC_SLIP39_ADVANCED_20, MNEMONIC_SLIP39_ADVANCED_33 +from ...input_flows import ( + InputFlowSlip39AdvancedRecovery, + InputFlowSlip39AdvancedRecoveryAbort, + InputFlowSlip39AdvancedRecoveryNoAbort, + InputFlowSlip39AdvancedRecoveryTwoSharesWarning, ) pytestmark = pytest.mark.skip_t1 @@ -46,27 +47,9 @@ VECTORS = ( def _test_secret( client: Client, shares: list[str], secret: str, click_info: bool = False ): - debug = client.debug - - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - # Proceed with recovery - yield from recovery_enter_shares( - debug, shares, groups=True, click_info=click_info - ) - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - # Proceed with recovery - yield from recovery_enter_shares_tr(debug, shares, groups=True) - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39AdvancedRecovery(client, shares, click_info=click_info) + client.set_input_flow(IF.get()) ret = device.recover( client, pin_protection=False, @@ -81,7 +64,7 @@ def _test_secret( assert client.features.pin_protection is False assert client.features.passphrase_protection is False assert client.features.backup_type is messages.BackupType.Slip39_Advanced - assert debug.state().mnemonic_secret.hex() == secret + assert client.debug.state().mnemonic_secret.hex() == secret @pytest.mark.parametrize("shares, secret", VECTORS) @@ -107,18 +90,9 @@ def test_extra_share_entered(client: Client): @pytest.mark.setup_client(uninitialized=True) def test_abort(client: Client): - debug = client.debug - - def input_flow(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - abort process - debug.press_no() - yield # Homescreen - confirm abort - debug.press_yes() - with client: - client.set_input_flow(input_flow) + IF = InputFlowSlip39AdvancedRecoveryAbort(client) + client.set_input_flow(IF.get()) with pytest.raises(exceptions.Cancelled): device.recover( client, pin_protection=False, label="label", show_tutorial=False @@ -129,35 +103,11 @@ def test_abort(client: Client): @pytest.mark.setup_client(uninitialized=True) def test_noabort(client: Client): - debug = client.debug - - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - abort process - debug.press_no() - yield # Homescreen - go back to process - debug.press_no() - yield from recovery_enter_shares( - debug, EXTRA_GROUP_SHARE + MNEMONIC_SLIP39_ADVANCED_20, groups=True - ) - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - abort process - debug.press_no() - yield # Homescreen - go back to process - debug.press_no() - yield from recovery_enter_shares_tr( - debug, EXTRA_GROUP_SHARE + MNEMONIC_SLIP39_ADVANCED_20, groups=True - ) - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39AdvancedRecoveryNoAbort( + client, EXTRA_GROUP_SHARE + MNEMONIC_SLIP39_ADVANCED_20 + ) + client.set_input_flow(IF.get()) device.recover(client, pin_protection=False, label="label", show_tutorial=False) client.init_device() assert client.features.initialized is True @@ -165,78 +115,17 @@ def test_noabort(client: Client): @pytest.mark.setup_client(uninitialized=True) def test_same_share(client: Client): - debug = client.debug # we choose the second share from the fixture because # the 1st is 1of1 and group threshold condition is reached first first_share = MNEMONIC_SLIP39_ADVANCED_20[1].split(" ") # second share is first 4 words of first second_share = MNEMONIC_SLIP39_ADVANCED_20[1].split(" ")[:4] - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input(str(len(first_share))) - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - for word in first_share: - debug.input(word) - - yield # Continue to next share - debug.press_yes() - yield # Homescreen - next share - debug.press_yes() - yield # Enter next share - for word in second_share: - debug.input(word) - - br = yield - assert br.code == messages.ButtonRequestType.Warning - - client.cancel() - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input(str(len(first_share))) - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - debug.press_yes() - yield # Enter first share - for word in first_share: - debug.input(word) - - yield # Continue to next share - debug.press_yes() - yield # Homescreen - next share - debug.press_yes() - yield # Homescreen - next share - debug.press_yes() - yield # Enter next share - for word in second_share: - debug.input(word) - - yield - br = yield - assert br.code == messages.ButtonRequestType.Warning - debug.press_right() - debug.press_yes() - yield - - client.cancel() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39AdvancedRecoveryTwoSharesWarning( + client, first_share, second_share + ) + client.set_input_flow(IF.get()) with pytest.raises(exceptions.Cancelled): device.recover( client, pin_protection=False, label="label", show_tutorial=False @@ -245,77 +134,16 @@ def test_same_share(client: Client): @pytest.mark.setup_client(uninitialized=True) def test_group_threshold_reached(client: Client): - debug = client.debug # first share in the fixture is 1of1 so we choose that first_share = MNEMONIC_SLIP39_ADVANCED_20[0].split(" ") # second share is first 3 words of first second_share = MNEMONIC_SLIP39_ADVANCED_20[0].split(" ")[:3] - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input(str(len(first_share))) - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - for word in first_share: - debug.input(word) - - yield # Continue to next share - debug.press_yes() - yield # Homescreen - next share - debug.press_yes() - yield # Enter next share - for word in second_share: - debug.input(word) - - br = yield - assert br.code == messages.ButtonRequestType.Warning - - client.cancel() - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input(str(len(first_share))) - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - debug.press_yes() - yield # Enter first share - for word in first_share: - debug.input(word) - - yield # Continue to next share - debug.press_yes() - yield # Homescreen - next share - debug.press_yes() - yield # Enter next share - debug.press_yes() - yield # Enter next share - for word in second_share: - debug.input(word) - - br = yield - br = yield - assert br.code == messages.ButtonRequestType.Warning - debug.press_right() - debug.press_yes() - yield - - client.cancel() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39AdvancedRecoveryTwoSharesWarning( + client, first_share, second_share + ) + client.set_input_flow(IF.get()) with pytest.raises(exceptions.Cancelled): device.recover( client, pin_protection=False, label="label", show_tutorial=False 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 d18cf247f7..31ecc8ad8a 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 @@ -20,11 +20,8 @@ from trezorlib import device, messages from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.exceptions import TrezorFailure -from ...common import ( - MNEMONIC_SLIP39_ADVANCED_20, - recovery_enter_shares, - recovery_enter_shares_tr, -) +from ...common import MNEMONIC_SLIP39_ADVANCED_20 +from ...input_flows import InputFlowSlip39AdvancedRecoveryDryRun pytestmark = pytest.mark.skip_t1 @@ -43,29 +40,11 @@ EXTRA_GROUP_SHARE = [ @pytest.mark.setup_client(mnemonic=MNEMONIC_SLIP39_ADVANCED_20, passphrase=False) def test_2of3_dryrun(client: Client): - debug = client.debug - - def input_flow_tt(): - yield # Confirm Dryrun - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares( - debug, EXTRA_GROUP_SHARE + MNEMONIC_SLIP39_ADVANCED_20, groups=True - ) - - def input_flow_tr(): - yield # Confirm Dryrun - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares_tr( - debug, EXTRA_GROUP_SHARE + MNEMONIC_SLIP39_ADVANCED_20, groups=True - ) - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39AdvancedRecoveryDryRun( + client, EXTRA_GROUP_SHARE + MNEMONIC_SLIP39_ADVANCED_20 + ) + client.set_input_flow(IF.get()) ret = device.recover( client, passphrase_protection=False, @@ -84,32 +63,14 @@ def test_2of3_dryrun(client: Client): @pytest.mark.setup_client(mnemonic=MNEMONIC_SLIP39_ADVANCED_20) def test_2of3_invalid_seed_dryrun(client: Client): - debug = client.debug - - def input_flow_tt(): - yield # Confirm Dryrun - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares( - debug, INVALID_SHARES_SLIP39_ADVANCED_20, groups=True - ) - - def input_flow_tr(): - yield # Confirm Dryrun - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares_tr( - debug, INVALID_SHARES_SLIP39_ADVANCED_20, groups=True - ) - # test fails because of different seed on device with client, pytest.raises( TrezorFailure, match=r"The seed does not match the one in the device" ): - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39AdvancedRecoveryDryRun( + client, INVALID_SHARES_SLIP39_ADVANCED_20 + ) + client.set_input_flow(IF.get()) device.recover( client, passphrase_protection=False, 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 9b2542a9c0..c566acf30d 100644 --- a/tests/device_tests/reset_recovery/test_recovery_slip39_basic.py +++ b/tests/device_tests/reset_recovery/test_recovery_slip39_basic.py @@ -22,8 +22,16 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client from ...common import ( MNEMONIC_SLIP39_BASIC_20_3of6, MNEMONIC_SLIP39_BASIC_20_3of6_SECRET, - recovery_enter_shares, - recovery_enter_shares_tr, +) +from ...input_flows import ( + InputFlowSlip39BasicRecovery, + InputFlowSlip39BasicRecoveryAbort, + InputFlowSlip39BasicRecoveryNoAbort, + InputFlowSlip39BasicRecoveryPIN, + InputFlowSlip39BasicRecoveryRetryFirst, + InputFlowSlip39BasicRecoveryRetrySecond, + InputFlowSlip39BasicRecoverySameShare, + InputFlowSlip39BasicRecoveryWrongNthWord, ) pytestmark = pytest.mark.skip_t1 @@ -50,25 +58,9 @@ VECTORS = ( @pytest.mark.setup_client(uninitialized=True) @pytest.mark.parametrize("shares, secret", VECTORS) def test_secret(client: Client, shares: list[str], secret: str): - debug = client.debug - - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares(debug, shares) - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares_tr(debug, shares) - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39BasicRecovery(client, shares) + client.set_input_flow(IF.get()) ret = device.recover( client, pin_protection=False, label="label", show_tutorial=False ) @@ -80,40 +72,16 @@ def test_secret(client: Client, shares: list[str], secret: str): assert client.features.backup_type is messages.BackupType.Slip39_Basic # Check mnemonic - assert debug.state().mnemonic_secret.hex() == secret + assert client.debug.state().mnemonic_secret.hex() == secret @pytest.mark.setup_client(uninitialized=True) def test_recover_with_pin_passphrase(client: Client): - debug = client.debug - - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - yield # Enter PIN - debug.input("654") - yield # Enter PIN again - debug.input("654") - # Proceed with recovery - yield from recovery_enter_shares(debug, MNEMONIC_SLIP39_BASIC_20_3of6) - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - yield # Enter PIN - debug.input("654") - yield # Please re-enter PIN to confirm - debug.press_yes() - yield # Enter PIN again - debug.input("654") - # Proceed with recovery - yield from recovery_enter_shares_tr(debug, MNEMONIC_SLIP39_BASIC_20_3of6) - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39BasicRecoveryPIN( + client, MNEMONIC_SLIP39_BASIC_20_3of6, "654" + ) + client.set_input_flow(IF.get()) ret = device.recover( client, pin_protection=True, @@ -122,7 +90,7 @@ def test_recover_with_pin_passphrase(client: Client): show_tutorial=False, ) - # Workflow succesfully ended + # Workflow successfully ended assert ret == messages.Success(message="Device recovered") assert client.features.pin_protection is True assert client.features.passphrase_protection is True @@ -131,18 +99,9 @@ def test_recover_with_pin_passphrase(client: Client): @pytest.mark.setup_client(uninitialized=True) def test_abort(client: Client): - debug = client.debug - - def input_flow(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - abort process - debug.press_no() - yield # Homescreen - confirm abort - debug.press_yes() - with client: - client.set_input_flow(input_flow) + IF = InputFlowSlip39BasicRecoveryAbort(client) + client.set_input_flow(IF.get()) with pytest.raises(exceptions.Cancelled): device.recover( client, pin_protection=False, label="label", show_tutorial=False @@ -153,31 +112,9 @@ def test_abort(client: Client): @pytest.mark.setup_client(uninitialized=True) def test_noabort(client: Client): - debug = client.debug - - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - abort process - debug.press_no() - yield # Homescreen - go back to process - debug.press_no() - yield from recovery_enter_shares(debug, MNEMONIC_SLIP39_BASIC_20_3of6) - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - abort process - debug.press_no() - yield # Homescreen - go back to process - debug.press_no() - yield from recovery_enter_shares_tr(debug, MNEMONIC_SLIP39_BASIC_20_3of6) - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39BasicRecoveryNoAbort(client, MNEMONIC_SLIP39_BASIC_20_3of6) + client.set_input_flow(IF.get()) device.recover(client, pin_protection=False, label="label", show_tutorial=False) client.init_device() assert client.features.initialized is True @@ -187,98 +124,9 @@ def test_noabort(client: Client): def test_ask_word_number(client: Client): if client.features.model == "R": pytest.skip("Flow is not working correctly for TR") - - debug = client.debug - - def input_flow_retry_first_tt(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input("20") - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - for _ in range(20): - debug.input("slush") - - br = yield # Invalid share - assert br.code == messages.ButtonRequestType.Warning - debug.press_yes() - - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input("33") - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - for _ in range(33): - debug.input("slush") - - br = yield # Invalid share - assert br.code == messages.ButtonRequestType.Warning - debug.press_yes() - - yield # Homescreen - debug.press_no() - yield # Confirm abort - debug.press_yes() - - def input_flow_retry_first_tr(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input("20") - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - debug.press_yes() - for _ in range(20): - debug.input("slush") - - yield - # assert br.code == messages.ButtonRequestType.Warning - debug.press_yes() - - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input("33") - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Homescreen - proceed to share entry - debug.press_yes() - yield - for _ in range(33): - debug.input("slush") - - yield - debug.press_yes() - - yield - debug.press_no() - - yield - debug.press_right() - - yield - debug.press_right() - - yield - debug.press_right() - - yield - debug.press_yes() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_retry_first_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_retry_first_tr) + IF = InputFlowSlip39BasicRecoveryRetryFirst(client) + client.set_input_flow(IF.get()) with pytest.raises(exceptions.Cancelled): device.recover( client, pin_protection=False, label="label", show_tutorial=False @@ -286,85 +134,11 @@ def test_ask_word_number(client: Client): client.init_device() assert client.features.initialized is False - def input_flow_retry_second_tt(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input("20") - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - share = MNEMONIC_SLIP39_BASIC_20_3of6[0].split(" ") - for word in share: - debug.input(word) - - yield # More shares needed - debug.press_yes() - - yield # Enter another share - share = share[:3] + ["slush"] * 17 - for word in share: - debug.input(word) - - br = yield # Invalid share - assert br.code == messages.ButtonRequestType.Warning - debug.press_yes() - - yield # Proceed to next share - share = MNEMONIC_SLIP39_BASIC_20_3of6[1].split(" ") - for word in share: - debug.input(word) - - yield # More shares needed - debug.press_no() - yield # Confirm abort - debug.press_yes() - - def input_flow_retry_second_tr(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input("20") - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - debug.press_yes() - yield # Enter first share - share = MNEMONIC_SLIP39_BASIC_20_3of6[0].split(" ") - for word in share: - debug.input(word) - - yield # More shares needed - debug.press_yes() - - yield # Enter another share - share = share[:3] + ["slush"] * 17 - for word in share: - debug.input(word) - - yield # Invalid share - # assert br.code == messages.ButtonRequestType.Warning - debug.press_yes() - - yield # Proceed to next share - share = MNEMONIC_SLIP39_BASIC_20_3of6[1].split(" ") - for word in share: - debug.input(word) - - yield # More shares needed - debug.press_no() - yield # Confirm abort - debug.press_yes() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_retry_second_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_retry_second_tr) + IF = InputFlowSlip39BasicRecoveryRetrySecond( + client, MNEMONIC_SLIP39_BASIC_20_3of6 + ) + client.set_input_flow(IF.get()) with pytest.raises(exceptions.Cancelled): device.recover( client, pin_protection=False, label="label", show_tutorial=False @@ -376,74 +150,10 @@ def test_ask_word_number(client: Client): @pytest.mark.setup_client(uninitialized=True) @pytest.mark.parametrize("nth_word", range(3)) def test_wrong_nth_word(client: Client, nth_word: int): - debug = client.debug share = MNEMONIC_SLIP39_BASIC_20_3of6[0].split(" ") - - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input(str(len(share))) - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - for word in share: - debug.input(word) - - yield # Continue to next share - debug.press_yes() - yield # Enter next share - for i, word in enumerate(share): - if i < nth_word: - debug.input(word) - else: - debug.input(share[-1]) - break - - br = yield - assert br.code == messages.ButtonRequestType.Warning - - client.cancel() - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input(str(len(share))) - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - debug.press_yes() - yield # Enter first share - for word in share: - debug.input(word) - - yield # Continue to next share - debug.press_yes() - yield # Enter next share - debug.press_yes() - yield # Enter next share - for i, word in enumerate(share): - if i < nth_word: - debug.input(word) - else: - debug.input(share[-1]) - break - - yield - # assert br.code == messages.ButtonRequestType.Warning - - client.cancel() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39BasicRecoveryWrongNthWord(client, share, nth_word) + client.set_input_flow(IF.get()) with pytest.raises(exceptions.Cancelled): device.recover( client, pin_protection=False, label="label", show_tutorial=False @@ -452,74 +162,12 @@ def test_wrong_nth_word(client: Client, nth_word: int): @pytest.mark.setup_client(uninitialized=True) def test_same_share(client: Client): - debug = client.debug 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] - - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input(str(len(first_share))) - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - for word in first_share: - debug.input(word) - - yield # Continue to next share - debug.press_yes() - yield # Enter next share - for word in second_share: - debug.input(word) - - br = yield - assert br.code == messages.ButtonRequestType.Warning - - # TODO: seems like the screenshots did not catch the WARNING - - client.cancel() - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - yield # Homescreen - start process - debug.press_yes() - yield # Enter number of words - debug.input(str(len(first_share))) - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Homescreen - proceed to share entry - debug.press_yes() - yield # Enter first share - for word in first_share: - debug.input(word) - - yield # Continue to next share - debug.press_yes() - yield # Continue to next share - debug.press_yes() - yield # Enter next share - for word in second_share: - debug.input(word) - - br = yield - br = yield - assert br.code == messages.ButtonRequestType.Warning - debug.press_right() - debug.press_yes() - yield - - client.cancel() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39BasicRecoverySameShare(client, first_share, second_share) + client.set_input_flow(IF.get()) with pytest.raises(exceptions.Cancelled): device.recover( client, pin_protection=False, label="label", show_tutorial=False @@ -528,29 +176,9 @@ def test_same_share(client: Client): @pytest.mark.setup_client(uninitialized=True) def test_1of1(client: Client): - debug = client.debug - - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - # Proceed with recovery - yield from recovery_enter_shares( - debug, MNEMONIC_SLIP39_BASIC_20_1of1, groups=False - ) - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - # Proceed with recovery - yield from recovery_enter_shares_tr( - debug, MNEMONIC_SLIP39_BASIC_20_1of1, groups=False - ) - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39BasicRecovery(client, MNEMONIC_SLIP39_BASIC_20_1of1) + client.set_input_flow(IF.get()) ret = device.recover( client, pin_protection=False, 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 6f17faf17a..f8dfc11cea 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 ...common import recovery_enter_shares, recovery_enter_shares_tr +from ...input_flows import InputFlowSlip39BasicRecovery pytestmark = pytest.mark.skip_t1 @@ -38,25 +38,9 @@ INVALID_SHARES_20_2of3 = [ @pytest.mark.setup_client(mnemonic=SHARES_20_2of3[0:2]) def test_2of3_dryrun(client: Client): - debug = client.debug - - def input_flow_tt(): - yield # Confirm Dryrun - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares(debug, SHARES_20_2of3[1:3]) - - def input_flow_tr(): - yield # Confirm Dryrun - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares_tr(debug, SHARES_20_2of3[1:3]) - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39BasicRecovery(client, SHARES_20_2of3[1:3]) + client.set_input_flow(IF.get()) ret = device.recover( client, passphrase_protection=False, @@ -75,28 +59,12 @@ def test_2of3_dryrun(client: Client): @pytest.mark.setup_client(mnemonic=SHARES_20_2of3[0:2]) def test_2of3_invalid_seed_dryrun(client: Client): - debug = client.debug - - def input_flow_tt(): - yield # Confirm Dryrun - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares(debug, INVALID_SHARES_20_2of3) - - def input_flow_tr(): - yield # Confirm Dryrun - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares_tr(debug, INVALID_SHARES_20_2of3) - # test fails because of different seed on device with client, pytest.raises( TrezorFailure, match=r"The seed does not match the one in the device" ): - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39BasicRecovery(client, INVALID_SHARES_20_2of3) + client.set_input_flow(IF.get()) device.recover( client, passphrase_protection=False, diff --git a/tests/device_tests/reset_recovery/test_reset_backup.py b/tests/device_tests/reset_recovery/test_reset_backup.py index eed06a9e8f..a95455f6c0 100644 --- a/tests/device_tests/reset_recovery/test_reset_backup.py +++ b/tests/device_tests/reset_recovery/test_reset_backup.py @@ -24,41 +24,19 @@ from trezorlib import device, messages from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.messages import BackupType, ButtonRequestType as B -from ...common import ( - EXTERNAL_ENTROPY, - click_through, - read_and_confirm_mnemonic, - read_and_confirm_mnemonic_tr, +from ...common import EXTERNAL_ENTROPY +from ...input_flows import ( + InputFlowBip39Backup, + InputFlowResetSkipBackup, + InputFlowSlip39AdvancedBackup, + InputFlowSlip39BasicBackup, ) -def backup_flow_bip39(client: Client): - mnemonic = None - - def input_flow(): - nonlocal mnemonic - - # 1. Confirm Reset - yield from click_through(client.debug, screens=1, code=B.ResetDevice) - - # mnemonic phrases - if client.debug.model == "R": - client.debug.watch_layout(True) - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - else: - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - - # confirm recovery seed check - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # confirm success - br = yield - assert br.code == B.Success - client.debug.press_yes() - +def backup_flow_bip39(client: Client) -> bytes: with client: + IF = InputFlowBip39Backup(client) + client.set_input_flow(IF.get()) client.set_expected_responses( [ *[ @@ -71,74 +49,16 @@ def backup_flow_bip39(client: Client): messages.Features, ] ) - client.set_input_flow(input_flow) device.backup(client) - return mnemonic.encode() + assert IF.mnemonic is not None + return IF.mnemonic.encode() def backup_flow_slip39_basic(client: Client): - mnemonics = [] - - def input_flow_tt(): - # 1. Checklist - # 2. Number of shares (5) - # 3. Checklist - # 4. Threshold (3) - # 5. Checklist - # 6. Confirm show seeds - yield from click_through(client.debug, screens=6, code=B.ResetDevice) - - # Mnemonic phrases - for _ in range(5): - # Phrase screen - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - mnemonics.append(mnemonic) - yield # Confirm continue to next - client.debug.press_yes() - - # Confirm backup - yield - client.debug.press_yes() - - def input_flow_tr(): - yield # Checklist - client.debug.press_yes() - yield # Number of shares info - client.debug.press_yes() - yield # Number of shares (5) - client.debug.input("5") - yield # Checklist - client.debug.press_yes() - yield # Threshold info - client.debug.press_yes() - yield # Threshold (3) - client.debug.input("3") - yield # Checklist - client.debug.press_yes() - yield # Confirm show seeds - client.debug.press_yes() - - # Mnemonic phrases - for _ in range(5): - # Phrase screen - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - mnemonics.append(mnemonic) - - br = yield # Confirm continue to next - assert br.code == B.Success - client.debug.press_yes() - - br = yield # Confirm backup - assert br.code == B.Success - client.debug.press_yes() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.debug.watch_layout(True) - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39BasicBackup(client, False) + client.set_input_flow(IF.get()) client.set_expected_responses( [messages.ButtonRequest(code=B.ResetDevice)] * 6 # intro screens + [ @@ -154,89 +74,15 @@ def backup_flow_slip39_basic(client: Client): ) device.backup(client) - groups = shamir.decode_mnemonics(mnemonics[:3]) + groups = shamir.decode_mnemonics(IF.mnemonics[:3]) ems = shamir.recover_ems(groups) return ems.ciphertext def backup_flow_slip39_advanced(client: Client): - mnemonics = [] - - def input_flow_tt(): - # 1. Confirm Reset - # 2. shares info - # 3. Set & Confirm number of groups - # 4. threshold info - # 5. Set & confirm group threshold value - # 6-15: for each of 5 groups: - # 1. Set & Confirm number of shares - # 2. Set & confirm share threshold value - # 16. Confirm show seeds - yield from click_through(client.debug, screens=16, code=B.ResetDevice) - - # show & confirm shares for all groups - for _ in range(5): - for _ in range(5): - # mnemonic phrases - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - mnemonics.append(mnemonic) - - # Confirm continue to next share - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # safety warning - br = yield - assert br.code == B.Success - client.debug.press_yes() - - def input_flow_tr(): - yield # Checklist - client.debug.press_yes() - yield # Set and confirm group count - client.debug.input("5") - yield # Checklist - client.debug.press_yes() - yield # Set and confirm group threshold - client.debug.input("3") - yield # Checklist - client.debug.press_yes() - for _ in range(5): # for each of 5 groups - yield # Number of shares info - client.debug.press_yes() - yield # Number of shares (5) - client.debug.input("5") - yield # Threshold info - client.debug.press_yes() - yield # Threshold (3) - client.debug.input("3") - yield # Confirm show seeds - client.debug.press_yes() - - # show & confirm shares for all groups - for _g in range(5): - for _h in range(5): - # mnemonic phrases - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - mnemonics.append(mnemonic) - - # Confirm continue to next share - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # safety warning - br = yield - assert br.code == B.Success - client.debug.press_yes() - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.debug.watch_layout(True) - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39AdvancedBackup(client, False) + client.set_input_flow(IF.get()) client.set_expected_responses( [messages.ButtonRequest(code=B.ResetDevice)] * 6 # intro screens + [ @@ -257,7 +103,7 @@ def backup_flow_slip39_advanced(client: Client): ) device.backup(client) - mnemonics = mnemonics[0:3] + mnemonics[5:8] + mnemonics[10:13] + mnemonics = IF.mnemonics[0:3] + IF.mnemonics[5:8] + IF.mnemonics[10:13] groups = shamir.decode_mnemonics(mnemonics) ems = shamir.recover_ems(groups) return ems.ciphertext @@ -309,19 +155,10 @@ def test_skip_backup_msg(client: Client, backup_type, backup_flow): @pytest.mark.parametrize("backup_type, backup_flow", VECTORS) @pytest.mark.setup_client(uninitialized=True) def test_skip_backup_manual(client: Client, backup_type, backup_flow): - def reset_skip_input_flow(): - yield # Confirm Recovery - client.debug.press_yes() - - yield # Skip Backup - client.debug.press_no() - - yield # Confirm skip backup - client.debug.press_no() - os_urandom = mock.Mock(return_value=EXTERNAL_ENTROPY) with mock.patch("os.urandom", os_urandom), client: - client.set_input_flow(reset_skip_input_flow) + IF = InputFlowResetSkipBackup(client) + client.set_input_flow(IF.get()) client.set_expected_responses( [ messages.ButtonRequest(code=B.ResetDevice), diff --git a/tests/device_tests/reset_recovery/test_reset_bip39_t2.py b/tests/device_tests/reset_recovery/test_reset_bip39_t2.py index 04a28b0b43..6ab7ba2be7 100644 --- a/tests/device_tests/reset_recovery/test_reset_bip39_t2.py +++ b/tests/device_tests/reset_recovery/test_reset_bip39_t2.py @@ -24,12 +24,11 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.exceptions import TrezorFailure from trezorlib.messages import ButtonRequestType as B -from ...common import ( - MNEMONIC12, - click_through, - generate_entropy, - read_and_confirm_mnemonic, - read_and_confirm_mnemonic_tr, +from ...common import MNEMONIC12, generate_entropy +from ...input_flows import ( + InputFlowBip39ResetBackup, + InputFlowBip39ResetFailedCheck, + InputFlowBip39ResetPIN, ) pytestmark = [pytest.mark.skip_t1] @@ -37,35 +36,11 @@ pytestmark = [pytest.mark.skip_t1] EXTERNAL_ENTROPY = b"zlutoucky kun upel divoke ody" * 2 -def reset_device(client: Client, strength): - mnemonic = None - - def input_flow(): - nonlocal mnemonic - # 1. Confirm Reset - # 2. Backup your seed - # 3. Confirm warning - yield from click_through(client.debug, screens=3, code=B.ResetDevice) - - # mnemonic phrases - if client.debug.model == "R": - client.debug.watch_layout(True) - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - else: - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - - # confirm recovery seed check - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # confirm success - br = yield - assert br.code == B.Success - client.debug.press_yes() - +def reset_device(client: Client, strength: int): os_urandom = mock.Mock(return_value=EXTERNAL_ENTROPY) with mock.patch("os.urandom", os_urandom), client: + IF = InputFlowBip39ResetBackup(client) + client.set_input_flow(IF.get()) client.set_expected_responses( [ messages.ButtonRequest(code=B.ResetDevice), @@ -81,7 +56,6 @@ def reset_device(client: Client, strength): messages.Features, ] ) - client.set_input_flow(input_flow) # No PIN, no passphrase, don't display random device.reset( @@ -101,7 +75,7 @@ def reset_device(client: Client, strength): expected_mnemonic = Mnemonic("english").to_mnemonic(entropy) # Compare that device generated proper mnemonic for given entropies - assert mnemonic == expected_mnemonic + assert IF.mnemonic == expected_mnemonic # Check if device is properly initialized resp = client.call_raw(messages.Initialize()) @@ -128,64 +102,12 @@ def test_reset_device_192(client: Client): @pytest.mark.setup_client(uninitialized=True) def test_reset_device_pin(client: Client): - mnemonic = None strength = 256 # 24 words - def input_flow(): - nonlocal mnemonic - - # Confirm Reset - br = yield - assert br.code == B.ResetDevice - client.debug.press_yes() - - # Enter new PIN - yield - client.debug.input("654") - - if client.debug.model == "R": - # Re-enter PIN - yield - client.debug.press_yes() - - # Confirm PIN - yield - client.debug.input("654") - - # Confirm entropy - br = yield - assert br.code == B.ResetDevice - client.debug.press_yes() - - # Backup your seed - br = yield - assert br.code == B.ResetDevice - client.debug.press_yes() - - # Confirm warning - br = yield - assert br.code == B.ResetDevice - client.debug.press_yes() - - # mnemonic phrases - if client.debug.model == "R": - client.debug.watch_layout(True) - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - else: - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - - # confirm recovery seed check - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # confirm success - br = yield - assert br.code == B.Success - client.debug.press_yes() - os_urandom = mock.Mock(return_value=EXTERNAL_ENTROPY) with mock.patch("os.urandom", os_urandom), client: + IF = InputFlowBip39ResetPIN(client) + client.set_input_flow(IF.get()) client.set_expected_responses( [ messages.ButtonRequest(code=B.ResetDevice), @@ -204,7 +126,6 @@ def test_reset_device_pin(client: Client): messages.Features, ] ) - client.set_input_flow(input_flow) # PIN, passphrase, display random device.reset( @@ -224,7 +145,7 @@ def test_reset_device_pin(client: Client): expected_mnemonic = Mnemonic("english").to_mnemonic(entropy) # Compare that device generated proper mnemonic for given entropies - assert mnemonic == expected_mnemonic + assert IF.mnemonic == expected_mnemonic # Check if device is properly initialized resp = client.call_raw(messages.Initialize()) @@ -236,51 +157,12 @@ def test_reset_device_pin(client: Client): @pytest.mark.setup_client(uninitialized=True) def test_reset_failed_check(client: Client): - mnemonic = None strength = 256 # 24 words - def input_flow(): - nonlocal mnemonic - # 1. Confirm Reset - # 2. Backup your seed - # 3. Confirm warning - yield from click_through(client.debug, screens=3, code=B.ResetDevice) - - # mnemonic phrases, wrong answer - if client.debug.model == "R": - client.debug.watch_layout(True) - mnemonic = yield from read_and_confirm_mnemonic_tr( - client.debug, choose_wrong=True - ) - else: - mnemonic = yield from read_and_confirm_mnemonic( - client.debug, choose_wrong=True - ) - - # warning screen - br = yield - assert br.code == B.ResetDevice - client.debug.press_yes() - - # mnemonic phrases - if client.debug.model == "R": - client.debug.watch_layout(True) - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - else: - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - - # confirm recovery seed check - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # confirm success - br = yield - assert br.code == B.Success - client.debug.press_yes() - os_urandom = mock.Mock(return_value=EXTERNAL_ENTROPY) with mock.patch("os.urandom", os_urandom), client: + IF = InputFlowBip39ResetFailedCheck(client) + client.set_input_flow(IF.get()) client.set_expected_responses( [ messages.ButtonRequest(code=B.ResetDevice), @@ -298,7 +180,6 @@ def test_reset_failed_check(client: Client): messages.Features, ] ) - client.set_input_flow(input_flow) # PIN, passphrase, display random device.reset( @@ -318,7 +199,7 @@ def test_reset_failed_check(client: Client): expected_mnemonic = Mnemonic("english").to_mnemonic(entropy) # Compare that device generated proper mnemonic for given entropies - assert mnemonic == expected_mnemonic + assert IF.mnemonic == expected_mnemonic # Check if device is properly initialized resp = client.call_raw(messages.Initialize()) 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 8b187884e6..950059bfa7 100644 --- a/tests/device_tests/reset_recovery/test_reset_recovery_bip39.py +++ b/tests/device_tests/reset_recovery/test_reset_recovery_bip39.py @@ -24,12 +24,8 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.messages import BackupType, ButtonRequestType as B from trezorlib.tools import parse_path -from ...common import ( - EXTERNAL_ENTROPY, - click_through, - read_and_confirm_mnemonic, - read_and_confirm_mnemonic_tr, -) +from ...common import EXTERNAL_ENTROPY +from ...input_flows import InputFlowBip39RecoveryNoPIN, InputFlowBip39ResetBackup @pytest.mark.skip_t1 @@ -45,53 +41,10 @@ def test_reset_recovery(client: Client): def reset(client: Client, strength: int = 128, skip_backup: bool = False) -> str: - mnemonic = None - - def input_flow_tt(): - nonlocal mnemonic - - # 1. Confirm Reset - # 2. Backup your seed - # 3. Confirm warning - yield from click_through(client.debug, screens=3, code=B.ResetDevice) - - # mnemonic phrases - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - - # confirm recovery seed check - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # confirm success - br = yield - assert br.code == B.Success - client.debug.press_yes() - - def input_flow_tr(): - nonlocal mnemonic - - # 1. Confirm Reset - # 2. Backup your seed - # 3. Confirm warning - yield from click_through(client.debug, screens=3, code=B.ResetDevice) - - # mnemonic phrases - client.debug.watch_layout(True) - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - - # confirm recovery seed check - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # confirm success - br = yield - assert br.code == B.Success - client.debug.press_yes() - os_urandom = mock.Mock(return_value=EXTERNAL_ENTROPY) with mock.patch("os.urandom", os_urandom), client: + IF = InputFlowBip39ResetBackup(client) + client.set_input_flow(IF.get()) client.set_expected_responses( [ messages.ButtonRequest(code=B.ResetDevice), @@ -107,10 +60,6 @@ def reset(client: Client, strength: int = 128, skip_backup: bool = False) -> str messages.Features, ] ) - if client.debug.model == "R": - client.set_input_flow(input_flow_tr) - elif client.debug.model == "T": - client.set_input_flow(input_flow_tt) # No PIN, no passphrase, don't display random device.reset( @@ -131,57 +80,16 @@ def reset(client: Client, strength: int = 128, skip_backup: bool = False) -> str assert client.features.pin_protection is False assert client.features.passphrase_protection is False - return mnemonic + assert IF.mnemonic is not None + return IF.mnemonic def recover(client: Client, mnemonic: str): - debug = client.debug words = mnemonic.split(" ") - - def input_flow_tt(): - yield # Confirm recovery - debug.press_yes() - yield # Homescreen - debug.press_yes() - - yield # Enter word count - debug.input(str(len(words))) - - yield # Homescreen - debug.press_yes() - yield # Enter words - for word in words: - debug.input(word) - - yield # confirm success - debug.press_yes() - - def input_flow_tr(): - yield # Confirm recovery - debug.press_yes() - yield # Homescreen - debug.press_yes() - - yield # Enter word count - debug.input(str(len(words))) - - yield # Homescreen - debug.press_yes() - yield # Homescreen - debug.press_yes() - yield # Enter words - for word in words: - debug.input(word) - - yield # confirm success - debug.press_yes() - yield - with client: - if client.debug.model == "R": - client.set_input_flow(input_flow_tr) - elif client.debug.model == "T": - client.set_input_flow(input_flow_tt) + IF = InputFlowBip39RecoveryNoPIN(client, words) + client.set_input_flow(IF.get()) + client.watch_layout() client.set_expected_responses( [ messages.ButtonRequest(code=B.ProtectCall), diff --git a/tests/device_tests/reset_recovery/test_reset_recovery_slip39_advanced.py b/tests/device_tests/reset_recovery/test_reset_recovery_slip39_advanced.py index 2a1c7ba585..aa74d30bd0 100644 --- a/tests/device_tests/reset_recovery/test_reset_recovery_slip39_advanced.py +++ b/tests/device_tests/reset_recovery/test_reset_recovery_slip39_advanced.py @@ -23,13 +23,10 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.messages import BackupType, ButtonRequestType as B from trezorlib.tools import parse_path -from ...common import ( - EXTERNAL_ENTROPY, - click_through, - read_and_confirm_mnemonic, - read_and_confirm_mnemonic_tr, - recovery_enter_shares, - recovery_enter_shares_tr, +from ...common import EXTERNAL_ENTROPY +from ...input_flows import ( + InputFlowSlip39AdvancedRecovery, + InputFlowSlip39AdvancedResetRecovery, ) @@ -63,85 +60,10 @@ def test_reset_recovery(client: Client): def reset(client: Client, strength: int = 128) -> list[str]: - all_mnemonics: list[str] = [] - - def input_flow_tt(): - # 1. Confirm Reset - # 2. Backup your seed - # 3. Confirm warning - # 4. shares info - # 5. Set & Confirm number of groups - # 6. threshold info - # 7. Set & confirm group threshold value - # 8-17: for each of 5 groups: - # 1. Set & Confirm number of shares - # 2. Set & confirm share threshold value - # 18. Confirm show seeds - yield from click_through(client.debug, screens=18, code=B.ResetDevice) - - # show & confirm shares for all groups - for _g in range(5): - for _h in range(5): - # mnemonic phrases - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - all_mnemonics.append(mnemonic) - - # Confirm continue to next share - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # safety warning - br = yield - assert br.code == B.Success - client.debug.press_yes() - - def input_flow_tr(): - yield # Wallet backup - client.debug.press_yes() - yield # Wallet creation - client.debug.press_yes() - yield # Checklist - client.debug.press_yes() - yield # Set and confirm group count - client.debug.input("5") - yield # Checklist - client.debug.press_yes() - yield # Set and confirm group threshold - client.debug.input("3") - yield # Checklist - client.debug.press_yes() - for _ in range(5): # for each of 5 groups - yield # Number of shares info - client.debug.press_yes() - yield # Number of shares (5) - client.debug.input("5") - yield # Threshold info - client.debug.press_yes() - yield # Threshold (3) - client.debug.input("3") - yield # Confirm show seeds - client.debug.press_yes() - - # show & confirm shares for all groups - for _g in range(5): - for _h in range(5): - # mnemonic phrases - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - all_mnemonics.append(mnemonic) - - # Confirm continue to next share - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # safety warning - br = yield - assert br.code == B.Success - client.debug.press_yes() - os_urandom = mock.Mock(return_value=EXTERNAL_ENTROPY) with mock.patch("os.urandom", os_urandom), client: + IF = InputFlowSlip39AdvancedResetRecovery(client, False) + client.set_input_flow(IF.get()) client.set_expected_responses( [ messages.ButtonRequest(code=B.ResetDevice), @@ -176,11 +98,6 @@ def reset(client: Client, strength: int = 128) -> list[str]: messages.Features, ] ) - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.debug.watch_layout(True) - client.set_input_flow(input_flow_tr) # No PIN, no passphrase, don't display random device.reset( @@ -201,29 +118,13 @@ def reset(client: Client, strength: int = 128) -> list[str]: assert client.features.pin_protection is False assert client.features.passphrase_protection is False - return all_mnemonics + return IF.mnemonics def recover(client: Client, shares: list[str]): - debug = client.debug - - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares(debug, shares, groups=True) - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares_tr(debug, shares, groups=True) - with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39AdvancedRecovery(client, shares, False) + client.set_input_flow(IF.get()) ret = device.recover( client, pin_protection=False, label="label", show_tutorial=False ) diff --git a/tests/device_tests/reset_recovery/test_reset_recovery_slip39_basic.py b/tests/device_tests/reset_recovery/test_reset_recovery_slip39_basic.py index 2142f268af..3aca6b9a02 100644 --- a/tests/device_tests/reset_recovery/test_reset_recovery_slip39_basic.py +++ b/tests/device_tests/reset_recovery/test_reset_recovery_slip39_basic.py @@ -24,12 +24,9 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.messages import BackupType, ButtonRequestType as B from trezorlib.tools import parse_path -from ...common import ( - click_through, - read_and_confirm_mnemonic, - read_and_confirm_mnemonic_tr, - recovery_enter_shares, - recovery_enter_shares_tr, +from ...input_flows import ( + InputFlowSlip39BasicRecovery, + InputFlowSlip39BasicResetRecovery, ) EXTERNAL_ENTROPY = b"zlutoucky kun upel divoke ody" * 2 @@ -54,72 +51,9 @@ def test_reset_recovery(client: Client): def reset(client: Client, strength: int = 128) -> list[str]: - all_mnemonics: list[str] = [] - - def input_flow_tt(): - # 1. Confirm Reset - # 2. Backup your seed - # 3. Confirm warning - # 4. shares info - # 5. Set & Confirm number of shares - # 6. threshold info - # 7. Set & confirm threshold value - # 8. Confirm show seeds - yield from click_through(client.debug, screens=8, code=B.ResetDevice) - - # show & confirm shares - for _ in range(5): - # mnemonic phrases - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - all_mnemonics.append(mnemonic) - - # Confirm continue to next share - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # safety warning - br = yield - assert br.code == B.Success - client.debug.press_yes() - - def input_flow_tr(): - yield # Confirm Reset - client.debug.press_yes() - yield # Backup your seed - client.debug.press_yes() - yield # Checklist - client.debug.press_yes() - yield # Number of shares info - client.debug.press_yes() - yield # Number of shares (5) - client.debug.input("5") - yield # Checklist - client.debug.press_yes() - yield # Threshold info - client.debug.press_yes() - yield # Threshold (3) - client.debug.input("3") - yield # Checklist - client.debug.press_yes() - yield # Confirm show seeds - client.debug.press_yes() - - # Mnemonic phrases - for _ in range(5): - # Phrase screen - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - all_mnemonics.append(mnemonic) - - br = yield # Confirm continue to next - assert br.code == B.Success - client.debug.press_yes() - - br = yield # Confirm backup - assert br.code == B.Success - client.debug.press_yes() - with client: + IF = InputFlowSlip39BasicResetRecovery(client) + client.set_input_flow(IF.get()) client.set_expected_responses( [ messages.ButtonRequest(code=B.ResetDevice), @@ -144,11 +78,6 @@ def reset(client: Client, strength: int = 128) -> list[str]: messages.Features, ] ) - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.debug.watch_layout(True) - client.set_input_flow(input_flow_tr) # No PIN, no passphrase, don't display random device.reset( @@ -169,29 +98,13 @@ def reset(client: Client, strength: int = 128) -> list[str]: assert client.features.pin_protection is False assert client.features.passphrase_protection is False - return all_mnemonics + return IF.mnemonics -def recover(client: Client, shares: list[str]) -> None: - debug = client.debug - - def input_flow_tt(): - yield # Confirm Recovery - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares(debug, shares) - - def input_flow_tr(): - yield # Confirm Recovery - debug.press_yes() - # run recovery flow - yield from recovery_enter_shares_tr(debug, shares) - +def recover(client: Client, shares: list[str]): with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39BasicRecovery(client, shares) + client.set_input_flow(IF.get()) ret = device.recover( client, pin_protection=False, label="label", show_tutorial=False ) diff --git a/tests/device_tests/reset_recovery/test_reset_slip39_advanced.py b/tests/device_tests/reset_recovery/test_reset_slip39_advanced.py index d2b59bbf0f..dd7dcbc465 100644 --- a/tests/device_tests/reset_recovery/test_reset_slip39_advanced.py +++ b/tests/device_tests/reset_recovery/test_reset_slip39_advanced.py @@ -24,12 +24,8 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.exceptions import TrezorFailure from trezorlib.messages import BackupType, ButtonRequestType as B -from ...common import ( - click_through, - generate_entropy, - read_and_confirm_mnemonic, - read_and_confirm_mnemonic_tr, -) +from ...common import generate_entropy +from ...input_flows import InputFlowSlip39AdvancedResetRecovery pytestmark = [pytest.mark.skip_t1] @@ -41,85 +37,11 @@ EXTERNAL_ENTROPY = b"zlutoucky kun upel divoke ody" * 2 def test_reset_device_slip39_advanced(client: Client): strength = 128 member_threshold = 3 - all_mnemonics = [] - - def input_flow_tt(): - # 1. Confirm Reset - # 2. Backup your seed - # 3. Confirm warning - # 4. shares info - # 5. Set & Confirm number of groups - # 6. threshold info - # 7. Set & confirm group threshold value - # 8-17: for each of 5 groups: - # 1. Set & Confirm number of shares - # 2. Set & confirm share threshold value - # 18. Confirm show seeds - yield from click_through(client.debug, screens=18, code=B.ResetDevice) - - # show & confirm shares for all groups - for _g in range(5): - for _h in range(5): - # mnemonic phrases - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - all_mnemonics.append(mnemonic) - - # Confirm continue to next share - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # safety warning - br = yield - assert br.code == B.Success - client.debug.press_yes() - - def input_flow_tr(): - yield # Wallet backup - client.debug.press_yes() - yield # Wallet creation - client.debug.press_yes() - yield # Checklist - client.debug.press_yes() - yield # Set and confirm group count - client.debug.input("5") - yield # Checklist - client.debug.press_yes() - yield # Set and confirm group threshold - client.debug.input("3") - yield # Checklist - client.debug.press_yes() - for _ in range(5): # for each of 5 groups - yield # Number of shares info - client.debug.press_yes() - yield # Number of shares (5) - client.debug.input("5") - yield # Threshold info - client.debug.press_yes() - yield # Threshold (3) - client.debug.input("3") - yield # Confirm show seeds - client.debug.press_yes() - - # show & confirm shares for all groups - for _g in range(5): - for _h in range(5): - # mnemonic phrases - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - all_mnemonics.append(mnemonic) - - # Confirm continue to next share - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # safety warning - br = yield - assert br.code == B.Success - client.debug.press_yes() os_urandom = mock.Mock(return_value=EXTERNAL_ENTROPY) with mock.patch("os.urandom", os_urandom), client: + IF = InputFlowSlip39AdvancedResetRecovery(client, False) + client.set_input_flow(IF.get()) client.set_expected_responses( [ messages.ButtonRequest(code=B.ResetDevice), @@ -154,11 +76,6 @@ def test_reset_device_slip39_advanced(client: Client): messages.Features, ] ) - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.debug.watch_layout(True) - client.set_input_flow(input_flow_tr) # No PIN, no passphrase, don't display random device.reset( @@ -178,7 +95,7 @@ def test_reset_device_slip39_advanced(client: Client): secret = generate_entropy(strength, internal_entropy, EXTERNAL_ENTROPY) # validate that all combinations will result in the correct master secret - validate_mnemonics(all_mnemonics, member_threshold, secret) + validate_mnemonics(IF.mnemonics, member_threshold, secret) # Check if device is properly initialized assert client.features.initialized is True diff --git a/tests/device_tests/reset_recovery/test_reset_slip39_basic.py b/tests/device_tests/reset_recovery/test_reset_slip39_basic.py index 08fc07913b..6743aaaf61 100644 --- a/tests/device_tests/reset_recovery/test_reset_slip39_basic.py +++ b/tests/device_tests/reset_recovery/test_reset_slip39_basic.py @@ -25,86 +25,19 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.exceptions import TrezorFailure from trezorlib.messages import BackupType, ButtonRequestType as B -from ...common import ( - EXTERNAL_ENTROPY, - click_through, - generate_entropy, - read_and_confirm_mnemonic, - read_and_confirm_mnemonic_tr, -) +from ...common import EXTERNAL_ENTROPY, generate_entropy +from ...input_flows import InputFlowSlip39BasicResetRecovery pytestmark = [pytest.mark.skip_t1] -def reset_device(client: Client, strength): +def reset_device(client: Client, strength: int): member_threshold = 3 - all_mnemonics = [] - - def input_flow_tt(): - # 1. Confirm Reset - # 2. Backup your seed - # 3. Confirm warning - # 4. shares info - # 5. Set & Confirm number of shares - # 6. threshold info - # 7. Set & confirm threshold value - # 8. Confirm show seeds - yield from click_through(client.debug, screens=8, code=B.ResetDevice) - - # show & confirm shares - for _ in range(5): - # mnemonic phrases - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - all_mnemonics.append(mnemonic) - - # Confirm continue to next share - br = yield - assert br.code == B.Success - client.debug.press_yes() - - # safety warning - br = yield - assert br.code == B.Success - client.debug.press_yes() - - def input_flow_tr(): - yield # Confirm Reset - client.debug.press_yes() - yield # Backup your seed - client.debug.press_yes() - yield # Checklist - client.debug.press_yes() - yield # Number of shares info - client.debug.press_yes() - yield # Number of shares (5) - client.debug.input("5") - yield # Checklist - client.debug.press_yes() - yield # Threshold info - client.debug.press_yes() - yield # Threshold (3) - client.debug.input("3") - yield # Checklist - client.debug.press_yes() - yield # Confirm show seeds - client.debug.press_yes() - - # Mnemonic phrases - for _ in range(5): - # Phrase screen - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - all_mnemonics.append(mnemonic) - - br = yield # Confirm continue to next - assert br.code == B.Success - client.debug.press_yes() - - br = yield # Confirm backup - assert br.code == B.Success - client.debug.press_yes() os_urandom = mock.Mock(return_value=EXTERNAL_ENTROPY) with mock.patch("os.urandom", os_urandom), client: + IF = InputFlowSlip39BasicResetRecovery(client) + client.set_input_flow(IF.get()) client.set_expected_responses( [ messages.ButtonRequest(code=B.ResetDevice), @@ -129,11 +62,6 @@ def reset_device(client: Client, strength): messages.Features, ] ) - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.debug.watch_layout(True) - client.set_input_flow(input_flow_tr) # No PIN, no passphrase, don't display random device.reset( @@ -153,7 +81,7 @@ def reset_device(client: Client, strength): secret = generate_entropy(strength, internal_entropy, EXTERNAL_ENTROPY) # validate that all combinations will result in the correct master secret - validate_mnemonics(all_mnemonics, member_threshold, secret) + validate_mnemonics(IF.mnemonics, member_threshold, secret) # Check if device is properly initialized assert client.features.initialized is True diff --git a/tests/device_tests/test_msg_backup_device.py b/tests/device_tests/test_msg_backup_device.py index 1583c4cd8a..70c3a9755d 100644 --- a/tests/device_tests/test_msg_backup_device.py +++ b/tests/device_tests/test_msg_backup_device.py @@ -27,41 +27,22 @@ from ..common import ( MNEMONIC12, MNEMONIC_SLIP39_ADVANCED_20, MNEMONIC_SLIP39_BASIC_20_3of6, - read_and_confirm_mnemonic, - read_and_confirm_mnemonic_tr, ) - - -def click_info_button(debug): - """Click Shamir backup info button and return back.""" - debug.press_info() - yield # Info screen with text - debug.press_yes() +from ..input_flows import ( + InputFlowBip39Backup, + InputFlowSlip39AdvancedBackup, + InputFlowSlip39BasicBackup, +) @pytest.mark.skip_t1 # TODO we want this for t1 too @pytest.mark.setup_client(needs_backup=True, mnemonic=MNEMONIC12) def test_backup_bip39(client: Client): assert client.features.needs_backup is True - mnemonic = None - - def input_flow(): - nonlocal mnemonic - yield # Confirm Backup - client.debug.press_yes() - # Mnemonic phrases - if client.debug.model == "R": - client.debug.watch_layout(True) - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - else: - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - yield # Confirm success - client.debug.press_yes() - yield # Backup is done - client.debug.press_yes() with client: - client.set_input_flow(input_flow) + IF = InputFlowBip39Backup(client) + client.set_input_flow(IF.get()) client.set_expected_responses( [ *[ @@ -76,7 +57,7 @@ def test_backup_bip39(client: Client): ) device.backup(client) - assert mnemonic == MNEMONIC12 + assert IF.mnemonic == MNEMONIC12 client.init_device() assert client.features.initialized is True assert client.features.needs_backup is False @@ -95,76 +76,10 @@ def test_backup_slip39_basic(client: Client, click_info: bool): pytest.skip("click_info not implemented on TR") assert client.features.needs_backup is True - mnemonics = [] - - def input_flow_tt(): - yield # Checklist - client.debug.press_yes() - if click_info: - yield from click_info_button(client.debug) - yield # Number of shares (5) - client.debug.press_yes() - yield # Checklist - client.debug.press_yes() - if click_info: - yield from click_info_button(client.debug) - yield # Threshold (3) - client.debug.press_yes() - yield # Checklist - client.debug.press_yes() - yield # Confirm show seeds - client.debug.press_yes() - - # Mnemonic phrases - for _ in range(5): - # Phrase screen - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - mnemonics.append(mnemonic) - br = yield # Confirm continue to next - assert br.code == B.Success - client.debug.press_yes() - - br = yield # Confirm backup - assert br.code == B.Success - client.debug.press_yes() - - def input_flow_tr(): - yield # Checklist - client.debug.press_yes() - yield # Number of shares info - client.debug.press_yes() - yield # Number of shares (5) - client.debug.input("5") - yield # Checklist - client.debug.press_yes() - yield # Threshold info - client.debug.press_yes() - yield # Threshold (3) - client.debug.input("3") - yield # Checklist - client.debug.press_yes() - yield # Confirm show seeds - client.debug.press_yes() - - # Mnemonic phrases - for _ in range(5): - # Phrase screen - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - mnemonics.append(mnemonic) - br = yield # Confirm continue to next - assert br.code == B.Success - client.debug.press_yes() - - br = yield # Confirm backup - assert br.code == B.Success - client.debug.press_yes() with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.debug.watch_layout(True) - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39BasicBackup(client, click_info) + client.set_input_flow(IF.get()) client.set_expected_responses( [messages.ButtonRequest(code=B.ResetDevice)] * (8 if click_info else 6) # intro screens (and optional info) @@ -189,7 +104,7 @@ def test_backup_slip39_basic(client: Client, click_info: bool): assert client.features.backup_type is messages.BackupType.Slip39_Basic expected_ms = shamir.combine_mnemonics(MNEMONIC_SLIP39_BASIC_20_3of6) - actual_ms = shamir.combine_mnemonics(mnemonics[:3]) + actual_ms = shamir.combine_mnemonics(IF.mnemonics[:3]) assert expected_ms == actual_ms @@ -203,88 +118,10 @@ def test_backup_slip39_advanced(client: Client, click_info: bool): pytest.skip("click_info not implemented on TR") assert client.features.needs_backup is True - mnemonics = [] - - def input_flow_tt(): - yield # Checklist - client.debug.press_yes() - if click_info: - yield from click_info_button(client.debug) - yield # Set and confirm group count - client.debug.press_yes() - yield # Checklist - client.debug.press_yes() - if click_info: - yield from click_info_button(client.debug) - yield # Set and confirm group threshold - client.debug.press_yes() - yield # Checklist - client.debug.press_yes() - for _ in range(5): # for each of 5 groups - if click_info: - yield from click_info_button(client.debug) - yield # Set & Confirm number of shares - client.debug.press_yes() - if click_info: - yield from click_info_button(client.debug) - yield # Set & confirm share threshold value - client.debug.press_yes() - yield # Confirm show seeds - client.debug.press_yes() - - # Mnemonic phrases - for _ in range(5): - for _ in range(5): - # Phrase screen - mnemonic = yield from read_and_confirm_mnemonic(client.debug) - mnemonics.append(mnemonic) - yield # Confirm continue to next - client.debug.press_yes() - - yield # Confirm backup - client.debug.press_yes() - - def input_flow_tr(): - yield # Checklist - client.debug.press_yes() - yield # Set and confirm group count - client.debug.input("5") - yield # Checklist - client.debug.press_yes() - yield # Set and confirm group threshold - client.debug.input("3") - yield # Checklist - client.debug.press_yes() - for _ in range(5): # for each of 5 groups - yield # Number of shares info - client.debug.press_yes() - yield # Number of shares (5) - client.debug.input("5") - yield # Threshold info - client.debug.press_yes() - yield # Threshold (3) - client.debug.input("3") - yield # Confirm show seeds - client.debug.press_yes() - - # Mnemonic phrases - for _ in range(5): - for _ in range(5): - # Phrase screen - mnemonic = yield from read_and_confirm_mnemonic_tr(client.debug) - mnemonics.append(mnemonic) - yield # Confirm continue to next - client.debug.press_yes() - - yield # Confirm backup - client.debug.press_yes() with client: - if client.features.model == "T": - client.set_input_flow(input_flow_tt) - elif client.features.model == "R": - client.debug.watch_layout(True) - client.set_input_flow(input_flow_tr) + IF = InputFlowSlip39AdvancedBackup(client, click_info) + client.set_input_flow(IF.get()) client.set_expected_responses( [messages.ButtonRequest(code=B.ResetDevice)] * (8 if click_info else 6) # intro screens (and optional info) @@ -317,7 +154,7 @@ def test_backup_slip39_advanced(client: Client, click_info: bool): expected_ms = shamir.combine_mnemonics(MNEMONIC_SLIP39_ADVANCED_20) actual_ms = shamir.combine_mnemonics( - mnemonics[:3] + mnemonics[5:8] + mnemonics[10:13] + IF.mnemonics[:3] + IF.mnemonics[5:8] + IF.mnemonics[10:13] ) assert expected_ms == actual_ms diff --git a/tests/device_tests/test_msg_change_wipe_code_t2.py b/tests/device_tests/test_msg_change_wipe_code_t2.py index 8b9f46ac62..e5571e95fa 100644 --- a/tests/device_tests/test_msg_change_wipe_code_t2.py +++ b/tests/device_tests/test_msg_change_wipe_code_t2.py @@ -21,6 +21,8 @@ from trezorlib.client import MAX_PIN_LENGTH, PASSPHRASE_TEST_PATH from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.exceptions import Cancelled, TrezorFailure +from ..input_flows import InputFlowNewCodeMismatch + PIN4 = "1234" WIPE_CODE4 = "4321" WIPE_CODE6 = "456789" @@ -95,24 +97,12 @@ def test_set_remove_wipe_code(client: Client): def test_set_wipe_code_mismatch(client: Client): - # Let's set a wipe code. - def input_flow(): - yield # do you want to set the wipe code? - client.debug.press_yes() - yield # enter new wipe code - client.debug.input(WIPE_CODE4) - yield # enter new wipe code again (but different) - client.debug.input(WIPE_CODE6) - - # failed retry - yield # enter new wipe code - client.cancel() - with client, pytest.raises(Cancelled): + IF = InputFlowNewCodeMismatch(client, WIPE_CODE4, WIPE_CODE6) + client.set_input_flow(IF.get()) client.set_expected_responses( [messages.ButtonRequest()] * 4 + [messages.Failure()] ) - client.set_input_flow(input_flow) device.change_wipe_code(client) diff --git a/tests/device_tests/test_msg_changepin_t2.py b/tests/device_tests/test_msg_changepin_t2.py index cf520d149d..7c0cac13e4 100644 --- a/tests/device_tests/test_msg_changepin_t2.py +++ b/tests/device_tests/test_msg_changepin_t2.py @@ -21,6 +21,12 @@ from trezorlib.client import MAX_PIN_LENGTH, PASSPHRASE_TEST_PATH from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.exceptions import Cancelled, TrezorFailure +from ..input_flows import ( + InputFlowCodeChangeFail, + InputFlowNewCodeMismatch, + InputFlowWrongPIN, +) + PIN4 = "1234" PIN60 = "789456" * 10 PIN_MAX = "".join(chr((i % 10) + ord("0")) for i in range(MAX_PIN_LENGTH)) @@ -28,7 +34,7 @@ PIN_MAX = "".join(chr((i % 10) + ord("0")) for i in range(MAX_PIN_LENGTH)) pytestmark = pytest.mark.skip_t1 -def _check_pin(client: Client, pin): +def _check_pin(client: Client, pin: str): client.lock() assert client.features.pin_protection is True assert client.features.unlocked is False @@ -116,22 +122,10 @@ def test_set_failed(client: Client): # Check that there's no PIN protection _check_no_pin(client) - # Let's set new PIN - def input_flow(): - yield # do you want to set pin? - client.debug.press_yes() - yield # enter new pin - client.debug.input(PIN4) - yield # enter new pin again (but different) - client.debug.input(PIN60) - - # failed retry - yield # enter new pin - client.cancel() - with client, pytest.raises(Cancelled): + IF = InputFlowNewCodeMismatch(client, PIN4, PIN60) + client.set_input_flow(IF.get()) client.set_expected_responses([messages.ButtonRequest] * 4 + [messages.Failure]) - client.set_input_flow(input_flow) device.change_pin(client) @@ -148,24 +142,10 @@ def test_change_failed(client: Client): # Check current PIN value _check_pin(client, PIN4) - # Let's set new PIN - def input_flow(): - yield # do you want to change pin? - client.debug.press_yes() - yield # enter current pin - client.debug.input(PIN4) - yield # enter new pin - client.debug.input("457891") - yield # enter new pin again (but different) - client.debug.input("381847") - - # failed retry - yield # enter current pin again - client.cancel() - with client, pytest.raises(Cancelled): + IF = InputFlowCodeChangeFail(client, PIN4, "457891", "381847") + client.set_input_flow(IF.get()) client.set_expected_responses([messages.ButtonRequest] * 5 + [messages.Failure]) - client.set_input_flow(input_flow) device.change_pin(client) @@ -182,18 +162,10 @@ def test_change_invalid_current(client: Client): # Check current PIN value _check_pin(client, PIN4) - # Let's set new PIN - def input_flow(): - yield # do you want to change pin? - client.debug.press_yes() - yield # enter wrong current pin - client.debug.input(PIN60) - yield - client.debug.press_no() - with client, pytest.raises(TrezorFailure): + IF = InputFlowWrongPIN(client, PIN60) + client.set_input_flow(IF.get()) client.set_expected_responses([messages.ButtonRequest] * 3 + [messages.Failure]) - client.set_input_flow(input_flow) device.change_pin(client) diff --git a/tests/device_tests/test_pin.py b/tests/device_tests/test_pin.py index abe4d99ce5..5096190977 100644 --- a/tests/device_tests/test_pin.py +++ b/tests/device_tests/test_pin.py @@ -22,7 +22,8 @@ from trezorlib import messages from trezorlib.debuglink import TrezorClientDebugLink as Client from trezorlib.exceptions import PinException -from ..common import get_test_address +from ..common import check_PIN_backoff_time, get_test_address +from ..input_flows import InputFlowPINBackoff PIN4 = "1234" BAD_PIN = "5678" @@ -79,13 +80,6 @@ def test_incorrect_pin_t2(client: Client): get_test_address(client) -def _check_backoff_time(attempts: int, start: float) -> None: - """Helper to assert the exponentially growing delay after incorrect PIN attempts""" - expected = (2**attempts) - 1 - got = round(time.time() - start, 2) - assert got >= expected - - @pytest.mark.skip_t2 @pytest.mark.skip_tr def test_exponential_backoff_t1(client: Client): @@ -94,21 +88,12 @@ def test_exponential_backoff_t1(client: Client): with client, pytest.raises(PinException): client.use_pin_sequence([BAD_PIN]) get_test_address(client) - _check_backoff_time(attempt, start) + check_PIN_backoff_time(attempt, start) @pytest.mark.skip_t1 def test_exponential_backoff_t2(client: Client): - def input_flow(): - """Inputting some bad PINs and finally the correct one""" - yield # PIN entry - for attempt in range(3): - start = time.time() - client.debug.input(BAD_PIN) - yield # PIN entry - _check_backoff_time(attempt, start) - client.debug.input(PIN4) - with client: - client.set_input_flow(input_flow) + IF = InputFlowPINBackoff(client, BAD_PIN, PIN4) + client.set_input_flow(IF.get()) get_test_address(client) diff --git a/tests/input_flows.py b/tests/input_flows.py new file mode 100644 index 0000000000..d65f6cace1 --- /dev/null +++ b/tests/input_flows.py @@ -0,0 +1,1738 @@ +""" +Central place for defining all input flows for the device tests. + +Each model has potentially its own input flow, and in most cases +we need to distinguish between them. Doing it at one place +offers a better overview of the differences and makes it easier +to maintain. The whole `device_tests` folder can then focus +only on the actual tests and data-assertions, not on the lower-level +input flow details. +""" + +import time +from typing import Callable, Generator, Optional + +from trezorlib import messages +from trezorlib.debuglink import ( + DebugLink, + LayoutContent, + TrezorClientDebugLink as Client, + multipage_content, +) + +from . import buttons +from .common import ( + check_PIN_backoff_time, + click_info_button, + click_through, + read_and_confirm_mnemonic, + recovery_enter_shares, +) + +GeneratorType = Generator[None, messages.ButtonRequest, None] + +B = messages.ButtonRequestType + + +class InputFlowBase: + def __init__(self, client: Client): + self.client = client + self.debug = client.debug + self.layout = client.debug.wait_layout + + def model(self) -> Optional[str]: + return self.client.features.model + + def get(self) -> Callable[[], GeneratorType]: + self.client.watch_layout(True) + + # There could be one common input flow for all models + if hasattr(self, "input_flow_common"): + return getattr(self, "input_flow_common") + elif self.model() == "T": + return self.input_flow_tt + elif self.model() == "R": + return self.input_flow_tr + else: + raise ValueError("Unknown model") + + def input_flow_tt(self) -> GeneratorType: + """Special for TT""" + raise NotImplementedError + + def input_flow_tr(self) -> GeneratorType: + """Special for TR""" + raise NotImplementedError + + +class InputFlowSetupDevicePINWIpeCode(InputFlowBase): + def __init__(self, client: Client, pin: str, wipe_code: str): + super().__init__(client) + self.pin = pin + self.wipe_code = wipe_code + + def input_flow_common(self) -> GeneratorType: + yield # do you want to set/change the wipe_code? + self.debug.press_yes() + yield # enter current pin + self.debug.input(self.pin) + yield # enter new wipe code + self.debug.input(self.wipe_code) + yield # enter new wipe code again + self.debug.input(self.wipe_code) + yield # success + self.debug.press_yes() + + +class InputFlowNewCodeMismatch(InputFlowBase): + def __init__(self, client: Client, first_code: str, second_code: str): + super().__init__(client) + self.first_code = first_code + self.second_code = second_code + + def input_flow_common(self) -> GeneratorType: + yield # do you want to set the PIN/wipe_code? + self.debug.press_yes() + yield # enter new PIN/wipe_code + self.debug.input(self.first_code) + yield # enter new PIN/wipe_code again (but different) + self.debug.input(self.second_code) + + # failed retry + yield # enter new PIN/wipe_code + self.client.cancel() + + +class InputFlowCodeChangeFail(InputFlowBase): + def __init__( + self, client: Client, current_pin: str, new_pin_1: str, new_pin_2: str + ): + super().__init__(client) + self.current_pin = current_pin + self.new_pin_1 = new_pin_1 + self.new_pin_2 = new_pin_2 + + def input_flow_common(self) -> GeneratorType: + 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 + self.debug.input(self.new_pin_1) + yield # enter new pin again (but different) + self.debug.input(self.new_pin_2) + + # failed retry + yield # enter current pin again + self.client.cancel() + + +class InputFlowWrongPIN(InputFlowBase): + def __init__(self, client: Client, wrong_pin: str): + super().__init__(client) + self.wrong_pin = wrong_pin + + def input_flow_common(self) -> GeneratorType: + yield # do you want to change pin? + self.debug.press_yes() + yield # enter wrong current pin + self.debug.input(self.wrong_pin) + yield + self.debug.press_no() + + +class InputFlowPINBackoff(InputFlowBase): + def __init__(self, client: Client, wrong_pin: str, good_pin: str): + super().__init__(client) + self.wrong_pin = wrong_pin + self.good_pin = good_pin + + def input_flow_common(self) -> GeneratorType: + """Inputting some bad PINs and finally the correct one""" + yield # PIN entry + for attempt in range(3): + start = time.time() + self.debug.input(self.wrong_pin) + yield # PIN entry + check_PIN_backoff_time(attempt, start) + self.debug.input(self.good_pin) + + +class InputFlowSignMessagePagination(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + self.message_read = "" + + def input_flow_tt(self) -> GeneratorType: + # collect screen contents into `message_read`. + # Using a helper debuglink function to assemble the final text. + layouts: list[LayoutContent] = [] + + br = yield # confirm address + self.debug.wait_layout() + self.debug.press_yes() + + br = yield + assert br.pages is not None + for i in range(br.pages): + layout = self.debug.wait_layout() + layouts.append(layout) + + if i < br.pages - 1: + self.debug.swipe_up() + + self.message_read = multipage_content(layouts) + + self.debug.press_yes() + + def input_flow_tr(self) -> GeneratorType: + # confirm address + yield + self.debug.press_yes() + + br = yield + # TODO: try load the message_read the same way as in model T + if br.pages is not None: + for i in range(br.pages): + if i < br.pages - 1: + self.debug.swipe_up() + self.debug.press_yes() + + +class InputFlowShowAddressQRCode(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + + def input_flow_tt(self) -> GeneratorType: + yield + self.debug.press_no() + yield + self.debug.press_yes() + + def input_flow_tr(self) -> GeneratorType: + yield + self.debug.press_right() + self.debug.press_yes() + + +class InputFlowShowMultisigXPUBs(InputFlowBase): + def __init__(self, client: Client, address: str, xpubs: list[str], index: int): + super().__init__(client) + self.address = address + self.xpubs = xpubs + self.index = index + + def input_flow_tt(self) -> GeneratorType: + yield # show address + layout = self.debug.wait_layout() # TODO: do not need to *wait* now? + assert layout.get_title() == "MULTISIG 2 OF 3" + assert layout.get_content().replace(" ", "") == self.address + + self.debug.press_no() + yield # show QR code + assert "Painter" in self.debug.wait_layout().text + + # Three xpub pages with the same testing logic + for xpub_num in range(3): + expected_title = f"XPUB #{xpub_num + 1} " + ( + "(yours)" if self.index == xpub_num else "(cosigner)" + ) + + self.debug.press_no() + yield # show XPUB#{xpub_num} + layout1 = self.debug.wait_layout() + assert layout1.get_title() == expected_title + self.debug.swipe_up() + + layout2 = self.debug.wait_layout() + assert layout2.get_title() == expected_title + content = (layout1.get_content() + layout2.get_content()).replace(" ", "") + assert content == self.xpubs[xpub_num] + + self.debug.press_yes() + + +class InputFlowPaymentRequestDetails(InputFlowBase): + def __init__(self, client: Client, outputs: list[messages.TxOutputType]): + super().__init__(client) + self.outputs = outputs + + def input_flow_tt(self) -> GeneratorType: + yield # request to see details + self.debug.wait_layout() + self.debug.press_info() + + yield # confirm first output + assert self.outputs[0].address[:16] in self.layout().text + self.debug.press_yes() + yield # confirm first output + self.debug.wait_layout() + self.debug.press_yes() + + yield # confirm second output + assert self.outputs[1].address[:16] in self.layout().text + self.debug.press_yes() + yield # confirm second output + self.debug.wait_layout() + self.debug.press_yes() + + yield # confirm transaction + self.debug.press_yes() + yield # confirm transaction + self.debug.press_yes() + + +class InputFlowSignTxHighFee(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + self.finished = False + + def go_through_all_screens(self, screens: list[B]) -> GeneratorType: + for expected in screens: + br = yield + assert br.code == expected + self.debug.press_yes() + + self.finished = True + + def input_flow_tt(self) -> GeneratorType: + screens = [ + B.ConfirmOutput, + B.ConfirmOutput, + B.FeeOverThreshold, + B.SignTx, + B.SignTx, + ] + yield from self.go_through_all_screens(screens) + + def input_flow_tr(self) -> GeneratorType: + screens = [ + B.ConfirmOutput, + B.FeeOverThreshold, + B.SignTx, + ] + yield from self.go_through_all_screens(screens) + + +def lock_time_input_flow_tt( + debug: DebugLink, layout_assert_func: Callable[[str], None] +) -> GeneratorType: + yield # confirm output + debug.wait_layout() + debug.press_yes() + yield # confirm output + debug.wait_layout() + debug.press_yes() + + yield # confirm locktime + layout_text = debug.wait_layout().text + layout_assert_func(layout_text) + debug.press_yes() + + yield # confirm transaction + debug.press_yes() + yield # confirm transaction + debug.press_yes() + + +def lock_time_input_flow_tr( + debug: DebugLink, layout_assert_func: Callable[[str], None] +) -> GeneratorType: + yield # confirm output + debug.wait_layout() + debug.swipe_up() + debug.wait_layout() + debug.press_yes() + + yield # confirm locktime + layout_text = debug.wait_layout().text + layout_assert_func(layout_text) + debug.press_yes() + + yield # confirm transaction + debug.press_yes() + + +class InputFlowLockTimeBlockHeight(InputFlowBase): + def __init__(self, client: Client, block_height: str): + super().__init__(client) + self.block_height = block_height + + def layout_assert_func(self, layout_text: str) -> None: + assert "blockheight" in layout_text + assert self.block_height in layout_text + + def input_flow_tt(self) -> GeneratorType: + yield from lock_time_input_flow_tt(self.debug, self.layout_assert_func) + + def input_flow_tr(self) -> GeneratorType: + yield from lock_time_input_flow_tr(self.debug, self.layout_assert_func) + + +class InputFlowLockTimeDatetime(InputFlowBase): + def __init__(self, client: Client, lock_time_str: str): + super().__init__(client) + self.lock_time_str = lock_time_str + + def layout_assert_func(self, layout_text: str) -> None: + assert self.lock_time_str in layout_text + + def input_flow_tt(self) -> GeneratorType: + yield from lock_time_input_flow_tt(self.debug, self.layout_assert_func) + + def input_flow_tr(self) -> GeneratorType: + yield from lock_time_input_flow_tr(self.debug, self.layout_assert_func) + + +class InputFlowEIP712ShowMore(InputFlowBase): + SHOW_MORE = (143, 167) + + def __init__(self, client: Client): + super().__init__(client) + self.same_for_all_models = True + + def _confirm_show_more(self) -> None: + """Model-specific, either clicks a screen or presses a button.""" + if self.model() == "T": + self.debug.click(self.SHOW_MORE) + elif self.model() == "R": + self.debug.press_yes() + + def input_flow_common(self) -> GeneratorType: + """Triggers show more wherever possible""" + yield # confirm domain + self.debug.wait_layout() + self._confirm_show_more() + + # confirm domain properties + for _ in range(4): + yield + self.debug.press_yes() + + yield # confirm message + self.debug.wait_layout() + self._confirm_show_more() + + yield # confirm message.from + self.debug.wait_layout() + self._confirm_show_more() + + # confirm message.from properties + for _ in range(2): + yield + self.debug.press_yes() + + yield # confirm message.to + self.debug.wait_layout() + self._confirm_show_more() + + # confirm message.to properties + for _ in range(2): + yield + self.debug.press_yes() + + yield # confirm message.contents + self.debug.press_yes() + + yield # confirm final hash + self.debug.press_yes() + + +class InputFlowEIP712Cancel(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + + def input_flow_common(self) -> GeneratorType: + """Clicks cancelling button""" + yield # confirm domain + self.debug.press_no() + + +class InputFlowEthereumSignTxSkip(InputFlowBase): + def __init__(self, client: Client, cancel: bool = False): + super().__init__(client) + self.cancel = cancel + + def input_flow_common(self) -> GeneratorType: + yield # confirm address + self.debug.press_yes() + yield # confirm amount + self.debug.wait_layout() + self.debug.press_yes() + yield # confirm data + if self.cancel: + self.debug.press_no() + else: + self.debug.press_yes() + yield # gas price + self.debug.press_yes() + yield # maximum fee + self.debug.press_yes() + yield # hold to confirm + self.debug.press_yes() + + +class InputFlowEthereumSignTxScrollDown(InputFlowBase): + SHOW_ALL = (143, 167) + + def __init__(self, client: Client, cancel: bool = False): + super().__init__(client) + self.cancel = cancel + + def input_flow_tt(self) -> GeneratorType: + yield # confirm address + self.debug.wait_layout() + self.debug.press_yes() + yield # confirm amount + self.debug.wait_layout() + self.debug.press_yes() + yield # confirm data + self.debug.wait_layout() + self.debug.click(self.SHOW_ALL) + + br = yield # paginated data + assert br.pages is not None + for i in range(br.pages): + self.debug.wait_layout() + if i < br.pages - 1: + self.debug.swipe_up() + + self.debug.press_yes() + yield # confirm data + if self.cancel: + self.debug.press_no() + else: + self.debug.press_yes() + yield # gas price + self.debug.press_yes() + yield # maximum fee + self.debug.press_yes() + yield # hold to confirm + self.debug.press_yes() + + def input_flow_tr(self) -> GeneratorType: + yield # confirm address + self.debug.wait_layout() + self.debug.press_yes() + yield # confirm amount + self.debug.wait_layout() + self.debug.swipe_up() + self.debug.press_yes() + + br = yield # paginated data + assert br.pages is not None + for i in range(br.pages): + self.debug.wait_layout() + if i < br.pages - 1: + self.debug.swipe_up() + + self.debug.press_yes() + yield # confirm data + if self.cancel: + self.debug.press_no() + else: + self.debug.press_yes() + yield # gas price + self.debug.press_yes() + yield # maximum fee + self.debug.press_yes() + yield # hold to confirm + self.debug.press_yes() + + +class InputFlowEthereumSignTxGoBack(InputFlowBase): + SHOW_ALL = (143, 167) + GO_BACK = (16, 220) + + def __init__(self, client: Client, cancel: bool = False): + super().__init__(client) + self.cancel = cancel + + def input_flow_tt(self) -> GeneratorType: + br = yield # confirm address + self.debug.wait_layout() + self.debug.press_yes() + br = yield # confirm amount + self.debug.wait_layout() + self.debug.press_yes() + br = yield # confirm data + self.debug.wait_layout() + self.debug.click(self.SHOW_ALL) + + br = yield # paginated data + assert br.pages is not None + for i in range(br.pages): + self.debug.wait_layout() + if i == 2: + self.debug.click(self.GO_BACK) + yield # confirm data + self.debug.wait_layout() + if self.cancel: + self.debug.press_no() + else: + self.debug.press_yes() + yield # confirm address + self.debug.wait_layout() + self.debug.press_yes() + yield # confirm amount + self.debug.wait_layout() + self.debug.press_yes() + yield # hold to confirm + self.debug.wait_layout() + self.debug.press_yes() + return + + elif i < br.pages - 1: + self.debug.swipe_up() + + +def get_mnemonic_and_confirm_success( + debug: DebugLink, +) -> Generator[None, "messages.ButtonRequest", str]: + # mnemonic phrases + mnemonic = yield from read_and_confirm_mnemonic(debug) + + br = yield # confirm recovery seed check + assert br.code == B.Success + debug.press_yes() + + br = yield # confirm success + assert br.code == B.Success + debug.press_yes() + + assert mnemonic is not None + return mnemonic + + +class InputFlowBip39Backup(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + self.mnemonic = None + + def input_flow_common(self) -> GeneratorType: + # 1. Confirm Reset + yield from click_through(self.debug, screens=1, code=B.ResetDevice) + + # mnemonic phrases and rest + self.mnemonic = yield from get_mnemonic_and_confirm_success(self.debug) + + +class InputFlowBip39ResetBackup(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + self.mnemonic = None + + # NOTE: same as above, just two more YES + def input_flow_common(self) -> GeneratorType: + # 1. Confirm Reset + # 2. Backup your seed + # 3. Confirm warning + yield from click_through(self.debug, screens=3, code=B.ResetDevice) + + # mnemonic phrases and rest + self.mnemonic = yield from get_mnemonic_and_confirm_success(self.debug) + + +class InputFlowBip39ResetPIN(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + self.mnemonic = None + + def input_flow_common(self) -> GeneratorType: + br = yield # Confirm Reset + assert br.code == B.ResetDevice + self.debug.press_yes() + + yield # Enter new PIN + self.debug.input("654") + + yield # Re-enter PIN + self.debug.press_yes() + + yield # Confirm PIN + self.debug.input("654") + + br = yield # Confirm entropy + assert br.code == B.ResetDevice + self.debug.press_yes() + + br = yield # Backup your seed + assert br.code == B.ResetDevice + self.debug.press_yes() + + br = yield # Confirm warning + assert br.code == B.ResetDevice + self.debug.press_yes() + + # mnemonic phrases + self.mnemonic = yield from read_and_confirm_mnemonic(self.debug) + + br = yield # confirm recovery seed check + assert br.code == B.Success + self.debug.press_yes() + + br = yield # confirm success + assert br.code == B.Success + self.debug.press_yes() + + +class InputFlowBip39ResetFailedCheck(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + self.mnemonic = None + + def input_flow_common(self) -> GeneratorType: + # 1. Confirm Reset + # 2. Backup your seed + # 3. Confirm warning + yield from click_through(self.debug, screens=3, code=B.ResetDevice) + + # mnemonic phrases, wrong answer + self.mnemonic = yield from read_and_confirm_mnemonic( + self.debug, choose_wrong=True + ) + + br = yield # warning screen + assert br.code == B.ResetDevice + self.debug.press_yes() + + # mnemonic phrases + self.mnemonic = yield from read_and_confirm_mnemonic(self.debug) + + br = yield # confirm recovery seed check + assert br.code == B.Success + self.debug.press_yes() + + br = yield # confirm success + assert br.code == B.Success + self.debug.press_yes() + + +def load_5_shares( + debug: DebugLink, +) -> Generator[None, "messages.ButtonRequest", list[str]]: + mnemonics: list[str] = [] + + for _ in range(5): + # Phrase screen + mnemonic = yield from read_and_confirm_mnemonic(debug) + assert mnemonic is not None + mnemonics.append(mnemonic) + + br = yield # Confirm continue to next + assert br.code == B.Success + debug.press_yes() + + return mnemonics + + +class InputFlowSlip39BasicBackup(InputFlowBase): + def __init__(self, client: Client, click_info: bool): + super().__init__(client) + self.mnemonics: list[str] = [] + self.click_info = click_info + + def input_flow_tt(self) -> GeneratorType: + yield # 1. Checklist + self.debug.press_yes() + if self.click_info: + yield from click_info_button(self.debug) + yield # 2. Number of shares (5) + self.debug.press_yes() + yield # 3. Checklist + self.debug.press_yes() + if self.click_info: + yield from click_info_button(self.debug) + yield # 4. Threshold (3) + self.debug.press_yes() + yield # 5. Checklist + self.debug.press_yes() + yield # 6. Confirm show seeds + self.debug.press_yes() + + # Mnemonic phrases + self.mnemonics = yield from load_5_shares(self.debug) + + br = yield # Confirm backup + assert br.code == B.Success + self.debug.press_yes() + + def input_flow_tr(self) -> GeneratorType: + yield # Checklist + self.debug.press_yes() + yield # Number of shares info + self.debug.press_yes() + yield # Number of shares (5) + self.debug.input("5") + yield # Checklist + self.debug.press_yes() + yield # Threshold info + self.debug.press_yes() + yield # Threshold (3) + self.debug.input("3") + yield # Checklist + self.debug.press_yes() + yield # Confirm show seeds + self.debug.press_yes() + + # Mnemonic phrases + self.mnemonics = yield from load_5_shares(self.debug) + + br = yield # Confirm backup + assert br.code == B.Success + self.debug.press_yes() + + +class InputFlowSlip39BasicResetRecovery(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + self.mnemonics: list[str] = [] + + def input_flow_tt(self) -> GeneratorType: + # 1. Confirm Reset + # 2. Backup your seed + # 3. Confirm warning + # 4. shares info + # 5. Set & Confirm number of shares + # 6. threshold info + # 7. Set & confirm threshold value + # 8. Confirm show seeds + yield from click_through(self.debug, screens=8, code=B.ResetDevice) + + # Mnemonic phrases + self.mnemonics = yield from load_5_shares(self.debug) + + br = yield # safety warning + assert br.code == B.Success + self.debug.press_yes() + + def input_flow_tr(self) -> GeneratorType: + yield # Confirm Reset + self.debug.press_yes() + yield # Backup your seed + self.debug.press_yes() + yield # Checklist + self.debug.press_yes() + yield # Number of shares info + self.debug.press_yes() + yield # Number of shares (5) + self.debug.input("5") + yield # Checklist + self.debug.press_yes() + yield # Threshold info + self.debug.press_yes() + yield # Threshold (3) + self.debug.input("3") + yield # Checklist + self.debug.press_yes() + yield # Confirm show seeds + self.debug.press_yes() + + # Mnemonic phrases + self.mnemonics = yield from load_5_shares(self.debug) + + br = yield # Confirm backup + assert br.code == B.Success + self.debug.press_yes() + + +def load_5_groups_5_shares( + debug: DebugLink, +) -> Generator[None, "messages.ButtonRequest", list[str]]: + mnemonics: list[str] = [] + + for _g in range(5): + for _s in range(5): + # Phrase screen + mnemonic = yield from read_and_confirm_mnemonic(debug) + assert mnemonic is not None + mnemonics.append(mnemonic) + br = yield # Confirm continue to next + assert br.code == B.Success + debug.press_yes() + + return mnemonics + + +class InputFlowSlip39AdvancedBackup(InputFlowBase): + def __init__(self, client: Client, click_info: bool): + super().__init__(client) + self.mnemonics: list[str] = [] + self.click_info = click_info + + def input_flow_tt(self) -> GeneratorType: + yield # 1. Checklist + self.debug.press_yes() + if self.click_info: + yield from click_info_button(self.debug) + yield # 2. Set and confirm group count + self.debug.press_yes() + yield # 3. Checklist + self.debug.press_yes() + if self.click_info: + yield from click_info_button(self.debug) + yield # 4. Set and confirm group threshold + self.debug.press_yes() + yield # 5. Checklist + self.debug.press_yes() + for _ in range(5): # for each of 5 groups + if self.click_info: + yield from click_info_button(self.debug) + yield # Set & Confirm number of shares + self.debug.press_yes() + if self.click_info: + yield from click_info_button(self.debug) + yield # Set & confirm share threshold value + self.debug.press_yes() + yield # Confirm show seeds + self.debug.press_yes() + + # Mnemonic phrases - show & confirm shares for all groups + self.mnemonics = yield from load_5_groups_5_shares(self.debug) + + br = yield # Confirm backup + assert br.code == B.Success + self.debug.press_yes() + + def input_flow_tr(self) -> GeneratorType: + yield # 1. Checklist + self.debug.press_yes() + yield # 2. Set and confirm group count + self.debug.input("5") + yield # 3. Checklist + self.debug.press_yes() + yield # 4. Set and confirm group threshold + self.debug.input("3") + yield # 5. Checklist + self.debug.press_yes() + for _ in range(5): # for each of 5 groups + yield # Number of shares info + self.debug.press_yes() + yield # Number of shares (5) + self.debug.input("5") + yield # Threshold info + self.debug.press_yes() + yield # Threshold (3) + self.debug.input("3") + yield # Confirm show seeds + self.debug.press_yes() + + # Mnemonic phrases - show & confirm shares for all groups + self.mnemonics = yield from load_5_groups_5_shares(self.debug) + + br = yield # Confirm backup + assert br.code == B.Success + self.debug.press_yes() + + +class InputFlowSlip39AdvancedResetRecovery(InputFlowBase): + def __init__(self, client: Client, click_info: bool): + super().__init__(client) + self.mnemonics: list[str] = [] + self.click_info = click_info + + def input_flow_tt(self) -> GeneratorType: + # 1. Confirm Reset + # 2. Backup your seed + # 3. Confirm warning + # 4. shares info + # 5. Set & Confirm number of groups + # 6. threshold info + # 7. Set & confirm group threshold value + # 8-17: for each of 5 groups: + # 1. Set & Confirm number of shares + # 2. Set & confirm share threshold value + # 18. Confirm show seeds + yield from click_through(self.debug, screens=18, code=B.ResetDevice) + + # Mnemonic phrases - show & confirm shares for all groups + self.mnemonics = yield from load_5_groups_5_shares(self.debug) + + br = yield # safety warning + assert br.code == B.Success + self.debug.press_yes() + + def input_flow_tr(self) -> GeneratorType: + yield # Wallet backup + self.debug.press_yes() + yield # Wallet creation + self.debug.press_yes() + yield # Checklist + self.debug.press_yes() + yield # Set and confirm group count + self.debug.input("5") + yield # Checklist + self.debug.press_yes() + yield # Set and confirm group threshold + self.debug.input("3") + yield # Checklist + self.debug.press_yes() + for _ in range(5): # for each of 5 groups + yield # Number of shares info + self.debug.press_yes() + yield # Number of shares (5) + self.debug.input("5") + yield # Threshold info + self.debug.press_yes() + yield # Threshold (3) + self.debug.input("3") + yield # Confirm show seeds + self.debug.press_yes() + + # Mnemonic phrases - show & confirm shares for all groups + self.mnemonics = yield from load_5_groups_5_shares(self.debug) + + br = yield # safety warning + assert br.code == B.Success + self.debug.press_yes() + + +def enter_recovery_seed_dry_run(debug: DebugLink, mnemonic: list[str]) -> GeneratorType: + yield + assert "check the recovery seed" in debug.wait_layout().get_content() + debug.click(buttons.OK) + + yield + assert "Select number of words" in debug.wait_layout().get_content() + debug.click(buttons.OK) + + yield + assert "SelectWordCount" in debug.wait_layout().text + # 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 recovery seed" in debug.wait_layout().get_content() + debug.click(buttons.OK) + + yield + for word in mnemonic: + assert debug.wait_layout().text == "< MnemonicKeyboard >" + debug.input(word) + + +class InputFlowBip39RecoveryDryRun(InputFlowBase): + def __init__(self, client: Client, mnemonic: list[str]): + super().__init__(client) + self.mnemonic = mnemonic + + def input_flow_tt(self) -> GeneratorType: + yield from enter_recovery_seed_dry_run(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.layout().text + self.debug.press_right() + + yield + assert "select the number of words" in self.layout().text + self.debug.press_yes() + + yield + yield + assert "NUMBER OF WORDS" in self.layout().text + word_options = (12, 18, 20, 24, 33) + index = word_options.index(len(self.mnemonic)) + for _ in range(index): + self.debug.press_right() + self.debug.input(str(len(self.mnemonic))) + + yield + assert "enter your recovery seed" in self.layout().text + self.debug.press_yes() + + yield + self.debug.press_right() + yield + for word in self.mnemonic: + assert "WORD" in self.layout().text + self.debug.input(word) + + self.debug.wait_layout() + self.debug.press_right() + yield + self.debug.press_yes() + yield + + +class InputFlowBip39RecoveryDryRunInvalid(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + + def input_flow_tt(self) -> GeneratorType: + mnemonic = ["stick"] * 12 + yield from enter_recovery_seed_dry_run(self.debug, mnemonic) + + br = yield + assert br.code == messages.ButtonRequestType.Warning + assert "invalid recovery seed" in self.layout().get_content() + self.debug.click(buttons.OK) + + yield # retry screen + assert "Select number of words" in self.layout().get_content() + self.debug.click(buttons.CANCEL) + + yield + assert "ABORT SEED CHECK" == self.layout().get_title() + self.debug.click(buttons.OK) + + def input_flow_tr(self) -> GeneratorType: + yield + assert "check the recovery seed" in self.layout().text + self.debug.press_right() + + yield + assert "select the number of words" in self.layout().text + self.debug.press_yes() + + yield + yield + assert "NUMBER OF WORDS" in self.layout().text + # select 12 words + self.debug.press_middle() + + yield + assert "enter your recovery seed" in self.layout().text + self.debug.press_yes() + + yield + assert "WORD ENTERING" in self.layout().text + self.debug.press_yes() + + yield + for _ in range(12): + assert "WORD" in self.layout().text + self.debug.input("stick") + + br = yield + assert br.code == messages.ButtonRequestType.Warning + assert "invalid recovery seed" in self.layout().text + self.debug.press_right() + + yield # retry screen + assert "select the number of words" in self.layout().text + self.debug.press_left() + + yield + assert "abort" in self.layout().text + self.debug.press_right() + + +def bip39_recovery_possible_pin( + debug: DebugLink, mnemonic: list[str], pin: Optional[str] +) -> GeneratorType: + yield + assert ( + "Do you really want to recover a wallet?" in debug.wait_layout().get_content() + ) + debug.press_yes() + + # PIN when requested + if pin is not None: + yield + assert debug.wait_layout().text == "< PinKeyboard >" + debug.input(pin) + + yield + assert debug.wait_layout().text == "< PinKeyboard >" + debug.input(pin) + + yield + assert "Select number of words" in debug.wait_layout().get_content() + debug.press_yes() + + yield + assert "SelectWordCount" in debug.wait_layout().text + debug.input(str(len(mnemonic))) + + yield + assert "Enter recovery seed" in debug.wait_layout().get_content() + debug.press_yes() + + yield + for word in mnemonic: + assert debug.wait_layout().text == "< MnemonicKeyboard >" + debug.input(word) + + yield + assert ( + "You have successfully recovered your wallet." + in debug.wait_layout().get_content() + ) + debug.press_yes() + + +class InputFlowBip39RecoveryPIN(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(self.debug, self.mnemonic, pin="654") + + def input_flow_tr(self) -> GeneratorType: + yield + assert "By continuing you agree" in self.layout().text + self.debug.press_right() + assert "trezor.io/tos" in self.layout().text + self.debug.press_yes() + + yield + assert "ENTER" in self.layout().text + self.debug.input("654") + + yield + assert "re-enter PIN to confirm" in self.layout().text + self.debug.press_right() + + yield + assert "ENTER" in self.layout().text + self.debug.input("654") + + yield + assert "select the number of words" in self.layout().text + self.debug.press_yes() + + yield + yield + assert "NUMBER OF WORDS" in self.layout().text + self.debug.input(str(len(self.mnemonic))) + + yield + assert "enter your recovery seed" in self.layout().text + self.debug.press_yes() + + yield + assert "WORD ENTERING" in self.layout().text + self.debug.press_right() + + yield + for word in self.mnemonic: + assert "WORD" in self.layout().text + self.debug.input(word) + + yield + assert "You have finished recovering your wallet." in self.layout().text + self.debug.press_yes() + + +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(self.debug, self.mnemonic, pin=None) + + def input_flow_tr(self) -> GeneratorType: + yield # Confirm recovery + self.debug.press_yes() + yield # Homescreen + self.debug.press_yes() + + yield # Enter word count + self.debug.input(str(len(self.mnemonic))) + + yield # Homescreen + self.debug.press_yes() + yield # Homescreen + self.debug.press_yes() + yield # Enter words + for word in self.mnemonic: + self.debug.input(word) + + yield # confirm success + self.debug.press_yes() + yield + + +class InputFlowSlip39AdvancedRecoveryDryRun(InputFlowBase): + def __init__(self, client: Client, shares: list[str]): + super().__init__(client) + self.shares = shares + + def input_flow_common(self) -> GeneratorType: + yield # Confirm Dryrun + self.debug.press_yes() + # run recovery flow + yield from recovery_enter_shares(self.debug, self.shares, groups=True) + + +class InputFlowSlip39AdvancedRecovery(InputFlowBase): + def __init__(self, client: Client, shares: list[str], click_info: bool): + super().__init__(client) + self.shares = shares + self.click_info = click_info + + def input_flow_common(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + # Proceed with recovery + yield from recovery_enter_shares( + self.debug, self.shares, groups=True, click_info=self.click_info + ) + + +class InputFlowSlip39AdvancedRecoveryAbort(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + + def input_flow_common(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + yield # Homescreen - abort process + self.debug.press_no() + yield # Homescreen - confirm abort + self.debug.press_yes() + + +class InputFlowSlip39AdvancedRecoveryNoAbort(InputFlowBase): + def __init__(self, client: Client, shares: list[str]): + super().__init__(client) + self.shares = shares + + def input_flow_common(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + yield # Homescreen - abort process + self.debug.press_no() + yield # Homescreen - go back to process + self.debug.press_no() + yield from recovery_enter_shares(self.debug, self.shares, groups=True) + + +class InputFlowSlip39AdvancedRecoveryTwoSharesWarning(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_tt(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + yield # Homescreen - start process + self.debug.press_yes() + yield # Enter number of words + self.debug.input(str(len(self.first_share))) + yield # Homescreen - proceed to share entry + self.debug.press_yes() + yield # Enter first share + for word in self.first_share: + self.debug.input(word) + + yield # Continue to next share + self.debug.press_yes() + yield # Homescreen - next share + self.debug.press_yes() + yield # Enter next share + for word in self.second_share: + self.debug.input(word) + + br = yield + assert br.code == messages.ButtonRequestType.Warning + + self.client.cancel() + + def input_flow_tr(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + yield # Homescreen - start process + self.debug.press_yes() + yield # Enter number of words + self.debug.input(str(len(self.first_share))) + yield # Homescreen - proceed to share entry + self.debug.press_yes() + yield # Enter first share + self.debug.press_yes() + yield # Enter first share + for word in self.first_share: + self.debug.input(word) + + yield # Continue to next share + self.debug.press_yes() + yield # Homescreen - next share + self.debug.press_yes() + yield # Homescreen - next share + self.debug.press_yes() + yield # Enter next share + for word in self.second_share: + self.debug.input(word) + + yield + br = yield + assert br.code == messages.ButtonRequestType.Warning + self.debug.press_right() + self.debug.press_yes() + yield + + self.client.cancel() + + +def slip39_recovery_possible_pin( + debug: DebugLink, shares: list[str], pin: Optional[str] +) -> GeneratorType: + yield # Confirm Recovery/Dryrun + debug.press_yes() + + if pin is not None: + yield # Enter PIN + debug.input(pin) + yield # Please re-enter PIN to confirm + debug.press_yes() + yield # Enter PIN again + debug.input(pin) + + # Proceed with recovery + yield from recovery_enter_shares(debug, shares) + + +class InputFlowSlip39BasicRecovery(InputFlowBase): + def __init__(self, client: Client, shares: list[str]): + super().__init__(client) + self.shares = shares + + def input_flow_common(self) -> GeneratorType: + yield from slip39_recovery_possible_pin(self.debug, self.shares, pin=None) + + +class InputFlowSlip39BasicRecoveryPIN(InputFlowBase): + def __init__(self, client: Client, shares: list[str], pin: str): + super().__init__(client) + self.shares = shares + self.pin = pin + + def input_flow_common(self) -> GeneratorType: + yield from slip39_recovery_possible_pin(self.debug, self.shares, pin=self.pin) + + +class InputFlowSlip39BasicRecoveryAbort(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + + def input_flow_common(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + yield # Homescreen - abort process + self.debug.press_no() + yield # Homescreen - confirm abort + self.debug.press_yes() + + +class InputFlowSlip39BasicRecoveryNoAbort(InputFlowBase): + def __init__(self, client: Client, shares: list[str]): + super().__init__(client) + self.shares = shares + + def input_flow_common(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + yield # Homescreen - abort process + self.debug.press_no() + yield # Homescreen - go back to process + self.debug.press_no() + # run recovery flow + yield from recovery_enter_shares(self.debug, self.shares) + + +def slip39_recovery_setup_and_first_share( + debug: DebugLink, first_share: list[str] +) -> GeneratorType: + yield # Homescreen - start process + debug.press_yes() + yield # Enter number of words + debug.input(str(len(first_share))) + yield # Homescreen - proceed to share entry + debug.press_yes() + yield # Enter first share + for word in first_share: + debug.input(word) + + +class InputFlowSlip39BasicRecoveryRetryFirst(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + + def input_flow_tt(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + + 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 + 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 + self.debug.press_yes() + + yield # Homescreen + self.debug.press_no() + yield # Confirm abort + self.debug.press_yes() + + def input_flow_tr(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + yield # Homescreen - start process + self.debug.press_yes() + yield # Enter number of words + self.debug.input("20") + yield # Homescreen - proceed to share entry + self.debug.press_yes() + yield # Enter first share + self.debug.press_yes() + for _ in range(20): + self.debug.input("slush") + + yield + # assert br.code == messages.ButtonRequestType.Warning + self.debug.press_yes() + + yield # Homescreen - start process + self.debug.press_yes() + yield # Enter number of words + self.debug.input("33") + yield # Homescreen - proceed to share entry + self.debug.press_yes() + yield # Homescreen - proceed to share entry + self.debug.press_yes() + yield + for _ in range(33): + self.debug.input("slush") + + yield + self.debug.press_yes() + + yield + self.debug.press_no() + + yield + self.debug.press_right() + + yield + self.debug.press_right() + + yield + self.debug.press_right() + + yield + self.debug.press_yes() + + +class InputFlowSlip39BasicRecoveryRetrySecond(InputFlowBase): + def __init__(self, client: Client, shares: list[str]): + super().__init__(client) + self.shares = shares + + def input_flow_tt(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + + # 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 + self.debug.press_yes() + + yield # Enter another share + invalid_share = first_share[:3] + ["slush"] * 17 + for word in invalid_share: + self.debug.input(word) + + br = yield # Invalid share + assert br.code == messages.ButtonRequestType.Warning + self.debug.press_yes() + + yield # Proceed to next share + second_share = self.shares[1].split(" ") + for word in second_share: + self.debug.input(word) + + yield # More shares needed + self.debug.press_no() + yield # Confirm abort + self.debug.press_yes() + + def input_flow_tr(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + yield # Homescreen - start process + self.debug.press_yes() + yield # Enter number of words + self.debug.input("20") + yield # Homescreen - proceed to share entry + self.debug.press_yes() + yield # Enter first share + self.debug.press_yes() + yield # Enter first share + share = self.shares[0].split(" ") + for word in share: + self.debug.input(word) + + yield # More shares needed + self.debug.press_yes() + + yield # Enter another share + share = share[:3] + ["slush"] * 17 + for word in share: + self.debug.input(word) + + yield # Invalid share + # assert br.code == messages.ButtonRequestType.Warning + self.debug.press_yes() + + yield # Proceed to next share + share = self.shares[1].split(" ") + for word in share: + self.debug.input(word) + + yield # More shares needed + self.debug.press_no() + yield # Confirm abort + self.debug.press_yes() + + +class InputFlowSlip39BasicRecoveryWrongNthWord(InputFlowBase): + def __init__(self, client: Client, share: list[str], nth_word: int): + super().__init__(client) + self.share = share + self.nth_word = nth_word + + def input_flow_tt(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + + # First complete share + yield from slip39_recovery_setup_and_first_share(self.debug, self.share) + + yield # Continue to next share + self.debug.press_yes() + yield # Enter next share + 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 + + self.client.cancel() + + def input_flow_tr(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + yield # Homescreen - start process + self.debug.press_yes() + yield # Enter number of words + self.debug.input(str(len(self.share))) + yield # Homescreen - proceed to share entry + self.debug.press_yes() + yield # Enter first share + self.debug.press_yes() + yield # Enter first share + for word in self.share: + self.debug.input(word) + + yield # Continue to next share + self.debug.press_yes() + yield # Enter next share + self.debug.press_yes() + yield # Enter next share + for i, word in enumerate(self.share): + if i < self.nth_word: + self.debug.input(word) + else: + self.debug.input(self.share[-1]) + break + + yield + # assert br.code == messages.ButtonRequestType.Warning + + self.client.cancel() + + +class InputFlowSlip39BasicRecoverySameShare(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_tt(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + + # First complete share + yield from slip39_recovery_setup_and_first_share(self.debug, self.first_share) + + yield # Continue to next share + self.debug.press_yes() + yield # Enter next share + for word in self.second_share: + self.debug.input(word) + + br = yield + assert br.code == messages.ButtonRequestType.Warning + + # TODO: seems like the screenshots did not catch the WARNING + + self.client.cancel() + + def input_flow_tr(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + yield # Homescreen - start process + self.debug.press_yes() + yield # Enter number of words + self.debug.input(str(len(self.first_share))) + yield # Homescreen - proceed to share entry + self.debug.press_yes() + yield # Homescreen - proceed to share entry + self.debug.press_yes() + yield # Enter first share + for word in self.first_share: + self.debug.input(word) + + yield # Continue to next share + self.debug.press_yes() + yield # Continue to next share + self.debug.press_yes() + yield # Enter next share + for word in self.second_share: + self.debug.input(word) + + br = yield + br = yield + assert br.code == messages.ButtonRequestType.Warning + self.debug.press_right() + self.debug.press_yes() + yield + + self.client.cancel() + + +class InputFlowResetSkipBackup(InputFlowBase): + def __init__(self, client: Client): + super().__init__(client) + + def input_flow_common(self) -> GeneratorType: + yield # Confirm Recovery + self.debug.press_yes() + yield # Skip Backup + self.debug.press_no() + yield # Confirm skip backup + self.debug.press_no() diff --git a/tests/persistence_tests/test_wipe_code.py b/tests/persistence_tests/test_wipe_code.py index 709eff37b5..0764250825 100644 --- a/tests/persistence_tests/test_wipe_code.py +++ b/tests/persistence_tests/test_wipe_code.py @@ -3,6 +3,7 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client from ..common import MNEMONIC12 from ..emulators import EmulatorWrapper +from ..input_flows import InputFlowSetupDevicePINWIpeCode from ..upgrade_tests import core_only, legacy_only PIN = "1234" @@ -26,24 +27,12 @@ def setup_device_core(client: Client, pin: str, wipe_code: str) -> None: client, MNEMONIC12, pin, passphrase_protection=False, label="WIPECODE" ) - def input_flow(): - yield # do you want to set/change the wipe_code? - client.debug.press_yes() - if pin is not None: - yield # enter current pin - client.debug.input(pin) - yield # enter new wipe code - client.debug.input(wipe_code) - yield # enter new wipe code again - client.debug.input(wipe_code) - yield # success - client.debug.press_yes() - with client: client.set_expected_responses( [messages.ButtonRequest()] * 5 + [messages.Success, messages.Features] ) - client.set_input_flow(input_flow) + IF = InputFlowSetupDevicePINWIpeCode(client, pin, wipe_code) + client.set_input_flow(IF.get()) device.change_wipe_code(client)