refactor(python): refactor THP-related exception types

[no changelog]

Author: romanz

python/src/trezorlib/cli/__init__.py

@@ -293,7 +293,7 @@ def session_context(
                     empty_passphrase=empty_passphrase,
                     must_resume=must_resume,
                 )
-        except exceptions.DeviceLockedException:
+        except exceptions.DeviceLocked:
             click.echo(
                 "Device is locked, enter a pin on the device.",
                 err=True,

python/src/trezorlib/debuglink.py

@@ -34,12 +34,7 @@
 
 from . import btc, mapping, messages, models, protobuf
 from .client import ProtocolVersion, TrezorClient
-from .exceptions import (
-    Cancelled,
-    DeviceLockedException,
-    TrezorFailure,
-    UnexpectedMessageError,
-)
+from .exceptions import Cancelled, TrezorFailure, UnexpectedMessageError
 from .log import DUMP_BYTES
 from .messages import DebugTouchEventType, DebugWaitType
 from .tools import parse_path
@@ -1317,14 +1312,7 @@ def get_pin(_msg: messages.PinMatrixRequest) -> str:
         self.pin_callback = get_pin
         self.button_callback = self.ui.button_request
 
-        try:
-            super().__init__(transport)
-        except DeviceLockedException:
-            LOG.debug("Locked device handling")
-            self.debug.input("")
-            self.debug.input(self.debug.encode_pin("1234"))
-            super().__init__(transport)
-
+        super().__init__(transport)
         self.sync_responses()
 
         # So that we can choose right screenshotting logic (T1 vs TT)

python/src/trezorlib/exceptions.py

@@ -111,13 +111,33 @@ class DerivationOnUninitaizedDeviceError(TrezorException):
     To communicate with uninitialized device, use seedless session instead."""
 
 
-class DeviceLockedException(TrezorException):
+class UnexpectedCodeEntryTagException(TrezorException):
     pass
 
 
-class UnexpectedCodeEntryTagException(TrezorException):
+class ThpError(TrezorException):
     pass
 
 
-class ThpError(TrezorException):
+class TransportBusy(ThpError):
+    pass
+
+
+class UnallocatedChannel(ThpError):
+    pass
+
+
+class DecryptionFailed(ThpError):
+    pass
+
+
+class InvalidData(ThpError):
+    pass
+
+
+class DeviceLocked(ThpError):
+    pass
+
+
+class ThpUnknownError(ThpError):
     pass

python/src/trezorlib/transport/thp/protocol_v2.py

@@ -231,11 +231,6 @@ def _send_handshake_init_request(self) -> None:
 
     def _read_handshake_init_response(self) -> bytes:
         header, payload = self._read_until_valid_crc_check()
-        if control_byte.is_error(header.ctrl_byte):
-            if payload == b"\x05":
-                raise exceptions.DeviceLockedException()
-            else:
-                raise exceptions.ThpError(_get_error_from_int(payload[0]))
 
         if not header.is_handshake_init_response():
             LOG.error("Received message is not a valid handshake init response message")
@@ -278,8 +273,6 @@ def _read_handshake_completion_response(self) -> int:
         header, data = self._read_until_valid_crc_check()
         if not header.is_handshake_comp_response():
             LOG.error("Received message is not a valid handshake completion response")
-            if control_byte.is_error(header.ctrl_byte):
-                raise exceptions.ThpError(_get_error_from_int(data[0]))
         trezor_state = self._noise.decrypt(bytes(data))
         assert trezor_state == b"\x00" or trezor_state == b"\x01"
         self._send_ack_bit(bit=1)
@@ -289,8 +282,6 @@ def _read_ack(self):
         header, payload = self._read_until_valid_crc_check()
         if not header.is_ack() or len(payload) > 0:
             LOG.error("Received message is not a valid ACK")
-            if control_byte.is_error(header.ctrl_byte):
-                raise exceptions.ThpError(_get_error_from_int(payload[0]))
 
     def _send_ack_bit(self, bit: int):
         if bit not in (0, 1):
@@ -336,8 +327,6 @@ def read_and_decrypt(
                 continue
             if control_byte.is_ack(header.ctrl_byte):
                 continue
-            if control_byte.is_error(header.ctrl_byte):
-                raise exceptions.ThpError(_get_error_from_int(raw_payload[0]))
             if not header.is_encrypted_transport():
                 LOG.error(
                     "Trying to decrypt not encrypted message! ("
@@ -392,6 +381,10 @@ def _read_until_valid_crc_check(
 
                 self.sync_bit_receive = 1 - self.sync_bit_receive
 
+            if control_byte.is_error(header.ctrl_byte):
+                code = payload[0]
+                raise _ERRORS_MAP.get(code) or exceptions.ThpUnknownError(code)
+
             return header, payload
 
     def _is_valid_channel_allocation_response(
@@ -420,16 +413,10 @@ def _is_valid_pong(
         return True
 
 
-def _get_error_from_int(error_code: int) -> str:
-    # TODO FIXME improve this (ThpErrorType)
-    if error_code == 1:
-        return "TRANSPORT BUSY"
-    if error_code == 2:
-        return "UNALLOCATED CHANNEL"
-    if error_code == 3:
-        return "DECRYPTION FAILED"
-    if error_code == 4:
-        return "INVALID DATA"
-    if error_code == 5:
-        return "DEVICE LOCKED"
-    raise Exception("Not Implemented error case")
+_ERRORS_MAP = {
+    1: exceptions.TransportBusy,
+    2: exceptions.UnallocatedChannel,
+    3: exceptions.DecryptionFailed,
+    4: exceptions.InvalidData,
+    5: exceptions.DeviceLocked,
+}

tests/device_tests/thp/test_multiple_hosts.py

@@ -31,10 +31,9 @@ def test_concurrent_handshakes(client: Client) -> None:
     # The second host starts handshake
     protocol_2._send_handshake_init_request()
 
-    # The second host should not be able to interrupt the first host's handshake
-    with pytest.raises(exceptions.ThpError) as e:
+    # The second host should not be able to interrupt the first host's handshake immediately
+    with pytest.raises(exceptions.TransportBusy):
         protocol_2._read_ack()
-    assert e.value.args[0] == "TRANSPORT BUSY"
 
     # The first host can complete handshake
     protocol_1._send_handshake_completion_request()

tests/device_tests/thp/test_pairing.py

@@ -406,9 +406,8 @@ def test_credential_phase(client: Client) -> None:
     protocol._noise.noise_protocol.cipher_state_encrypt.n = 250
 
     protocol._send_message(ButtonAck())
-    with pytest.raises(exceptions.ThpError) as e:
+    with pytest.raises(exceptions.DecryptionFailed):
         protocol.read(1)
-    assert e.value.args[0] == "DECRYPTION FAILED"
 
     # Connect using credential with confirmation and ask for autoconnect credential.
     protocol = prepare_protocol_for_handshake(client)
@@ -457,9 +456,8 @@ def test_credential_phase(client: Client) -> None:
     protocol._noise.noise_protocol.cipher_state_encrypt.n = 100
 
     protocol._send_message(ButtonAck())
-    with pytest.raises(exceptions.ThpError) as e:
+    with pytest.raises(exceptions.DecryptionFailed):
         protocol.read(1)
-    assert e.value.args[0] == "DECRYPTION FAILED"
 
     # Connect using autoconnect credential - should work the same as above
     protocol = prepare_protocol_for_handshake(client)