You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
trezor-firmware/core/src/apps/management/recovery_device/homescreen.py

274 lines
9.3 KiB

from typing import TYPE_CHECKING
import storage.device as storage_device
import storage.recovery as storage_recovery
import storage.recovery_shares as storage_recovery_shares
from trezor import TR, wire
from trezor.messages import Success
from .. import backup_types
from . import layout, recover
if TYPE_CHECKING:
from trezor.enums import BackupType
async def recovery_homescreen() -> None:
from trezor import workflow
from apps.homescreen import homescreen
if not storage_recovery.is_in_progress():
workflow.set_default(homescreen)
return
await recovery_process()
async def recovery_process() -> Success:
import storage
from trezor.enums import MessageType
wire.message_handler.AVOID_RESTARTING_FOR = (
MessageType.Initialize,
MessageType.GetFeatures,
)
try:
return await _continue_recovery_process()
except recover.RecoveryAborted:
dry_run = storage_recovery.is_dry_run()
if dry_run:
storage_recovery.end_progress()
else:
storage.wipe()
raise wire.ActionCancelled
async def _continue_recovery_process() -> Success:
from trezor import utils
from trezor.errors import MnemonicError
# gather the current recovery state from storage
dry_run = storage_recovery.is_dry_run()
word_count, backup_type = recover.load_slip39_state()
# 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:
assert word_count is not None
# If we continue recovery, show starting screen with word count immediately.
await _request_share_first_screen(word_count)
secret = None
while secret is None:
if is_first_step:
# If we are starting recovery, ask for word count first...
# _request_word_count
# For TT, just continuing straight to word count keyboard
if utils.INTERNAL_MODEL == "T2B1":
await layout.homescreen_dialog(
TR.buttons__continue, TR.recovery__num_of_words
)
# ask for the number of words
word_count = await layout.request_word_count(dry_run)
# ...and only then show the starting screen with word count.
await _request_share_first_screen(word_count)
assert word_count is not None
# ask for mnemonic words one by one
words = await layout.request_mnemonic(word_count, backup_type)
# if they were invalid or some checks failed we continue and request them again
if not words:
continue
try:
secret, backup_type = await _process_words(words)
# If _process_words succeeded, we now have both backup_type (from
# its result) and word_count (from _request_word_count earlier), which means
# that the first step is complete.
is_first_step = False
except MnemonicError:
await layout.show_invalid_mnemonic(word_count)
assert backup_type is not None
if dry_run:
result = await _finish_recovery_dry_run(secret, backup_type)
else:
result = await _finish_recovery(secret, backup_type)
return result
async def _finish_recovery_dry_run(secret: bytes, backup_type: BackupType) -> Success:
from trezor import utils
from trezor.crypto.hashlib import sha256
from apps.common import mnemonic
if backup_type is None:
raise RuntimeError
digest_input = sha256(secret).digest()
stored = mnemonic.get_secret()
digest_stored = sha256(stored).digest()
result = utils.consteq(digest_stored, digest_input)
is_slip39 = backup_types.is_slip39_backup_type(backup_type)
# Check that the identifier and iteration exponent match as well
if is_slip39:
result &= (
storage_device.get_slip39_identifier()
== storage_recovery.get_slip39_identifier()
)
result &= (
storage_device.get_slip39_iteration_exponent()
== storage_recovery.get_slip39_iteration_exponent()
)
storage_recovery.end_progress()
await layout.show_dry_run_result(result, is_slip39)
if result:
return Success(message="The seed is valid and matches the one in the device")
else:
raise wire.ProcessError("The seed does not match the one in the device")
async def _finish_recovery(secret: bytes, backup_type: BackupType) -> Success:
from trezor.enums import BackupType
from trezor.ui.layouts import show_success
if backup_type is None:
raise RuntimeError
storage_device.store_mnemonic_secret(
secret, backup_type, needs_backup=False, no_backup=False
)
if backup_type in (BackupType.Slip39_Basic, BackupType.Slip39_Advanced):
identifier = storage_recovery.get_slip39_identifier()
exponent = storage_recovery.get_slip39_iteration_exponent()
if identifier is None or exponent is None:
# Identifier and exponent need to be stored in storage at this point
raise RuntimeError
storage_device.set_slip39_identifier(identifier)
storage_device.set_slip39_iteration_exponent(exponent)
storage_recovery.end_progress()
await show_success("success_recovery", TR.recovery__wallet_recovered)
return Success(message="Device recovered")
async def _process_words(words: str) -> tuple[bytes | None, BackupType]:
word_count = len(words.split(" "))
is_slip39 = backup_types.is_slip39_word_count(word_count)
share = None
if not is_slip39: # BIP-39
secret: bytes | None = recover.process_bip39(words)
else:
secret, share = recover.process_slip39(words)
backup_type = backup_types.infer_backup_type(is_slip39, share)
if secret is None: # SLIP-39
assert share is not None
if share.group_count and share.group_count > 1:
await layout.show_group_share_success(share.index, share.group_index)
await _request_share_next_screen()
return secret, backup_type
async def _request_share_first_screen(word_count: int) -> None:
if backup_types.is_slip39_word_count(word_count):
remaining = storage_recovery.fetch_slip39_remaining_shares()
if remaining:
await _request_share_next_screen()
else:
await layout.homescreen_dialog(
TR.buttons__enter_share,
TR.recovery__enter_any_share,
TR.recovery__word_count_template.format(word_count),
show_info=True,
)
else: # BIP-39
await layout.homescreen_dialog(
TR.buttons__continue,
TR.recovery__enter_backup,
TR.recovery__word_count_template.format(word_count),
show_info=True,
)
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:
# 'remaining' should be stored at this point
raise RuntimeError
if group_count > 1:
await layout.homescreen_dialog(
TR.buttons__enter,
TR.recovery__more_shares_needed,
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
# TODO: consider kwargs in format here
entered = TR.recovery__x_of_y_entered_template.format(
already_entered_shares, overall_needed
)
needed = strings.format_plural(
TR.recovery__x_more_shares_needed_template_plural,
still_needed_shares,
TR.plurals__x_shares_needed,
)
await layout.homescreen_dialog(TR.buttons__enter_share, entered, needed)
async def _show_remaining_groups_and_shares() -> None:
"""
Show info dialog for Slip39 Advanced - what shares are to be entered.
"""
from trezor.crypto import slip39
shares_remaining = storage_recovery.fetch_slip39_remaining_shares()
# should be stored at this point
assert shares_remaining
groups = set()
first_entered_index = -1
for i, group_count in enumerate(shares_remaining):
if group_count < slip39.MAX_SHARE_COUNT:
first_entered_index = i
share = None
for index, remaining in enumerate(shares_remaining):
if 0 <= remaining < slip39.MAX_SHARE_COUNT:
m = storage_recovery_shares.fetch_group(index)[0]
if not share:
share = slip39.decode_mnemonic(m)
identifier = m.split(" ")[0:3]
groups.add((remaining, tuple(identifier)))
elif remaining == slip39.MAX_SHARE_COUNT: # no shares yet
identifier = storage_recovery_shares.fetch_group(first_entered_index)[
0
].split(" ")[0:2]
groups.add((remaining, tuple(identifier)))
assert share # share needs to be set
return await layout.show_remaining_shares(
groups, shares_remaining, share.group_threshold
)