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

253 lines
8.9 KiB

import storage
import storage.device
import storage.recovery
import storage.recovery_shares
from trezor import strings, utils, wire, workflow
from trezor.crypto import slip39
from trezor.crypto.hashlib import sha256
from trezor.errors import MnemonicError
from trezor.messages import BackupType
from trezor.messages.Success import Success
from trezor.ui.layouts import show_success
from apps.common import mnemonic
from apps.homescreen.homescreen import homescreen
from .. import backup_types
from . import layout, recover
if False:
from trezor.messages.ResetDevice import EnumTypeBackupType
async def recovery_homescreen() -> None:
if not storage.recovery.is_in_progress():
workflow.set_default(homescreen)
return
# recovery process does not communicate on the wire
ctx = wire.DUMMY_CONTEXT
await recovery_process(ctx)
async def recovery_process(ctx: wire.GenericContext) -> Success:
try:
return await _continue_recovery_process(ctx)
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(ctx: wire.GenericContext) -> Success:
# 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(ctx, word_count)
secret = None
while secret is None:
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)
assert word_count is not None
# ask for mnemonic words one by one
words = await layout.request_mnemonic(ctx, 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(ctx, 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(ctx, word_count)
assert backup_type is not None
if dry_run:
result = await _finish_recovery_dry_run(ctx, secret, backup_type)
else:
result = await _finish_recovery(ctx, secret, backup_type)
return result
async def _finish_recovery_dry_run(
ctx: wire.GenericContext, secret: bytes, backup_type: EnumTypeBackupType
) -> Success:
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(ctx, 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(
ctx: wire.GenericContext, secret: bytes, backup_type: EnumTypeBackupType
) -> 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(
ctx, "success_recovery", "You have successfully recovered your wallet."
)
return Success(message="Device recovered")
async def _request_word_count(ctx: wire.GenericContext, dry_run: bool) -> int:
homepage = layout.RecoveryHomescreen("Select number of words")
await layout.homescreen_dialog(ctx, homepage, "Select")
# ask for the number of words
return await layout.request_word_count(ctx, dry_run)
async def _process_words(
ctx: wire.GenericContext, words: str
) -> tuple[bytes | None, EnumTypeBackupType]:
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(ctx, share.index, share.group_index)
await _request_share_next_screen(ctx)
return secret, backup_type
async def _request_share_first_screen(
ctx: wire.GenericContext, 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(ctx)
else:
content = layout.RecoveryHomescreen(
"Enter any share", "(%d words)" % word_count
)
await layout.homescreen_dialog(ctx, content, "Enter share")
else: # BIP-39
content = layout.RecoveryHomescreen(
"Enter recovery seed", "(%d words)" % word_count
)
await layout.homescreen_dialog(ctx, content, "Enter seed")
async def _request_share_next_screen(ctx: wire.GenericContext) -> None:
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:
content = layout.RecoveryHomescreen("More shares needed")
await layout.homescreen_dialog(
ctx, content, "Enter", _show_remaining_groups_and_shares
)
else:
text = strings.format_plural("{count} more {plural}", remaining[0], "share")
content = layout.RecoveryHomescreen(text, "needed to enter")
await layout.homescreen_dialog(ctx, content, "Enter share")
async def _show_remaining_groups_and_shares(ctx: wire.GenericContext) -> None:
"""
Show info dialog for Slip39 Advanced - what shares are to be entered.
"""
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 in range(len(shares_remaining)):
if shares_remaining[i] < 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(
ctx, groups, shares_remaining, share.group_threshold
)