fixup! fixup! wip: single packet decryption (not finished) [no changelog]

core/src/trezor/wire/thp/__init__.py

@@ -44,6 +44,10 @@ class ThpInvalidDataError(ThpError):
     pass
 
 
+class ThpDeviceLockedError(ThpError):
+    pass
+
+
 class ThpUnallocatedSessionError(ThpError):
 
     def __init__(self, session_id: int) -> None:
@@ -55,6 +59,7 @@ class ThpErrorType(IntEnum):
     UNALLOCATED_CHANNEL = 2
     DECRYPTION_FAILED = 3
     INVALID_DATA = 4
+    DEVICE_LOCKED = 5
 
 
 class ChannelState(IntEnum):

core/src/trezor/wire/thp/received_message_handler.py

@@ -15,7 +15,7 @@ from storage.cache_thp import (
     update_channel_last_used,
     update_session_last_used,
 )
-from trezor import log, loop, protobuf, utils
+from trezor import config, log, loop, protobuf, utils
 from trezor.enums import FailureType
 from trezor.messages import Failure
 from trezor.wire.thp import memory_manager
@@ -35,6 +35,7 @@ from . import (
     ThpErrorType,
     ThpInvalidDataError,
     ThpUnallocatedSessionError,
+    ThpDeviceLockedError,
 )
 from . import alternating_bit_protocol as ABP
 from . import (
@@ -139,6 +140,9 @@ async def handle_received_message(
     except ThpInvalidDataError:
         await ctx.write_error(ThpErrorType.INVALID_DATA)
         ctx.clear()
+    except ThpDeviceLockedError:
+        await ctx.write_error(ThpErrorType.DEVICE_LOCKED)
+
     if __debug__ and utils.ALLOW_DEBUG_MESSAGES:
         log.debug(__name__, "handle_received_message - end")
 
@@ -226,6 +230,9 @@ async def _handle_state_TH1(
     if not payload_length == PUBKEY_LENGTH + CHECKSUM_LENGTH:
         raise ThpError("Message received is not a valid handshake init request!")
 
+    if not config.is_unlocked():
+        raise ThpDeviceLockedError
+
     ctx.handshake = Handshake()
 
     buffer = memory_manager.get_existing_read_buffer(ctx.get_channel_id_int())
@@ -272,6 +279,9 @@ async def _handle_state_TH2(ctx: Channel, message_length: int, ctrl_byte: int) -
     if ctx.handshake is None:
         raise Exception("Handshake object is not prepared. Retry handshake.")
 
+    if not config.is_unlocked():
+        raise ThpDeviceLockedError
+
     buffer = memory_manager.get_existing_read_buffer(ctx.get_channel_id_int())
     # if buffer is BufferError:
     # pass  # TODO handle

python/src/trezorlib/cli/__init__.py

@@ -262,6 +262,7 @@ def with_session(
         def function_with_session(
             obj: TrezorConnection, *args: "P.args", **kwargs: "P.kwargs"
         ) -> "R":
+            try:
                 if management:
                     session = obj.get_management_session()
                 else:
@@ -271,8 +272,13 @@ def with_session(
                         empty_passphrase=empty_passphrase,
                         must_resume=must_resume,
                     )
-            try:
+
                 return func(session, *args, **kwargs)
+            except exceptions.DeviceLockedException:
+                click.echo(
+                    "Device is locked, enter a pin on the device.",
+                    err=True,
+                )
             finally:
                 pass
                 # TODO try end session if not resumed

python/src/trezorlib/exceptions.py

@@ -55,3 +55,7 @@ class Cancelled(TrezorException):
 
 class OutdatedFirmwareError(TrezorException):
     pass
+
+
+class DeviceLockedException(TrezorException):
+    pass

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

@@ -173,7 +173,13 @@ class ProtocolV2(ProtocolAndChannel):
         header, payload = self._read_until_valid_crc_check()
         self._send_ack_0()
 
+        if header.ctrl_byte == 0x42:
+            if payload == b"\x05":
+                raise exceptions.DeviceLockedException()
+
         if not header.is_handshake_init_response():
+            LOG.debug("Received message is not a valid handshake init response message")
+
             click.echo(
                 "Received message is not a valid handshake init response message",
                 err=True,