diff --git a/core/src/trezor/wire/protocol_common.py b/core/src/trezor/wire/protocol_common.py index 93b77e617..8d091d8e5 100644 --- a/core/src/trezor/wire/protocol_common.py +++ b/core/src/trezor/wire/protocol_common.py @@ -2,19 +2,30 @@ from trezor import protobuf class Message: + def __init__( + self, + message_data: bytes, + ) -> None: + self.data = message_data + + def to_bytes(self): + return self.data + + +class MessageWithType(Message): def __init__( self, message_type: int, message_data: bytes, ) -> None: self.type = message_type - self.data = message_data + super().__init__(message_data) def to_bytes(self): return self.type.to_bytes(2, "big") + self.data -class MessageWithId(Message): +class MessageWithId(MessageWithType): def __init__( self, message_type: int, diff --git a/core/src/trezor/wire/thp/thp_messages.py b/core/src/trezor/wire/thp/thp_messages.py index 12b4649dd..1bcd87382 100644 --- a/core/src/trezor/wire/thp/thp_messages.py +++ b/core/src/trezor/wire/thp/thp_messages.py @@ -60,7 +60,7 @@ _ERROR_UNALLOCATED_SESSION = ( def get_device_properties() -> Message: - return Message(1000, _ENCODED_PROTOBUF_DEVICE_PROPERTIES) + return Message(_ENCODED_PROTOBUF_DEVICE_PROPERTIES) def get_channel_allocation_response(nonce: bytes, new_cid: int) -> bytes: