trezor-firmware/tests/click_tests/reset.py

from typing import TYPE_CHECKING

from shamir_mnemonic import shamir

from trezorlib import messages

from .. import buttons

if TYPE_CHECKING:
    from trezorlib.debuglink import DebugLink


def confirm_wait(debug: "DebugLink", title: str) -> None:
    layout = debug.wait_layout()
    assert title.upper() in layout.get_title()
    debug.click(buttons.OK, wait=True)


def confirm_read(debug: "DebugLink", title: str) -> None:
    layout = debug.read_layout()
    if title == "Caution":
        assert "OK, I UNDERSTAND" in layout.text
    elif title == "Success":
        assert any(
            text in layout.get_content() for text in ("success", "finished", "done")
        )
    else:
        assert title.upper() in layout.get_title()
    debug.click(buttons.OK, wait=True)


def set_selection(debug: "DebugLink", button: tuple[int, int], diff: int) -> None:
    layout = debug.read_layout()
    assert "NumberInputDialog" in layout.text
    for _ in range(diff):
        debug.click(button, wait=False)
    debug.click(buttons.OK, wait=True)


def read_words(debug: "DebugLink", is_advanced: bool = False) -> list[str]:
    words: list[str] = []
    layout = debug.read_layout()
    if is_advanced:
        assert layout.get_title().startswith("GROUP")
    else:
        assert layout.get_title().startswith("RECOVERY SHARE #")

    # Swiping through all the page and loading the words
    for _ in range(layout.get_page_count() - 1):
        words.extend(layout.get_seed_words())
        layout = debug.input(swipe=messages.DebugSwipeDirection.UP, wait=True)
    words.extend(layout.get_seed_words())

    debug.press_yes()

    return words


def confirm_words(debug: "DebugLink", words: list[str]) -> None:
    layout = debug.wait_layout()
    assert "Select word" in layout.text
    for _ in range(3):
        # "Select word 3 of 20"
        #              ^
        word_pos = int(layout.get_content().split()[2])
        # Unifying both the buttons and words to lowercase
        btn_texts = [text.lower() for text in layout.get_button_texts()]
        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)


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