scanner_protobuf_c.py

"""
scanner_protobuf_c.py - protobuf-c 扫描器 (C 语言)
识别 ProtobufCMessageDescriptor 结构体 (magic: 0x28AAEEF9)
支持 32/64-bit ELF/PE
"""

import struct
from binary_reader import BinaryReader
from proto_generator import WELL_KNOWN_PROTOS

# protobuf-c magic
PROTOBUF_C_MESSAGE_DESCRIPTOR_MAGIC = 0x28AAEEF9
PROTOBUF_C_ENUM_DESCRIPTOR_MAGIC = 0x114315AF

PROTOBUF_C_MAGIC_BYTES_LE = struct.pack('<I', PROTOBUF_C_MESSAGE_DESCRIPTOR_MAGIC)
PROTOBUF_C_ENUM_MAGIC_BYTES_LE = struct.pack('<I', PROTOBUF_C_ENUM_DESCRIPTOR_MAGIC)


class ProtobufCType:
    INT32 = 0
    SINT32 = 1
    SFIXED32 = 2
    INT64 = 3
    SINT64 = 4
    SFIXED64 = 5
    UINT32 = 6
    FIXED32 = 7
    UINT64 = 8
    FIXED64 = 9
    FLOAT = 10
    DOUBLE = 11
    BOOL = 12
    ENUM = 13
    STRING = 14
    BYTES = 15
    MESSAGE = 16

    NAMES = {
        0: 'int32', 1: 'sint32', 2: 'sfixed32',
        3: 'int64', 4: 'sint64', 5: 'sfixed64',
        6: 'uint32', 7: 'fixed32',
        8: 'uint64', 9: 'fixed64',
        10: 'float', 11: 'double', 12: 'bool',
        13: 'enum', 14: 'string', 15: 'bytes', 16: 'message',
    }


class ProtobufCLabel:
    REQUIRED = 0
    OPTIONAL = 1
    REPEATED = 2
    NONE = 3  # proto3

    NAMES = {0: 'required', 1: 'optional', 2: 'repeated', 3: ''}


class ProtobufCFieldFlag:
    PACKED = 1 << 0
    DEPRECATED = 1 << 1
    ONEOF = 1 << 2


class ProtobufCFieldDescriptorInfo:
    """解析后的 ProtobufCFieldDescriptor"""
    def __init__(self, reader: BinaryReader, offset: int):
        self.reader = reader
        self.file_offset = offset
        e = '>' if reader.endian == 'big' else '<'
        p = reader.ptr_size

        if p == 8:
            # 64-bit layout (0x48 = 72 bytes):
            # ptr name(0), u32 id(8), u32 label(12), u32 type(16),
            # u32 quantifier_offset(20), u32 offset(24), [4-byte pad to align ptr],
            # ptr descriptor(32), ptr default_value(40),
            # u32 flags(48), u32 reserved_flags(52), ptr reserved2(56), ptr reserved3(64)
            fmt = f'{e}QIIIII4xQQIIQQ'
            size = struct.calcsize(fmt)  # 72
            data = reader.read_bytes(offset, size)
            if len(data) < size:
                data = data.ljust(size, b'\x00')
            vals = struct.unpack(fmt, data[:size])
            self.name_ptr = vals[0]
            self.id = vals[1]
            self.label = vals[2]
            self.type = vals[3]
            self.quantifier_offset = vals[4]
            self.offset = vals[5]
            self.descriptor_ptr = vals[6]
            self.default_value_ptr = vals[7]
            self.flags = vals[8]
            self.reserved_flags = vals[9]
            self.reserved2 = vals[10]
            self.reserved3 = vals[11]
            self.struct_size = size
        else:
            # 32-bit layout (48 bytes, no alignment padding):
            # ptr name, u32 id, u32 label, u32 type, u32 quantifier_offset,
            # u32 offset, ptr descriptor, ptr default_value,
            # u32 flags, u32 reserved_flags, ptr reserved2, ptr reserved3
            fmt = f'{e}IIIIIIIIIIII'
            size = struct.calcsize(fmt)  # 48
            data = reader.read_bytes(offset, size)
            if len(data) < size:
                data = data.ljust(size, b'\x00')
            vals = struct.unpack(fmt, data[:size])
            self.name_ptr = vals[0]
            self.id = vals[1]
            self.label = vals[2]
            self.type = vals[3]
            self.quantifier_offset = vals[4]
            self.offset = vals[5]
            self.descriptor_ptr = vals[6]
            self.default_value_ptr = vals[7]
            self.flags = vals[8]
            self.reserved_flags = vals[9]
            self.reserved2 = vals[10]
            self.reserved3 = vals[11]
            self.struct_size = size

        # 读取字段名
        self.name = reader.read_cstring_va(self.name_ptr) if self.name_ptr else ''

    def is_valid(self) -> bool:
        return (
            bool(self.name)
            and 0 < self.id < 0x20000000
            and 0 <= self.type <= 16
            and 0 <= self.label <= 3
        )

    @property
    def type_name(self) -> str:
        return ProtobufCType.NAMES.get(self.type, f'unknown({self.type})')

    @property
    def label_name(self) -> str:
        return ProtobufCLabel.NAMES.get(self.label, '')


