mirror of
https://github.com/trezor/trezor-firmware.git
synced 2025-07-09 16:18:10 +00:00
37 lines
1.2 KiB
Python
37 lines
1.2 KiB
Python
from trezor.messages import MessageType
|
|
|
|
if __debug__:
|
|
from trezor import log
|
|
|
|
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):
|
|
"""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(MessageType):
|
|
# Modules contain internal variables that may cause exception here.
|
|
# No Message begins with underscore so it's safe to skip those.
|
|
if msg_name[0] == "_":
|
|
continue
|
|
type_to_name[getattr(MessageType, msg_name)] = msg_name
|