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,