trezor-firmware/ci/hardware_tests/device/device.py

import datetime
import sys
import time
from subprocess import run


class Device:
    def __init__(self, *, uhub_location=None, hub_vendor=None, device_port=None):
        self.uhub_location = uhub_location
        self.hub_vendor = hub_vendor
        self.device_port = device_port

    @staticmethod
    def log(msg):
        print(msg, flush=True, file=sys.stderr)

    def run_trezorctl(self, cmd: str, **kwargs):
        full_cmd = "trezorctl "
        full_cmd += cmd
        self.log(f"[software/trezorctl] Running '{full_cmd}'")
        return run(full_cmd, shell=True, check=True, **kwargs)

    def check_model(self, model=None):
        res = self.run_trezorctl("list", capture_output=True, text=True)
        self.log(res.stdout)
        self.log(res.stderr)
        self.run_trezorctl("get-features | grep version")
        lines = res.stdout.splitlines()
        if len(lines) != 1:
            raise RuntimeError(f"{len(lines)} trezors connected")
        if model and model not in lines[0]:
            raise RuntimeError(f"invalid trezor model connected (expected {model})")
        return lines[0].split()[0]

    def reboot(self):
        self.power_off()
        self.power_on()

    def _hub(self):
        if self.hub_vendor:
            return f"--vendor {self.hub_vendor}"
        else:
            return f"-l {self.uhub_location}"

    def power_on(self):
        self.now()
        self.log("[hardware/usb] Turning power on...")
        run(
            f"uhubctl {self._hub()} -p {self.device_port} -a on",
            shell=True,
            check=True,
        )
        self.wait(3)

    def power_off(self):
        self.now()
        self.log("[hardware/usb] Turning power off...")
        run(
            f"uhubctl {self._hub()} -p {self.device_port} -r 5 -a off",
            shell=True,
            check=True,
        )
        self.wait(3)

    def touch(self, location, action):
        raise NotImplementedError

    @staticmethod
    def wait(seconds):
        Device.now()
        Device.log(f"[software] Waiting for {seconds} seconds...")
        time.sleep(seconds)

    @staticmethod
    def now():
        Device.log(f"\n[timestamp] {datetime.datetime.now()}")