Added registering custom protobuf messages by application.

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

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