protobufc-ida-analyze.py

import os
import tempfile
from binascii import b2a_hex
from functools import partial

import ida_struct
import idaapi
from ida_bytes import get_dword, get_word, get_byte, get_flags, is_code
from ida_nalt import get_imagebase
from idc import get_strlit_contents, get_qword, parse_decl, apply_type, TINFO_GUESSED, find_imm, SEARCH_DOWN, BADADDR, \
    PT_FILE, find_binary, get_type
from idaapi import auto_wait
from struct import pack


class ProtobufMagic(object):
    PROTOBUF_C__SERVICE_DESCRIPTOR_MAGIC = 0x14159bc3
    PROTOBUF_C__MESSAGE_DESCRIPTOR_MAGIC = 0x28aaeef9
    PROTOBUF_C__ENUM_DESCRIPTOR_MAGIC = 0x114315af


class ProtobufCLable(object):
    PROTOBUF_C_LABEL_REQUIRED = 0x0
    PROTOBUF_C_LABEL_OPTIONAL = 0x1
    PROTOBUF_C_LABEL_REPEATED = 0x2
    PROTOBUF_C_LABEL_NONE = 0x3


class ProtobufCType(object):
    PROTOBUF_C_TYPE_INT32 = 0x0
    PROTOBUF_C_TYPE_SINT32 = 0x1
    PROTOBUF_C_TYPE_SFIXED32 = 0x2
    PROTOBUF_C_TYPE_INT64 = 0x3
    PROTOBUF_C_TYPE_SINT64 = 0x4
    PROTOBUF_C_TYPE_SFIXED64 = 0x5
    PROTOBUF_C_TYPE_UINT32 = 0x6
    PROTOBUF_C_TYPE_FIXED32 = 0x7
    PROTOBUF_C_TYPE_UINT64 = 0x8
    PROTOBUF_C_TYPE_FIXED64 = 0x9
    PROTOBUF_C_TYPE_FLOAT = 0xA
    PROTOBUF_C_TYPE_DOUBLE = 0xB
    PROTOBUF_C_TYPE_BOOL = 0xC
    PROTOBUF_C_TYPE_ENUM = 0xD
    PROTOBUF_C_TYPE_STRING = 0xE
    PROTOBUF_C_TYPE_BYTES = 0xF
    PROTOBUF_C_TYPE_MESSAGE = 0x10


def debug(*args):
    print(args)


def getStructSize(type_name):
    struct_id = ida_struct.get_struc_id(type_name)
    size = ida_struct.get_struc_size(struct_id)
    assert size != 0, 'isnt struct'
    return size


def setType(addr, type_str):
    info = parse_decl(type_str, 0)
    return apply_type(addr, info, TINFO_GUESSED)


def getFieldIntVal(addr, type_name, field_name):
    off, size = getFieldOffsetAndSize(type_name, field_name)
    if size == 4:
        return get_dword(addr + off)
    elif size == 8:
        return get_qword(addr + off)
    elif size == 2:
        return get_word(addr + off)
    elif size == 1:
        return get_byte(addr + off)
    else:
        raise NotImplementedError('the field size cant resolve...')


def getFieldOffsetAndSize(type_name, field_name):
    struct_id = ida_struct.get_struc_id(type_name)
    struct = ida_struct.get_struc(struct_id)
    member = ida_struct.get_member_by_name(struct, field_name)
    return member.soff, member.eoff - member.soff


def setStructArr(addr, struct_name, arr_len):
    """set addr is struct arr"""
    # setType(addr, 'ProtobufCEnumValue dummy[%d]' % size)
    size = getStructSize(struct_name)
    for i in range(arr_len):
        setType(addr + i * size, struct_name)


def parseFieldDescriptor(addr):
    STRUCT_NAME = 'ProtobufCFieldDescriptor'
    if get_type(addr) == STRUCT_NAME:
        return True
    get_int_val = partial(getFieldIntVal, addr, STRUCT_NAME)
    setType(addr, STRUCT_NAME)
    type_ = get_int_val('type')
    descriptor_addr = get_int_val('descriptor')
    if type_ == ProtobufCType.PROTOBUF_C_TYPE_MESSAGE:
        parseMessageDescriptor(descriptor_addr)
    elif type_ == ProtobufCType.PROTOBUF_C_TYPE_ENUM:
        parseEnumDescriptor(addr)


