From e96a9e8d3959c7604c4128b8ea0cd655b0f46d7e Mon Sep 17 00:00:00 2001 From: matejcik Date: Wed, 26 Aug 2020 16:01:58 +0200 Subject: [PATCH] python/trezorctl: make use of EndSession --- python/src/trezorlib/cli/__init__.py | 51 ++++++++++++++++++++++----- python/src/trezorlib/cli/trezorctl.py | 36 +++++++++++-------- python/src/trezorlib/client.py | 12 +++++-- 3 files changed, 74 insertions(+), 25 deletions(-) diff --git a/python/src/trezorlib/cli/__init__.py b/python/src/trezorlib/cli/__init__.py index a56eda5ee..622081f47 100644 --- a/python/src/trezorlib/cli/__init__.py +++ b/python/src/trezorlib/cli/__init__.py @@ -16,6 +16,7 @@ import functools import sys +from contextlib import contextmanager import click @@ -61,25 +62,59 @@ class TrezorConnection: ui = self.get_ui() return TrezorClient(transport, ui=ui, session_id=self.session_id) + @contextmanager + def client_context(self): + """Get a client instance as a context manager. Handle errors in a manner + appropriate for end-users. -def with_client(func): - @click.pass_obj - @functools.wraps(func) - def trezorctl_command_with_client(obj, *args, **kwargs): + Usage: + >>> with obj.client_context() as client: + >>> do_your_actions_here() + """ try: - client = obj.get_client() + client = self.get_client() except Exception: click.echo("Failed to find a Trezor device.") - if obj.path is not None: - click.echo("Using path: {}".format(obj.path)) + if self.path is not None: + click.echo("Using path: {}".format(self.path)) sys.exit(1) try: - return func(client, *args, **kwargs) + yield client except exceptions.Cancelled: + # handle cancel action click.echo("Action was cancelled.") sys.exit(1) except exceptions.TrezorException as e: + # handle any Trezor-sent exceptions as user-readable raise click.ClickException(str(e)) from e + # other exceptions may cause a traceback + + +def with_client(func): + """Wrap a Click command in `with obj.client_context() as client`. + + Sessions are handled transparently. The user is warned when session did not resume + cleanly. The session is closed after the command completes - unless the session + was resumed, in which case it should remain open. + """ + + @click.pass_obj + @functools.wraps(func) + def trezorctl_command_with_client(obj, *args, **kwargs): + with obj.client_context() as client: + session_was_resumed = obj.session_id == client.session_id + if not session_was_resumed and obj.session_id is not None: + # tried to resume but failed + click.echo("Warning: failed to resume session.", err=True) + + try: + return func(client, *args, **kwargs) + finally: + if not session_was_resumed: + try: + client.end_session() + except Exception: + pass return trezorctl_command_with_client diff --git a/python/src/trezorlib/cli/trezorctl.py b/python/src/trezorlib/cli/trezorctl.py index b29e66f3e..966f230a4 100755 --- a/python/src/trezorlib/cli/trezorctl.py +++ b/python/src/trezorlib/cli/trezorctl.py @@ -159,7 +159,7 @@ def configure_logging(verbose: int): def cli(ctx, path, verbose, is_json, passphrase_on_host, session_id): configure_logging(verbose) - if session_id: + if session_id is not None: try: session_id = bytes.fromhex(session_id) except ValueError: @@ -216,6 +216,11 @@ def list_devices(no_resolve): for transport in enumerate_devices(): client = TrezorClient(transport, ui=ui.ClickUI()) click.echo("{} - {}".format(transport, format_device_name(client.features))) + try: + # firmwares <2.3.4 do not recognize EndSession + client.end_session() + except Exception: + pass @cli.command() @@ -241,8 +246,8 @@ def ping(client, message, button_protection): @cli.command() -@with_client -def get_session(client): +@click.pass_obj +def get_session(obj): """Get a session ID for subsequent commands. Unlocks Trezor with a passphrase and returns a session ID. Use this session ID with @@ -252,17 +257,20 @@ def get_session(client): The session ID is valid until another client starts using Trezor, until the next get-session call, or until Trezor is disconnected. """ - from ..btc import get_address - from ..client import PASSPHRASE_TEST_PATH - - if client.features.model == "1" and client.version < (1, 9, 0): - raise click.ClickException("Upgrade your firmware to enable session support.") - - get_address(client, "Testnet", PASSPHRASE_TEST_PATH) - if client.session_id is None: - raise click.ClickException("Passphrase not enabled or firmware too old.") - else: - return client.session_id.hex() + # make sure session is not resumed + obj.session_id = None + + with obj.client_context() as client: + if client.features.model == "1" and client.version < (1, 9, 0): + raise click.ClickException( + "Upgrade your firmware to enable session support." + ) + + client.ensure_unlocked() + if client.session_id is None: + raise click.ClickException("Passphrase not enabled or firmware too old.") + else: + return client.session_id.hex() @cli.command() diff --git a/python/src/trezorlib/client.py b/python/src/trezorlib/client.py index be1d2f663..b9bcac450 100644 --- a/python/src/trezorlib/client.py +++ b/python/src/trezorlib/client.py @@ -88,9 +88,9 @@ class TrezorClient: LOG.info("creating client instance for device: {}".format(transport.get_path())) self.transport = transport self.ui = ui - self.session_id = session_id self.session_counter = 0 - self.init_device() + self.session_id = session_id + self.init_device(session_id=session_id) def open(self): if self.session_counter == 0: @@ -294,6 +294,11 @@ class TrezorClient: if not isinstance(resp, messages.Features): raise exceptions.TrezorException("Unexpected response to Initialize") + if resp.session_id == self.session_id: + LOG.info("Successfully resumed session") + elif session_id is not None: + LOG.info("Failed to resume session") + # TT < 2.3.0 compatibility: # _refresh_features will clear out the session_id field. We want this function # to return its value, so that callers can rely on it being either a valid @@ -385,6 +390,7 @@ class TrezorClient: The session will become invalid until `init_device()` is called again. If passphrase is enabled, further actions will prompt for it again. """ + # since: 2.3.4, 1.9.4 self.call(messages.EndSession()) self.session_id = None @@ -400,4 +406,4 @@ class TrezorClient: # call LockDevice manually to save one refresh_features() call self.call(messages.LockDevice()) self.end_session() - self.init_device() + self.init_device(new_session=True)