From 82d4b5c065ec26e7ad4b176d37a7852eee4c9bcc Mon Sep 17 00:00:00 2001 From: matejcik Date: Mon, 20 Nov 2023 11:58:33 +0100 Subject: [PATCH] refactor(core): improve recovery_enter_share --- core/embed/rust/librust_qstr.h | 1 + core/embed/rust/src/ui/model_tr/layout.rs | 6 +-- core/embed/rust/src/ui/model_tt/layout.rs | 1 + core/mocks/generated/trezorui2.pyi | 3 +- .../management/recovery_device/homescreen.py | 27 ++---------- .../apps/management/recovery_device/layout.py | 41 ++++++++++++++++++- core/src/trezor/ui/layouts/tr/recovery.py | 6 +-- core/src/trezor/ui/layouts/tt/recovery.py | 4 +- 8 files changed, 56 insertions(+), 33 deletions(-) diff --git a/core/embed/rust/librust_qstr.h b/core/embed/rust/librust_qstr.h index fa94d7086d..c5420ae021 100644 --- a/core/embed/rust/librust_qstr.h +++ b/core/embed/rust/librust_qstr.h @@ -122,6 +122,7 @@ static void _librust_qstrs(void) { MP_QSTR_show_homescreen; MP_QSTR_show_info; MP_QSTR_show_info_with_cancel; + MP_QSTR_show_instructions; MP_QSTR_show_lockscreen; MP_QSTR_show_mismatch; MP_QSTR_show_passphrase; diff --git a/core/embed/rust/src/ui/model_tr/layout.rs b/core/embed/rust/src/ui/model_tr/layout.rs index 7f1163ed19..643f834316 100644 --- a/core/embed/rust/src/ui/model_tr/layout.rs +++ b/core/embed/rust/src/ui/model_tr/layout.rs @@ -1460,11 +1460,11 @@ extern "C" fn new_confirm_recovery(n_args: usize, args: *const Obj, kwargs: *mut let description: StrBuffer = kwargs.get(Qstr::MP_QSTR_description)?.try_into()?; let button: StrBuffer = kwargs.get(Qstr::MP_QSTR_button)?.try_into()?; let dry_run: bool = kwargs.get(Qstr::MP_QSTR_dry_run)?.try_into()?; - let show_info: bool = kwargs.get(Qstr::MP_QSTR_show_info)?.try_into()?; + let show_instructions: bool = kwargs.get(Qstr::MP_QSTR_show_instructions)?.try_into()?; let mut paragraphs = ParagraphVecShort::new(); paragraphs.add(Paragraph::new(&theme::TEXT_NORMAL, description)); - if show_info { + if show_instructions { let first = "You'll only have to select the first 2-4 letters of each word."; let second = "Position of the cursor will change between entries for enhanced security."; @@ -1994,7 +1994,7 @@ pub static mp_module_trezorui2: Module = obj_module! { /// button: str, /// dry_run: bool, /// info_button: bool, # unused on TR - /// show_info: bool, + /// show_instructions: bool, /// ) -> LayoutObj[UiResult]: /// """Device recovery homescreen.""" Qstr::MP_QSTR_confirm_recovery => obj_fn_kw!(0, new_confirm_recovery).as_obj(), diff --git a/core/embed/rust/src/ui/model_tt/layout.rs b/core/embed/rust/src/ui/model_tt/layout.rs index d41dca6330..6a30543fbe 100644 --- a/core/embed/rust/src/ui/model_tt/layout.rs +++ b/core/embed/rust/src/ui/model_tt/layout.rs @@ -2043,6 +2043,7 @@ pub static mp_module_trezorui2: Module = obj_module! { /// button: str, /// dry_run: bool, /// info_button: bool = False, + /// show_instructions: bool = False, # unused on TT /// ) -> LayoutObj[UiResult]: /// """Device recovery homescreen.""" Qstr::MP_QSTR_confirm_recovery => obj_fn_kw!(0, new_confirm_recovery).as_obj(), diff --git a/core/mocks/generated/trezorui2.pyi b/core/mocks/generated/trezorui2.pyi index 9229cdd61b..eb8e63ae12 100644 --- a/core/mocks/generated/trezorui2.pyi +++ b/core/mocks/generated/trezorui2.pyi @@ -371,7 +371,7 @@ def confirm_recovery( button: str, dry_run: bool, info_button: bool, # unused on TR - show_info: bool, + show_instructions: bool, ) -> LayoutObj[UiResult]: """Device recovery homescreen.""" @@ -883,6 +883,7 @@ def confirm_recovery( button: str, dry_run: bool, info_button: bool = False, + show_instructions: bool = False, # unused on TT ) -> LayoutObj[UiResult]: """Device recovery homescreen.""" diff --git a/core/src/apps/management/recovery_device/homescreen.py b/core/src/apps/management/recovery_device/homescreen.py index ef44dfe592..337ba275c0 100644 --- a/core/src/apps/management/recovery_device/homescreen.py +++ b/core/src/apps/management/recovery_device/homescreen.py @@ -188,12 +188,7 @@ async def _request_share_first_screen(word_count: int) -> None: if remaining: await _request_share_next_screen() else: - await layout.homescreen_dialog( - "Enter share", - "Enter any share", - f"({word_count} words)", - show_info=True, - ) + await layout.enter_share(word_count=word_count) else: # BIP-39 await layout.homescreen_dialog( "Continue", @@ -204,8 +199,6 @@ async def _request_share_first_screen(word_count: int) -> None: async def _request_share_next_screen() -> None: - from trezor import strings - remaining = storage_recovery.fetch_slip39_remaining_shares() group_count = storage_recovery.get_slip39_group_count() if not remaining: @@ -213,22 +206,10 @@ async def _request_share_next_screen() -> None: raise RuntimeError if group_count > 1: - await layout.homescreen_dialog( - "Enter", - "More shares needed", - info_func=_show_remaining_groups_and_shares, - ) + await layout.enter_share(info_func=_show_remaining_groups_and_shares) else: - still_needed_shares = remaining[0] - already_entered_shares = len(storage_recovery_shares.fetch_group(0)) - overall_needed = still_needed_shares + already_entered_shares - entered = ( - f"{already_entered_shares} of {overall_needed} shares entered successfully." - ) - needed = strings.format_plural( - "{count} more {plural} needed.", still_needed_shares, "share" - ) - await layout.homescreen_dialog("Enter share", entered, needed) + entered = len(storage_recovery_shares.fetch_group(0)) + await layout.enter_share(entered_remaining=(entered, remaining[0])) async def _show_remaining_groups_and_shares() -> None: diff --git a/core/src/apps/management/recovery_device/layout.py b/core/src/apps/management/recovery_device/layout.py index 97415a221c..76ba2e9dd1 100644 --- a/core/src/apps/management/recovery_device/layout.py +++ b/core/src/apps/management/recovery_device/layout.py @@ -12,7 +12,7 @@ from trezor.ui.layouts.recovery import ( # noqa: F401 from .. import backup_types if TYPE_CHECKING: - from typing import Callable + from typing import Awaitable, Callable from trezor.enums import BackupType @@ -120,6 +120,45 @@ async def show_invalid_mnemonic(word_count: int) -> None: ) +def enter_share( + word_count: int | None = None, + entered_remaining: tuple[int, int] | None = None, + info_func: Callable | None = None, +) -> Awaitable[None]: + from trezor import strings + + show_instructions = False + + if word_count is not None: + # First-time entry. Show instructions and word count. + text = "Enter any share" + subtext = f"({word_count} words)" + show_instructions = True + + elif entered_remaining is not None: + # Basic Shamir. There is only one group, we report entered/remaining count. + entered, remaining = entered_remaining + total = entered + remaining + text = f"{entered} of {total} shares entered successfully." + subtext = strings.format_plural( + "{count} more {plural} needed.", remaining, "share" + ) + + else: + # SuperShamir. We cannot easily show entered/remaining across groups, + # the caller provided an info_func that has the details. + text = "More shares needed." + subtext = None + + return homescreen_dialog( + "Enter share", + text, + subtext, + info_func, + show_instructions, + ) + + async def homescreen_dialog( button_label: str, text: str, diff --git a/core/src/trezor/ui/layouts/tr/recovery.py b/core/src/trezor/ui/layouts/tr/recovery.py index 8d50d4c31b..54de8c88b3 100644 --- a/core/src/trezor/ui/layouts/tr/recovery.py +++ b/core/src/trezor/ui/layouts/tr/recovery.py @@ -67,7 +67,7 @@ async def continue_recovery( subtext: str | None, info_func: Callable | None, dry_run: bool, - show_info: bool = False, + show_instructions: bool = False, ) -> bool: # TODO: implement info_func? # There is very limited space on the screen @@ -75,7 +75,7 @@ async def continue_recovery( # Never showing info for dry-run, user already saw it and it is disturbing if dry_run: - show_info = False + show_instructions = False if subtext: text += f"\n\n{subtext}" @@ -86,7 +86,7 @@ async def continue_recovery( button=button_label.upper(), info_button=False, dry_run=dry_run, - show_info=show_info, # type: ignore [No parameter named "show_info"] + show_instructions=show_instructions, ) result = await interact( homepage, diff --git a/core/src/trezor/ui/layouts/tt/recovery.py b/core/src/trezor/ui/layouts/tt/recovery.py index cd3fec45c1..7fd231612b 100644 --- a/core/src/trezor/ui/layouts/tt/recovery.py +++ b/core/src/trezor/ui/layouts/tt/recovery.py @@ -106,9 +106,9 @@ async def continue_recovery( subtext: str | None, info_func: Callable | None, dry_run: bool, - show_info: bool = False, # unused on TT + show_instructions: bool = False, ) -> bool: - if show_info: + if show_instructions: # Show this just one-time description = "You'll only have to select the first 2-4 letters of each word." else: