From 2731742d340429fff7aa1ef401357d1b98d3dff5 Mon Sep 17 00:00:00 2001 From: Jan Pochyla Date: Thu, 5 Apr 2018 12:34:25 +0200 Subject: [PATCH] src/trezor/messages: allow registration of custom messages --- src/trezor/messages/__init__.py | 35 +++++++++++++++++++++++++-------- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/src/trezor/messages/__init__.py b/src/trezor/messages/__init__.py index dac057b09..1bde9d28e 100644 --- a/src/trezor/messages/__init__.py +++ b/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