1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-12-05 14:08:39 +00:00
trezor-firmware/tests/common.py
2024-11-28 13:59:17 +01:00

385 lines
14 KiB
Python

# This file is part of the Trezor project.
#
# Copyright (C) 2012-2019 SatoshiLabs and contributors
#
# This library is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version 3
# as published by the Free Software Foundation.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the License along with this library.
# If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
from __future__ import annotations
import json
import re
import time
from pathlib import Path
from typing import TYPE_CHECKING, Any, Generator, Optional
from unittest import mock
import pytest
from trezorlib import btc, messages, models, tools
from trezorlib.debuglink import LayoutType
from . import buttons
if TYPE_CHECKING:
from _pytest.mark.structures import MarkDecorator
from trezorlib.debuglink import DebugLink
from trezorlib.messages import ButtonRequest
from trezorlib.transport.session import Session
PRIVATE_KEYS_DEV = [byte * 32 for byte in (b"\xdd", b"\xde", b"\xdf")]
BRGeneratorType = Generator[None, messages.ButtonRequest, None]
# fmt: off
# 1 2 3 4 5 6 7 8 9 10 11 12
MNEMONIC12 = "alcohol woman abuse must during monitor noble actual mixed trade anger aisle"
MNEMONIC_SLIP39_BASIC_20_3of6 = [
"extra extend academic bishop cricket bundle tofu goat apart victim enlarge program behavior permit course armed jerky faint language modern",
"extra extend academic acne away best indicate impact square oasis prospect painting voting guest either argue username racism enemy eclipse",
"extra extend academic arcade born dive legal hush gross briefing talent drug much home firefly toxic analysis idea umbrella slice",
]
MNEMONIC_SLIP39_BASIC_20_3of6_SECRET = "491b795b80fc21ccdf466c0fbc98c8fc"
MNEMONIC_SLIP39_BASIC_EXT_20_2of3 = [
"enemy favorite academic acid cowboy phrase havoc level response walnut budget painting inside trash adjust froth kitchen learn tidy punish",
"enemy favorite academic always academic sniff script carpet romp kind promise scatter center unfair training emphasis evening belong fake enforce",
]
MNEMONIC_SLIP39_BASIC_EXT_20_2of3_SECRET = "644c905b0c4da21692f06fff3ed8b3e1"
# Shamir shares (128 bits, 2 groups from 1 of 1, 1 of 1, 3 of 5, 2 of 6)
MNEMONIC_SLIP39_ADVANCED_20 = [
"eraser senior beard romp adorn nuclear spill corner cradle style ancient family general leader ambition exchange unusual garlic promise voice",
"eraser senior ceramic snake clay various huge numb argue hesitate auction category timber browser greatest hanger petition script leaf pickup",
"eraser senior ceramic shaft dynamic become junior wrist silver peasant force math alto coal amazing segment yelp velvet image paces",
"eraser senior ceramic round column hawk trust auction smug shame alive greatest sheriff living perfect corner chest sled fumes adequate",
]
# Shamir shares (256 bits, 2 groups from 1 of 1, 1 of 1, 3 of 5, 2 of 6):
MNEMONIC_SLIP39_ADVANCED_33 = [
"wildlife deal beard romp alcohol space mild usual clothes union nuclear testify course research heat listen task location thank hospital slice smell failure fawn helpful priest ambition average recover lecture process dough stadium",
"wildlife deal acrobat romp anxiety axis starting require metric flexible geology game drove editor edge screw helpful have huge holy making pitch unknown carve holiday numb glasses survive already tenant adapt goat fangs",
]
MNEMONIC_SLIP39_CUSTOM_1of1 = ["tolerate flexible academic academic average dwarf square home promise aspect temple cluster roster forward hand unfair tenant emperor ceramic element forget perfect knit adapt review usual formal receiver typical pleasure duke yield party"]
MNEMONIC_SLIP39_CUSTOM_SECRET = "3439316237393562383066633231636364663436366330666263393863386663"
MNEMONIC_SLIP39_SINGLE_EXT_20 = ["academic again academic academic academic academic academic academic academic academic academic academic academic academic academic academic academic pecan provide remember"]
# External entropy mocked as received from trezorlib.
EXTERNAL_ENTROPY = b"zlutoucky kun upel divoke ody" * 2
# fmt: on
TEST_ADDRESS_N = tools.parse_path("m/44h/1h/0h/0/0")
COMMON_FIXTURES_DIR = (
Path(__file__).resolve().parent.parent / "common" / "tests" / "fixtures"
)
# So that all the random things are consistent
MOCK_OS_URANDOM = mock.Mock(return_value=EXTERNAL_ENTROPY)
WITH_MOCK_URANDOM = mock.patch("os.urandom", MOCK_OS_URANDOM)
def parametrize_using_common_fixtures(*paths: str) -> "MarkDecorator":
fixtures = []
for path in paths:
fixtures.append(json.loads((COMMON_FIXTURES_DIR / path).read_text()))
tests = []
for fixture in fixtures:
for test in fixture["tests"]:
test_id = test.get("name")
if not test_id:
test_id = test.get("description")
if test_id is not None:
test_id = test_id.lower().replace(" ", "_")
skip_models = test.get("skip_models", [])
skiplist = []
# TODO: genericify this
for skip_model in skip_models:
if skip_model in ("t1", "t1b1"):
skiplist.append(models.T1B1)
if skip_model in ("t2", "t2t1"):
skiplist.append(models.T2T1)
if skip_model in ("tr", "t2b1"):
skiplist.append(models.T2B1)
skiplist.append(models.T3B1)
if skip_model == "t3t1":
skiplist.append(models.T3T1)
if skiplist:
skip_marks = [pytest.mark.models(skip=skiplist)]
else:
skip_marks = []
tests.append(
pytest.param(
test["parameters"],
test["result"],
marks=[
pytest.mark.setup_client(
passphrase=fixture["setup"]["passphrase"],
mnemonic=fixture["setup"]["mnemonic"],
)
]
+ skip_marks,
id=test_id,
)
)
return pytest.mark.parametrize("parameters, result", tests)
def generate_entropy(
strength: int, internal_entropy: bytes, external_entropy: bytes
) -> bytes:
"""
strength - length of produced seed. One of 128, 192, 256
random - binary stream of random data from external HRNG
"""
import hashlib
if strength not in (128, 192, 256):
raise ValueError("Invalid strength")
if not internal_entropy:
raise ValueError("Internal entropy is not provided")
if len(internal_entropy) < 32:
raise ValueError("Internal entropy too short")
if not external_entropy:
raise ValueError("External entropy is not provided")
if len(external_entropy) < 32:
raise ValueError("External entropy too short")
entropy = hashlib.sha256(internal_entropy + external_entropy).digest()
entropy_stripped = entropy[: strength // 8]
if len(entropy_stripped) * 8 != strength:
raise ValueError("Entropy length mismatch")
return entropy_stripped
def click_through(
debug: "DebugLink", screens: int, code: Optional[messages.ButtonRequestType] = None
) -> Generator[None, "ButtonRequest", None]:
"""Click through N dialog screens.
For use in an input flow function.
Example:
def input_flow():
# 1. Confirm reset
# 2. Backup your seed
# 3. Confirm warning
# 4. Shares info
yield from click_through(client.debug, screens=4, code=ButtonRequestType.ResetDevice)
"""
for _ in range(screens):
received = yield
if code is not None:
assert received.code == code
debug.press_yes()
def read_and_confirm_mnemonic(
debug: "DebugLink", choose_wrong: bool = False
) -> Generator[None, "ButtonRequest", Optional[str]]:
"""Read a given number of mnemonic words from the screen and answer
confirmation questions.
Return the full mnemonic or None if `choose_wrong` is True.
For use in an input flow function.
Example:
def input_flow():
yield from click_through(client.debug, screens=3)
mnemonic = yield from read_and_confirm_mnemonic(client.debug)
"""
if debug.layout_type is LayoutType.TT:
mnemonic = yield from read_mnemonic_from_screen_tt(debug)
elif debug.layout_type is LayoutType.TR:
mnemonic = yield from read_mnemonic_from_screen_tr(debug)
elif debug.layout_type is LayoutType.Mercury:
mnemonic = yield from read_mnemonic_from_screen_mercury(debug)
else:
raise ValueError(f"Unknown model: {debug.layout_type}")
if not check_share(debug, mnemonic, choose_wrong):
return None
return " ".join(mnemonic)
def read_mnemonic_from_screen_tt(
debug: "DebugLink",
) -> Generator[None, "ButtonRequest", list[str]]:
mnemonic: list[str] = []
br = yield
assert br.pages is not None
for i in range(br.pages):
words = debug.read_layout().seed_words()
mnemonic.extend(words)
# Not swiping on the last page
if i < br.pages - 1:
debug.swipe_up()
debug.press_yes()
return mnemonic
def read_mnemonic_from_screen_tr(
debug: "DebugLink",
) -> Generator[None, "ButtonRequest", list[str]]:
mnemonic: list[str] = []
yield # write down all 12 words in order
debug.press_yes()
br = yield
assert br.pages is not None
for _ in range(br.pages - 1):
layout = debug.read_layout()
words = layout.seed_words()
mnemonic.extend(words)
debug.press_right()
debug.press_yes()
yield # Select correct words...
debug.press_right()
return mnemonic
def read_mnemonic_from_screen_mercury(
debug: "DebugLink",
) -> Generator[None, "ButtonRequest", list[str]]:
mnemonic: list[str] = []
br = yield
assert br.pages is not None
debug.read_layout()
debug.swipe_up()
for _ in range(br.pages - 2):
words = debug.read_layout().seed_words()
mnemonic.extend(words)
debug.swipe_up()
debug.read_layout()
debug.press_yes()
return mnemonic
def check_share(
debug: "DebugLink", mnemonic: list[str], choose_wrong: bool = False
) -> bool:
"""
Given the mnemonic word list, proceed with the backup check:
three rounds of `Select word X of Y` choices.
"""
re_num_of_word = r"\d+"
for _ in range(3):
if debug.layout_type is LayoutType.TT:
# T2T1 has position as the first number in the text
word_pos_match = re.search(
re_num_of_word, debug.read_layout().text_content()
)
elif debug.layout_type is LayoutType.TR:
# other models have the instruction in the title/subtitle
word_pos_match = re.search(re_num_of_word, debug.read_layout().title())
elif debug.layout_type is LayoutType.Mercury:
word_pos_match = re.search(re_num_of_word, debug.read_layout().subtitle())
else:
word_pos_match = None
assert word_pos_match is not None
word_pos = int(word_pos_match.group(0))
index = word_pos - 1
if choose_wrong:
debug.input(mnemonic[(index + 1) % len(mnemonic)])
return False
else:
debug.input(mnemonic[index])
return True
def click_info_button_tt(debug: "DebugLink") -> Generator[Any, Any, ButtonRequest]:
"""Click Shamir backup info button and return back."""
debug.press_info()
debug.press_yes()
return (yield)
def click_info_button_mercury(debug: "DebugLink"):
"""Click Shamir backup info button and return back."""
layout = debug.click(buttons.CORNER_BUTTON)
assert "VerticalMenu" in layout.all_components()
debug.click(buttons.VERTICAL_MENU[0])
debug.click(buttons.CORNER_BUTTON)
debug.click(buttons.CORNER_BUTTON)
def check_pin_backoff_time(attempts: int, start: float) -> None:
"""Helper to assert the exponentially growing delay after incorrect PIN attempts"""
expected = (2**attempts) - 1
got = round(time.time() - start, 2)
assert got >= expected
def get_test_address(session: "Session") -> str:
"""Fetch a testnet address on a fixed path. Useful to make a pin/passphrase
protected call, or to identify the root secret (seed+passphrase)"""
return btc.get_address(session, "Testnet", TEST_ADDRESS_N)
def compact_size(n: int) -> bytes:
if n < 253:
return n.to_bytes(1, "little")
elif n < 0x1_0000:
return bytes([253]) + n.to_bytes(2, "little")
elif n < 0x1_0000_0000:
return bytes([254]) + n.to_bytes(4, "little")
else:
return bytes([255]) + n.to_bytes(8, "little")
def get_text_possible_pagination(debug: "DebugLink", br: messages.ButtonRequest) -> str:
text = debug.read_layout().text_content()
if br.pages is not None:
for _ in range(br.pages - 1):
debug.swipe_up()
text += " "
text += debug.read_layout().text_content()
return text
def swipe_if_necessary(
debug: "DebugLink", br_code: messages.ButtonRequestType | None = None
) -> BRGeneratorType:
br = yield
if br_code is not None:
assert br.code == br_code
swipe_till_the_end(debug, br)
def swipe_till_the_end(debug: "DebugLink", br: messages.ButtonRequest) -> None:
if br.pages is not None:
for _ in range(br.pages - 1):
debug.swipe_up()
def is_core(session: "Session") -> bool:
return session.model is not models.T1B1