1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-11-17 21:22:10 +00:00
trezor-firmware/tests/click_tests/reset.py
Martin Milata 73a7223e7b test(core): add T3T1 support
[no changelog]
2024-03-12 20:55:23 +00:00

136 lines
4.8 KiB
Python

import re
from typing import TYPE_CHECKING
from shamir_mnemonic import shamir # type: ignore
from trezorlib import messages, models
from .. import buttons
from .. import translations as TR
if TYPE_CHECKING:
from trezorlib.debuglink import DebugLink
def confirm_new_wallet(debug: "DebugLink") -> None:
TR.assert_equals_multiple(
debug.read_layout().title(),
["reset__title_create_wallet", "reset__title_create_wallet_shamir"],
)
if debug.model in (models.T2T1, models.T3T1):
debug.click(buttons.OK, wait=True)
elif debug.model in (models.T2B1,):
debug.press_right(wait=True)
debug.press_right(wait=True)
def confirm_read(debug: "DebugLink", middle_r: bool = False) -> None:
if debug.model in (models.T2T1, models.T3T1):
debug.click(buttons.OK, wait=True)
elif debug.model in (models.T2B1,):
page_count = debug.read_layout().page_count()
if page_count > 1:
for _ in range(page_count - 1):
debug.press_right(wait=True)
if middle_r:
debug.press_middle(wait=True)
else:
debug.press_right(wait=True)
def set_selection(debug: "DebugLink", button: tuple[int, int], diff: int) -> None:
if debug.model in (models.T2T1, models.T3T1):
assert "NumberInputDialog" in debug.read_layout().all_components()
for _ in range(diff):
debug.click(button)
debug.click(buttons.OK, wait=True)
elif debug.model in (models.T2B1,):
layout = debug.read_layout()
if layout.title() in TR.translate(
"reset__title_number_of_shares"
) + TR.translate("words__title_threshold"):
# Special info screens
layout = debug.press_right(wait=True)
assert "NumberInput" in layout.all_components()
if button == buttons.RESET_MINUS:
for _ in range(diff):
debug.press_left(wait=True)
else:
for _ in range(diff):
debug.press_right(wait=True)
debug.press_middle(wait=True)
def read_words(
debug: "DebugLink", backup_type: messages.BackupType, do_htc: bool = True
) -> list[str]:
words: list[str] = []
if debug.model in (models.T2B1,):
debug.press_right(wait=True)
# Swiping through all the pages and loading the words
layout = debug.read_layout()
for _ in range(layout.page_count() - 1):
words.extend(layout.seed_words())
layout = debug.swipe_up(wait=True)
assert layout is not None
if debug.model in (models.T2T1, models.T3T1):
words.extend(layout.seed_words())
# There is hold-to-confirm button
if do_htc:
if debug.model in (models.T2T1, models.T3T1):
debug.click_hold(buttons.OK, hold_ms=1500)
elif debug.model in (models.T2B1,):
debug.press_right_htc(1200)
else:
# It would take a very long time to test 16-of-16 with doing 1500 ms HTC after
# each word set
debug.press_yes()
return words
def confirm_words(debug: "DebugLink", words: list[str]) -> None:
layout = debug.wait_layout()
if debug.model in (models.T2T1, models.T3T1):
TR.assert_template(layout.text_content(), "reset__select_word_x_of_y_template")
for _ in range(3):
# "Select word 3 of 20"
# ^
word_pos_match = re.search(r"\d+", debug.wait_layout().text_content())
assert word_pos_match is not None
word_pos = int(word_pos_match.group(0))
# Unifying both the buttons and words to lowercase
btn_texts = [
text.lower() for text in layout.tt_check_seed_button_contents()
]
wanted_word = words[word_pos - 1].lower()
button_pos = btn_texts.index(wanted_word)
layout = debug.click(buttons.RESET_WORD_CHECK[button_pos], wait=True)
elif debug.model in (models.T2B1,):
TR.assert_in(layout.text_content(), "reset__select_correct_word")
layout = debug.press_right(wait=True)
for _ in range(3):
# "SELECT 2ND WORD"
# ^
word_pos_match = re.search(r"\d+", layout.title())
assert word_pos_match is not None
word_pos = int(word_pos_match.group(0))
wanted_word = words[word_pos - 1].lower()
while not layout.get_middle_choice() == wanted_word:
layout = debug.press_right(wait=True)
layout = debug.press_middle(wait=True)
def validate_mnemonics(mnemonics: list[str], expected_ems: bytes) -> None:
# We expect these combinations to recreate the secret properly
# In case of click tests the mnemonics are always XofX so no need for combinations
groups = shamir.decode_mnemonics(mnemonics)
ems = shamir.recover_ems(groups)
assert expected_ems == ems.ciphertext