def parseEnumDescriptor(addr):
    STRUCT_NAME = 'ProtobufCEnumDescriptor'
    if get_type(addr) == STRUCT_NAME:
        return True
    get_int_val = partial(getFieldIntVal, addr, STRUCT_NAME)
    assert get_int_val('magic') == ProtobufMagic.PROTOBUF_C__ENUM_DESCRIPTOR_MAGIC
    setType(addr, STRUCT_NAME)
    name = get_strlit_contents(get_int_val('name'))
    package_name = get_strlit_contents(get_int_val('package_name'))
    debug('the enum name:"{}" package_name:"{}"'.format(name, package_name))
    # set fields...
    n_values = get_int_val('n_values')
    values_addr = get_int_val('values')
    setStructArr(values_addr, 'ProtobufCEnumValue', n_values)
    n_value_names = get_int_val('n_value_names')
    values_by_name_addr = get_int_val('values_by_name')
    setStructArr(values_by_name_addr, 'ProtobufCEnumValueIndex', n_value_names)
    n_value_ranges = get_int_val('n_value_ranges')
    value_ranges_addr = get_int_val('value_ranges')
    setStructArr(value_ranges_addr, 'ProtobufCIntRange', n_value_ranges)


def parseMessageDescriptor(addr):
    STRUCT_NAME = 'ProtobufCMessageDescriptor'
    if get_type(addr) == STRUCT_NAME:
        return True
    get_int_val = partial(getFieldIntVal, addr, STRUCT_NAME)
    assert get_int_val('magic') == ProtobufMagic.PROTOBUF_C__MESSAGE_DESCRIPTOR_MAGIC
    if not setType(addr, STRUCT_NAME):
        debug('set addr "{}" type "{}" failed'.format(hex(addr), STRUCT_NAME))
        return False
    n_fields = get_int_val('n_fields')
    fields_addr = get_int_val('fields')
    FIELD_DESCRIPTOR_SIZE = getStructSize('ProtobufCFieldDescriptor')
    # setStructArr(fields_addr, 'ProtobufCFieldDescriptor', n_fields)
    for i in range(n_fields):
        addr = FIELD_DESCRIPTOR_SIZE * i + fields_addr
        parseFieldDescriptor(addr)


def isMessage(addr):
    flag = get_flags(addr)
    if is_code(flag):
        return False
    try:
        # name is readable
        get_strlit_contents(addr + 0x08)
        # reserved is 0
        assert get_qword(addr + 0x60) == get_qword(addr + 0x68) == get_qword(addr + 0x70) == 0
    except:
        return False
    else:
        return True


def autoParseAllMessage():
    start_addr = get_imagebase()
    magic = ProtobufMagic.PROTOBUF_C__MESSAGE_DESCRIPTOR_MAGIC
    magic = b2a_hex(pack('>I', magic)).decode()
    while True:
        addr = find_binary(start_addr, SEARCH_DOWN, magic)
        start_addr = addr + 1
        if addr == BADADDR:
            debug('scan finish...')
            break
        debug('find addr {}'.format(hex(addr)))
        if not isMessage(addr):
            continue
        try:
            parseMessageDescriptor(addr)
        except Exception as e:
            debug('addr {} err: {}'.format(hex(addr), e))
        auto_wait()


