cwb_markers.py v0.2.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_new_size_arg(ea, typename=""):
    cur_ea = ea

    # Search backwards for a call to new
    found_new_call = False
    for i in range(50):
        cur_ea = idc.PrevHead(cur_ea)
        if idc.GetMnem(cur_ea) == "call":
            dname = demangle_name(idc.GetOpnd(cur_ea, 0))
            if dname == "operator new(unsigned __int64)":
                found_new_call = True
                break
    
    if not found_new_call:
        return -1

    # Found the new call, now search for the (r||e)cx register argument.
    for i in range(20):
        cur_ea = idc.PrevHead(cur_ea)
        if idc.GetMnem(cur_ea) == "mov" and idc.GetOpnd(cur_ea, 0).endswith('cx'):
            if idc.GetOpType(cur_ea, 1) == idc.o_imm:
                return idc.GetOperandValue(cur_ea, 1)
            else:
                print("[CWT!] ERROR: got non-immediate operand for `new` argument. Can't parse this.")
                print("[CWT!] \tAt:{:X}, for starting ea: {:X}, type:{}".format(cur_ea, ea, typename))
                return -1
    

    return -1


def get_new_size_arg_from_first_xref_ea(ea, typename=""):
    cur_ea = -1

    # Get ea of first xref.
    for xref in idautils.XrefsTo(ea, 0):
        cur_ea = xref.frm
        break
    
    if cur_ea == -1:
        return -1

    return get_new_size_arg(cur_ea, typename)

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)
        if idc.GetMnem(cur_ea) == "call":
            dname = demangle_name(idc.GetOpnd(cur_ea, 0))
            if 'j_j_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 get_exception_size(typename):
    # Go over each name until we get the RTTI type descriptor,
    # then go over it's xrefs until we find a comment in some RTTI that ida parses, but can't demangle the name of,
    # then grab the parsed size.
    for ea, name in idautils.Names():
        dname = idc.Demangle(name, idc.GetLongPrm(idc.INF_SHORT_DN))
        if dname != None and typename in dname and 'RTTI Type Descriptor' in dname:
            for xref in idautils.XrefsTo(ea, 0):
                size_ea = xref.frm+16
                if(Comment(size_ea) == "size of thrown object"):
                    return ida_bytes.get_dword(size_ea)
    
    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.
    for typename, fields in type_to_info.iteritems():
        # First try to the size via `free` in vtable[0] dtor.
        size = get_free_size_arg_from_vtable_ea(fields['vtable_ea'], typename)
        
        # If the size wasn't gotten via `free`, try to get it via the `new` call before ctor.
        if size == -1:
            size = get_new_size_arg_from_first_xref_ea(fields['ctor_ea'], typename)

        # If the size still wasn't gotten, the ctor might be inlined:
        if size == -1:
            size = get_new_size_arg_from_first_xref_ea(fields['vtable_ea'])

        # Or maybe it's an exception after all.
        if size == -1:
            size = get_exception_size(typename)

        if size == -1:
            # Aww.
            print("[CWT_MARKERS!] Failed to get size for {}".format(typename))

        type_to_info[typename]['size'] = size

    # 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']))


main()