class ProtobufCMessageDescriptorInfo:
    """解析后的 ProtobufCMessageDescriptor"""
    def __init__(self, reader: BinaryReader, offset: int):
        self.reader = reader
        self.file_offset = offset
        e = '>' if reader.endian == 'big' else '<'
        p = reader.ptr_size

        if p == 8:
            # 64-bit layout:
            # u32 magic, pad4, ptr name, ptr short_name, ptr c_name, ptr package_name,
            # u64 sizeof_message, u32 n_fields, pad4, ptr fields, ptr fields_sorted_by_name,
            # u32 n_field_ranges, pad4, ptr field_ranges, ptr message_init,
            # ptr reserved1, ptr reserved2, ptr reserved3
            fmt = f'{e}I4xQQQQQI4xQQI4xQQQQQ'
            size = struct.calcsize(fmt)
            data = reader.read_bytes(offset, size)
            if len(data) < size:
                data = data.ljust(size, b'\x00')
            vals = struct.unpack(fmt, data[:size])
            self.magic = vals[0]
            self.name_ptr = vals[1]
            self.short_name_ptr = vals[2]
            self.c_name_ptr = vals[3]
            self.package_name_ptr = vals[4]
            self.sizeof_message = vals[5]
            self.n_fields = vals[6]
            self.fields_ptr = vals[7]
            self.fields_sorted_ptr = vals[8]
            self.n_field_ranges = vals[9]
            self.field_ranges_ptr = vals[10]
            self.message_init_ptr = vals[11]
            self.struct_size = size
        else:
            # 32-bit layout:
            # u32 magic, ptr name, ptr short_name, ptr c_name, ptr package_name,
            # u32 sizeof_message, u32 n_fields, ptr fields, ptr fields_sorted_by_name,
            # u32 n_field_ranges, ptr field_ranges, ptr message_init,
            # ptr reserved1, ptr reserved2, ptr reserved3
            fmt = f'{e}IIIIIIIIIIIIIII'
            size = struct.calcsize(fmt)
            data = reader.read_bytes(offset, size)
            if len(data) < size:
                data = data.ljust(size, b'\x00')
            vals = struct.unpack(fmt, data[:size])
            self.magic = vals[0]
            self.name_ptr = vals[1]
            self.short_name_ptr = vals[2]
            self.c_name_ptr = vals[3]
            self.package_name_ptr = vals[4]
            self.sizeof_message = vals[5]
            self.n_fields = vals[6]
            self.fields_ptr = vals[7]
            self.fields_sorted_ptr = vals[8]
            self.n_field_ranges = vals[9]
            self.field_ranges_ptr = vals[10]
            self.message_init_ptr = vals[11]
            self.struct_size = size

        # 读取字符串
        self.name = reader.read_cstring_va(self.name_ptr) if self.name_ptr else ''
        self.short_name = reader.read_cstring_va(self.short_name_ptr) if self.short_name_ptr else ''
        self.c_name = reader.read_cstring_va(self.c_name_ptr) if self.c_name_ptr else ''
        self.package_name = reader.read_cstring_va(self.package_name_ptr) if self.package_name_ptr else ''

        # 解析字段
        self.fields = []
        if self.fields_ptr and 0 < self.n_fields < 10000:
            self._parse_fields()

    def _parse_fields(self):
        fields_offset = self.reader.va_to_offset(self.fields_ptr)
        # 计算单个 FieldDescriptor 的大小
        dummy = ProtobufCFieldDescriptorInfo(self.reader, fields_offset)
        field_size = dummy.struct_size

        for i in range(self.n_fields):
            fo = fields_offset + i * field_size
            field = ProtobufCFieldDescriptorInfo(self.reader, fo)
            if field.is_valid():
                self.fields.append(field)
            else:
                break

    def is_valid(self) -> bool:
        return (
            self.magic == PROTOBUF_C_MESSAGE_DESCRIPTOR_MAGIC
            and bool(self.name)
            and self.n_fields < 10000
        )


