feat(python): add trezorctl possibility to record screen changes

python/.changelog.d/2547.added

@@ -0,0 +1 @@
+Add possibility to save emulator screenshots.

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 <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
-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()

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:

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"