From 653ed4a67bfd7e5386d498c5152a2c87c4937fab Mon Sep 17 00:00:00 2001 From: slush Date: Sun, 17 Dec 2017 03:17:37 +0100 Subject: [PATCH] Added registering custom protobuf messages by application. --- trezorlib/client.py | 5 +++++ trezorlib/mapping.py | 16 ++++++++++++---- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/trezorlib/client.py b/trezorlib/client.py index 71ed664f3..ba0d4c125 100644 --- a/trezorlib/client.py +++ b/trezorlib/client.py @@ -32,6 +32,7 @@ from mnemonic import Mnemonic from . import messages as proto from . import tools +from . import mapping from .coins import coins_slip44 from .debuglink import DebugLink @@ -204,6 +205,10 @@ class BaseClient(object): raise CallException(msg.code, msg.message) + def register_message(self, msg): + '''Allow application to register custom protobuf message type''' + mapping.register_message(msg) + class VerboseWireMixin(object): def call_raw(self, msg): diff --git a/trezorlib/mapping.py b/trezorlib/mapping.py index 96b538a21..776d21a6e 100644 --- a/trezorlib/mapping.py +++ b/trezorlib/mapping.py @@ -32,13 +32,21 @@ def build_map(): try: msg_class = getattr(messages, msg_name) except AttributeError: - raise raise ValueError("Implementation of protobuf message '%s' is missing" % msg_name) - wire_type = getattr(messages.MessageType, 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) - map_type_to_class[wire_type] = msg_class - map_class_to_type[msg_class] = wire_type + 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):