def main():
parser = ArgumentParser()
parser.add_argument("-m", dest="mnemonic", help="Set mnemonic", type=str)
parser.add_argument("-p", dest="pin", help="Set pin", type=str)
parser.add_argument("--passphrase",
dest="passphrase",
help="Enable passphrase",
action="store_true")
parser.add_argument("--no-passphrase",
dest="passphrase",
help="Enable passphrase",
action="store_false")
parser.set_defaults(passphrase=True)
args = parser.parse_args()
# Setup link
wirelink = get_device()
client = TrezorClientDebugLink(wirelink)
client.open()
device.wipe(client)
debuglink.load_device_by_mnemonic(client,
mnemonic=args.mnemonic,
pin=args.pin,
passphrase_protection=args.passphrase,
label='test')
print(args.mnemonic)
print(client.features)
client.close()
def apply_settings(
language: Optional[str] = None,
label: Optional[str] = None,
use_passphrase: Optional[bool] = None,
homescreen: Optional[str] = None,
auto_lock_delay_ms: Optional[int] = None,
display_rotation: Optional[int] = None,
passphrase_always_on_device: Optional[bool] = None,
safety_checks: Optional[int] = None,
) -> None:
"""Forwards settings fields to be applied on a device.
NOTE: does not handle the experimental_features argument,
seems that it is not yet supported in latest trezorlib
"""
# Homescreen needs to be bytes object, so if there,
# it should be encoded from the received string
homescreen_bytes = homescreen.encode() if homescreen else None
client = TrezorClientDebugLink(get_device())
client.open()
time.sleep(SLEEP)
device.apply_settings(
client,
label=label,
language=language,
use_passphrase=use_passphrase,
homescreen=homescreen_bytes,
passphrase_always_on_device=passphrase_always_on_device,
auto_lock_delay_ms=auto_lock_delay_ms,
display_rotation=display_rotation,
safety_checks=safety_checks,
)
client.close()
def allow_unsafe() -> None:
client = TrezorClientDebugLink(get_device())
client.open()
time.sleep(SLEEP)
# T1 does not support PromptAlways
if client.features.major_version == 1:
safety_checks = messages.SafetyCheckLevel.PromptTemporarily
else:
safety_checks = messages.SafetyCheckLevel.PromptAlways
# Some older devices do not support safety checks, so we know
# the command will fail with a specific error message
# T1 supports safety checks from 1.10.1 and T2 from 2.3.2
try:
device.apply_settings(client, safety_checks=safety_checks)
except TrezorFailure as err:
# Catching only specific error message, otherwise reraising the exception
if "No setting provided" in str(err):
log(
f"Could not allow unsafe. Device does not support safety checks. Err: {err}",
"red",
)
else:
raise
client.close()
def client():
wirelink = get_device()
client = TrezorClientDebugLink(wirelink)
wipe_device(client)
client.open()
yield client
client.close()
def allow_unsafe():
client = TrezorClientDebugLink(get_device())
# ignore for Legacy firmware, there is no such setting
if client.features.major_version == 1:
return
client.open()
time.sleep(SLEEP)
device.apply_settings(client, safety_checks=1) # TODO
client.close()
def apply_settings(passphrase_always_on_device=None):
client = TrezorClientDebugLink(get_device())
client.open()
time.sleep(SLEEP)
device.apply_settings(
client,
passphrase_always_on_device=bool(passphrase_always_on_device),
)
client.close()
class TrezorTest:
# 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"
mnemonic18 = "owner little vague addict embark decide pink prosper true fork panda embody mixture exchange choose canoe electric jewel"
mnemonic24 = "dignity pass list indicate nasty swamp pool script soccer toe leaf photo multiply desk host tomato cradle drill spread actor shine dismiss champion exotic"
mnemonic_all = " ".join(["all"] * 12)
# fmt: on
pin4 = "1234"
pin6 = "789456"
pin8 = "45678978"
def setup_method(self, method):
wirelink = conftest.get_device()
self.client = TrezorClientDebugLink(wirelink)
# self.client.set_buttonwait(3)
device.wipe(self.client)
self.client.open()
def teardown_method(self, method):
self.client.close()
def _setup_mnemonic(self, mnemonic=None, pin="", passphrase=False):
if mnemonic is None:
mnemonic = TrezorTest.mnemonic12
debuglink.load_device_by_mnemonic(
self.client,
mnemonic=mnemonic,
pin=pin,
passphrase_protection=passphrase,
label="test",
language="english",
)
if conftest.TREZOR_VERSION == 1:
# remove cached PIN (introduced via load_device)
self.client.clear_session()
if conftest.TREZOR_VERSION > 1 and passphrase:
device.apply_settings(self.client,
passphrase_source=PASSPHRASE_ON_HOST)
def setup_mnemonic_allallall(self):
self._setup_mnemonic(mnemonic=TrezorTest.mnemonic_all)
def setup_mnemonic_nopin_nopassphrase(self):
self._setup_mnemonic()
def setup_mnemonic_nopin_passphrase(self):
self._setup_mnemonic(passphrase=True)
def setup_mnemonic_pin_nopassphrase(self):
self._setup_mnemonic(pin=TrezorTest.pin4)
def setup_mnemonic_pin_passphrase(self):
self._setup_mnemonic(pin=TrezorTest.pin4, passphrase=True)
class EmulatorWrapper:
def __init__(self, gen, tag, storage=None):
self.gen = gen
self.tag = tag
self.workdir = tempfile.TemporaryDirectory()
if storage:
open(self._storage_file(), "wb").write(storage)
def __enter__(self):
if self.tag.startswith("/"): # full path+filename provided
args = [self.tag]
else: # only gen+tag provided
args = ["%s/trezor-emu-%s-%s" % (BINDIR, self.gen, self.tag)]
env = ENV
if self.gen == "core":
args += ["-m", "main"]
env["TREZOR_PROFILE_DIR"] = self.workdir.name
self.process = subprocess.Popen(args,
cwd=self.workdir.name,
env=ENV,
stdout=open(os.devnull, "w"))
# wait until emulator is started
while True:
try:
self.transport = get_transport("udp:127.0.0.1:21324")
except TransportException:
time.sleep(0.1)
continue
break
self.client = TrezorClientDebugLink(self.transport)
self.client.open()
return self
def __exit__(self, exc_type, exc_value, traceback):
self.client.close()
self.process.terminate()
try:
self.process.wait(1)
except subprocess.TimeoutExpired:
self.process.kill()
self.workdir.cleanup()
def _storage_file(self):
if self.gen == "legacy":
return self.workdir.name + "/emulator.img"
elif self.gen == "core":
return self.workdir.name + "/trezor.flash"
else:
raise ValueError("Unknown gen")
def storage(self):
return open(self._storage_file(), "rb").read()
def client():
wirelink = get_device()
client = TrezorClientDebugLink(wirelink)
wipe_device(client)
client.transport.session_begin()
yield client
client.transport.session_end()
# XXX debuglink session must also be closed
# client.close accomplishes that for now; going forward, there should
# also be proper session handling for debuglink
client.close()
def setup_device(mnemonic,
pin,
passphrase_protection,
label,
needs_backup=None):
# TODO:
# - check if device is acquired otherwise throws "wrong previous session" from bridge
client = TrezorClientDebugLink(get_device())
client.open()
debuglink.load_device(client,
mnemonic,
pin,
passphrase_protection,
label,
needs_backup=bool(needs_backup))
client.close()
def setup_device(
mnemonic: str,
pin: str,
passphrase_protection: bool,
label: str,
needs_backup: bool = False,
) -> None:
# TODO: check if device is acquired, otherwise throws
# "wrong previous session" from bridge
client = TrezorClientDebugLink(get_device())
client.open()
time.sleep(SLEEP)
debuglink.load_device(
client,
mnemonic,
pin,
passphrase_protection,
label,
needs_backup=needs_backup,
)
client.close()
class EmulatorWrapper:
def __init__(self, gen, tag=None, executable=None, storage=None):
self.gen = gen
self.tag = tag
if executable is not None:
self.executable = executable
elif tag is not None:
self.executable = filename_from_tag(gen, tag)
else:
self.executable = LOCAL_BUILD_PATHS[gen]
if not os.path.exists(self.executable):
raise ValueError(
f"emulator executable not found: {self.executable}")
self.workdir = tempfile.TemporaryDirectory()
if storage:
open(self._storage_file(), "wb").write(storage)
with gzip.open(SD_CARD_GZ, "rb") as gz:
with open(self.workdir.name + "/trezor.sdcard", "wb") as sd:
sd.write(gz.read())
self.client = None
def _get_params_core(self):
env = ENV.copy()
args = [self.executable, "-m", "main"]
# for firmware 2.1.2 and newer
env["TREZOR_PROFILE_DIR"] = self.workdir.name
# for firmware 2.1.1 and older
env["TREZOR_PROFILE"] = self.workdir.name
if self.executable == LOCAL_BUILD_PATHS["core"]:
cwd = ROOT + "/core/src"
else:
cwd = self.workdir.name
return env, args, cwd
def _get_params_legacy(self):
env = ENV.copy()
args = [self.executable]
cwd = self.workdir.name
return env, args, cwd
def _get_params(self):
if self.gen == "core":
return self._get_params_core()
elif self.gen == "legacy":
return self._get_params_legacy()
else:
raise ValueError("Unknown gen")
def start(self):
env, args, cwd = self._get_params()
self.process = subprocess.Popen(args,
cwd=cwd,
env=env,
stdout=open(os.devnull, "w"))
# wait until emulator is listening
transport = UdpTransport("127.0.0.1:21324")
transport.open()
for _ in range(300):
if transport._ping():
break
if self.process.poll() is not None:
self._cleanup()
raise RuntimeError("Emulator proces died")
time.sleep(0.1)
else:
# could not connect after 300 attempts * 0.1s = 30s of waiting
self._cleanup()
raise RuntimeError("Can't connect to emulator")
transport.close()
self.client = TrezorClientDebugLink(transport)
self.client.open()
check_version(self.tag, self.client.version)
def stop(self):
if self.client:
self.client.close()
self.process.terminate()
try:
self.process.wait(1)
except subprocess.TimeoutExpired:
self.process.kill()
def restart(self):
self.stop()
self.start()
def __enter__(self):
self.start()
return self
def __exit__(self, exc_type, exc_value, traceback):
self._cleanup()
def _cleanup(self):
self.stop()
self.workdir.cleanup()
def _storage_file(self):
if self.gen == "legacy":
return self.workdir.name + "/emulator.img"
elif self.gen == "core":
return self.workdir.name + "/trezor.flash"
else:
raise ValueError("Unknown gen")
def storage(self):
return open(self._storage_file(), "rb").read()
def wipe_device():
client = TrezorClientDebugLink(get_device())
client.open()
wipe(client)
client.close()
def reset_device():
client = TrezorClientDebugLink(get_device())
client.open()
reset(client, skip_backup=True, pin_protection=False)
client.close()
class EmulatorWrapper:
def __init__(self, gen, tag=None, executable=None, storage=None):
self.gen = gen
self.tag = tag
if executable is not None:
self.executable = executable
elif tag is not None:
self.executable = filename_from_tag(gen, tag)
else:
self.executable = LOCAL_BUILD_PATHS[gen]
if not os.path.exists(self.executable):
raise ValueError(
f"emulator executable not found: {self.executable}")
self.workdir = tempfile.TemporaryDirectory()
if storage:
open(self._storage_file(), "wb").write(storage)
self.client = None
def __enter__(self):
args = [self.executable]
env = ENV
if self.gen == "core":
args += ["-m", "main"]
# for firmware 2.1.2 and newer
env["TREZOR_PROFILE_DIR"] = self.workdir.name
# for firmware 2.1.1 and older
env["TREZOR_PROFILE"] = self.workdir.name
self.process = subprocess.Popen(args,
cwd=self.workdir.name,
env=env,
stdout=open(os.devnull, "w"))
# wait until emulator is listening
for _ in range(300):
try:
time.sleep(0.1)
transport = get_transport("udp:127.0.0.1:21324")
break
except TransportException:
pass
if self.process.poll() is not None:
self._cleanup()
raise RuntimeError("Emulator proces died")
else:
# could not connect after 300 attempts * 0.1s = 30s of waiting
self._cleanup()
raise RuntimeError("Can't connect to emulator")
self.client = TrezorClientDebugLink(transport)
self.client.open()
check_version(self.tag, self.client.version)
return self
def __exit__(self, exc_type, exc_value, traceback):
self._cleanup()
return False
def _cleanup(self):
if self.client:
self.client.close()
self.process.terminate()
try:
self.process.wait(1)
except subprocess.TimeoutExpired:
self.process.kill()
self.workdir.cleanup()
def _storage_file(self):
if self.gen == "legacy":
return self.workdir.name + "/emulator.img"
elif self.gen == "core":
return self.workdir.name + "/trezor.flash"
else:
raise ValueError("Unknown gen")
def storage(self):
return open(self._storage_file(), "rb").read()
def reset_device() -> None:
client = TrezorClientDebugLink(get_device())
client.open()
time.sleep(SLEEP)
device.reset(client, skip_backup=True, pin_protection=False)
client.close()
def wipe_device() -> None:
client = TrezorClientDebugLink(get_device())
client.open()
time.sleep(SLEEP)
device.wipe(client)
client.close()
class Emulator:
STORAGE_FILENAME = None
def __init__(self,
executable,
profile_dir,
*,
logfile=None,
storage=None,
headless=False,
debug=True,
extra_args=()):
self.executable = Path(executable).resolve()
if not executable.exists():
raise ValueError("emulator executable not found: {}".format(
self.executable))
self.profile_dir = Path(profile_dir).resolve()
if not self.profile_dir.exists():
self.profile_dir.mkdir(parents=True)
elif not self.profile_dir.is_dir():
raise ValueError("profile_dir is not a directory")
self.workdir = self.profile_dir
self.storage = self.profile_dir / self.STORAGE_FILENAME
if storage:
self.storage.write_bytes(storage)
if logfile:
self.logfile = logfile
else:
self.logfile = self.profile_dir / "trezor.log"
self.client = None
self.process = None
self.port = 21324
self.headless = headless
self.debug = debug
self.extra_args = list(extra_args)
def make_args(self):
return []
def make_env(self):
return os.environ.copy()
def _get_transport(self):
return UdpTransport("127.0.0.1:{}".format(self.port))
def wait_until_ready(self, timeout=30):
transport = self._get_transport()
transport.open()
start = time.monotonic()
try:
while True:
if transport._ping():
break
if self.process.poll() is not None:
raise RuntimeError("Emulator proces died")
elapsed = time.monotonic() - start
if elapsed >= timeout:
raise TimeoutError("Can't connect to emulator")
time.sleep(0.1)
finally:
transport.close()
def wait(self, timeout=None):
ret = self.process.wait(timeout=None)
self.stop()
return ret
def launch_process(self):
args = self.make_args()
env = self.make_env()
if hasattr(self.logfile, "write"):
output = self.logfile
else:
output = open(self.logfile, "w")
return subprocess.Popen(
[self.executable] + args + self.extra_args,
cwd=self.workdir,
stdout=output,
stderr=subprocess.STDOUT,
env=env,
)
def start(self):
if self.process:
if self.process.poll() is not None:
# process has died, stop and start again
self.stop()
else:
# process is running, no need to start again
return
self.process = self.launch_process()
try:
self.wait_until_ready()
except TimeoutError:
# Assuming that after the default 30-second timeout, the process is stuck
self.process.kill()
raise
(self.profile_dir /
"trezor.pid").write_text(str(self.process.pid) + "\n")
(self.profile_dir / "trezor.port").write_text(str(self.port) + "\n")
transport = self._get_transport()
self.client = TrezorClientDebugLink(transport,
auto_interact=self.debug)
self.client.open()
def stop(self):
if self.client:
self.client.close()
self.client = None
if self.process:
self.process.terminate()
try:
self.process.wait(1)
except subprocess.TimeoutExpired:
self.process.kill()
_rm_f(self.profile_dir / "trezor.pid")
_rm_f(self.profile_dir / "trezor.port")
self.process = None
def restart(self):
self.stop()
self.start()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
self.stop()
def get_storage(self):
return self.storage.read_bytes()
def client(request: pytest.FixtureRequest,
_raw_client: Client) -> Generator[Client, None, None]:
"""Client fixture.
Every test function that requires a client instance will get it from here.
If we can't connect to a debuggable device, the test will fail.
If 'skip_t2' is used and TT is connected, the test is skipped. Vice versa with T1
and 'skip_t1'.
The client instance is wiped and preconfigured with "all all all..." mnemonic, no
password and no pin. It is possible to customize this with the `setup_client`
marker.
To specify a custom mnemonic and/or custom pin and/or enable passphrase:
@pytest.mark.setup_client(mnemonic=MY_MNEMONIC, pin="9999", passphrase=True)
To receive a client instance that was not initialized:
@pytest.mark.setup_client(uninitialized=True)
"""
if request.node.get_closest_marker(
"skip_t2") and _raw_client.features.model == "T":
pytest.skip("Test excluded on Trezor T")
if request.node.get_closest_marker(
"skip_t1") and _raw_client.features.model == "1":
pytest.skip("Test excluded on Trezor 1")
sd_marker = request.node.get_closest_marker("sd_card")
if sd_marker and not _raw_client.features.sd_card_present:
raise RuntimeError("This test requires SD card.\n"
"To skip all such tests, run:\n"
" pytest -m 'not sd_card' <test path>")
test_ui = request.config.getoption("ui")
_raw_client.reset_debug_features()
_raw_client.open()
try:
_raw_client.init_device()
except Exception:
request.session.shouldstop = "Failed to communicate with Trezor"
pytest.fail("Failed to communicate with Trezor")
if test_ui:
# we need to reseed before the wipe
_raw_client.debug.reseed(0)
if sd_marker:
should_format = sd_marker.kwargs.get("formatted", True)
_raw_client.debug.erase_sd_card(format=should_format)
wipe_device(_raw_client)
setup_params = dict(
uninitialized=False,
mnemonic=" ".join(["all"] * 12),
pin=None,
passphrase=False,
needs_backup=False,
no_backup=False,
)
marker = request.node.get_closest_marker("setup_client")
if marker:
setup_params.update(marker.kwargs)
use_passphrase = setup_params["passphrase"] is True or isinstance(
setup_params["passphrase"], str)
if not setup_params["uninitialized"]:
debuglink.load_device(
_raw_client,
mnemonic=setup_params["mnemonic"],
pin=setup_params["pin"],
passphrase_protection=use_passphrase,
label="test",
language="en-US",
needs_backup=setup_params["needs_backup"],
no_backup=setup_params["no_backup"],
)
if _raw_client.features.model == "T":
apply_settings(_raw_client, experimental_features=True)
if use_passphrase and isinstance(setup_params["passphrase"], str):
_raw_client.use_passphrase(setup_params["passphrase"])
_raw_client.clear_session()
if test_ui:
with ui_tests.screen_recording(_raw_client, request):
yield _raw_client
else:
yield _raw_client
_raw_client.close()
