src/trezor/wire: allow passing MessageType classes to register()

src/trezor/wire/__init__.py

@@ -13,6 +13,8 @@ workflow_handlers = {}
 
 def register(mtype, handler, *args):
     '''Register `handler` to get scheduled after `mtype` message is received.'''
+    if isinstance(mtype, type) and issubclass(mtype, protobuf.MessageType):
+        mtype = mtype.MESSAGE_WIRE_TYPE
     if mtype in workflow_handlers:
         raise KeyError
     workflow_handlers[mtype] = (handler, args)