trezor-firmware/tools/pb2py

#!/usr/bin/env python3
# Converts Google's protobuf python definitions of TREZOR wire messages
# to plain-python objects as used in TREZOR Core and python-trezor

import argparse
import importlib
import logging
import os
import sys
from collections import namedtuple

ProtoField = namedtuple('ProtoField', 'name, number, proto_type, py_type, repeated, required, orig')


def parse_field(number, field):
    FIELD_TYPES = {
        field.TYPE_UINT64:  ('p.UVarintType', 'int'),
        field.TYPE_UINT32:  ('p.UVarintType', 'int'),
        field.TYPE_ENUM:    ('p.UVarintType', 'int'),
        field.TYPE_SINT32:  ('p.SVarintType', 'int'),
        field.TYPE_SINT64:  ('p.SVarintType', 'int'),
        field.TYPE_STRING:  ('p.UnicodeType', 'str'),
        field.TYPE_BOOL:    ('p.BoolType', 'bool'),
        field.TYPE_BYTES:   ('p.BytesType', 'bytes'),
    }
    repeated = (field.label == field.LABEL_REPEATED)
    required = (field.label == field.LABEL_REQUIRED)
    if field.type == field.TYPE_MESSAGE:
        proto_type = py_type = field.message_type.name
    else:
        try:
            proto_type, py_type = FIELD_TYPES[field.type]
        except KeyError:
            raise ValueError("Unknown field type %d for field %s" % (field.type, field.name)) from None

    if repeated:
        py_type = "List[%s]" % py_type

    return ProtoField(
        name=field.name,
        number=number,
        proto_type=proto_type,
        py_type=py_type,
        repeated=repeated,
        required=required,
        orig=field,
    )


def import_pb2(name):
    return importlib.import_module("%s_pb2" % name)


def create_message_import(name):
    return "from .%s import %s" % (name, name)


def remove_from_start(s, prefix):
    if s.startswith(prefix):
        return s[len(prefix):]
    else:
        return s


def process_message_imports(descriptor):
    imports = set()

    for field in descriptor.fields:
        if field.type == field.TYPE_MESSAGE:
            imports.add(field.message_type.name)

    for name in sorted(imports):
        yield create_message_import(name)


def create_init_method(fields):
    yield         "    def __init__("
    yield         "        self,"
    for field in fields[:-1]:
        yield     "        %s: %s = None," % (field.name, field.py_type)
    # last field must not have a traling comma
    yield         "        %s: %s = None" % (fields[-1].name, fields[-1].py_type)
    yield         "    ) -> None:"

    for field in fields:
        if field.repeated:
            yield "        self.{0} = {0} if {0} is not None else []".format(field.name)
        else:
            yield "        self.{0} = {0}".format(field.name)


def process_message(descriptor, protobuf_module, msg_id, is_upy):
    logging.debug("Processing message %s", descriptor.name)

    if is_upy:
        yield "import protobuf as p"
    else:
        yield "from .. import protobuf as p"

    fields = list(parse_field(number, field)
                  for number, field
                  in descriptor.fields_by_number.items())

    if any(field.repeated for field in fields):
        yield "if __debug__:"
        yield "    try:"
        yield "        from typing import List"
        yield "    except ImportError:"
        yield "        List = None"

    yield from process_message_imports(descriptor)

    yield ""
    yield ""
    yield "class %s(p.MessageType):" % descriptor.name

    if msg_id is not None:
        yield "    MESSAGE_WIRE_TYPE = %d" % msg_id

    if fields:
        yield "    FIELDS = {"
        for field in fields:
            comments = []
            if field.required:
                comments.append('required')
            if field.orig.has_default_value:
                comments.append("default=%s" % repr(field.orig.default_value))

            if comments:
                comment = "  # %s" % ' '.join(comments)
            else:
                comment = ''

            if field.repeated:
                flags = 'p.FLAG_REPEATED'
            else:
                flags = '0'

            yield "        %d: ('%s', %s, %s),%s" % (field.number, field.name, field.proto_type, flags, comment)

        yield "    }"
        yield ""
        yield from create_init_method(fields)

    if not fields and not msg_id:
        yield "    pass"


def process_enum(descriptor, is_upy):
    logging.debug("Processing enum %s", descriptor.name)

    for name, value in descriptor.values_by_name.items():
        # Remove type name from the beginning of the constant
        # For example "PinMatrixRequestType_Current" -> "Current"
        enum_prefix = descriptor.name
        name = remove_from_start(name, "%s_" % enum_prefix)

        # If type ends with *Type, but constant use type name without *Type, remove it too :)
        # For example "ButtonRequestType & ButtonRequest_Other" => "Other"
        if enum_prefix.endswith("Type"):
            enum_prefix, _ = enum_prefix.rsplit("Type", 1)
            name = remove_from_start(name, "%s_" % enum_prefix)

        yield "%s = %s" % (name, value.number)


def process_file(descriptor, protobuf_module, genpath, modlist, is_upy):
    logging.info("Processing module %s", descriptor.name)

    msg_types = import_pb2('messages').MessageType

    for name, message_descriptor in sorted(descriptor.message_types_by_name.items()):
        # Find message type for given class
        try:
            msg_id = msg_types.Value("MessageType_%s" % name)
        except ValueError:
            msg_id = None

        out = process_message(message_descriptor, protobuf_module, msg_id, is_upy)
        write_to_file(genpath, name, out)
        if modlist:
            modlist.write(create_message_import(name) + "\n")

    for name, enum_descriptor in descriptor.enum_types_by_name.items():
        out = process_enum(enum_descriptor, is_upy)
        write_to_file(genpath, name, out)
        if modlist:
            modlist.write("from . import %s\n" % name)


def write_to_file(genpath, t, out):
    # Write generated sourcecode to given file
    with open(os.path.join(genpath, "%s.py" % t), 'w') as f:
        f.write("# Automatically generated by pb2py\n")
        for line in out:
            f.write(line + "\n")


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)

    parser = argparse.ArgumentParser()
    parser.add_argument('module', help="Name of module to generate")
    parser.add_argument('genpath', help="Directory for generated source code")
    parser.add_argument('-P', '--protobuf-module', default="protobuf", help="Name of protobuf module")
    parser.add_argument('-l', '--modlist', type=argparse.FileType('a'), help="Generate list of modules")
    parser.add_argument('-p', '--protopath', type=str, help="Path to search for pregenerated Google's python sources")
    parser.add_argument('-m', '--micropython', action='store_true', help="Use micropython-favoured source code")
    args = parser.parse_args()

    if args.protopath:
        sys.path.append(args.protopath)

    # This must be done after sys.path.append
    module = import_pb2(args.module)

    process_file(module.DESCRIPTOR, args.protobuf_module, args.genpath, args.modlist, args.micropython)