mirror of
https://github.com/trezor/trezor-firmware.git
synced 2024-12-30 01:58:11 +00:00
155 lines
4.9 KiB
Python
155 lines
4.9 KiB
Python
from typing import TYPE_CHECKING
|
|
|
|
from .. import buttons
|
|
from .common import go_next
|
|
|
|
if TYPE_CHECKING:
|
|
from trezorlib.debuglink import DebugLink, LayoutContent
|
|
|
|
|
|
def enter_word(
|
|
debug: "DebugLink", word: str, is_slip39: bool = False
|
|
) -> "LayoutContent":
|
|
if debug.model == "T":
|
|
typed_word = word[:4]
|
|
for coords in buttons.type_word(typed_word, is_slip39=is_slip39):
|
|
debug.click(coords)
|
|
|
|
return debug.click(buttons.CONFIRM_WORD, wait=True)
|
|
elif debug.model == "R":
|
|
letter_index = 0
|
|
layout = debug.read_layout()
|
|
|
|
# Letter choices
|
|
while layout.find_values_by_key("letter_choices"):
|
|
letter = word[letter_index]
|
|
while not layout.get_middle_choice() == letter:
|
|
layout = debug.press_right(wait=True)
|
|
|
|
layout = debug.press_middle(wait=True)
|
|
letter_index += 1
|
|
|
|
# Word choices
|
|
while not layout.get_middle_choice() == word:
|
|
layout = debug.press_right(wait=True)
|
|
|
|
return debug.press_middle(wait=True)
|
|
else:
|
|
raise ValueError("Unknown model")
|
|
|
|
|
|
def confirm_recovery(debug: "DebugLink") -> None:
|
|
layout = debug.wait_layout()
|
|
if debug.model == "T":
|
|
assert layout.title().startswith(("RECOVER WALLET", "BACKUP CHECK"))
|
|
debug.click(buttons.OK, wait=True)
|
|
elif debug.model == "R":
|
|
assert layout.title() == "RECOVER WALLET"
|
|
debug.press_right(wait=True)
|
|
debug.press_right()
|
|
|
|
|
|
def select_number_of_words(
|
|
debug: "DebugLink", num_of_words: int = 20, wait: bool = True
|
|
) -> None:
|
|
if wait:
|
|
debug.wait_layout()
|
|
if debug.model == "T":
|
|
assert "number of words" in debug.read_layout().text_content()
|
|
assert debug.read_layout().title() in (
|
|
"BACKUP CHECK",
|
|
"RECOVER WALLET",
|
|
)
|
|
|
|
# click the number
|
|
word_option_offset = 6
|
|
word_options = (12, 18, 20, 24, 33)
|
|
index = word_option_offset + word_options.index(
|
|
num_of_words
|
|
) # raises if num of words is invalid
|
|
coords = buttons.grid34(index % 3, index // 3)
|
|
layout = debug.click(coords, wait=True)
|
|
elif debug.model == "R":
|
|
assert "number of words" in debug.read_layout().text_content()
|
|
layout = debug.press_right(wait=True)
|
|
|
|
assert layout.title() == "NUMBER OF WORDS"
|
|
|
|
# navigate to the number and confirm it
|
|
word_options = (12, 18, 20, 24, 33)
|
|
index = word_options.index(num_of_words)
|
|
for _ in range(index):
|
|
debug.press_right(wait=True)
|
|
layout = debug.press_middle(wait=True)
|
|
else:
|
|
raise ValueError("Unknown model")
|
|
|
|
if num_of_words in (20, 33):
|
|
assert "Enter any share" in layout.text_content()
|
|
else:
|
|
assert "Enter your backup" in layout.text_content()
|
|
|
|
|
|
def enter_share(
|
|
debug: "DebugLink", share: str, is_first: bool = True
|
|
) -> "LayoutContent":
|
|
if debug.model == "T":
|
|
layout = debug.click(buttons.OK, wait=True)
|
|
|
|
assert layout.main_component() == "MnemonicKeyboard"
|
|
for word in share.split(" "):
|
|
layout = enter_word(debug, word, is_slip39=True)
|
|
|
|
return layout
|
|
elif debug.model == "R":
|
|
assert "RECOVER WALLET" in debug.wait_layout().title()
|
|
layout = debug.press_right(wait=True)
|
|
if is_first:
|
|
# Word entering info
|
|
debug.press_right()
|
|
layout = debug.press_right(wait=True)
|
|
assert "MnemonicKeyboard" in layout.all_components()
|
|
|
|
for word in share.split(" "):
|
|
layout = enter_word(debug, word, is_slip39=True)
|
|
|
|
return layout
|
|
else:
|
|
raise ValueError("Unknown model")
|
|
|
|
|
|
def enter_shares(debug: "DebugLink", shares: list[str]) -> None:
|
|
layout = debug.read_layout()
|
|
expected_text = "Enter any share"
|
|
for index, share in enumerate(shares):
|
|
assert expected_text in layout.text_content()
|
|
layout = enter_share(debug, share, is_first=index == 0)
|
|
expected_text = f"{index + 1} of {len(shares)} shares entered"
|
|
|
|
assert "Wallet recovered successfully" in layout.text_content()
|
|
|
|
|
|
def enter_seed(debug: "DebugLink", seed_words: list[str]) -> None:
|
|
assert "Enter" in debug.read_layout().text_content()
|
|
if debug.model == "T":
|
|
layout = debug.click(buttons.OK, wait=True)
|
|
assert layout.main_component() == "MnemonicKeyboard"
|
|
elif debug.model == "R":
|
|
layout = debug.press_right(wait=True)
|
|
assert "RECOVER WALLET" in layout.title()
|
|
debug.press_right()
|
|
|
|
layout = debug.press_right(wait=True)
|
|
assert "MnemonicKeyboard" in layout.all_components()
|
|
|
|
for word in seed_words:
|
|
layout = enter_word(debug, word, is_slip39=False)
|
|
|
|
assert "Wallet recovered successfully" in layout.text_content() # type: ignore
|
|
|
|
|
|
def finalize(debug: "DebugLink") -> None:
|
|
layout = go_next(debug, wait=True)
|
|
assert layout is not None
|
|
assert layout.main_component() == "Homescreen"
|