From 41551ffffa6252eeaa42aa386c3bd24685634f0b Mon Sep 17 00:00:00 2001 From: M1nd3r Date: Tue, 26 Mar 2024 15:37:19 +0100 Subject: [PATCH] Fix imports in thp_v1.py --- core/src/trezor/wire/protocol.py | 7 +++---- core/src/trezor/wire/thp/channel_context.py | 14 ++++++++------ core/src/trezor/wire/thp/thp_messages.py | 2 -- core/src/trezor/wire/thp_v1.py | 14 +++++++------- core/tests/test_trezor.wire.thp_v1.py | 2 +- 5 files changed, 19 insertions(+), 20 deletions(-) diff --git a/core/src/trezor/wire/protocol.py b/core/src/trezor/wire/protocol.py index fca86c730..63b3ef313 100644 --- a/core/src/trezor/wire/protocol.py +++ b/core/src/trezor/wire/protocol.py @@ -1,7 +1,7 @@ from typing import TYPE_CHECKING # pyright: ignore[reportShadowedImports] from trezor import utils -from trezor.wire import codec_v1, thp_v1 +from trezor.wire import codec_v1 from trezor.wire.protocol_common import MessageWithId if TYPE_CHECKING: @@ -10,13 +10,12 @@ if TYPE_CHECKING: async def read_message(iface: WireInterface, buffer: utils.BufferType) -> MessageWithId: if utils.USE_THP: - return await thp_v1.read_message(iface, buffer) + raise Exception("THP protocol should be used instead") return await codec_v1.read_message(iface, buffer) async def write_message(iface: WireInterface, message: MessageWithId) -> None: if utils.USE_THP: - await thp_v1.write_message_with_sync_control(iface, message) - return + raise Exception("THP protocol should be used instead") await codec_v1.write_message(iface, message.type, message.data) return diff --git a/core/src/trezor/wire/thp/channel_context.py b/core/src/trezor/wire/thp/channel_context.py index ba2414cff..3423f3f52 100644 --- a/core/src/trezor/wire/thp/channel_context.py +++ b/core/src/trezor/wire/thp/channel_context.py @@ -13,11 +13,9 @@ from . import thp_session as THP from .checksum import CHECKSUM_LENGTH from .thp_messages import ( ACK_MESSAGE, - CONT_DATA_OFFSET, CONTINUATION_PACKET, ENCRYPTED_TRANSPORT, HANDSHAKE_INIT, - INIT_DATA_OFFSET, ) from .thp_session import ThpError @@ -30,8 +28,12 @@ _MOCK_INTERFACE_HID = b"\x00" _PUBKEY_LENGTH = const(32) -_REPORT_LENGTH = const(64) -_MAX_PAYLOAD_LEN = const(60000) +INIT_DATA_OFFSET = const(5) +CONT_DATA_OFFSET = const(3) + + +REPORT_LENGTH = const(64) +MAX_PAYLOAD_LEN = const(60000) class ChannelContext(Context): @@ -267,7 +269,7 @@ def _encode_iface(iface: WireInterface) -> bytes: def _get_buffer_for_payload( - payload_length: int, existing_buffer: utils.BufferType, max_length=_MAX_PAYLOAD_LEN + payload_length: int, existing_buffer: utils.BufferType, max_length=MAX_PAYLOAD_LEN ) -> utils.BufferType: if payload_length > max_length: raise ThpError("Message too large") @@ -276,7 +278,7 @@ def _get_buffer_for_payload( try: payload: utils.BufferType = bytearray(payload_length) except MemoryError: - payload = bytearray(_REPORT_LENGTH) + payload = bytearray(REPORT_LENGTH) raise ThpError("Message too large") return payload diff --git a/core/src/trezor/wire/thp/thp_messages.py b/core/src/trezor/wire/thp/thp_messages.py index b4a764997..0b30746e8 100644 --- a/core/src/trezor/wire/thp/thp_messages.py +++ b/core/src/trezor/wire/thp/thp_messages.py @@ -12,8 +12,6 @@ HANDSHAKE_INIT = 0x00 ACK_MESSAGE = 0x20 _ERROR = 0x41 _CHANNEL_ALLOCATION_RES = 0x40 -INIT_DATA_OFFSET = const(5) -CONT_DATA_OFFSET = const(3) class InitHeader: diff --git a/core/src/trezor/wire/thp_v1.py b/core/src/trezor/wire/thp_v1.py index 56e5f2aa0..fb5193eab 100644 --- a/core/src/trezor/wire/thp_v1.py +++ b/core/src/trezor/wire/thp_v1.py @@ -9,18 +9,18 @@ from .protocol_common import MessageWithId from .thp import ChannelState, ack_handler, checksum, thp_messages from .thp import thp_session as THP from .thp.channel_context import ( - _MAX_PAYLOAD_LEN, - _REPORT_LENGTH, + CONT_DATA_OFFSET, + INIT_DATA_OFFSET, + MAX_PAYLOAD_LEN, + REPORT_LENGTH, ChannelContext, load_cached_channels, ) from .thp.checksum import CHECKSUM_LENGTH from .thp.thp_messages import ( CODEC_V1, - CONT_DATA_OFFSET, CONTINUATION_PACKET, ENCRYPTED_TRANSPORT, - INIT_DATA_OFFSET, InitHeader, InterruptingInitPacket, ) @@ -192,7 +192,7 @@ def _get_loop_wait_read(iface: WireInterface): def _get_buffer_for_payload( - payload_length: int, existing_buffer: utils.BufferType, max_length=_MAX_PAYLOAD_LEN + payload_length: int, existing_buffer: utils.BufferType, max_length=MAX_PAYLOAD_LEN ) -> utils.BufferType: if payload_length > max_length: raise ThpError("Message too large") @@ -201,7 +201,7 @@ def _get_buffer_for_payload( try: payload: utils.BufferType = bytearray(payload_length) except MemoryError: - payload = bytearray(_REPORT_LENGTH) + payload = bytearray(REPORT_LENGTH) raise ThpError("Message too large") return payload @@ -302,7 +302,7 @@ async def write_to_wire( payload_length = len(payload) # prepare the report buffer with header data - report = bytearray(_REPORT_LENGTH) + report = bytearray(REPORT_LENGTH) header.pack_to_buffer(report) # write initial report diff --git a/core/tests/test_trezor.wire.thp_v1.py b/core/tests/test_trezor.wire.thp_v1.py index a2ffb2aec..728deb5cc 100644 --- a/core/tests/test_trezor.wire.thp_v1.py +++ b/core/tests/test_trezor.wire.thp_v1.py @@ -322,7 +322,7 @@ class TestWireTrezorHostProtocolV1(unittest.TestCase): # ensure that a message this big won't fit into memory # Note: this control is changed, because THP has only 2 byte length field - self.assertTrue(message_size > thp_v1._MAX_PAYLOAD_LEN) + self.assertTrue(message_size > thp_v1.MAX_PAYLOAD_LEN) # self.assertRaises(MemoryError, bytearray, message_size) header = make_header(PLAINTEXT_1, COMMON_CID, message_size) packet = header + MESSAGE_TYPE_BYTES + (b"\x00" * INIT_MESSAGE_DATA_LENGTH)