diff --git a/core/src/apps/management/recovery_device/homescreen.py b/core/src/apps/management/recovery_device/homescreen.py index 3826dff420..4d4252b7cf 100644 --- a/core/src/apps/management/recovery_device/homescreen.py +++ b/core/src/apps/management/recovery_device/homescreen.py @@ -46,13 +46,23 @@ async def _continue_recovery_process(ctx: wire.Context) -> Success: # gather the current recovery state from storage dry_run = storage.recovery.is_dry_run() word_count, backup_type = recover.load_slip39_state() - if word_count: + + # Both word_count and backup_type are derived from the same data. Both will be + # either set or unset. We use 'backup_type is None' to detect status of both. + # The following variable indicates that we are (re)starting the first recovery step, + # which includes word count selection. + is_first_step = backup_type is None + + if not is_first_step: + # If we continue recovery, show starting screen with word count immediately. await _request_share_first_screen(ctx, word_count) secret = None while secret is None: - if not word_count: # the first run, prompt word count from the user + if is_first_step: + # If we are starting recovery, ask for word count first... word_count = await _request_word_count(ctx, dry_run) + # ...and only then show the starting screen with word count. await _request_share_first_screen(ctx, word_count) # ask for mnemonic words one by one @@ -63,7 +73,10 @@ async def _continue_recovery_process(ctx: wire.Context) -> Success: continue try: - secret, word_count, backup_type = await _process_words(ctx, words) + secret, backup_type = await _process_words(ctx, words) + # If _process_words succeeded, backup_type will be set. + # Otherwise we are still in "first step". + is_first_step = backup_type is None except MnemonicError: await layout.show_invalid_mnemonic(ctx, word_count) @@ -143,7 +156,6 @@ async def _request_word_count(ctx: wire.Context, dry_run: bool) -> int: async def _process_words( ctx: wire.Context, words: str ) -> Tuple[Optional[bytes], EnumTypeBackupType, int]: - word_count = len(words.split(" ")) is_slip39 = backup_types.is_slip39_word_count(word_count) @@ -159,7 +171,7 @@ async def _process_words( await layout.show_group_share_success(ctx, share.index, share.group_index) await _request_share_next_screen(ctx) - return secret, word_count, backup_type + return secret, backup_type async def _request_share_first_screen(ctx: wire.Context, word_count: int) -> None: diff --git a/tests/device_tests/test_msg_recoverydevice_slip39_basic.py b/tests/device_tests/test_msg_recoverydevice_slip39_basic.py index d62769a56f..eacc852fa2 100644 --- a/tests/device_tests/test_msg_recoverydevice_slip39_basic.py +++ b/tests/device_tests/test_msg_recoverydevice_slip39_basic.py @@ -133,6 +133,97 @@ def test_noabort(client): assert client.features.initialized is True +@pytest.mark.setup_client(uninitialized=True) +def test_ask_word_number(client): + debug = client.debug + + def input_flow_retry_first(): + 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") + + code = yield # Invalid share + assert 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") + + code = yield # Invalid share + assert code == messages.ButtonRequestType.Warning + debug.press_yes() + + yield # Homescreen + debug.press_no() + yield # Confirm abort + debug.press_yes() + + with client: + client.set_input_flow(input_flow_retry_first) + with pytest.raises(exceptions.Cancelled): + device.recover(client, pin_protection=False, label="label") + client.init_device() + assert client.features.initialized is False + + def input_flow_retry_second(): + 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) + + code = yield # Invalid share + assert 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: + client.set_input_flow(input_flow_retry_second) + with pytest.raises(exceptions.Cancelled): + device.recover(client, pin_protection=False, label="label") + client.init_device() + assert client.features.initialized is False + + @pytest.mark.setup_client(uninitialized=True) @pytest.mark.parametrize("nth_word", range(3)) def test_wrong_nth_word(client, nth_word):