From 46d1605ef19fbbd10a9f1500701cd08a67a2520c Mon Sep 17 00:00:00 2001
From: M1nd3r <petrsedlacek.km@seznam.cz>
Date: Wed, 18 Dec 2024 15:55:51 +0100
Subject: [PATCH] fixup! fixup! wip: single packet decryption (not finished)
 [no changelog]

---
 core/src/trezor/wire/thp/__init__.py          |  5 ++++
 .../wire/thp/received_message_handler.py      | 12 +++++++++-
 python/src/trezorlib/cli/__init__.py          | 24 ++++++++++++-------
 python/src/trezorlib/exceptions.py            |  4 ++++
 .../trezorlib/transport/thp/protocol_v2.py    |  6 +++++
 5 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/core/src/trezor/wire/thp/__init__.py b/core/src/trezor/wire/thp/__init__.py
index 76288c1f9b..3250b5aad9 100644
--- a/core/src/trezor/wire/thp/__init__.py
+++ b/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):
diff --git a/core/src/trezor/wire/thp/received_message_handler.py b/core/src/trezor/wire/thp/received_message_handler.py
index 1ee3b581e1..6403315568 100644
--- a/core/src/trezor/wire/thp/received_message_handler.py
+++ b/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
diff --git a/python/src/trezorlib/cli/__init__.py b/python/src/trezorlib/cli/__init__.py
index 18253a06ec..462a265c50 100644
--- a/python/src/trezorlib/cli/__init__.py
+++ b/python/src/trezorlib/cli/__init__.py
@@ -262,17 +262,23 @@ def with_session(
         def function_with_session(
             obj: TrezorConnection, *args: "P.args", **kwargs: "P.kwargs"
         ) -> "R":
-            if management:
-                session = obj.get_management_session()
-            else:
-                # TODO try (sys.exit ve finally)
-                session = obj.get_session(
-                    derive_cardano=derive_cardano,
-                    empty_passphrase=empty_passphrase,
-                    must_resume=must_resume,
-                )
             try:
+                if management:
+                    session = obj.get_management_session()
+                else:
+                    # TODO try (sys.exit ve finally)
+                    session = obj.get_session(
+                        derive_cardano=derive_cardano,
+                        empty_passphrase=empty_passphrase,
+                        must_resume=must_resume,
+                    )
+
                 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
diff --git a/python/src/trezorlib/exceptions.py b/python/src/trezorlib/exceptions.py
index fd7133d12f..f56a82db71 100644
--- a/python/src/trezorlib/exceptions.py
+++ b/python/src/trezorlib/exceptions.py
@@ -55,3 +55,7 @@ class Cancelled(TrezorException):
 
 class OutdatedFirmwareError(TrezorException):
     pass
+
+
+class DeviceLockedException(TrezorException):
+    pass
diff --git a/python/src/trezorlib/transport/thp/protocol_v2.py b/python/src/trezorlib/transport/thp/protocol_v2.py
index 07ff2cadd4..2637a2e064 100644
--- a/python/src/trezorlib/transport/thp/protocol_v2.py
+++ b/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,