cwb_markers.py v0.1.0

cwb_markers.py

import idautils
import idc
import ida_bytes
from pprint import pprint, pformat


def demangle_name(name):
    return idc.Demangle(name, idc.GetLongPrm(idc.INF_SHORT_DN)) or name

def get_primary_vtables():
    vtables = {}
    for ea, name in idautils.Names():
        dname = idc.Demangle(name, idc.GetLongPrm(idc.INF_SHORT_DN))
        if dname != None and 'table' in dname:
            if not name.endswith('_0'): # Skip secondary vtables with the _0 suffix.
                vtables[dname] = ea
    return vtables

def get_free_size_arg_from_vtable_ea(vtable_ea, typename):
    # Read the dtor from vtable[0].
    dtor_ea = ida_bytes.get_qword(vtable_ea)
    cur_ea = dtor_ea

    # Search forwards for a call to free.
    found_free_call = False
    func_end_ea = idc.GetFunctionAttr(dtor_ea, idc.FUNCATTR_END)
    while cur_ea < func_end_ea:
        cur_ea = idc.NextHead(cur_ea)
        print(cur_ea)
        if idc.GetMnem(cur_ea) == "call":
            dname = demangle_name(idc.GetOpnd(cur_ea, 0))
            if 'free' in dname:
                found_free_call = True
                break
    
    if not found_free_call:
        return -1

    # Now go back and find the size argument passed in (r||e)dx
    while cur_ea > dtor_ea:
        cur_ea = idc.PrevHead(cur_ea)
        if idc.GetMnem(cur_ea) == "mov" and idc.GetOpnd(cur_ea, 0).endswith('dx'):
            if idc.GetOpType(cur_ea, 1) == idc.o_imm:
                return idc.GetOperandValue(cur_ea, 1)
            else:
                print("[CWT_MARKERS!] ERROR: got non-immediate operand for `free` argument. Can't parse this.")
                print("[CWT_MARKERS!] \tAt:{:X}, for dtor: {:X}, type:{}".format(cur_ea, dtor_ea, typename))
                return -1

    return -1
    
def main():
    type_to_info = {}

    # Parse out the type name from the vftable symbol.
    for key, value in get_primary_vtables().iteritems():
        typename = key
        if typename.endswith("::`vftable'"):
            typename = typename[:-11]
        if typename.startswith("const "):
            typename = typename[6:]
        
        type_to_info[typename] = {'vtable_ea' :value}

    # Get the type ctor function by finding the first xref.
    for typename, fields in type_to_info.iteritems():
        for xref in idautils.XrefsTo(fields['vtable_ea'], 0):
            type_to_info[typename]['ctor_ea'] = idc.GetFunctionAttr(xref.frm, idc.FUNCATTR_START)
            break

    # Try to get class size via `free` in vtable[0] dtor.
    for typename, fields in type_to_info.iteritems():
        type_to_info[typename]['size'] = get_free_size_arg_from_vtable_ea(fields['vtable_ea'], typename)

    # Print the classes and their sizes
    for typename in sorted(type_to_info):
        fields = type_to_info[typename]
        if typename.startswith("cube") or typename.startswith("plasma") or typename.startswith("gfx"):
            name_pad = ' ' * (70 - len(typename))
            print("Class: {}, {} Size: 0x{:X}".format(typename, name_pad, fields['size']))

    #pprint(type_to_info)

main()