stuff to take advantage of classinformer, label all functions and find likely constructors for all RTTI types

random IDAPython functions.py

import idaapi
import idautils
import ida_funcs
import ida_allins
import ida_ua

import ctypes as C


NETNODE_NAME = '$ClassInformer_node'
NN_DATA_TAG = 'A'
NN_TABLE_TAG = 'S'


NIDX_VERSION = 0
NIDX_COUNT = 1

CI_FLAG_BCD_NOTVISIBLE          = 0x01
CI_FLAG_BCD_AMBIGUOUS           = 0x02
CI_FLAG_BCD_PRIVORPROTINCOMPOBJ = 0x04
CI_FLAG_BCD_PRIVORPROTBASE      = 0x08
CI_FLAG_BCD_VBOFCONTOBJ         = 0x10
CI_FLAG_BCD_NONPOLYMORPHIC      = 0x20
CI_FLAG_BCD_HASPCHD             = 0x40

CI_FLAG_CHD_MULTINH   = 0x01 # Multiple inheritance
CI_FLAG_CHD_VIRTINH   = 0x02 # Virtual inheritance
CI_FLAG_CHD_AMBIGUOUS = 0x04 # Ambiguous inheritance

if idaapi.get_inf_structure().is_64bit():
    ea_t = C.c_uint64
    ptr_sz = 8
    read_ptr = idaapi.get_qword
else:
    ea_t = C.c_uint32
    ptr_sz = 4
    read_ptr = idaapi.get_dword


class TBLENTRY(C.Structure):
    _fields_ = [
        ('vft', ea_t),
        ('methods', C.c_uint16),
        ('flags', C.c_uint16),
        ('strSize', C.c_uint16),
        ('str', C.c_char * (idaapi.MAXSPECSIZE - (C.sizeof(C.c_uint32) + C.sizeof(C.c_uint16) * 3))),
    ]

    def __repr__(self):
        return f'TBLENTRY(\n\tvft={self.vft:X}\n\tmethods={self.methods}\n\tflags={self.flags:X}\n\tstrSize={self.strSize}\n\tstr={self.str}\n\t)'

    @property
    def class_name(self):
        return self.str.partition(b'@')[0].decode('utf-8')


    @property
    def class_name_sanitized(self):
        cname = self.str.partition(b'@')[0].decode('utf-8')
        cname = cname.replace('<','[').replace('>',']').replace('*','').replace(',','&').replace(' ','')
        return cname
    


ci_netnode = idaapi.netnode(NETNODE_NAME, 0, False)

def ci_count():
    return ci_netnode.altval_idx8(NIDX_COUNT, NN_DATA_TAG)

def ci_get_table_idx(idx):
    data = ci_netnode.supval(idx, NN_TABLE_TAG)
    if data is None:
        return None
    pad_sz = 0x400 - len(data)
    if pad_sz > 0:
        data += b'\x00' * pad_sz
    return TBLENTRY.from_buffer_copy(data)

def ci_all_entries():
    for i in range(ci_count()):
        yield ci_get_table_idx(i)

def ci_vtable_fns(entry):
    cname = entry.class_name
    for i in range(entry.methods):
        vtf_ea_ea = entry.vft + i * ptr_sz
        vtf_ea = read_ptr(vtf_ea_ea)
        vtf_name = idaapi.get_name(vtf_ea)
        yield vtf_ea, vtf_name


def prefix_vt_fns():
    for e in ci_all_entries():
        cname = e.class_name_sanitized
        for ea, name in ci_vtable_fns(e):
            if cname not in name:
                new_name = cname + '::' + name
                print(f'rewrite ea:{ea:X} from {name} to {new_name}')
                idaapi.set_name(ea, new_name)




def trace_thisptr(ea):
    fn = ida_funcs.get_func(ea)
    if not fn:
        return []
    thisptr_regs = [idautils.procregs.ecx.reg]
    for isn_ea in Heads(fn.start_ea, ea):
        insn = idaapi.insn_t()
        length = idaapi.decode_insn(insn, isn_ea)
        if insn.itype != ida_allins.NN_mov:
            continue
        
        if insn.Op2.reg in thisptr_regs:
            thisptr_regs.append(insn.Op1.reg)
        elif insn.Op1.reg in thisptr_regs and insn.Op2.reg not in thisptr_regs:
            thisptr_regs.remove(insn.Op1.reg)
    return thisptr_regs


def find_entry_ctors(entry):
    for ref in idautils.XrefsTo(entry.vft):
        insn = idaapi.insn_t()
        length = idaapi.decode_insn(insn, ref.frm)

        if insn.itype != ida_allins.NN_mov:
            continue

        this_regs = trace_thisptr(ref.frm)

        if insn.Op1.reg in this_regs:
            fn = ida_funcs.get_func(ref.frm)
            yield fn.start_ea, idaapi.get_name(fn.start_ea)
             

def write_known_ctors():
    for e in ci_all_entries():
        cname = e.class_name_sanitized
        for ea, name in find_entry_ctors(e):
            if cname not in name:
                resolution_str = ''
                new_name = cname + '::' + 'ctor' + resolution_str
                print(f'rewrite ea:{ea:X} from {name} to {new_name}')
                while idaapi.set_name(ea, new_name) == 0:
                    if resolution_str == '':
                        resolution_str = 1
                    else:
                        resolution_str += 1
                    print(f'rewrite ea:{ea:X} from {name} to {new_name}')
                    new_name = cname + '::' + 'ctor' + str(resolution_str)