feat(python): ScriptUI for trezorctl, reference client implementation

python/.changelog.d/2023.added

@@ -0,0 +1 @@
+Add ScriptUI for trezorctl, spawned by --script option

python/src/trezorlib/cli/__init__.py

@@ -23,10 +23,13 @@ import click
 
 from .. import exceptions
 from ..client import TrezorClient
-from ..transport import Transport, get_transport
+from ..transport import get_transport
-from ..ui import ClickUI
+from ..ui import ClickUI, ScriptUI
 
 if TYPE_CHECKING:
+    from ..transport import Transport
+    from ..ui import TrezorClientUI
+
     # Needed to enforce a return value from decorators
     # More details: https://www.python.org/dev/peps/pep-0612/
     from typing import TypeVar
@@ -50,13 +53,18 @@ class ChoiceType(click.Choice):
 
 class TrezorConnection:
     def __init__(
-        self, path: str, session_id: Optional[bytes], passphrase_on_host: bool
+        self,
+        path: str,
+        session_id: Optional[bytes],
+        passphrase_on_host: bool,
+        script: bool,
     ) -> None:
         self.path = path
         self.session_id = session_id
         self.passphrase_on_host = passphrase_on_host
+        self.script = script
 
-    def get_transport(self) -> Transport:
+    def get_transport(self) -> "Transport":
         try:
             # look for transport without prefix search
             return get_transport(self.path, prefix_search=False)
@@ -68,7 +76,13 @@ class TrezorConnection:
         # if this fails, we want the exception to bubble up to the caller
         return get_transport(self.path, prefix_search=True)
 
-    def get_ui(self) -> ClickUI:
+    def get_ui(self) -> "TrezorClientUI":
+        if self.script:
+            # It is alright to return just the class object instead of instance,
+            # as the ScriptUI class object itself is the implementation of TrezorClientUI
+            # (ScriptUI is just a set of staticmethods)
+            return ScriptUI  # type: ignore [Expression of type "Type[ScriptUI]" cannot be assigned to return type "TrezorClientUI"]
+        else:
             return ClickUI(passphrase_on_host=self.passphrase_on_host)
 
     def get_client(self) -> TrezorClient:

python/src/trezorlib/cli/trezorctl.py

@@ -155,6 +155,12 @@ def configure_logging(verbose: int) -> None:
     is_flag=True,
     help="Enter passphrase on host.",
 )
