mirror of
https://github.com/trezor/trezor-firmware.git
synced 2025-02-07 13:12:41 +00:00
9a330f3475
Protobuf encoding now happens in TrezorClient, and transports get encoded blobs to (chunkify and) send. This is a better design because transports don't need to know about protobuf. It also lays groundwork for sending raw bytes feature (#116) This commit also removes all vestiges of ProtocolV2 which was never used and will probably need to be redesigned from the ground up anyway. The code is still ready for protocol flexibility.
79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
# This file is part of the Trezor project.
|
|
#
|
|
# Copyright (C) 2012-2019 SatoshiLabs and contributors
|
|
#
|
|
# This library is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Lesser General Public License version 3
|
|
# as published by the Free Software Foundation.
|
|
#
|
|
# This library is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Lesser General Public License for more details.
|
|
#
|
|
# You should have received a copy of the License along with this library.
|
|
# If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
|
|
|
|
import io
|
|
from typing import Tuple
|
|
|
|
from . import messages, protobuf
|
|
|
|
map_type_to_class = {}
|
|
map_class_to_type = {}
|
|
|
|
|
|
def build_map():
|
|
for msg_name in dir(messages.MessageType):
|
|
if msg_name.startswith("__"):
|
|
continue
|
|
|
|
try:
|
|
msg_class = getattr(messages, msg_name)
|
|
except AttributeError:
|
|
raise ValueError(
|
|
"Implementation of protobuf message '%s' is missing" % msg_name
|
|
)
|
|
|
|
if msg_class.MESSAGE_WIRE_TYPE != getattr(messages.MessageType, msg_name):
|
|
raise ValueError(
|
|
"Inconsistent wire type and MessageType record for '%s'" % msg_class
|
|
)
|
|
|
|
register_message(msg_class)
|
|
|
|
|
|
def register_message(msg_class):
|
|
if msg_class.MESSAGE_WIRE_TYPE in map_type_to_class:
|
|
raise Exception(
|
|
"Message for wire type %s is already registered by %s"
|
|
% (msg_class.MESSAGE_WIRE_TYPE, get_class(msg_class.MESSAGE_WIRE_TYPE))
|
|
)
|
|
|
|
map_class_to_type[msg_class] = msg_class.MESSAGE_WIRE_TYPE
|
|
map_type_to_class[msg_class.MESSAGE_WIRE_TYPE] = msg_class
|
|
|
|
|
|
def get_type(msg):
|
|
return map_class_to_type[msg.__class__]
|
|
|
|
|
|
def get_class(t):
|
|
return map_type_to_class[t]
|
|
|
|
|
|
def encode(msg: protobuf.MessageType) -> Tuple[int, bytes]:
|
|
message_type = msg.MESSAGE_WIRE_TYPE
|
|
buf = io.BytesIO()
|
|
protobuf.dump_message(buf, msg)
|
|
return message_type, buf.getvalue()
|
|
|
|
|
|
def decode(message_type: int, message_bytes: bytes) -> protobuf.MessageType:
|
|
cls = get_class(message_type)
|
|
buf = io.BytesIO(message_bytes)
|
|
return protobuf.load_message(buf, cls)
|
|
|
|
|
|
build_map()
|