mirror of
https://github.com/trezor/trezor-firmware.git
synced 2024-11-22 07:28:10 +00:00
301 lines
9.7 KiB
Python
301 lines
9.7 KiB
Python
from typing import TYPE_CHECKING
|
|
|
|
from trezorlib.debuglink import LayoutType
|
|
|
|
from .. import buttons
|
|
from .. import translations as TR
|
|
from .common import get_possible_btn_texts, go_next
|
|
|
|
if TYPE_CHECKING:
|
|
from trezorlib.debuglink import DebugLink, LayoutContent
|
|
|
|
|
|
DELETE_BTN_TEXTS = get_possible_btn_texts("inputs__delete") + get_possible_btn_texts(
|
|
"inputs__previous"
|
|
)
|
|
|
|
|
|
def enter_word(
|
|
debug: "DebugLink", word: str, is_slip39: bool = False
|
|
) -> "LayoutContent":
|
|
if debug.layout_type in (LayoutType.TT, LayoutType.Mercury):
|
|
typed_word = word[:4]
|
|
for coords in buttons.type_word(typed_word, is_slip39=is_slip39):
|
|
debug.click(coords)
|
|
if debug.layout_type is LayoutType.Mercury and not is_slip39 and len(word) > 4:
|
|
# T3T1 (mercury) BIP39 keyboard allows to "confirm" only if the word is fully written, you need to click the word to auto-complete
|
|
debug.click(buttons.CONFIRM_WORD)
|
|
return debug.click(buttons.CONFIRM_WORD)
|
|
elif debug.layout_type is LayoutType.TR:
|
|
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()
|
|
|
|
layout = debug.press_middle()
|
|
letter_index += 1
|
|
|
|
# Word choices
|
|
while not layout.get_middle_choice() == word:
|
|
layout = debug.press_right()
|
|
|
|
return debug.press_middle()
|
|
else:
|
|
raise ValueError("Unknown model")
|
|
|
|
|
|
def confirm_recovery(debug: "DebugLink", title: str = "recovery__title") -> None:
|
|
layout = debug.read_layout()
|
|
TR.assert_equals(layout.title(), title)
|
|
if debug.layout_type is LayoutType.TT:
|
|
debug.click(buttons.OK)
|
|
elif debug.layout_type is LayoutType.Mercury:
|
|
debug.swipe_up()
|
|
elif debug.layout_type is LayoutType.TR:
|
|
for _ in range(layout.page_count()):
|
|
debug.press_right()
|
|
|
|
|
|
def select_number_of_words(
|
|
debug: "DebugLink",
|
|
num_of_words: int = 20,
|
|
unlock_repeated_backup=False,
|
|
) -> None:
|
|
def select_tt() -> "LayoutContent":
|
|
# click the button from ValuePad
|
|
if unlock_repeated_backup:
|
|
coords_map = {20: buttons.grid34(0, 2), 33: buttons.grid34(1, 2)}
|
|
else:
|
|
coords_map = {
|
|
12: buttons.grid34(0, 2),
|
|
18: buttons.grid34(1, 2),
|
|
20: buttons.grid34(2, 2),
|
|
24: buttons.grid34(0, 3),
|
|
33: buttons.grid34(1, 3),
|
|
}
|
|
coords = coords_map.get(num_of_words)
|
|
if coords is None:
|
|
raise ValueError("Invalid num_of_words")
|
|
return debug.click(coords)
|
|
|
|
def select_tr() -> "LayoutContent":
|
|
# navigate to the number and confirm it
|
|
word_options = (20, 33) if unlock_repeated_backup else (12, 18, 20, 24, 33)
|
|
index = word_options.index(num_of_words)
|
|
for _ in range(index):
|
|
debug.press_right()
|
|
return debug.press_middle()
|
|
|
|
def select_mercury() -> "LayoutContent":
|
|
# click the button from ValuePad
|
|
if unlock_repeated_backup:
|
|
coords_map = {20: buttons.MERCURY_NO, 33: buttons.MERCURY_YES}
|
|
else:
|
|
coords_map = {
|
|
12: buttons.grid34(0, 1),
|
|
18: buttons.grid34(2, 1),
|
|
20: buttons.grid34(0, 2),
|
|
24: buttons.grid34(2, 2),
|
|
33: buttons.grid34(1, 3),
|
|
}
|
|
coords = coords_map.get(num_of_words)
|
|
if coords is None:
|
|
raise ValueError("Invalid num_of_words")
|
|
return debug.click(coords)
|
|
|
|
if debug.layout_type is LayoutType.TT:
|
|
TR.assert_equals(debug.read_layout().text_content(), "recovery__num_of_words")
|
|
layout = select_tt()
|
|
elif debug.layout_type is LayoutType.TR:
|
|
layout = debug.press_right()
|
|
TR.assert_equals(layout.title(), "word_count__title")
|
|
layout = select_tr()
|
|
elif debug.layout_type is LayoutType.Mercury:
|
|
layout = select_mercury()
|
|
else:
|
|
raise ValueError("Unknown model")
|
|
|
|
if unlock_repeated_backup:
|
|
if debug.layout_type is LayoutType.TR:
|
|
TR.assert_in(layout.text_content(), "recovery__enter_backup")
|
|
else:
|
|
TR.assert_in_multiple(
|
|
layout.text_content(),
|
|
["recovery__only_first_n_letters", "recovery__enter_each_word"],
|
|
)
|
|
elif num_of_words in (20, 33):
|
|
TR.assert_in_multiple(
|
|
layout.text_content(),
|
|
[
|
|
"recovery__enter_backup",
|
|
"recovery__enter_any_share",
|
|
"recovery__only_first_n_letters",
|
|
"recovery__enter_each_word",
|
|
],
|
|
)
|
|
else: # BIP-39
|
|
TR.assert_in_multiple(
|
|
layout.text_content(),
|
|
[
|
|
"recovery__enter_backup",
|
|
"recovery__only_first_n_letters",
|
|
"recovery__enter_each_word",
|
|
],
|
|
)
|
|
|
|
|
|
def enter_share(
|
|
debug: "DebugLink",
|
|
share: str,
|
|
is_first: bool = True,
|
|
before_title: str = "recovery__title_recover",
|
|
) -> "LayoutContent":
|
|
if debug.layout_type is LayoutType.TR:
|
|
TR.assert_in(debug.read_layout().title(), before_title)
|
|
layout = debug.read_layout()
|
|
for _ in range(layout.page_count()):
|
|
layout = debug.press_right()
|
|
elif debug.layout_type is LayoutType.Mercury:
|
|
layout = debug.swipe_up()
|
|
else:
|
|
TR.assert_in(debug.read_layout().title(), before_title)
|
|
layout = debug.click(buttons.OK)
|
|
|
|
assert "MnemonicKeyboard" in layout.all_components()
|
|
|
|
for word in share.split(" "):
|
|
layout = enter_word(debug, word, is_slip39=True)
|
|
|
|
return layout
|
|
|
|
|
|
def enter_shares(
|
|
debug: "DebugLink",
|
|
shares: list[str],
|
|
enter_share_before_title: str = "recovery__title_recover",
|
|
text: str = "recovery__enter_any_share",
|
|
after_layout_text: str = "recovery__wallet_recovered",
|
|
) -> None:
|
|
TR.assert_in_multiple(
|
|
debug.read_layout().text_content(),
|
|
[
|
|
"recovery__enter_backup",
|
|
"recovery__enter_any_share",
|
|
"recovery__only_first_n_letters",
|
|
"recovery__enter_each_word",
|
|
text,
|
|
],
|
|
)
|
|
for index, share in enumerate(shares):
|
|
enter_share(
|
|
debug, share, is_first=index == 0, before_title=enter_share_before_title
|
|
)
|
|
if index < len(shares) - 1:
|
|
# FIXME: when ui-t3t1 done for shamir, we want to check the template below
|
|
TR.assert_in(debug.read_layout().title(), enter_share_before_title)
|
|
# TR.assert_in(
|
|
# debug.read_layout().text_content(),
|
|
# "recovery__x_of_y_entered_template",
|
|
# template=(index + 1, len(shares)),
|
|
# )
|
|
|
|
TR.assert_in(debug.read_layout().text_content(), after_layout_text)
|
|
|
|
|
|
def enter_seed(
|
|
debug: "DebugLink",
|
|
seed_words: list[str],
|
|
is_slip39=False,
|
|
prepare_layout_text: str = "recovery__enter_backup",
|
|
after_layout_text: str = "recovery__wallet_recovered",
|
|
) -> None:
|
|
prepare_enter_seed(debug, prepare_layout_text)
|
|
|
|
for word in seed_words:
|
|
enter_word(debug, word, is_slip39=is_slip39)
|
|
|
|
TR.assert_in(debug.read_layout().text_content(), after_layout_text)
|
|
|
|
|
|
def enter_seed_previous_correct(
|
|
debug: "DebugLink", seed_words: list[str], bad_indexes: dict[int, str]
|
|
) -> None:
|
|
prepare_enter_seed(debug)
|
|
|
|
i = 0
|
|
go_back = False
|
|
bad_word = ""
|
|
while True:
|
|
assert i >= 0
|
|
|
|
if i >= len(seed_words):
|
|
break
|
|
|
|
if go_back:
|
|
go_back = False
|
|
if debug.layout_type is LayoutType.TT:
|
|
debug.swipe_right()
|
|
for _ in range(len(bad_word)):
|
|
debug.click(buttons.RECOVERY_DELETE)
|
|
elif debug.layout_type is LayoutType.TR:
|
|
layout = debug.read_layout()
|
|
|
|
while layout.get_middle_choice() not in DELETE_BTN_TEXTS:
|
|
layout = debug.press_right()
|
|
layout = debug.press_middle()
|
|
|
|
for _ in range(len(bad_word)):
|
|
while layout.get_middle_choice() not in DELETE_BTN_TEXTS:
|
|
layout = debug.press_left()
|
|
layout = debug.press_middle()
|
|
elif debug.layout_type is LayoutType.Mercury:
|
|
debug.click(buttons.RECOVERY_DELETE) # Top-left
|
|
for _ in range(len(bad_word)):
|
|
debug.click(buttons.RECOVERY_DELETE)
|
|
continue
|
|
|
|
if i in bad_indexes:
|
|
word = bad_indexes.pop(i)
|
|
bad_word = word
|
|
go_back = True
|
|
else:
|
|
word = seed_words[i]
|
|
i += 1
|
|
layout = enter_word(debug, word, is_slip39=False)
|
|
|
|
# TR.assert_in(debug.read_layout().text_content(), "recovery__wallet_recovered")
|
|
|
|
|
|
def prepare_enter_seed(
|
|
debug: "DebugLink", layout_text: str = "recovery__enter_backup"
|
|
) -> None:
|
|
TR.assert_in_multiple(
|
|
debug.read_layout().text_content(),
|
|
[
|
|
"recovery__enter_backup",
|
|
"recovery__only_first_n_letters",
|
|
"recovery__enter_each_word",
|
|
layout_text,
|
|
],
|
|
)
|
|
if debug.layout_type is LayoutType.TT:
|
|
debug.click(buttons.OK)
|
|
elif debug.layout_type is LayoutType.Mercury:
|
|
debug.swipe_up()
|
|
debug.swipe_up()
|
|
elif debug.layout_type is LayoutType.TR:
|
|
debug.press_right()
|
|
debug.press_right()
|
|
layout = debug.press_right()
|
|
assert "MnemonicKeyboard" in layout.all_components()
|
|
|
|
|
|
def finalize(debug: "DebugLink") -> None:
|
|
layout = go_next(debug)
|
|
assert layout is not None
|
|
assert layout.main_component() == "Homescreen"
|