protobuf: add enum value validation (#363)

common/protob/pb2py

@@ -10,13 +10,12 @@ import shutil
 import subprocess
 import sys
 import tempfile
-from collections import namedtuple
+from collections import namedtuple, defaultdict
+
+import attr
 
 from google.protobuf import descriptor_pb2
 
-ProtoField = namedtuple(
-    "ProtoField", "name, number, proto_type, py_type, repeated, required, orig"
-)
 
 AUTO_HEADER = "# Automatically generated by pb2py\n"
 
@@ -24,7 +23,7 @@ AUTO_HEADER = "# Automatically generated by pb2py\n"
 FIELD_TYPES = {
     descriptor_pb2.FieldDescriptorProto.TYPE_UINT64:  ('p.UVarintType', 'int'),
     descriptor_pb2.FieldDescriptorProto.TYPE_UINT32:  ('p.UVarintType', 'int'),
-    descriptor_pb2.FieldDescriptorProto.TYPE_ENUM:    ('p.UVarintType', 'int'),
+#    descriptor_pb2.FieldDescriptorProto.TYPE_ENUM:    ('p.UVarintType', 'int'),
     descriptor_pb2.FieldDescriptorProto.TYPE_SINT32:  ('p.SVarintType', 'int'),
     descriptor_pb2.FieldDescriptorProto.TYPE_SINT64:  ('p.SVarintType', 'int'),
     descriptor_pb2.FieldDescriptorProto.TYPE_STRING:  ('p.UnicodeType', 'str'),
@@ -42,6 +41,53 @@ PROTOC_PREFIX = os.path.dirname(os.path.dirname(PROTOC))
 PROTOC_INCLUDE = os.path.join(PROTOC_PREFIX, "include")
 
 
+@attr.s
+class ProtoField:
+    name = attr.ib()
+    number = attr.ib()
+    orig = attr.ib()
+    repeated = attr.ib()
+    required = attr.ib()
+    type_name = attr.ib()
+    proto_type = attr.ib()
+    py_type = attr.ib()
+
+    @classmethod
+    def from_field(cls, descriptor, field):
+        repeated = field.label == field.LABEL_REPEATED
+        required = field.label == field.LABEL_REQUIRED
+        # ignore package path
+        type_name = field.type_name.rsplit(".")[-1]
+
+        if field.type == field.TYPE_MESSAGE:
+            proto_type = py_type = type_name
+        elif field.type == field.TYPE_ENUM:
+            valuestr = ", ".join(str(v) for v in descriptor.enum_types[type_name])
+            proto_type = 'p.EnumType("{}", ({}))'.format(type_name, valuestr)
+            py_type = "EnumType" + type_name
+        else:
+            try:
+                proto_type, py_type = FIELD_TYPES[field.type]
+            except KeyError:
+                raise ValueError(
+                    "Unknown field type {} for field {}".format(field.type, field.name)
+                ) from None
+
+        if repeated:
+            py_type = "List[{}]".format(py_type)
+
+        return cls(
+            name=field.name,
+            number=field.number,
+            orig=field,
+            repeated=repeated,
+            required=required,
+            type_name=type_name,
+            proto_type=proto_type,
+            py_type=py_type,
+        )
+
+
 def protoc(files, additional_includes=()):
     """Compile code with protoc and return the data."""
     include_dirs = set()
@@ -113,6 +159,7 @@ class Descriptor:
         # find messages and enums
         self.messages = []
         self.enums = []
+        self.enum_types = defaultdict(set)
         for file in self.files:
             self.messages += file.message_type
             self.enums += file.enum_type
@@ -136,9 +183,7 @@ class Descriptor:
     def find_message_types(self, message_type):
         message_types = {}
         try:
-            message_type_enum = next(
-                enum for enum in self.enums if enum.name == message_type
-            )
+            message_type_enum = next(e for e in self.enums if e.name == message_type)
             for value in message_type_enum.value:
                 name = strip_leader(value.name, message_type)
                 message_types[name] = value.number
@@ -151,38 +196,10 @@ class Descriptor:
 
         return message_types
 
-    def parse_field(self, field):
-        repeated = field.label == field.LABEL_REPEATED
-        required = field.label == field.LABEL_REQUIRED
-        if field.type == field.TYPE_MESSAGE:
-            # ignore package path
-            type_name = field.type_name.rsplit(".")[-1]
-            proto_type = py_type = type_name
-        else:
-            try:
-                proto_type, py_type = FIELD_TYPES[field.type]
-            except KeyError:
-                raise ValueError(
-                    "Unknown field type {} for field {}".format(field.type, field.name)
-                ) from None
-
-        if repeated:
-            py_type = "List[{}]".format(py_type)
-
-        return ProtoField(
-            name=field.name,
-            number=field.number,
-            proto_type=proto_type,
-            py_type=py_type,
-            repeated=repeated,
-            required=required,
-            orig=field,
-        )
-
     def create_message_import(self, name):
         return "from .{0} import {0}".format(name)
 
-    def process_message_imports(self, fields):
+    def process_subtype_imports(self, fields):
         imports = set(
             field.proto_type
             for field in fields
@@ -251,16 +268,26 @@ class Descriptor:
         # "from .. import protobuf as p"
         yield self.protobuf_import + " as p"
 
-        fields = [self.parse_field(field) for field in message.field]
+        fields = [ProtoField.from_field(self, field) for field in message.field]
 
-        yield from self.process_message_imports(fields)
+        yield from self.process_subtype_imports(fields)
 
         yield ""
         yield "if __debug__:"
         yield "    try:"
         yield "        from typing import Dict, List, Optional"
+        yield "        from typing_extensions import Literal  # noqa: F401"
+
+        all_enums = [field for field in fields if field.type_name in self.enum_types]
+        for field in all_enums:
+            allowed_values = self.enum_types[field.type_name]
+            valuestr = ", ".join(str(v) for v in sorted(allowed_values))
+            yield "        {} = Literal[{}]".format(field.py_type, valuestr)
+
         yield "    except ImportError:"
         yield "        Dict, List, Optional = None, None, None  # type: ignore"
+        for field in all_enums:
+            yield "        {} = None  # type: ignore".format(field.py_type)
 
         yield ""
         yield ""
@@ -294,6 +321,7 @@ class Descriptor:
                 enum_prefix, _ = enum_prefix.rsplit("Type", 1)
                 name = strip_leader(name, enum_prefix)
 
+            self.enum_types[enum.name].add(value.number)
             yield "{} = {}".format(name, value.number)
 
     def process_messages(self, messages):
@@ -325,8 +353,8 @@ class Descriptor:
 
     def write_classes(self, out_dir, init_py=True):
         self.out_dir = out_dir
-        self.process_messages(self.messages)
         self.process_enums(self.enums)
+        self.process_messages(self.messages)
         if init_py:
             self.write_init_py()
 

core/src/protobuf.py

@@ -121,6 +121,19 @@ class BoolType:
     WIRE_TYPE = 0
 
 
+class EnumType:
+    WIRE_TYPE = 0
+
+    def __init__(self, _, enum_values):
+        self.enum_values = enum_values
+
+    def validate(self, fvalue: int) -> int:
+        if fvalue in self.enum_values:
+            return fvalue
+        else:
+            raise TypeError("Invalid enum value")
+
+
 class BytesType:
     WIRE_TYPE = 2
 
@@ -210,6 +223,8 @@ async def load_message(
             fvalue = uint_to_sint(ivalue)
         elif ftype is BoolType:
             fvalue = bool(ivalue)
+        elif isinstance(ftype, EnumType):
+            fvalue = ftype.validate(ivalue)
         elif ftype is BytesType:
             fvalue = bytearray(ivalue)
             await reader.areadinto(fvalue)
@@ -258,7 +273,7 @@ async def dump_message(
             repvalue[0] = fvalue
             fvalue = repvalue
 
-        if issubclass(ftype, MessageType):
+        if isinstance(ftype, type) and issubclass(ftype, MessageType):
             ffields = ftype.get_fields()
         else:
             ffields = None
@@ -275,6 +290,9 @@ async def dump_message(
             elif ftype is BoolType:
                 await dump_uvarint(writer, int(svalue))
 
+            elif isinstance(ftype, EnumType):
+                await dump_uvarint(writer, svalue)
+
             elif ftype is BytesType:
                 if isinstance(svalue, list):
                     await dump_uvarint(writer, _count_bytes_list(svalue))
@@ -332,6 +350,10 @@ def count_message(msg: MessageType, fields: Dict = None) -> int:
             for svalue in fvalue:
                 nbytes += count_uvarint(int(svalue))
 
+        elif isinstance(ftype, EnumType):
+            for svalue in fvalue:
+                nbytes += count_uvarint(svalue)
+
         elif ftype is BytesType:
             for svalue in fvalue:
                 if isinstance(svalue, list):

core/src/trezor/messages/Address.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/ApplyFlags.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/ApplySettings.py

@@ -5,8 +5,11 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
+        EnumTypePassphraseSourceType = Literal[0, 1, 2]
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
+        EnumTypePassphraseSourceType = None  # type: ignore
 
 
 class ApplySettings(p.MessageType):
@@ -18,7 +21,7 @@ class ApplySettings(p.MessageType):
         label: str = None,
         use_passphrase: bool = None,
         homescreen: bytes = None,
-        passphrase_source: int = None,
+        passphrase_source: EnumTypePassphraseSourceType = None,
         auto_lock_delay_ms: int = None,
         display_rotation: int = None,
     ) -> None:
@@ -37,7 +40,7 @@ class ApplySettings(p.MessageType):
             2: ('label', p.UnicodeType, 0),
             3: ('use_passphrase', p.BoolType, 0),
             4: ('homescreen', p.BytesType, 0),
-            5: ('passphrase_source', p.UVarintType, 0),
+            5: ('passphrase_source', p.EnumType("PassphraseSourceType", (0, 1, 2)), 0),
             6: ('auto_lock_delay_ms', p.UVarintType, 0),
             7: ('display_rotation', p.UVarintType, 0),
         }

core/src/trezor/messages/BackupDevice.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/BinanceAddress.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/BinanceCancelMsg.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/BinanceCoin.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/BinanceGetAddress.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/BinanceGetPublicKey.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/BinanceInputOutput.py

@@ -7,6 +7,7 @@ from .BinanceCoin import BinanceCoin
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/BinanceOrderMsg.py

@@ -5,8 +5,15 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
+        EnumTypeBinanceOrderType = Literal[0, 1, 2, 3]
+        EnumTypeBinanceOrderSide = Literal[0, 1, 2]
+        EnumTypeBinanceTimeInForce = Literal[0, 1, 2, 3]
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
+        EnumTypeBinanceOrderType = None  # type: ignore
+        EnumTypeBinanceOrderSide = None  # type: ignore
+        EnumTypeBinanceTimeInForce = None  # type: ignore
 
 
 class BinanceOrderMsg(p.MessageType):
@@ -15,13 +22,13 @@ class BinanceOrderMsg(p.MessageType):
     def __init__(
         self,
         id: str = None,
-        ordertype: int = None,
+        ordertype: EnumTypeBinanceOrderType = None,
         price: int = None,
         quantity: int = None,
         sender: str = None,
-        side: int = None,
+        side: EnumTypeBinanceOrderSide = None,
         symbol: str = None,
-        timeinforce: int = None,
+        timeinforce: EnumTypeBinanceTimeInForce = None,
     ) -> None:
         self.id = id
         self.ordertype = ordertype
@@ -36,11 +43,11 @@ class BinanceOrderMsg(p.MessageType):
     def get_fields(cls) -> Dict:
         return {
             1: ('id', p.UnicodeType, 0),
-            2: ('ordertype', p.UVarintType, 0),
+            2: ('ordertype', p.EnumType("BinanceOrderType", (0, 1, 2, 3)), 0),
             3: ('price', p.SVarintType, 0),
             4: ('quantity', p.SVarintType, 0),
             5: ('sender', p.UnicodeType, 0),
-            6: ('side', p.UVarintType, 0),
+            6: ('side', p.EnumType("BinanceOrderSide", (0, 1, 2)), 0),
             7: ('symbol', p.UnicodeType, 0),
-            8: ('timeinforce', p.UVarintType, 0),
+            8: ('timeinforce', p.EnumType("BinanceTimeInForce", (0, 1, 2, 3)), 0),
         }

core/src/trezor/messages/BinancePublicKey.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/BinanceSignTx.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/BinanceSignedTx.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/BinanceTransferMsg.py

@@ -7,6 +7,7 @@ from .BinanceInputOutput import BinanceInputOutput
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/BinanceTxRequest.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/ButtonAck.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/ButtonRequest.py

@@ -5,8 +5,11 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
+        EnumTypeButtonRequestType = Literal[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
+        EnumTypeButtonRequestType = None  # type: ignore
 
 
 class ButtonRequest(p.MessageType):
@@ -14,7 +17,7 @@ class ButtonRequest(p.MessageType):
 
     def __init__(
         self,
-        code: int = None,
+        code: EnumTypeButtonRequestType = None,
         data: str = None,
     ) -> None:
         self.code = code
@@ -23,6 +26,6 @@ class ButtonRequest(p.MessageType):
     @classmethod
     def get_fields(cls) -> Dict:
         return {
-            1: ('code', p.UVarintType, 0),
+            1: ('code', p.EnumType("ButtonRequestType", (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18)), 0),
             2: ('data', p.UnicodeType, 0),
         }

core/src/trezor/messages/Cancel.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CardanoAddress.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CardanoGetAddress.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CardanoGetPublicKey.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CardanoPublicKey.py

@@ -7,6 +7,7 @@ from .HDNodeType import HDNodeType
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CardanoSignTx.py

@@ -8,6 +8,7 @@ from .CardanoTxOutputType import CardanoTxOutputType
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CardanoSignedTx.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CardanoTxAck.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CardanoTxInputType.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CardanoTxOutputType.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CardanoTxRequest.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/ChangePin.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CipherKeyValue.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/CipheredKeyValue.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/ClearSession.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/DebugLinkDecision.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/DebugLinkGetState.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/DebugLinkState.py

@@ -7,6 +7,7 @@ from .HDNodeType import HDNodeType
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/DebugMoneroDiagAck.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/DebugMoneroDiagRequest.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/ECDHSessionKey.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/Entropy.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EntropyAck.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EntropyRequest.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionBuyRam.py

@@ -7,6 +7,7 @@ from .EosAsset import EosAsset
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionBuyRamBytes.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionCommon.py

@@ -7,6 +7,7 @@ from .EosPermissionLevel import EosPermissionLevel
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionDelegate.py

@@ -7,6 +7,7 @@ from .EosAsset import EosAsset
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionDeleteAuth.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionLinkAuth.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionNewAccount.py

@@ -7,6 +7,7 @@ from .EosAuthorization import EosAuthorization
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionRefund.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionSellRam.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionTransfer.py

@@ -7,6 +7,7 @@ from .EosAsset import EosAsset
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionUndelegate.py

@@ -7,6 +7,7 @@ from .EosAsset import EosAsset
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionUnknown.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionUnlinkAuth.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionUpdateAuth.py

@@ -7,6 +7,7 @@ from .EosAuthorization import EosAuthorization
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosActionVoteProducer.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosAsset.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosAuthorization.py

@@ -9,6 +9,7 @@ from .EosAuthorizationWait import EosAuthorizationWait
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosAuthorizationAccount.py

@@ -7,6 +7,7 @@ from .EosPermissionLevel import EosPermissionLevel
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosAuthorizationKey.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosAuthorizationWait.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosGetPublicKey.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosPermissionLevel.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosPublicKey.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosSignTx.py

@@ -7,6 +7,7 @@ from .EosTxHeader import EosTxHeader
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosSignedTx.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosTxActionAck.py

@@ -21,6 +21,7 @@ from .EosActionVoteProducer import EosActionVoteProducer
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosTxActionRequest.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EosTxHeader.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EthereumAddress.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EthereumGetAddress.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EthereumGetPublicKey.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EthereumMessageSignature.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EthereumPublicKey.py

@@ -7,6 +7,7 @@ from .HDNodeType import HDNodeType
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EthereumSignMessage.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EthereumSignTx.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EthereumTxAck.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EthereumTxRequest.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/EthereumVerifyMessage.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/Failure.py

@@ -5,8 +5,11 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
+        EnumTypeFailureType = Literal[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 99]
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
+        EnumTypeFailureType = None  # type: ignore
 
 
 class Failure(p.MessageType):
@@ -14,7 +17,7 @@ class Failure(p.MessageType):
 
     def __init__(
         self,
-        code: int = None,
+        code: EnumTypeFailureType = None,
         message: str = None,
     ) -> None:
         self.code = code
@@ -23,6 +26,6 @@ class Failure(p.MessageType):
     @classmethod
     def get_fields(cls) -> Dict:
         return {
-            1: ('code', p.UVarintType, 0),
+            1: ('code', p.EnumType("FailureType", (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 99)), 0),
             2: ('message', p.UnicodeType, 0),
         }

core/src/trezor/messages/Features.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/GetAddress.py

@@ -7,8 +7,11 @@ from .MultisigRedeemScriptType import MultisigRedeemScriptType
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
+        EnumTypeInputScriptType = Literal[0, 1, 2, 3, 4]
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
+        EnumTypeInputScriptType = None  # type: ignore
 
 
 class GetAddress(p.MessageType):
@@ -20,7 +23,7 @@ class GetAddress(p.MessageType):
         coin_name: str = None,
         show_display: bool = None,
         multisig: MultisigRedeemScriptType = None,
-        script_type: int = None,
+        script_type: EnumTypeInputScriptType = None,
     ) -> None:
         self.address_n = address_n if address_n is not None else []
         self.coin_name = coin_name
@@ -35,5 +38,5 @@ class GetAddress(p.MessageType):
             2: ('coin_name', p.UnicodeType, 0),  # default=Bitcoin
             3: ('show_display', p.BoolType, 0),
             4: ('multisig', MultisigRedeemScriptType, 0),
-            5: ('script_type', p.UVarintType, 0),  # default=SPENDADDRESS
+            5: ('script_type', p.EnumType("InputScriptType", (0, 1, 2, 3, 4)), 0),  # default=SPENDADDRESS
         }

core/src/trezor/messages/GetECDHSessionKey.py

@@ -7,6 +7,7 @@ from .IdentityType import IdentityType
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/GetEntropy.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/GetFeatures.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/GetPublicKey.py

@@ -5,8 +5,11 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
+        EnumTypeInputScriptType = Literal[0, 1, 2, 3, 4]
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
+        EnumTypeInputScriptType = None  # type: ignore
 
 
 class GetPublicKey(p.MessageType):
@@ -18,7 +21,7 @@ class GetPublicKey(p.MessageType):
         ecdsa_curve_name: str = None,
         show_display: bool = None,
         coin_name: str = None,
-        script_type: int = None,
+        script_type: EnumTypeInputScriptType = None,
     ) -> None:
         self.address_n = address_n if address_n is not None else []
         self.ecdsa_curve_name = ecdsa_curve_name
@@ -33,5 +36,5 @@ class GetPublicKey(p.MessageType):
             2: ('ecdsa_curve_name', p.UnicodeType, 0),
             3: ('show_display', p.BoolType, 0),
             4: ('coin_name', p.UnicodeType, 0),  # default=Bitcoin
-            5: ('script_type', p.UVarintType, 0),  # default=SPENDADDRESS
+            5: ('script_type', p.EnumType("InputScriptType", (0, 1, 2, 3, 4)), 0),  # default=SPENDADDRESS
         }

core/src/trezor/messages/HDNodePathType.py

@@ -7,6 +7,7 @@ from .HDNodeType import HDNodeType
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/HDNodeType.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/IdentityType.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/Initialize.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/LiskAddress.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/LiskDelegateType.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/LiskGetAddress.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/LiskGetPublicKey.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/LiskMessageSignature.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/LiskMultisignatureType.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore
 

core/src/trezor/messages/LiskPublicKey.py

@@ -5,6 +5,7 @@ import protobuf as p
 if __debug__:
     try:
         from typing import Dict, List, Optional
+        from typing_extensions import Literal  # noqa: F401
     except ImportError:
         Dict, List, Optional = None, None, None  # type: ignore