class ProtobufCEnumDescriptorInfo:
    """解析后的 ProtobufCEnumDescriptor"""
    def __init__(self, reader: BinaryReader, offset: int):
        self.reader = reader
        self.file_offset = offset
        e = '>' if reader.endian == 'big' else '<'
        p = reader.ptr_size

        if p == 8:
            fmt = f'{e}I4xQQQQI4xQI4xQI4xQ'
            size = struct.calcsize(fmt)
            data = reader.read_bytes(offset, size)
            if len(data) < size:
                data = data.ljust(size, b'\x00')
            vals = struct.unpack(fmt, data[:size])
            self.magic = vals[0]
            self.name_ptr = vals[1]
            self.short_name_ptr = vals[2]
            self.c_name_ptr = vals[3]
            self.package_name_ptr = vals[4]
            self.n_values = vals[5]
            self.values_ptr = vals[6]
            self.n_value_names = vals[7]
            self.values_by_name_ptr = vals[8]
            self.n_value_ranges = vals[9]
            self.value_ranges_ptr = vals[10]
        else:
            fmt = f'{e}IIIIIIIIIII'
            size = struct.calcsize(fmt)
            data = reader.read_bytes(offset, size)
            if len(data) < size:
                data = data.ljust(size, b'\x00')
            vals = struct.unpack(fmt, data[:size])
            self.magic = vals[0]
            self.name_ptr = vals[1]
            self.short_name_ptr = vals[2]
            self.c_name_ptr = vals[3]
            self.package_name_ptr = vals[4]
            self.n_values = vals[5]
            self.values_ptr = vals[6]
            self.n_value_names = vals[7]
            self.values_by_name_ptr = vals[8]
            self.n_value_ranges = vals[9]
            self.value_ranges_ptr = vals[10]

        self.name = reader.read_cstring_va(self.name_ptr) if self.name_ptr else ''
        self.short_name = reader.read_cstring_va(self.short_name_ptr) if self.short_name_ptr else ''
        self.package_name = reader.read_cstring_va(self.package_name_ptr) if self.package_name_ptr else ''

        # 解析 enum values: ProtobufCEnumValue = { const char *name; const char *c_name; int value; }
        self.values = []
        if self.values_ptr and 0 < self.n_values < 10000:
            self._parse_values()

    def _parse_values(self):
        vo = self.reader.va_to_offset(self.values_ptr)
        p = self.reader.ptr_size
        e = '>' if self.reader.endian == 'big' else '<'

        if p == 8:
            entry_size = 8 + 8 + 4 + 4  # name_ptr, c_name_ptr, value, pad
            for i in range(self.n_values):
                off = vo + i * entry_size
                name_ptr = struct.unpack_from(f'{e}Q', self.reader.data, off)[0]
                c_name_ptr = struct.unpack_from(f'{e}Q', self.reader.data, off + 8)[0]
                value = struct.unpack_from(f'{e}i', self.reader.data, off + 16)[0]
                name = self.reader.read_cstring_va(name_ptr) if name_ptr else ''
                self.values.append((name, value))
        else:
            entry_size = 4 + 4 + 4  # name_ptr, c_name_ptr, value
            for i in range(self.n_values):
                off = vo + i * entry_size
                name_ptr = struct.unpack_from(f'{e}I', self.reader.data, off)[0]
                c_name_ptr = struct.unpack_from(f'{e}I', self.reader.data, off + 4)[0]
                value = struct.unpack_from(f'{e}i', self.reader.data, off + 8)[0]
                name = self.reader.read_cstring_va(name_ptr) if name_ptr else ''
                self.values.append((name, value))


def scan_protobuf_c(reader: BinaryReader) -> list:
    """
    扫描 protobuf-c 的 ProtobufCMessageDescriptor 结构体
    返回: [ProtobufCMessageDescriptorInfo, ...]
    """
    results = []
    magic_bytes = PROTOBUF_C_MAGIC_BYTES_LE if reader.endian == 'little' else struct.pack('>I', PROTOBUF_C_MESSAGE_DESCRIPTOR_MAGIC)

    offsets = reader.search(magic_bytes)
    for off in offsets:
        try:
            desc = ProtobufCMessageDescriptorInfo(reader, off)
            if desc.is_valid():
                results.append(desc)
        except Exception:
            continue
    return results


