diff --git a/python/.changelog.d/2547.added b/python/.changelog.d/2547.added new file mode 100644 index 0000000000..f15539097b --- /dev/null +++ b/python/.changelog.d/2547.added @@ -0,0 +1 @@ +Add possibility to save emulator screenshots. diff --git a/python/src/trezorlib/cli/debug.py b/python/src/trezorlib/cli/debug.py index 285bc52844..741aff2d57 100644 --- a/python/src/trezorlib/cli/debug.py +++ b/python/src/trezorlib/cli/debug.py @@ -14,11 +14,12 @@ # You should have received a copy of the License along with this library. # If not, see . -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Union import click from .. import mapping, messages, protobuf +from ..debuglink import TrezorClientDebugLink, record_screen if TYPE_CHECKING: from . import TrezorConnection @@ -74,3 +75,26 @@ def send_bytes( click.echo(protobuf.format_message(msg)) except Exception as e: click.echo(f"Could not parse response: {e}") + + +@cli.command() +@click.argument("directory", required=False) +@click.option("-s", "--stop", is_flag=True, help="Stop the recording") +@click.pass_obj +def record(obj: "TrezorConnection", directory: Union[str, None], stop: bool) -> None: + """Record screen changes into a specified directory. + + Recording can be stopped with `-s / --stop` option. + """ + record_screen_from_connection(obj, None if stop else directory) + + +def record_screen_from_connection( + obj: "TrezorConnection", directory: Union[str, None] +) -> None: + """Record screen helper to transform TrezorConnection into TrezorClientDebugLink.""" + transport = obj.get_transport() + debug_client = TrezorClientDebugLink(transport, auto_interact=False) + debug_client.open() + record_screen(debug_client, directory, report_func=click.echo) + debug_client.close() diff --git a/python/src/trezorlib/cli/trezorctl.py b/python/src/trezorlib/cli/trezorctl.py index 555818ecca..b820f28d96 100755 --- a/python/src/trezorlib/cli/trezorctl.py +++ b/python/src/trezorlib/cli/trezorctl.py @@ -186,6 +186,11 @@ def configure_logging(verbose: int) -> None: help="Resume given session ID.", default=os.environ.get("TREZOR_SESSION_ID"), ) +@click.option( + "-r", + "--record", + help="Record screen changes into a specified directory.", +) @click.version_option(version=__version__) @click.pass_context def cli_main( @@ -196,6 +201,7 @@ def cli_main( passphrase_on_host: bool, script: bool, session_id: Optional[str], + record: Optional[str], ) -> None: configure_logging(verbose) @@ -208,6 +214,10 @@ def cli_main( ctx.obj = TrezorConnection(path, bytes_session_id, passphrase_on_host, script) + # Optionally record the screen into a specified directory. + if record: + debug.record_screen_from_connection(ctx.obj, record) + # Creating a cli function that has the right types for future usage cli = cast(TrezorctlGroup, cli_main) @@ -241,6 +251,19 @@ def print_result(res: Any, is_json: bool, script: bool, **kwargs: Any) -> None: click.echo(res) +@cli.set_result_callback() +@click.pass_obj +def stop_recording_action(obj: TrezorConnection, *args: Any, **kwargs: Any) -> None: + """Stop recording screen changes when the recording was started by `cli_main`. + + (When user used the `-r / --record` option of `trezorctl` command.) + + It allows for isolating screen directories only for specific actions/commands. + """ + if kwargs.get("record"): + debug.record_screen_from_connection(obj, None) + + def format_device_name(features: messages.Features) -> str: model = features.model or "1" if features.bootloader_mode: diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py index 103205b505..187c663035 100644 --- a/python/src/trezorlib/debuglink.py +++ b/python/src/trezorlib/debuglink.py @@ -18,6 +18,7 @@ import logging import textwrap from collections import namedtuple from copy import deepcopy +from datetime import datetime from enum import IntEnum from itertools import zip_longest from pathlib import Path @@ -779,3 +780,51 @@ def self_test(client: "TrezorClient") -> protobuf.MessageType: payload=b"\x00\xFF\x55\xAA\x66\x99\x33\xCCABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!\x00\xFF\x55\xAA\x66\x99\x33\xCC" ) ) + + +def record_screen( + debug_client: "TrezorClientDebugLink", + directory: Union[str, None], + report_func: Union[Callable[[str], None], None] = None, +) -> None: + """Record screen changes into a specified directory. + + Passing `None` as `directory` stops the recording. + + Creates subdirectories inside a specified directory, one for each session + (for each new call of this function). + (So that older screenshots are not overwritten by new ones.) + + Is available only for emulators, hardware devices are not capable of that. + """ + + def get_session_screenshot_dir(directory: Path) -> Path: + """Create and return screenshot dir for the current session, according to datetime.""" + session_dir = directory / datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + session_dir.mkdir(parents=True, exist_ok=True) + return session_dir + + if not _is_emulator(debug_client): + raise RuntimeError("Recording is only supported on emulator.") + + if directory is None: + debug_client.debug.stop_recording() + if report_func is not None: + report_func("Recording stopped.") + else: + # Transforming the directory into an absolute path, + # because emulator demands it + abs_directory = Path(directory).resolve() + # Creating the dir when it does not exist yet + if not abs_directory.exists(): + abs_directory.mkdir(parents=True, exist_ok=True) + # Getting a new screenshot dir for the current session + current_session_dir = get_session_screenshot_dir(abs_directory) + debug_client.debug.start_recording(str(current_session_dir)) + if report_func is not None: + report_func(f"Recording started into {current_session_dir}.") + + +def _is_emulator(debug_client: "TrezorClientDebugLink") -> bool: + """Check if we are connected to emulator, in contrast to hardware device.""" + return debug_client.features.fw_vendor == "EMULATOR"