-
-
Save jon1scr/ade603b6672da9b0625025b37b2a014a to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from idautils import * | |
from idaapi import * | |
from ida_name import * | |
from idc import * | |
from ida_hexrays import * | |
from ida_frame import * | |
from ida_struct import * | |
INHERITANCE = { | |
# Auto Object base classes. | |
'KAutoObject' : None, | |
'KAutoObjectWithList' : 'KAutoObject', | |
'KAutoObjectWithListAllocatorAdapter' : 'KAutoObjectWithList', | |
'KSynchronizationObject' : 'KAutoObjectWithList', | |
'KSynchronizationObjectAllocatorAdapter' : 'KSynchronizationObject', | |
'KReadableEvent' : 'KSynchronizationObject', | |
'KReadableEventAllocatorAdapter' : 'KReadableEvent', | |
'KDebugBase' : 'KSynchronizationObject', | |
'KDebugBaseAllocatorAdapter' : 'KDebugBase', | |
# Auto Object final classes | |
'KClientPort' : 'KSynchronizationObject', | |
'KClientSession' : 'KAutoObjectWithListAllocatorAdapter', | |
'KCodeMemory' : 'KAutoObjectWithListAllocatorAdapter', | |
'KDebug' : 'KDebugBaseAllocatorAdapter', | |
'KDeviceAddressSpace' : 'KAutoObjectWithListAllocatorAdapter', | |
'KEvent' : 'KAutoObjectWithListAllocatorAdapter', | |
'KInterruptEvent' : 'KReadableEventAllocatorAdapter', | |
'KLightClientSession' : 'KAutoObjectWithListAllocatorAdapter', | |
'KLightServerSession' : 'KAutoObjectWithListAllocatorAdapter', | |
'KLightSession' : 'KAutoObjectWithListAllocatorAdapter', | |
'KPort' : 'KAutoObjectWithListAllocatorAdapter', | |
'KProcess' : 'KSynchronizationObjectAllocatorAdapter', | |
'KResourceLimit' : 'KAutoObjectWithListAllocatorAdapter', | |
'KServerPort' : 'KSynchronizationObject', | |
'KServerSession' : 'KSynchronizationObject', | |
'KSession' : 'KAutoObjectWithListAllocatorAdapter', | |
'KSessionRequest' : 'KAutoObject', | |
'KSharedMemory' : 'KAutoObjectWithListAllocatorAdapter', | |
'KThread' : 'KSynchronizationObjectAllocatorAdapter', | |
'KTransferMemory' : 'KAutoObjectWithListAllocatorAdapter', | |
'KWritableEvent' : 'KAutoObjectWithListAllocatorAdapter', | |
# Fake helper class | |
'IdObjectHelper' : 'KAutoObjectWithList' | |
} | |
CLASS_TOKENS = { | |
'KAutoObject' : 0x0, | |
'KClientSession' : 0xD00, | |
'KResourceLimit' : 0x2500, | |
'KLightSession' : 0x4500, | |
'KPort' : 0x8500, | |
'KSession' : 0x1900, | |
'KSharedMemory' : 0x2900, | |
'KEvent' : 0x4900, | |
'KWritableEvent' : 0x8900, | |
'KLightClientSession' : 0x3100, | |
'KLightServerSession' : 0x5100, | |
'KTransferMemory' : 0x9100, | |
'KDeviceAddressSpace' : 0x6100, | |
'KSessionRequest' : 0xA100, | |
'KCodeMemory' : 0xC100, | |
'KSynchronizationObject' : 0x1, | |
'KDebug' : 0xB01, | |
'KThread' : 0x1301, | |
'KServerPort' : 0x2301, | |
'KServerSession' : 0x4301, | |
'KClientPort' : 0x8301, | |
'KProcess' : 0x1501, | |
'KReadableEvent' : 0x3, | |
'KInterruptEvent' : 0x703, | |
} | |
INVERSE_CLASS_TOKENS = {v : k for k,v in CLASS_TOKENS.items()} | |
def MakeClassFunction(ret, name): | |
return (name, lambda cn: '%s %s::%s(%s *__hidden this)' % (ret, cn, name, cn), lambda cn: '%s::%s' % (cn, name)) | |
def MakeClassFunctionWithArgs(ret, name, *args): | |
if len(args) == 0: | |
return MakeClassFunction(ret, name) | |
else: | |
return (name, lambda cn: '%s %s::%s(%s *__hidden this, %s)' % (ret, cn, name, cn, ', '.join(args)), lambda cn: '%s::%s' % (cn, name)) | |
VTABLES = { | |
# Auto Object base classes. | |
'KAutoObject' : [ | |
( | |
'(base object destructor)', | |
lambda cn: 'void %s_base_destructor(%s *this);' % (cn, cn), | |
lambda cn: '_ZN%d%sD2Ev' % (len(cn), cn) | |
), | |
( | |
'(deleting destructor)', | |
lambda cn: 'void %s_deleting_destructor(%s *this);' % (cn, cn), | |
lambda cn: '_ZN%d%sD0Ev' % (len(cn), cn) | |
), | |
MakeClassFunction('void', 'Destroy'), | |
MakeClassFunction('void', 'Finalize'), | |
MakeClassFunction('KProcess *', 'GetOwnerProcess'), | |
MakeClassFunction('KTypeObj', 'GetTypeObj'), | |
MakeClassFunction('const char *', 'GetName'), | |
], | |
'KAutoObjectWithList': [ | |
MakeClassFunction('_QWORD', 'GetId'), | |
], | |
'KSynchronizationObject': [ | |
MakeClassFunction('void', 'OnFinalizeSynchronization'), | |
MakeClassFunction('bool', 'IsSignaled'), | |
MakeClassFunction('void', 'DumpWaiters'), | |
], | |
'KReadableEvent': [ | |
MakeClassFunction('Result', 'Signal'), | |
MakeClassFunction('Result', 'Clear'), | |
MakeClassFunction('Result', 'Reset'), | |
], | |
'KAutoObjectWithListAllocatorAdapter': [ | |
MakeClassFunction('bool', 'IsInitialized'), | |
MakeClassFunction('uintptr_t', 'GetPostFinalizeArgument'), | |
], | |
'KSynchronizationObjectAllocatorAdapter': [ | |
MakeClassFunction('bool', 'IsInitialized'), | |
MakeClassFunction('uintptr_t', 'GetPostFinalizeArgument'), | |
], | |
'KReadableEventAllocatorAdapter': [ | |
MakeClassFunction('bool', 'IsInitialized'), | |
MakeClassFunction('uintptr_t', 'GetPostFinalizeArgument'), | |
], | |
'KDebugBase' : [ | |
MakeClassFunctionWithArgs('Result', 'GetThreadContextImpl', 'ThreadContext *ctx', 'KThread *thread', 'uint32_t flags'), | |
MakeClassFunctionWithArgs('Result', 'SetThreadContextImpl', 'const ThreadContext *ctx', 'KThread *thread', 'uint32_t flags'), | |
], | |
'KDebugBaseAllocatorAdapter': [ | |
MakeClassFunction('bool', 'IsInitialized'), | |
MakeClassFunction('uintptr_t', 'GetPostFinalizeArgument'), | |
], | |
# Auto Object final classes | |
'KThread' : [ | |
MakeClassFunction('void', 'OnTimer'), # Really from multiple inheritance KTimerTask | |
MakeClassFunction('void', 'DoTask'), # Really from multiple inhertiance KWorkerTask | |
] | |
} | |
def GetVT(cls): | |
assert cls in INHERITANCE | |
funcs = [] | |
if INHERITANCE[cls] in INHERITANCE: | |
funcs += GetVT(INHERITANCE[cls]) | |
if cls in VTABLES: | |
funcs += VTABLES[cls] | |
return funcs | |
LABELED = { | |
} | |
NAMES = { | |
} | |
seg_mapping = {idc.get_segm_name(x): (idc.get_segm_start(x), idc.get_segm_end(x)) for x in Segments()} | |
text_start, text_end = seg_mapping['.text'] | |
ro_start, ro_end = seg_mapping['.rodata'] | |
def IsInText(ea): | |
return text_start <= ea and ea < text_end | |
def IsAncestorOf(cls, other): | |
if cls == other: | |
return True | |
elif INHERITANCE[other] in INHERITANCE: | |
return IsAncestorOf(cls, INHERITANCE[other]) | |
else: | |
return False | |
def TestCommonAncestor(ancestor, classes): | |
for cls in classes: | |
if not IsAncestorOf(ancestor, cls): | |
return False | |
return True | |
def GetShallowestCommonAncestorOfFunction(cls, others, func_id): | |
if TestCommonAncestor(cls, others): | |
for (vt_func_id, _, _) in GetVT(cls): | |
if vt_func_id == func_id: | |
return cls | |
if INHERITANCE[cls] in INHERITANCE: | |
ancestor = GetShallowestCommonAncestorOfFunction(INHERITANCE[cls], others, func_id) | |
if ancestor != None: | |
return ancestor | |
return None | |
def ApplyFunction(func_ea, get_type, get_name, cls): | |
func_type = get_type(cls) | |
func_name = get_name(cls) | |
idc.set_name(func_ea, func_name, SN_CHECK) | |
idc.SetType(func_ea, func_type) | |
if func_ea not in LABELED: | |
LABELED[func_ea] = [cls] | |
else: | |
LABELED[func_ea].append(cls) | |
NAMES[func_name] = func_ea | |
#print 'Labeled %s (%x)' % (func_name, func_ea) | |
def ResolveAncestryConflict(cls, func_ea, func_id, get_type, get_name): | |
common_ancestor = GetShallowestCommonAncestorOfFunction(cls, LABELED[func_ea], func_id) | |
assert common_ancestor != None | |
# Simple case, common ancestor name not in use. | |
anc_name = get_name(common_ancestor) | |
if anc_name not in NAMES or NAMES[anc_name] == func_ea: | |
ApplyFunction(func_ea, get_type, get_name, common_ancestor) | |
else: | |
print common_ancestor | |
print anc_name | |
print NAMES | |
print '%X' % NAMES[anc_name] | |
print '%X' % func_ea | |
# Common ancestor name is in use. | |
# TODO: Resolve more complicated ancestry conflict. | |
assert False | |
def ApplyVirtualTable(cls, vt, vt_ea): | |
for i in xrange(len(vt)): | |
assert IsInText(ida_bytes.get_64bit(vt_ea + 8 * i)) | |
for i, (func_id, get_func_type, get_func_name) in enumerate(vt): | |
func_ea = ida_bytes.get_64bit(vt_ea + 8 * i) | |
if func_ea not in LABELED: | |
ApplyFunction(func_ea, get_func_type, get_func_name, cls) | |
else: | |
ResolveAncestryConflict(cls, func_ea, func_id, get_func_type, get_func_name) | |
def ProcessClass(cls): | |
if INHERITANCE[cls] in INHERITANCE: | |
ProcessClass(INHERITANCE[cls]) | |
vt = GetVT(cls) | |
ea = get_name_ea(BADADDR, '%s::vt' % cls) | |
if ea != BADADDR: | |
# print '%s: %x' % (cls, ea) | |
ApplyVirtualTable(cls, vt, ea) | |
def Disassemble(head): | |
disasm = GetDisasm(head).lower().lstrip().rstrip().replace(',',' ') | |
if ';' in disasm: | |
disasm = disasm[:disasm.index(';')] | |
return disasm.split() | |
def IsGetTypeObj(disasms): | |
if len(disasms) == 3: | |
if len(disasms[0]) != 3: | |
return False | |
if disasms[0][0] != 'adrl': | |
return False | |
if disasms[0][1] != 'x0': | |
return False | |
if len(disasms[1]) != 3: | |
return False | |
if disasms[1][0] != 'mov': | |
return False | |
if disasms[1][1] == 'w1': | |
if disasms[1][2].startswith('#'): | |
pass | |
elif disasms[1][2] == 'wzr': | |
pass | |
else: | |
return False | |
elif disasms[1][1] == 'x1': | |
if disasms[1][2].startswith('#'): | |
pass | |
elif disasms[1][2] == 'xzr': | |
pass | |
else: | |
return False | |
if len(disasms[2]) != 1: | |
return False | |
return disasms[2][0] == 'ret' | |
elif len(disasms) == 5: | |
if len(disasms[0]) != 3: | |
return False | |
if disasms[0][0] != 'adrp': | |
return False | |
if disasms[0][1] != 'x8': | |
return False | |
if len(disasms[1]) != 4: | |
return False | |
if disasms[1][0] != 'ldr': | |
return False | |
if disasms[1][1] != 'x8': | |
return False | |
if disasms[1][2] != '[x8': | |
return False | |
if disasms[1][3] != '%soff]' % disasms[0][2]: | |
return False | |
if len(disasms[2]) != 3: | |
return False | |
if disasms[2][0] != 'ldr': | |
return False | |
if disasms[2][1] != 'x0': | |
return False | |
if disasms[2][2] != '[x8]': | |
return False | |
if len(disasms[3]) != 3: | |
return False | |
if disasms[3][0] != 'mov': | |
return False | |
if disasms[3][1] == 'w1': | |
if disasms[3][2].startswith('#'): | |
pass | |
elif disasms[3][2] == 'wzr': | |
pass | |
else: | |
return False | |
elif disasms[3][1] == 'x1': | |
if disasms[3][2].startswith('#'): | |
pass | |
elif disasms[3][2] == 'xzr': | |
pass | |
else: | |
return False | |
else: | |
return False | |
if len(disasms[4]) != 1: | |
return False | |
return disasms[4][0] == 'ret' | |
ID_OBJECT_HELPER_VT = None | |
def GetVtableAddress(get_type_obj_ea): | |
global ID_OBJECT_HELPER_VT | |
candidates = [] | |
ofs = ro_start & ~7 | |
while ofs < ro_end: | |
val = ida_bytes.get_64bit(ofs) | |
if val == get_type_obj_ea: | |
candidates.append(ofs - 0x28) | |
ofs += 8 | |
if len(candidates) == 1: | |
return candidates[0] | |
elif len(candidates) == 2: | |
# KAutoObject | |
assert ID_OBJECT_HELPER_VT is None | |
if ida_bytes.get_64bit(candidates[0] + 0x38) == 0: | |
ID_OBJECT_HELPER_VT = candidates[1] | |
return candidates[0] | |
else: | |
assert ida_bytes.get_64bit(candidates[1] + 0x38) == 0 | |
ID_OBJECT_HELPER_VT = candidates[0] | |
return candidates[1] | |
return None | |
# Identify vtables | |
TYPE_OBJS = {} | |
INVERSE_TYPE_OBJS = {} | |
for segea in [text_start]: | |
for funcea in Functions(segea, get_segm_end(segea)): | |
chunks = [chunk for chunk in Chunks(funcea)] | |
if len(chunks) != 1: | |
continue | |
startea, endea = chunks[0] | |
heads = [head for head in Heads(startea, endea)] | |
if len(heads) != 3 and len(heads) != 5: | |
continue | |
disasms = [Disassemble(head) for head in heads] | |
if not IsGetTypeObj(disasms): | |
continue | |
token = disasms[-2][2][1:] | |
token = int(token, 0) if token != 'zr' else 0 | |
assert token not in INVERSE_TYPE_OBJS | |
vt_ea = GetVtableAddress(funcea) | |
assert vt_ea is not None | |
TYPE_OBJS[vt_ea] = token | |
INVERSE_TYPE_OBJS[token] = vt_ea | |
assert set(INVERSE_TYPE_OBJS.keys()) == set(INVERSE_CLASS_TOKENS.keys()) | |
assert ID_OBJECT_HELPER_VT is not None | |
for token in INVERSE_CLASS_TOKENS.keys(): | |
vt_name = '%s::vt' % INVERSE_CLASS_TOKENS[token] | |
vt_ea = INVERSE_TYPE_OBJS[token] | |
cur_ea = get_name_ea(BADADDR, vt_name) | |
if cur_ea != BADADDR: | |
idc.set_name(cur_ea, '', SN_CHECK) | |
idc.set_name(vt_ea, vt_name, SN_CHECK) | |
print 'Found vtable for %s at 0x%x' % (INVERSE_CLASS_TOKENS[token], vt_ea) | |
idc.set_name(ID_OBJECT_HELPER_VT, 'IdObjectHelper::vt', SN_CHECK) | |
# Label identified vtables | |
for cls in INHERITANCE.keys(): | |
ProcessClass(cls) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment