import bitkey_pb2 as proto map_type_to_class = { 0: proto.Initialize, 1: proto.Ping, 2: proto.Success, 3: proto.Failure, 4: proto.GetUUID, 5: proto.UUID, 6: proto.OtpRequest, 7: proto.OtpAck, 8: proto.OtpCancel, 9: proto.GetEntropy, 10: proto.Entropy, 11: proto.GetMasterPublicKey, 12: proto.MasterPublicKey, 13: proto.LoadDevice, 14: proto.ResetDevice, 15: proto.SignTx, 16: proto.SignedTx, 17: proto.Features, 18: proto.PinRequest, 19: proto.PinAck, 20: proto.PinCancel, } map_class_to_type = {} def get_type(msg): return map_class_to_type[msg.__class__] def get_class(t): return map_type_to_class[t] def build_index(): for k, v in map_type_to_class.items(): map_class_to_type[v] = k def check_missing(): from google.protobuf import reflection types = [ proto.__dict__[item] for item in dir(proto) if issubclass(proto.__dict__[item].__class__, reflection.GeneratedProtocolMessageType) ] missing = list(set(types) - set(map_type_to_class.values())) if len(missing): raise Exception("Following protobuf messages are not defined in mapping: %s" % missing) check_missing() build_index()