core: introduce emulator runner (fixes #466)

core/emu.py

@@ -0,0 +1,268 @@
+#!/usr/bin/env python3
+import gzip
+import os
+import platform
+import signal
+import subprocess
+import sys
+import tempfile
+import time
+from pathlib import Path
+
+import click
+
+import trezorlib.debuglink
+import trezorlib.device
+from trezorlib._internal.emulator import CoreEmulator
+
+try:
+    import inotify.adapters
+except ImportError:
+    inotify = None
+
+
+HERE = Path(__file__).parent.resolve()
+MICROPYTHON = HERE / "build" / "unix" / "micropython"
+SRC_DIR = HERE / "src"
+SD_CARD_GZ = HERE / "trezor.sdcard.gz"
+
+PROFILING_WRAPPER = HERE / "prof" / "prof.py"
+
+PROFILE_BASE = Path.home() / ".trezoremu"
+
+
+def run_command_with_emulator(emulator, command):
+    with emulator:
+        # first start the subprocess
+        process = subprocess.Popen(command)
+        # After the subprocess is started, ignore SIGINT in parent
+        # (so that we don't need to handle KeyboardInterrupts)
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
+        # SIGINTs will be sent to all children by the OS, so we should be able to safely
+        # wait for their exit.
+        return process.wait()
+
+
+def run_emulator(emulator):
+    with emulator:
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
+        return emulator.wait()
+
+
+def watch_emulator(emulator):
+    watch = inotify.adapters.InotifyTree(str(SRC_DIR))
+    try:
+        for _, type_names, _, _ in watch.event_gen(yield_nones=False):
+            if "IN_CLOSE_WRITE" in type_names:
+                click.echo("Restarting...")
+                emulator.restart()
+    except KeyboardInterrupt:
+        emulator.stop()
+    return 0
+
+
+def run_debugger(emulator):
+    os.chdir(emulator.workdir)
+    env = emulator.make_env()
+    if platform.system() == "Darwin":
+        env["PATH"] = "/usr/bin"
+        os.execvpe(
+            "lldb",
+            ["lldb", "-f", emulator.executable, "--"] + emulator.make_args(),
+            env,
+        )
+    else:
+        os.execvpe(
+            "gdb", ["gdb", "--args", emulator.executable] + emulator.make_args(), env
+        )
+
+
+@click.command(context_settings=dict(ignore_unknown_options=True))
+# fmt: off
+@click.option("-a", "--disable-animation", is_flag=True, default=os.environ.get("TREZOR_DISABLE_ANIMATION") == "1", help="Disable animation")
+@click.option("-c", "--command", "run_command", is_flag=True, help="Run command while emulator is running")
+@click.option("-d", "--production", is_flag=True, default=os.environ.get("PYOPT") == "1", help="Production mode (debuglink disabled)")
+@click.option("-D", "--debugger", is_flag=True, help="Run emulator in debugger (gdb/lldb)")
+@click.option("--executable", type=click.Path(exists=True, dir_okay=False), default=os.environ.get("MICROPYTHON"), help="Alternate emulator executable")
+@click.option("-g", "--profiling", is_flag=True, default=os.environ.get("TREZOR_PROFILING"), help="Run with profiler wrapper")
+@click.option("-h", "--headless", is_flag=True, help="Headless mode (no display)")
+@click.option("--heap-size", metavar="SIZE", default="20M", help="Configure heap size")
+@click.option("--main", help="Path to python main file")
+@click.option("--mnemonic", "mnemonics", multiple=True, help="Initialize device with given mnemonic. Specify multiple times for Shamir shares.")
+@click.option("--log-memory", is_flag=True, default=os.environ.get("TREZOR_LOG_MEMORY") == "1", help="Print memory usage after workflows")
+@click.option("-o", "--output", type=click.File("w"), default="-", help="Redirect emulator output to file")
+@click.option("-p", "--profile", metavar="NAME", help="Profile name or path")
+@click.option("-P", "--port", metavar="PORT", type=int, default=int(os.environ.get("TREZOR_UDP_PORT", 0)) or None, help="UDP port number")
+@click.option("-q", "--quiet", is_flag=True, help="Silence emulator output")
+@click.option("-s", "--slip0014", is_flag=True, help="Initialize device with SLIP-14 seed (all all all...)")
+@click.option("-t", "--temporary-profile", is_flag=True, help="Create an empty temporary profile")
+@click.option("-w", "--watch", is_flag=True, help="Restart emulator if sources change")
+@click.option("-X", "--extra-arg", "extra_args", multiple=True, help="Extra argument to pass to micropython")
+# fmt: on
+@click.argument("command", nargs=-1, type=click.UNPROCESSED)
+def cli(
+    disable_animation,
+    run_command,
+    production,
+    debugger,
+    executable,
+    profiling,
+    headless,
+    heap_size,
+    main,
+    mnemonics,
+    log_memory,
+    profile,
+    port,
+    output,
+    quiet,
+    slip0014,
+    temporary_profile,
+    watch,
+    extra_args,
+    command,
+):
+    """Run the trezor-core emulator.
+
+    If -c is specified, extra arguments are treated as a command that is executed with
+    the running emulator. This command can access the following environment variables:
+    
+    \b
+    TREZOR_PROFILE_DIR - path to storage directory
+    TREZOR_PATH - trezorlib connection string
+    TREZOR_UDP_PORT - UDP port on which the emulator listens
+    TREZOR_FIDO2_UDP_PORT - UDP port for FIDO2
+
+    By default, emulator output goes to stdout. If silenced with -q, it is redirected
+    to $TREZOR_PROFILE_DIR/trezor.log. You can also specify a custom path with -o.
+    """
+    if executable:
+        executable = Path(executable)
+    else:
+        executable = MICROPYTHON
+
+    if command and not run_command:
+        raise click.ClickException("Extra arguments found. Did you mean to use -c?")
+
+    if watch and (command or debugger or frozen):
+        raise click.ClickException("Cannot use -w together with -c or -D or -F")
+
+    if watch and inotify is None:
+        raise click.ClickException("inotify module is missing, install with pip")
+
+    if main and profiling:
+        raise click.ClickException("Cannot use --main and -g together")
+
+    if slip0014 and mnemonics:
+        raise click.ClickException("Cannot use -s and --mnemonic together")
+
+    if slip0014:
+        mnemonics = [" ".join(["all"] * 12)]
+
+    if mnemonics and debugger:
+        raise click.ClickException("Cannot load mnemonics when running in debugger")
+
+    if mnemonics and production:
+        raise click.ClickException("Cannot load mnemonics in production mode")
+
+    if profiling:
+        main_args = [str(PROFILING_WRAPPER)]
+    elif main:
+        main_args = [main]
+    else:
+        main_args = ["-m", "main"]
+
+    if profile and temporary_profile:
+        raise click.ClickException("Cannot use -p and -t together")
+
+    tempdir = None
+    if profile:
+        if "/" in profile:
+            profile_dir = Path(profile)
+        else:
+            profile_dir = PROFILE_BASE / profile
+
+    elif temporary_profile:
+        tempdir = tempfile.TemporaryDirectory(prefix="trezor-emulator-")
+        profile_dir = Path(tempdir.name)
+        # unpack empty SD card for faster start-up
+        with gzip.open(SD_CARD_GZ, "rb") as gz:
+            (profile_dir / "trezor.sdcard").write_bytes(gz.read())
+
+    elif "TREZOR_PROFILE_DIR" in os.environ:
+        profile_dir = Path(os.environ["TREZOR_PROFILE_DIR"])
+
+    else:
+        profile_dir = Path("/var/tmp")
+
+    if quiet:
+        output = None
+
+    emulator = CoreEmulator(
+        executable,
+        profile_dir,
+        logfile=output,
+        port=port,
+        headless=headless,
+        debug=not production,
+        extra_args=extra_args,
+        main_args=main_args,
+        heap_size=heap_size,
+        disable_animation=disable_animation,
+        workdir=SRC_DIR,
+    )
+
+    emulator_env = dict(
+        TREZOR_PATH=f"udp:127.0.0.1:{emulator.port}",
+        TREZOR_PROFILE_DIR=str(profile_dir.resolve()),
+        TREZOR_UDP_PORT=str(emulator.port),
+        TREZOR_FIDO2_UDP_PORT=str(emulator.port + 2),
+    )
+    os.environ.update(emulator_env)
+    for k, v in emulator_env.items():
+        click.echo(f"{k}={v}")
+
+    if log_memory:
+        os.environ["TREZOR_LOG_MEMORY"] = "1"
+
+    if debugger:
+        run_debugger(emulator)
+        raise RuntimeError("run_debugger should not return")
+
+    click.echo("Waiting for emulator to come up... ", err=True)
+    start = time.monotonic()
+    emulator.start()
+    end = time.monotonic()
+    click.echo(f"Emulator ready after {end - start:.3f} seconds", err=True)
+
+    if mnemonics:
+        if slip0014:
+            label = "SLIP-0014"
+        elif profile:
+            label = profile_dir.name
+        else:
+            label = "Emulator"
+
+        trezorlib.device.wipe(emulator.client)
+        trezorlib.debuglink.load_device(
+            emulator.client,
+            mnemonics,
+            pin=None,
+            passphrase_protection=False,
+            label=label,
+        )
+
+    if run_command:
+        ret = run_command_with_emulator(emulator, command)
+    elif watch:
+        ret = watch_emulator(emulator)
+    else:
+        ret = run_emulator(emulator)
+
+    if tempdir is not None:
+        tempdir.cleanup()
+    sys.exit(ret)
+
+
+if __name__ == "__main__":
+    cli()

python/src/trezorlib/_internal/emulator.py

@@ -0,0 +1,237 @@
+# 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>.
+
+import os
+import subprocess
+import time
+from pathlib import Path
+
+from trezorlib.debuglink import TrezorClientDebugLink
+from trezorlib.transport.udp import UdpTransport
+
+
+def _rm_f(path):
+    try:
+        path.unlink()
+    except FileNotFoundError:
+        pass
+
+
+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 RuntimeError("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()
+        self.wait_until_ready()
+
+        (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()
+
+
+class CoreEmulator(Emulator):
+    STORAGE_FILENAME = "trezor.flash"
+
+    def __init__(
+        self,
+        *args,
+        port=None,
+        main_args=("-m", "main"),
+        workdir=None,
+        sdcard=None,
+        disable_animation=True,
+        heap_size="20M",
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        if workdir is not None:
+            self.workdir = Path(workdir).resolve()
+
+        self.sdcard = self.profile_dir / "trezor.sdcard"
+        if sdcard is not None:
+            self.sdcard.write_bytes(sdcard)
+
+        if port:
+            self.port = port
+        self.disable_animation = disable_animation
+        self.main_args = list(main_args)
+        self.heap_size = heap_size
+
+    def make_env(self):
+        env = super().make_env()
+        env.update(
+            TREZOR_PROFILE_DIR=str(self.profile_dir),
+            TREZOR_PROFILE=str(self.profile_dir),
+            TREZOR_UDP_PORT=str(self.port),
+        )
+        if self.headless:
+            env["SDL_VIDEODRIVER"] = "dummy"
+        if self.disable_animation:
+            env["TREZOR_DISABLE_FADE"] = "1"
+            env["TREZOR_DISABLE_ANIMATION"] = "1"
+
+        return env
+
+    def make_args(self):
+        pyopt = "-O0" if self.debug else "-O1"
+        return (
+            [pyopt, "-X", "heapsize={}".format(self.heap_size)]
+            + self.main_args
+            + self.extra_args
+        )
+
+
+class LegacyEmulator(Emulator):
+    STORAGE_FILENAME = "emulator.img"
+
+    def make_env(self):
+        env = super().make_env()
+        if self.headless:
+            env["SDL_VIDEODRIVER"] = "dummy"
+        return env