def scan_protobuf_c_enums(reader: BinaryReader) -> list:
    """扫描 protobuf-c 的 ProtobufCEnumDescriptor"""
    results = []
    magic_bytes = PROTOBUF_C_ENUM_MAGIC_BYTES_LE if reader.endian == 'little' else struct.pack('>I', PROTOBUF_C_ENUM_DESCRIPTOR_MAGIC)

    offsets = reader.search(magic_bytes)
    for off in offsets:
        try:
            desc = ProtobufCEnumDescriptorInfo(reader, off)
            if desc.name:
                results.append(desc)
        except Exception:
            continue
    return results


def protobuf_c_to_proto(messages: list, enums: list = None) -> str:
    """从 protobuf-c 扫描结果生成 .proto 内容"""
    if not messages:
        return ''

    # 按 package 分组
    packages = {}
    for msg in messages:
        pkg = msg.package_name or ''
        packages.setdefault(pkg, []).append(msg)

    enum_map = {}
    if enums:
        for e in enums:
            enum_map[e.name] = e

    output_parts = []
    for pkg, msgs in packages.items():
        lines = ['syntax = "proto3";', '']
        if pkg:
            lines.append(f'package {pkg};')
            lines.append('')

        # enum definitions referenced by fields
        emitted_enums = set()
        for msg in msgs:
            for field in msg.fields:
                if field.type == ProtobufCType.ENUM and field.descriptor_ptr:
                    # 尝试查找该 enum
                    for en in (enums or []):
                        if en.name not in emitted_enums:
                            emitted_enums.add(en.name)
                            lines.append(f'enum {en.short_name} {{')
                            for vname, vnum in en.values:
                                lines.append(f'  {vname} = {vnum};')
                            lines.append('}')
                            lines.append('')

        for msg in msgs:
            lines.append(f'message {msg.short_name} {{')
            for field in msg.fields:
                label = ProtobufCLabel.NAMES.get(field.label, '')
                type_name = ProtobufCType.NAMES.get(field.type, f'unknown_{field.type}')

                # 对于 MESSAGE 和 ENUM 类型，尝试解析引用的描述符名称
                if field.type == ProtobufCType.MESSAGE and field.descriptor_ptr:
                    ref_name = _resolve_message_ref(reader=msg.reader, ptr=field.descriptor_ptr, messages=messages)
                    if ref_name:
                        type_name = ref_name
                elif field.type == ProtobufCType.ENUM and field.descriptor_ptr:
                    ref_name = _resolve_enum_ref(reader=msg.reader, ptr=field.descriptor_ptr, enums=enums or [])
                    if ref_name:
                        type_name = ref_name

                if label:
                    lines.append(f'  {label} {type_name} {field.name} = {field.id};')
                else:
                    lines.append(f'  {type_name} {field.name} = {field.id};')
            lines.append('}')
            lines.append('')

        output_parts.append('\n'.join(lines))

    return '\n'.join(output_parts)


def _resolve_message_ref(reader: BinaryReader, ptr: int, messages: list) -> str | None:
    """尝试通过 descriptor 指针解析引用的 message 名称"""
    for msg in messages:
        msg_va = reader.offset_to_va(msg.file_offset)
        if msg_va == ptr:
            return msg.short_name
    # fallback: 尝试读取被引用 descriptor 的 short_name
    try:
        off = reader.va_to_offset(ptr)
        magic = reader.read_u32(off)
        if magic == PROTOBUF_C_MESSAGE_DESCRIPTOR_MAGIC:
            desc = ProtobufCMessageDescriptorInfo(reader, off)
            if desc.is_valid():
                return desc.short_name
    except Exception:
        pass
    return None


def _resolve_enum_ref(reader: BinaryReader, ptr: int, enums: list) -> str | None:
    """尝试通过 descriptor 指针解析引用的 enum 名称"""
    for en in enums:
        en_va = reader.offset_to_va(en.file_offset)
        if en_va == ptr:
            return en.short_name
    try:
        off = reader.va_to_offset(ptr)
        magic = reader.read_u32(off)
        if magic == PROTOBUF_C_ENUM_DESCRIPTOR_MAGIC:
            desc = ProtobufCEnumDescriptorInfo(reader, off)
            if desc.name:
                return desc.short_name
    except Exception:
        pass
    return None