src/trezor/messages: allow registration of custom messages

src/trezor/messages/__init__.py

@@ -1,13 +1,32 @@
-from . import wire_types
+from trezor.messages import wire_types
 
+if __debug__:
+    from trezor import log
 
-def get_type_name(wire_type):
-    for name in dir(wire_types):
-        if getattr(wire_types, name) == wire_type:
-            return name
+type_to_name = {}  # int -> string, reverse table of wire_type mapping
+registered = {}  # int -> class, dynamically registered message types
+
+
+def register(msg_type):
+    '''Register custom message type in runtime.'''
+    if __debug__:
+        log.debug(__name__, 'register %s', msg_type)
+    registered[msg_type.MESSAGE_WIRE_TYPE] = msg_type
 
 
 def get_type(wire_type):
-    name = get_type_name(wire_type)
-    module = __import__('trezor.messages.%s' % name, None, None, (name, ), 0)
-    return getattr(module, name)
+    '''Get message class for handling given wire_type.'''
+    if wire_type in registered:
+        # message class is explicitly registered
+        msg_type = registered[wire_type]
+    else:
+        # import message class from trezor.messages dynamically
+        name = type_to_name[wire_type]
+        module = __import__('trezor.messages.%s' % name, None, None, (name, ), 0)
+        msg_type = getattr(module, name)
+    return msg_type
+
+
+# build reverse table of wire types
+for msg_name in dir(wire_types):
+    type_to_name[getattr(wire_types, msg_name)] = msg_name