refactor(python): improve trezorctl error handling

Following #5613.

[no changelog]

Author: romanz

python/src/trezorlib/cli/__init__.py

@@ -238,28 +238,15 @@ def get_seedless_session(self) -> Session:
         seedless_session = client.get_seedless_session()
         return seedless_session
 
-    @contextmanager
-    def client_context(self):
-        """Get a client instance as a context manager. Handle errors in a manner
-        appropriate for end-users.
-
-        Usage:
-        >>> with obj.client_context() as client:
-        >>>     do_your_actions_here()
-        """
+    def _connection_context(self, connect_fn: t.Callable[[], t.Any]):
         try:
-            client = self.get_client()
-        except transport.DeviceIsBusy:
-            click.echo("Device is in use by another process.")
-            sys.exit(1)
-        except Exception:
-            click.echo("Failed to find a Trezor device.")
-            if self.path is not None:
-                click.echo(f"Using path: {self.path}")
+            conn = connect_fn()
+        except Exception as e:
+            self._print_exception(e, "Failed to connect")
             sys.exit(1)
 
         try:
-            yield client
+            yield conn
         except exceptions.Cancelled:
             # handle cancel action
             click.echo("Action was cancelled.")
@@ -269,6 +256,17 @@ def client_context(self):
             raise click.ClickException(str(e)) from e
             # other exceptions may cause a traceback
 
+    @contextmanager
+    def client_context(self):
+        """Get a client instance as a context manager. Handle errors in a manner
+        appropriate for end-users.
+
+        Usage:
+        >>> with obj.client_context() as client:
+        >>>     do_your_actions_here()
+        """
+        yield from self._connection_context(self.get_client)
+
     @contextmanager
     def session_context(
         self,
@@ -277,55 +275,25 @@ def session_context(
         seedless: bool = False,
         must_resume: bool = False,
     ):
-        """Get a session instance as a context manager. Handle errors in a manner
-        appropriate for end-users.
-
-        Usage:
-        >>> with obj.session_context() as session:
-        >>>     do_your_actions_here()
-        """
-        try:
-            if seedless:
-                session = self.get_seedless_session()
-            else:
-                session = self.get_session(
-                    derive_cardano=derive_cardano,
-                    empty_passphrase=empty_passphrase,
-                    must_resume=must_resume,
-                )
-        except exceptions.DeviceLocked:
-            click.echo(
-                "Device is locked, enter a pin on the device.",
-                err=True,
+        yield from self._connection_context(
+            self.get_seedless_session
+            if seedless
+            else lambda: self.get_session(
+                derive_cardano=derive_cardano,
+                empty_passphrase=empty_passphrase,
+                must_resume=must_resume,
             )
-            sys.exit(1)
-        except transport.DeviceIsBusy:
-            click.echo("Device is in use by another process.")
-            sys.exit(1)
-        except exceptions.UnexpectedCodeEntryTagException:
-            click.echo("Entered Code is invalid.")
-            sys.exit(1)
-        except exceptions.FailedSessionResumption:
-            sys.exit(1)
-        except exceptions.DerivationOnUninitaizedDeviceError:
-            click.echo("Device is not initialized.")
-            sys.exit(1)
-        except Exception:
-            click.echo("Failed to find a Trezor device.")
-            if self.path is not None:
-                click.echo(f"Using path: {self.path}")
-            sys.exit(1)
+        )
 
-        try:
-            yield session
-        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 _print_exception(self, exc: Exception, message: str):
+        LOG.debug(message, exc_info=True)
+        message = f"{message}: {exc.__class__.__name__}"
+        if description := str(exc):
+            message = f"{message} ({description})"
+
+        click.echo(message)
+        if self.path is not None:
+            click.echo(f"Using path: {self.path}")
 
 
 def with_session(

python/src/trezorlib/client.py

@@ -172,20 +172,14 @@ def _handle_code_entry(self, session: SessionV2) -> None:
         sha_ctx = sha256(cpace.shared_secret)
         tag = sha_ctx.digest()
 
-        try:
-            secret_msg = session.call(
-                messages.ThpCodeEntryCpaceHostTag(
-                    cpace_host_public_key=cpace.host_public_key,
-                    tag=tag,
-                ),
-                expect=messages.ThpCodeEntrySecret,
-                skip_firmware_version_check=True,
-            )
-        except exceptions.TrezorFailure as e:
-            if e.message == "Unexpected Code Entry Tag":
-                raise exceptions.UnexpectedCodeEntryTagException
-            else:
-                raise e
+        secret_msg = session.call(
+            messages.ThpCodeEntryCpaceHostTag(
+                cpace_host_public_key=cpace.host_public_key,
+                tag=tag,
+            ),
+            expect=messages.ThpCodeEntrySecret,
+            skip_firmware_version_check=True,
+        )
 
         # Check `commitment` and `code`
         assert secret_msg.secret is not None
@@ -218,14 +212,7 @@ def get_session(
     ) -> Session:
         """
         Returns a new session.
-
-        In the case of seed derivation, the function will fail if the device is not initialized.
         """
-        if self.features.initialized is False and passphrase is not SEEDLESS:
-            raise exceptions.DerivationOnUninitaizedDeviceError(
-                "Calling uninitialized device with a passphrase. Call get_seedless_session instead."
-            )
-
         if isinstance(self.protocol, ProtocolV1Channel):
             from .transport.session import SessionV1, derive_seed
 

python/src/trezorlib/exceptions.py

@@ -105,16 +105,6 @@ class InvalidSessionError(TrezorException):
     Raised when Trezor returns unexpected PassphraseRequest"""
 
 
-class DerivationOnUninitaizedDeviceError(TrezorException):
-    """Tried to derive seed on uninitialized device.
-
-    To communicate with uninitialized device, use seedless session instead."""
-
-
-class UnexpectedCodeEntryTagException(TrezorException):
-    pass
-
-
 class ThpError(TrezorException):
     pass