python/trezorctl: make use of EndSession

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

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()

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)