diff --git a/python/.changelog.d/2023.added b/python/.changelog.d/2023.added new file mode 100644 index 000000000..052d16562 --- /dev/null +++ b/python/.changelog.d/2023.added @@ -0,0 +1 @@ +Add ScriptUI for trezorctl, spawned by --script option diff --git a/python/src/trezorlib/cli/__init__.py b/python/src/trezorlib/cli/__init__.py index 6213be026..5f41a8673 100644 --- a/python/src/trezorlib/cli/__init__.py +++ b/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 ..ui import ClickUI +from ..transport import get_transport +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,8 +76,14 @@ 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: - return ClickUI(passphrase_on_host=self.passphrase_on_host) + 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: transport = self.get_transport() diff --git a/python/src/trezorlib/cli/trezorctl.py b/python/src/trezorlib/cli/trezorctl.py index 69a769032..0394f1e6a 100755 --- a/python/src/trezorlib/cli/trezorctl.py +++ b/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: diff --git a/python/src/trezorlib/ui.py b/python/src/trezorlib/ui.py index 6deb0a3c5..9d955e6b8 100644 --- a/python/src/trezorlib/ui.py +++ b/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]: diff --git a/python/tools/trezorctl_script_client.py b/python/tools/trezorctl_script_client.py new file mode 100644 index 000000000..72aa3287b --- /dev/null +++ b/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()