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):