+@click.option(
+    "-S",
+    "--script",
+    is_flag=True,
+    help="Use UI for usage in scripts.",
+)
 @click.option(
     "-s",
     "--session-id",
@@ -170,6 +176,7 @@ def cli_main(
     verbose: int,
     is_json: bool,
     passphrase_on_host: bool,
+    script: bool,
     session_id: Optional[str],
 ) -> None:
     configure_logging(verbose)
@@ -181,7 +188,7 @@ def cli_main(
         except ValueError:
             raise click.ClickException(f"Not a valid session id: {session_id}")
 
-    ctx.obj = TrezorConnection(path, bytes_session_id, passphrase_on_host)
+    ctx.obj = TrezorConnection(path, bytes_session_id, passphrase_on_host, script)
 
 
 # Creating a cli function that has the right types for future usage
@@ -189,10 +196,14 @@ cli = cast(TrezorctlGroup, cli_main)
 
 
 @cli.resultcallback()
-def print_result(res: Any, is_json: bool, **kwargs: Any) -> None:
+def print_result(res: Any, is_json: bool, script: bool, **kwargs: Any) -> None:
     if is_json:
         if isinstance(res, protobuf.MessageType):
-            click.echo(json.dumps({res.__class__.__name__: res.__dict__}))
+            res = protobuf.to_dict(res, hexlify_bytes=True)
+
+        # No newlines for scripts, pretty-print for users
+        if script:
+            click.echo(json.dumps(res))
         else:
             click.echo(json.dumps(res, sort_keys=True, indent=4))
     else:

python/src/trezorlib/ui.py

@@ -172,6 +172,56 @@ class ClickUI:
                 raise Cancelled from None
 
 
+class ScriptUI:
+    """Interface to be used by scripts, not directly by user.
+
+    Communicates with a client application using print() and input().
+
+    Lot of `ClickUI` logic is outsourced to the client application, which
+    is responsible for supplying the PIN and passphrase.
+
+    Reference client implementation can be found under `tools/trezorctl_script_client.py`.
+    """
+
+    @staticmethod
+    def button_request(br: messages.ButtonRequest) -> None:
+        # TODO: send name={br.name} when it will be supported
+        code = br.code.name if br.code else None
+        print(f"?BUTTON code={code} pages={br.pages}")
+
+    @staticmethod
+    def get_pin(code: Optional[PinMatrixRequestType] = None) -> str:
+        if code is None:
+            print("?PIN")
+        else:
+            print(f"?PIN code={code.name}")
+
+        pin = input()
+        if pin == "CANCEL":
+            raise Cancelled from None
+        elif not pin.startswith(":"):
+            raise RuntimeError("Sent PIN must start with ':'")
+        else:
+            return pin[1:]
+
+    @staticmethod
+    def get_passphrase(available_on_device: bool) -> Union[str, object]:
+        if available_on_device:
+            print("?PASSPHRASE available_on_device")
+        else:
+            print("?PASSPHRASE")
+
+        passphrase = input()
+        if passphrase == "CANCEL":
+            raise Cancelled from None
+        elif passphrase == "ON_DEVICE":
+            return PASSPHRASE_ON_DEVICE
+        elif not passphrase.startswith(":"):
+            raise RuntimeError("Sent passphrase must start with ':'")
+        else:
+            return passphrase[1:]
+
+
 def mnemonic_words(
     expand: bool = False, language: str = "english"
 ) -> Callable[[WordRequestType], str]:

python/tools/trezorctl_script_client.py

@@ -0,0 +1,134 @@
+"""
+Reference client implementation consuming trezorctl's script interface
+(ScriptUI class) available by using `--script` flag in any trezorctl command.
+
+Function `get_address()` is showing the communication with ScriptUI
+on a specific example
+"""
+
+import os
+import subprocess
+from typing import Dict, List, Optional, Tuple, Union
+
+import click
+
+
+def parse_args_from_line(line: str) -> Tuple[str, Dict[str, Union[str, bool]]]:
+    # ?PIN code=123
+    # ?PASSPHRASE available_on_device
+    command, *args = line.split(" ")
+    result: Dict[str, Union[str, bool]] = {}
+    for arg in args:
+        if "=" in arg:
+            key, value = arg.split("=")
+            result[key] = value
+        else:
+            result[arg] = True
+    return command, result
+
+
+def get_pin_from_user(code: Optional[str] = None) -> str:
+    # ?PIN
+    # ?PIN code=Current
+    while True:
+        try:
+            pin = click.prompt(
+                f"Enter PIN (code: {code})",
+                hide_input=True,
+                default="",
+                show_default=False,
+            )
+        except click.Abort:
+            return "CANCEL"
+        if not all(c in "123456789" for c in pin):
+            click.echo("PIN must only be numbers 1-9")
+            continue
+        return ":" + pin
+
+
+def show_button_request(
+    code: Optional[str] = None, pages: Optional[str] = None, name: Optional[str] = None
+) -> None:
+    # ?BUTTON code=Other
+    # ?BUTTON code=SignTx pages=2
+    # ?BUTTON code=ProtectCall name=confirm_set_pin
+    print(f"Please confirm action on Trezor (code={code} name={name} pages={pages})")
+
+
+def get_passphrase_from_user(available_on_device: bool = False) -> str:
+    # ?PASSPHRASE
+    # ?PASSPHRASE available_on_device
+    if available_on_device:
+        if click.confirm("Enter passphrase on device?", default=True):
+            return "ON_DEVICE"
+
+    env_passphrase = os.getenv("PASSPHRASE")
+    if env_passphrase:
+        if click.confirm("Use env PASSPHRASE?", default=False):
+            return ":" + env_passphrase
+
+    while True:
+        try:
+            passphrase = click.prompt("Enter passphrase", hide_input=True, default="")
+        except click.Abort:
+            return "CANCEL"
+
+        passphrase2 = click.prompt(
+            "Enter passphrase again", hide_input=True, default=""
+        )
+        if passphrase != passphrase2:
+            click.echo("Passphrases do not match")
+            continue
+        return ":" + passphrase
+
+
+def get_address() -> str:
+    args = """
+        trezorctl --script get-address -n "m/49'/0'/0'/0/0"
+    """.strip()
+    p = subprocess.Popen(  # type: ignore [No overloads for "__new__" match the provided arguments]
+        args,
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        text=True,
+        shell=True,
+        bufsize=0,
+    )
+
+    assert p.stdout is not None
+    assert p.stdin is not None
+
+    text_result: List[str] = []
+    while True:
+        line = p.stdout.readline().strip()
+        if not line:
+            break
+
+        if line.startswith("?"):
+            command, args = parse_args_from_line(line)
+            if command == "?PIN":
+                response = get_pin_from_user(**args)
+                p.stdin.write(response + "\n")
+            elif command == "?PASSPHRASE":
+                response = get_passphrase_from_user(**args)
+                p.stdin.write(response + "\n")
+            elif command == "?BUTTON":
+                show_button_request(**args)
+            else:
+                print("Unrecognized script command:", line)
+
+        text_result.append(line)
+        print(line)
+
+    address = text_result[-1]
+    print("Address:", address)
+    return address
+
+
+def clear_session_to_enable_pin():
+    os.system("trezorctl clear-session")
+
+
+if __name__ == "__main__":
+    get_address()
+    clear_session_to_enable_pin()