IDA script to rename vftable automatically

ida_name_vftable.py

from idaapi import *
from idautils import *
from idc import *
from ida_typeinf import *
import re

IS64 = get_inf_structure().is_64bit()
print('64bit mode' if IS64 else '32bit mode')

sub_regex = re.compile(r'sub_[0-9A-F]+')

def get_iter():
    return 8 if IS64 else 4

def dereference(current_ea):
    return get_qword(current_ea) if IS64 else get_dword(current_ea)

def get_default_prototype(local_type):
    return '__int64 __fastcall PROTO({} *);'.format(local_type) if IS64 else 'unsigned int __thiscall PROTO({} *);'.format(local_type)

def get_vtbl_prototype(local_type):
    return '__int64 (*__fastcall PROTO)({} *);'.format(local_type) if IS64 else 'unsigned int (*__thiscall PROTO)({} *);'.format(local_type)

# generate the object type associated to the vftable
# if the type already exists it is untouched
def generate_local_type(class_name, nfuncs):
    name = '{}_t'.format(class_name)
    vtbl_struct_name = None

    # try to create struct
    struct_id = idaapi.get_struc_id(name)
    if struct_id != idaapi.BADADDR:
        print('structure {} already exists'.format(name))
        return name, vtbl_struct_name
    
    obj_struct_id = add_struc(0, name, 0)

    print('Added type {} to database'.format(name))

    # create vtable definition
    proto = get_vtbl_prototype(name)
    
    vtbl_struct_name = '{}_vtbl'.format(class_name.upper())
    vtable_struct_id = add_struc(0, vtbl_struct_name, 0)
    
    offset = get_iter()
    idatype = idaapi.FF_QWORD if IS64 else idaapi.FF_DWORD
    for i in range(nfuncs):
        add_struc_member(vtable_struct_id, 'Vfunc{}'.format(i), i * offset, idatype, -1, offset)
        idc.SetType(idc.get_member_id(vtable_struct_id, i * offset), '{}'.format(proto))

    add_struc_member(obj_struct_id, '__vftable', 0, idatype, -1, offset)
    idc.SetType(idc.get_member_id(obj_struct_id, 0), '{} *'.format(vtbl_struct_name))

    print('Added type {} to database'.format(vtbl_struct_name))

    return name, vtbl_struct_name

def process_vfunc(class_name, local_type, current_ea, index):
    vfunc_ea = dereference(current_ea)
    vfunc_name = get_name(vfunc_ea)

    # only rename auto generated subroutines (prevent to rename user defined subroutines names)
    # nullsub are ommited on purpose
    if sub_regex.match(vfunc_name) is not None:
        new_vfunc_name = '{}::Vfunc{}'.format(class_name, str(index))
        set_name(vfunc_ea, new_vfunc_name)
        print('> renamed func at {} to {} (old name: {})'.format(hex(vfunc_ea), new_vfunc_name, vfunc_name))

        # if idaapi.decompile(vfunc_ea) is None:
        #     print('Failed to decompile {}'.format(hex(vfunc_ea)))

        # guess function type
        type_info = tinfo_t()
        if not idaapi.get_tinfo(type_info, vfunc_ea):
            if not idaapi.guess_tinfo(type_info, vfunc_ea):
                print('Failed to guess info for {}'.format(hex(vfunc_ea)))

        # change function first parameter type
        func_data = func_type_data_t()
        type_info.get_func_details(func_data)

        if func_data.size() >= 1:
            thisobj_type_info = tinfo_t()
            thisobj_type_info.get_named_type(get_idati(), local_type)

            thisptr_type_info = tinfo_t()
            thisptr_type_info.create_ptr(thisobj_type_info)

            func_data[0].type = thisptr_type_info

            function_type_info = tinfo_t()
            function_type_info.create_func(func_data)
            apply_tinfo(vfunc_ea, function_type_info, TINFO_DEFINITE)

            print('> Updated prototype for {} ({})'.format(new_vfunc_name, hex(vfunc_ea)))
        else:
            default_prototype = get_default_prototype(local_type)
            idc.SetType(vfunc_ea, '') # remove any user defined prototype
            idc.SetType(vfunc_ea, default_prototype)

            print('> Applied default prototype for {} ({})'.format(new_vfunc_name, hex(vfunc_ea)))



def process(class_name, local_type, addresses):
    for ea in addresses:
        process_vfunc(class_name, local_type, ea, addresses.index(ea))

def run():
    start_address = ask_long(0, 'vftable start address')
    if start_address is None or start_address == BADADDR or start_address <= 0:
        print('invalid start address')
        return

    end_address = ask_long(0, 'vftable end address')
    if end_address is None or end_address == BADADDR or end_address <= 0:
        print('invalid end address')
        return

    if start_address >= end_address:
        print('invalid start and end address')
        return
        
    class_name = ask_str('', 0, 'class name')
    if class_name is None or class_name == '':
        print('invalid class name')
        return

    offset = get_iter()
    addresses = list(range(start_address, end_address + offset, offset))

    # generate vtbl local type
    local_type, local_vtbl_type = generate_local_type(class_name, len(addresses))

    # process
    process(class_name, local_type, addresses)

    # rename vftable offset
    set_name(start_address, 's_' + class_name + '_vtbl')
    if local_vtbl_type is not None:
        idc.SetType(start_address, '{} *'.format(local_vtbl_type))

if __name__ == '__main__':
    run()