def importProtobufCHeaders():
    """load protobuf-c structs"""
    if ida_struct.get_struc_id('ProtobufCMessageDescriptor') != BADADDR:
        return True
    headers = '''
#define __int8 char
#define __int16 short
#define __int32 int
#define __int64 long long

/* 82 */
enum $64AFC66EE04482A03C9A353016AE4D0A
{
  PROTOBUF_C_LABEL_REQUIRED = 0x0,
  PROTOBUF_C_LABEL_OPTIONAL = 0x1,
  PROTOBUF_C_LABEL_REPEATED = 0x2,
  PROTOBUF_C_LABEL_NONE = 0x3,
};

/* 83 */
typedef enum $64AFC66EE04482A03C9A353016AE4D0A ProtobufCLabel;

/* 84 */
enum $7463941E82CC51592461BDB53E7E6651
{
  PROTOBUF_C_TYPE_INT32 = 0x0,
  PROTOBUF_C_TYPE_SINT32 = 0x1,
  PROTOBUF_C_TYPE_SFIXED32 = 0x2,
  PROTOBUF_C_TYPE_INT64 = 0x3,
  PROTOBUF_C_TYPE_SINT64 = 0x4,
  PROTOBUF_C_TYPE_SFIXED64 = 0x5,
  PROTOBUF_C_TYPE_UINT32 = 0x6,
  PROTOBUF_C_TYPE_FIXED32 = 0x7,
  PROTOBUF_C_TYPE_UINT64 = 0x8,
  PROTOBUF_C_TYPE_FIXED64 = 0x9,
  PROTOBUF_C_TYPE_FLOAT = 0xA,
  PROTOBUF_C_TYPE_DOUBLE = 0xB,
  PROTOBUF_C_TYPE_BOOL = 0xC,
  PROTOBUF_C_TYPE_ENUM = 0xD,
  PROTOBUF_C_TYPE_STRING = 0xE,
  PROTOBUF_C_TYPE_BYTES = 0xF,
  PROTOBUF_C_TYPE_MESSAGE = 0x10,
};

/* 85 */
typedef enum $7463941E82CC51592461BDB53E7E6651 ProtobufCType;

/* 5 */
typedef unsigned int uint32_t;

/* 86 */
struct ProtobufCFieldDescriptor
{
  const char *name;
  uint32_t id;
  ProtobufCLabel label;
  ProtobufCType type;
  unsigned int quantifier_offset;
  unsigned int offset;
  const void *descriptor;
  const void *default_value;
  uint32_t flags;
  unsigned int reserved_flags;
  void *reserved2;
  void *reserved3;
};

/* 87 */
struct ProtobufCIntRange
{
  int start_value;
  unsigned int orig_index;
};

/* 38 */
typedef unsigned __int64 size_t;

/* 88 */
struct ProtobufCMessageDescriptor
{
  uint32_t magic;
  const char *name;
  const char *short_name;
  const char *c_name;
  const char *package_name;
  size_t sizeof_message;
  unsigned int n_fields;
  const ProtobufCFieldDescriptor *fields;
  const unsigned int *fields_sorted_by_name;
  unsigned int n_field_ranges;
  const ProtobufCIntRange *field_ranges;
  void *message_init;
  void *reserved1;
  void *reserved2;
  void *reserved3;
};

/* 89 */
struct ProtobufCEnumValue
{
  const char *name;
  const char *c_name;
  int value;
};

/* 90 */
struct ProtobufCEnumValueIndex
{
  const char *name;
  unsigned int index;
};

/* 91 */
struct ProtobufCEnumDescriptor
{
  uint32_t magic;
  const char *name;
  const char *short_name;
  const char *c_name;
  const char *package_name;
  unsigned int n_values;
  const ProtobufCEnumValue *values;
  unsigned int n_value_names;
  const ProtobufCEnumValueIndex *values_by_name;
  unsigned int n_value_ranges;
  const ProtobufCIntRange *value_ranges;
  void *reserved1;
  void *reserved2;
  void *reserved3;
  void *reserved4;
};

/* 92 */
enum $7482D6500CDB2FFD70D784F4AD59D116
{
  PROTOBUF_C_WIRE_TYPE_VARINT = 0x0,
  PROTOBUF_C_WIRE_TYPE_64BIT = 0x1,
  PROTOBUF_C_WIRE_TYPE_LENGTH_PREFIXED = 0x2,
  PROTOBUF_C_WIRE_TYPE_32BIT = 0x5,
};

/* 93 */
typedef enum $7482D6500CDB2FFD70D784F4AD59D116 ProtobufCWireType;

/* 3 */
typedef unsigned __int8 uint8_t;

/* 94 */
struct ProtobufCMessageUnknownField
{
  uint32_t tag;
  ProtobufCWireType wire_type;
  size_t len;
  uint8_t *data;
};

/* 95 */
struct ProtobufCMessage
{
  const ProtobufCMessageDescriptor *descriptor;
  unsigned int n_unknown_fields;
  ProtobufCMessageUnknownField *unknown_fields;
};
'''
    return idaapi.idc_parse_types(headers, 0) == 0


def main():
    # import headers
    if not importProtobufCHeaders():
        debug('loads protobufc structs failed...')
    else:
        auto_wait()
        autoParseAllMessage()


if __name__ == '__main__':
    main()