Replace HexRaysCodeXplorer and HexraysInvertIf with HexRaysPyTools

This commit is contained in:
ecx86
2018-08-05 01:08:36 -04:00
parent 035b4ae117
commit 454ae02a36
22 changed files with 5639 additions and 0 deletions

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,583 @@
import logging
import idaapi
import idc
from Core.Helper import to_hex, get_member_name, get_func_argument_info, get_funcs_calling_address, is_imported_ea
logger = logging.getLogger(__name__)
SETTING_START_FROM_CURRENT_EXPR = True
def decompile_function(address):
try:
cfunc = idaapi.decompile(address)
if cfunc:
return cfunc
except idaapi.DecompilationFailure:
pass
logger.warn("IDA failed to decompile function at 0x{address:08X}".format(address=address))
class ScanObject(object):
def __init__(self):
self.ea = idaapi.BADADDR
self.name = None
self.tinfo = None
self.id = 0
@staticmethod
def create(cfunc, arg):
# Creates object suitable for scaning either from cexpr_t or ctree_item_t
if isinstance(arg, idaapi.ctree_item_t):
lvar = arg.get_lvar()
if lvar:
index = list(cfunc.get_lvars()).index(lvar)
result = VariableObject(lvar, index)
if arg.e:
result.ea = ScanObject.get_expression_address(cfunc, arg.e)
return result
cexpr = arg.e
else:
cexpr = arg
if cexpr.op == idaapi.cot_var:
lvar = cfunc.get_lvars()[cexpr.v.idx]
result = VariableObject(lvar, cexpr.v.idx)
result.ea = ScanObject.get_expression_address(cfunc, cexpr)
return result
elif cexpr.op == idaapi.cot_memptr:
t = cexpr.x.type.get_pointed_object()
result = StructPtrObject(t.dstr(), cexpr.m)
result.name = get_member_name(t, cexpr.m)
elif cexpr.op == idaapi.cot_memref:
t = cexpr.x.type
result = StructRefObject(t.dstr(), cexpr.m)
result.name = get_member_name(t, cexpr.m)
elif cexpr.op == idaapi.cot_obj:
result = GlobalVariableObject(cexpr.obj_ea)
result.name = idaapi.get_short_name(cexpr.obj_ea)
else:
return
result.tinfo = cexpr.type
result.ea = ScanObject.get_expression_address(cfunc, cexpr)
return result
@staticmethod
def get_expression_address(cfunc, cexpr):
expr = cexpr
while expr and expr.ea == idaapi.BADADDR:
expr = expr.to_specific_type
expr = cfunc.body.find_parent_of(expr)
assert expr is not None
return expr.ea
def __hash__(self):
return hash((self.id, self.name))
def __eq__(self, other):
return self.id == other.id and self.name == other.name
def __repr__(self):
return self.name
SO_LOCAL_VARIABLE = 1 # cexpr.op == idaapi.cot_var
SO_STRUCT_POINTER = 2 # cexpr.op == idaapi.cot_memptr
SO_STRUCT_REFERENCE = 3 # cexpr.op == idaapi.cot_memref
SO_GLOBAL_OBJECT = 4 # cexpr.op == idaapi.cot_obj
SO_CALL_ARGUMENT = 5 # cexpr.op == idaapi.cot_call
SO_MEMORY_ALLOCATOR = 6
SO_RETURNED_OBJECT = 7
class VariableObject(ScanObject):
# Represents `var` expression
def __init__(self, lvar, index):
super(VariableObject, self).__init__()
self.lvar = lvar
self.tinfo = lvar.type()
self.name = lvar.name
self.index = index
self.id = SO_LOCAL_VARIABLE
def is_target(self, cexpr):
return cexpr.op == idaapi.cot_var and cexpr.v.idx == self.index
class StructPtrObject(ScanObject):
# Represents `x->m` expression
def __init__(self, struct_name, offset):
super(StructPtrObject, self).__init__()
self.struct_name = struct_name
self.offset = offset
self.id = SO_STRUCT_POINTER
def is_target(self, cexpr):
return cexpr.op == idaapi.cot_memptr and cexpr.m == self.offset and \
cexpr.x.type.get_pointed_object().dstr() == self.struct_name
class StructRefObject(ScanObject):
# Represents `x.m` expression
def __init__(self, struct_name, offset):
super(StructRefObject, self).__init__()
self.__struct_name = struct_name
self.__offset = offset
self.id = SO_STRUCT_REFERENCE
def is_target(self, cexpr):
return cexpr.op == idaapi.cot_memref and cexpr.m == self.__offset and cexpr.x.type.dstr() == self.__struct_name
class GlobalVariableObject(ScanObject):
# Represents global object
def __init__(self, object_address):
super(GlobalVariableObject, self).__init__()
self.obj_ea = object_address
self.id = SO_GLOBAL_OBJECT
def is_target(self, cexpr):
return cexpr.op == idaapi.cot_obj and self.obj_ea == cexpr.obj_ea
class CallArgObject(ScanObject):
# Represents call of function and specified argument
def __init__(self, func_address, arg_idx):
super(CallArgObject, self).__init__()
self.__func_ea = func_address
self.__arg_idx = arg_idx
self.id = SO_CALL_ARGUMENT
def is_target(self, cexpr):
return cexpr.op == idaapi.cot_call and cexpr.x.obj_ea == self.__func_ea
def create_scan_obj(self, cfunc, cexpr):
e = cexpr.a[self.__arg_idx]
while e.op in (idaapi.cot_cast, idaapi.cot_ref, idaapi.cot_add, idaapi.cot_sub, idaapi.cot_idx):
e = e.x
return ScanObject.create(cfunc, e)
@staticmethod
def create(cfunc, arg_idx):
result = CallArgObject(cfunc.entry_ea, arg_idx)
result.name = cfunc.get_lvars()[arg_idx].name
result.tinfo = cfunc.type
return result
def __repr__(self):
return "{}"
class ReturnedObject(ScanObject):
# Represents value returned by function
def __init__(self, func_address):
super(ReturnedObject, self).__init__()
self.__func_ea = func_address
self.id = SO_RETURNED_OBJECT
def is_target(self, cexpr):
return cexpr.op == idaapi.cot_call and cexpr.x.obj_ea == self.__func_ea
class MemoryAllocationObject(ScanObject):
# Represents `operator new()' or `malloc'
def __init__(self, name, size):
super(MemoryAllocationObject, self).__init__()
self.name = name
self.size = size
self.id = SO_MEMORY_ALLOCATOR
@staticmethod
def create(cfunc, cexpr):
if cexpr.op == idaapi.cot_call:
e = cexpr
elif cexpr.op == idaapi.cot_cast and cexpr.x.op == idaapi.cot_call:
e = cexpr.x
else:
return
func_name = idaapi.get_short_name(e.x.obj_ea)
if "malloc" in func_name or "operator new" in func_name:
carg = e.a[0]
if carg.op == idaapi.cot_num:
size = carg.numval()
else:
size = 0
result = MemoryAllocationObject(func_name, size)
result.ea = ScanObject.get_expression_address(cfunc, e)
return result
ASSIGNMENT_RIGHT = 1
ASSIGNMENT_LEFT = 2
class ObjectVisitor(idaapi.ctree_parentee_t):
def __init__(self, cfunc, obj, data, skip_until_object):
super(ObjectVisitor, self).__init__()
self._cfunc = cfunc
self._objects = [obj]
self._init_obj = obj
self._data = data
self._start_ea = obj.ea
self._skip = skip_until_object if self._start_ea != idaapi.BADADDR else False
self.crippled = False
def process(self):
self.apply_to(self._cfunc.body, None)
def set_callbacks(self, manipulate=None):
if manipulate:
self.__manipulate = manipulate.__get__(self, ObjectDownwardsVisitor)
def _get_line(self):
for p in reversed(self.parents):
if not p.is_expr():
return idaapi.tag_remove(p.print1(self._cfunc.__ref__()))
AssertionError("Parent instruction is not found")
def _manipulate(self, cexpr, obj):
"""
Method called for every object having assignment relationship with starter object. This method should be
reimplemented in order to do something useful
:param cexpr: idaapi_cexpr_t
:param id: one of the SO_* constants
:return: None
"""
self.__manipulate(cexpr, obj)
def __manipulate(self, cexpr, obj):
logger.debug("Expression {} at {} Id - {}".format(cexpr.opname, to_hex(self._find_asm_address(cexpr)), obj.id))
class ObjectDownwardsVisitor(ObjectVisitor):
def __init__(self, cfunc, obj, data=None, skip_until_object=False):
super(ObjectDownwardsVisitor, self).__init__(cfunc, obj, data, skip_until_object)
self.cv_flags |= idaapi.CV_POST
def visit_expr(self, cexpr):
if self._skip:
if self._is_initial_object(cexpr):
self._skip = False
else:
return 0
if cexpr.op != idaapi.cot_asg:
return 0
x_cexpr = cexpr.x
if cexpr.y.op == idaapi.cot_cast:
y_cexpr = cexpr.y.x
else:
y_cexpr = cexpr.y
for obj in self._objects:
if obj.is_target(x_cexpr):
if self.__is_object_overwritten(x_cexpr, obj, y_cexpr):
logger.info("Removed object {} from scanning at {}".format(
obj, to_hex(self._find_asm_address(x_cexpr))))
self._objects.remove(obj)
return 0
elif obj.is_target(y_cexpr):
new_obj = ScanObject.create(self._cfunc, x_cexpr)
if new_obj:
self._objects.append(new_obj)
return 0
return 0
def leave_expr(self, cexpr):
if self._skip:
return 0
for obj in self._objects:
if obj.is_target(cexpr) and obj.id != SO_RETURNED_OBJECT:
self._manipulate(cexpr, obj)
return 0
return 0
def _is_initial_object(self, cexpr):
if cexpr.op == idaapi.cot_asg:
cexpr = cexpr.y
if cexpr.op == idaapi.cot_cast:
cexpr = cexpr.x
return self._init_obj.is_target(cexpr) and self._find_asm_address(cexpr) == self._start_ea
def __is_object_overwritten(self, x_cexpr, obj, y_cexpr):
if len(self._objects) < 2:
return False
if y_cexpr.op == idaapi.cot_cast:
e = y_cexpr.x
else:
e = y_cexpr
if e.op != idaapi.cot_call or len(e.a) == 0:
return True
for obj in self._objects:
if obj.is_target(e. a[0]):
return False
return True
class ObjectUpwardsVisitor(ObjectVisitor):
STAGE_PREPARE = 1
STAGE_PARSING = 2
def __init__(self, cfunc, obj, data=None, skip_after_object=False):
super(ObjectUpwardsVisitor, self).__init__(cfunc, obj, data, skip_after_object)
self._stage = self.STAGE_PREPARE
self._tree = {}
self._call_obj = obj if obj.id == SO_CALL_ARGUMENT else None
def visit_expr(self, cexpr):
if self._stage == self.STAGE_PARSING:
return 0
if self._call_obj and self._call_obj.is_target(cexpr):
obj = self._call_obj.create_scan_obj(self._cfunc, cexpr)
if obj:
self._objects.append(obj)
return 0
if cexpr.op != idaapi.cot_asg:
return 0
x_cexpr = cexpr.x
if cexpr.y.op == idaapi.cot_cast:
y_cexpr = cexpr.y.x
else:
y_cexpr = cexpr.y
obj_left = ScanObject.create(self._cfunc, x_cexpr)
obj_right = ScanObject.create(self._cfunc, y_cexpr)
if obj_left and obj_right:
self.__add_object_assignment(obj_left, obj_right)
if self._skip and self._is_initial_object(cexpr):
return 1
return 0
def leave_expr(self, cexpr):
if self._stage == self.STAGE_PREPARE:
return 0
if self._skip and self._is_initial_object(cexpr):
self._manipulate(cexpr, self._init_obj)
return 1
for obj in self._objects:
if obj.is_target(cexpr):
self._manipulate(cexpr, obj)
return 0
return 0
def process(self):
self._stage = self.STAGE_PREPARE
self.cv_flags &= ~idaapi.CV_POST
super(ObjectUpwardsVisitor, self).process()
self._stage = self.STAGE_PARSING
self.cv_flags |= idaapi.CV_POST
self.__prepare()
super(ObjectUpwardsVisitor, self).process()
def _is_initial_object(self, cexpr):
return self._init_obj.is_target(cexpr) and self._find_asm_address(cexpr) == self._start_ea
def __add_object_assignment(self, from_obj, to_obj):
if from_obj in self._tree:
self._tree[from_obj].add(to_obj)
else:
self._tree[from_obj] = {to_obj}
def __prepare(self):
result = set()
todo = set(self._objects)
while todo:
obj = todo.pop()
result.add(obj)
if obj.id == SO_CALL_ARGUMENT or obj not in self._tree:
continue
o = self._tree[obj]
todo |= o - result
result |= o
self._objects = list(result)
self._tree.clear()
class RecursiveObjectVisitor(ObjectVisitor):
def __init__(self, cfunc, obj, data=None, skip_until_object=False, visited=None):
super(RecursiveObjectVisitor, self).__init__(cfunc, obj, data, skip_until_object)
self._visited = visited if visited else set()
self._new_for_visit = set()
self.crippled = False
self._arg_idx = -1
self._debug_scan_tree = {}
self.__debug_scan_tree_root = idc.Name(self._cfunc.entry_ea)
self.__debug_message = []
def visit_expr(self, cexpr):
return super(RecursiveObjectVisitor, self).visit_expr(cexpr)
def set_callbacks(self, manipulate=None, start=None, start_iteration=None, finish=None, finish_iteration=None):
super(RecursiveObjectVisitor, self).set_callbacks(manipulate)
if start:
self._start = start.__get__(self, RecursiveObjectDownwardsVisitor)
if start_iteration:
self._start_iteration = start_iteration.__get__(self, RecursiveObjectDownwardsVisitor)
if finish:
self._finish = finish.__get__(self, RecursiveObjectDownwardsVisitor)
if finish_iteration:
self._finish_iteration = finish_iteration.__get__(self, RecursiveObjectDownwardsVisitor)
def prepare_new_scan(self, cfunc, arg_idx, obj, skip=False):
self._cfunc = cfunc
self._arg_idx = arg_idx
self._objects = [obj]
self._init_obj = obj
self._skip = False
self.crippled = self.__is_func_crippled()
def process(self):
self._start()
self._recursive_process()
self._finish()
self.dump_scan_tree()
def dump_scan_tree(self):
self.__prepare_debug_message()
logger.info("{}\n---------------".format("\n".join(self.__debug_message)))
def __prepare_debug_message(self, key=None, level=1):
if key is None:
key = (self.__debug_scan_tree_root, -1)
self.__debug_message.append("--- Scan Tree---\n{}".format(self.__debug_scan_tree_root))
if key in self._debug_scan_tree:
for func_name, arg_idx in self._debug_scan_tree[key]:
prefix = " | " * (level - 1) + " |_ "
self.__debug_message.append("{}{} (idx: {})".format(prefix, func_name, arg_idx))
self.__prepare_debug_message((func_name, arg_idx), level + 1)
def _recursive_process(self):
self._start_iteration()
super(RecursiveObjectVisitor, self).process()
self._finish_iteration()
def _manipulate(self, cexpr, obj):
self._check_call(cexpr)
super(RecursiveObjectVisitor, self)._manipulate(cexpr, obj)
def _check_call(self, cexpr):
raise NotImplemented
def _add_visit(self, func_ea, arg_idx):
if (func_ea, arg_idx) not in self._visited:
self._visited.add((func_ea, arg_idx))
self._new_for_visit.add((func_ea, arg_idx))
return True
return False
def _add_scan_tree_info(self, func_ea, arg_idx):
head_node = (idc.Name(self._cfunc.entry_ea), self._arg_idx)
tail_node = (idc.Name(func_ea), arg_idx)
if head_node in self._debug_scan_tree:
self._debug_scan_tree[head_node].add(tail_node)
else:
self._debug_scan_tree[head_node] = {tail_node}
def _start(self):
""" Called at the beginning of visiting """
pass
def _start_iteration(self):
""" Called every time new function visiting started """
pass
def _finish(self):
""" Called after all visiting happened """
pass
def _finish_iteration(self):
""" Called every time new function visiting finished """
pass
def __is_func_crippled(self):
# Check if function is just call to another function
b = self._cfunc.body.cblock
if b.size() == 1:
e = b.at(0)
return e.op == idaapi.cit_return or (e.op == idaapi.cit_expr and e.cexpr.op == idaapi.cot_call)
return False
class RecursiveObjectDownwardsVisitor(RecursiveObjectVisitor, ObjectDownwardsVisitor):
def __init__(self, cfunc, obj, data=None, skip_until_object=False, visited=None):
super(RecursiveObjectDownwardsVisitor, self).__init__(cfunc, obj, data, skip_until_object, visited)
def _check_call(self, cexpr):
parent = self.parent_expr()
grandparent = self.parents.at(self.parents.size() - 2)
if parent.op == idaapi.cot_call:
call_cexpr = parent
arg_cexpr = cexpr
elif parent.op == idaapi.cot_cast and grandparent.op == idaapi.cot_call:
call_cexpr = grandparent.cexpr
arg_cexpr = parent
else:
return
idx, _ = get_func_argument_info(call_cexpr, arg_cexpr)
func_ea = call_cexpr.x.obj_ea
if func_ea == idaapi.BADADDR:
return
if self._add_visit(func_ea, idx):
self._add_scan_tree_info(func_ea, idx)
def _recursive_process(self):
super(RecursiveObjectDownwardsVisitor, self)._recursive_process()
while self._new_for_visit:
func_ea, arg_idx = self._new_for_visit.pop()
if is_imported_ea(func_ea):
continue
cfunc = decompile_function(func_ea)
if cfunc:
assert arg_idx < len(cfunc.get_lvars()), "Wrong argument at func {}".format(to_hex(func_ea))
obj = VariableObject(cfunc.get_lvars()[arg_idx], arg_idx)
self.prepare_new_scan(cfunc, arg_idx, obj)
self._recursive_process()
class RecursiveObjectUpwardsVisitor(RecursiveObjectVisitor, ObjectUpwardsVisitor):
def __init__(self, cfunc, obj, data=None, skip_after_object=False, visited=None):
super(RecursiveObjectUpwardsVisitor, self).__init__(cfunc, obj, data, skip_after_object, visited)
def prepare_new_scan(self, cfunc, arg_idx, obj, skip=False):
super(RecursiveObjectUpwardsVisitor, self).prepare_new_scan(cfunc, arg_idx, obj, skip)
self._call_obj = obj if obj.id == SO_CALL_ARGUMENT else None
def _check_call(self, cexpr):
if cexpr.op == idaapi.cot_var and self._cfunc.get_lvars()[cexpr.v.idx].is_arg_var:
func_ea = self._cfunc.entry_ea
arg_idx = cexpr.v.idx
if self._add_visit(func_ea, arg_idx):
for callee_ea in get_funcs_calling_address(func_ea):
self._add_scan_tree_info(callee_ea, arg_idx)
def _recursive_process(self):
super(RecursiveObjectUpwardsVisitor, self)._recursive_process()
while self._new_for_visit:
new_visit = list(self._new_for_visit)
self._new_for_visit.clear()
for func_ea, arg_idx in new_visit:
funcs = get_funcs_calling_address(func_ea)
obj = CallArgObject.create(idaapi.decompile(func_ea), arg_idx)
for callee_ea in funcs:
cfunc = decompile_function(callee_ea)
if cfunc:
self.prepare_new_scan(cfunc, arg_idx, obj, False)
super(RecursiveObjectUpwardsVisitor, self)._recursive_process()

View File

@@ -0,0 +1,62 @@
# Information about explored functions
import re
import idaapi
import idautils
import idc
imported_ea = set()
demangled_names = {}
touched_functions = set()
temporary_structure = None
def init_imported_ea(*args):
def imp_cb(ea, name, ord):
imported_ea.add(ea)
# True -> Continue enumeration
# False -> Stop enumeration
return True
print "[Info] Collecting information about imports"
imported_ea.clear()
nimps = idaapi.get_import_module_qty()
for i in xrange(0, nimps):
name = idaapi.get_import_module_name(i)
if not name:
print "[Warning] Failed to get import module name for #%d" % i
continue
# print "Walking-> %s" % name
idaapi.enum_import_names(i, imp_cb)
print "[Info] Done..."
def init_demangled_names(*args):
"""
Creates dictionary of demangled names => address, that will be used further at double click on methods got from
symbols.
"""
demangled_names.clear()
for address, name in idautils.Names():
short_name = idc.Demangle(name, idc.GetLongPrm(idc.INF_SHORT_DN))
if short_name:
demangled_names[short_name.split('(')[0]] = address - idaapi.get_imagebase()
# Names can have templates and should be transformed before creating local type
name = re.sub(r'[<>]', '_t_', name)
# Thunk functions with name like "[thunk]:CWarmupHostProvider::Release`adjustor{8}'"
result = re.search(r"(\[thunk\]:)?([^`]*)(.*\{(\d+)}.*)?", short_name)
name, adjustor = result.group(2), result.group(4)
if adjustor:
demangled_names[name + "_adj_" + adjustor] = address - idaapi.get_imagebase()
print "[DEBUG] Demangled names have been initialized"
def reset_touched_functions(*args):
touched_functions = set()

View File

@@ -0,0 +1,536 @@
import idc
import idaapi
import HexRaysPyTools.Forms
import Helper
# from PySide import QtGui, QtCore
from HexRaysPyTools.Cute import *
all_virtual_functions = {} # name -> VirtualMethod
all_virtual_tables = {} # ordinal -> VirtualTable
class VirtualMethod(object):
def __init__(self, tinfo, name, parent):
self.tinfo = tinfo
self.tinfo_modified = False
self.name = name
self.class_name = None
self.name_modified = False
self.parents = [parent]
self.base_address = Helper.get_virtual_func_address(name)
if self.base_address:
self.base_address -= idaapi.get_imagebase()
self.rowcount = 0
self.children = []
@staticmethod
def create(tinfo, name, parent):
result = all_virtual_functions.get(name)
if result:
result.parents.append(parent)
return result
result = VirtualMethod(tinfo, name, parent)
all_virtual_functions[name] = result
return result
def update(self, name, tinfo):
self.name = name
self.tinfo = tinfo
self.name_modified = False
self.tinfo_modified = False
self.base_address = idc.LocByName(self.name)
if self.base_address != idaapi.BADADDR:
self.base_address -= idaapi.get_imagebase()
def data(self, column):
if column == 0:
return self.name
elif column == 1:
return self.tinfo.get_pointed_object().dstr()
elif column == 2:
return "0x{0:08X}".format(self.address) if self.address else None
def setData(self, column, value):
if column == 0:
if idaapi.isident(value) and self.name != value:
self.name = value
self.name_modified = True
for parent in self.parents:
parent.modified = True
return True
elif column == 1:
tinfo = idaapi.tinfo_t()
split = value.split('(')
if len(split) == 2:
value = split[0] + ' ' + self.name + '(' + split[1] + ';'
if idaapi.parse_decl2(idaapi.cvar.idati, value, '', tinfo, idaapi.PT_TYP):
if tinfo.is_func():
tinfo.create_ptr(tinfo)
if tinfo.dstr() != self.tinfo.dstr():
self.tinfo = tinfo
self.tinfo_modified = True
for parent in self.parents:
parent.modified = True
return True
return False
def font(self, column):
if column == 0 and self.name_modified:
return QtGui.QFont("Consolas", 10, italic=True)
elif column == 1 and self.tinfo_modified:
return QtGui.QFont("Consolas", 10, italic=True)
return QtGui.QFont("Consolas", 10, 0)
def flags(self, column):
if column != 2:
if self.address:
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsEditable
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled
@property
def color(self):
# return QtGui.QBrush(QtGui.QColor("#ffffb3"))
return QtGui.QColor("#fefbd8")
@property
def tooltip(self):
return None
@property
def address(self):
return self.base_address + idaapi.get_imagebase() if self.base_address else None
def set_first_argument_type(self, name):
func_data = idaapi.func_type_data_t()
func_tinfo = self.tinfo.get_pointed_object()
class_tinfo = idaapi.tinfo_t()
if func_tinfo.get_func_details(func_data) and func_tinfo.get_nargs() and \
class_tinfo.get_named_type(idaapi.cvar.idati, name):
class_tinfo.create_ptr(class_tinfo)
first_arg_tinfo = func_data[0].type
if (first_arg_tinfo.is_ptr() and first_arg_tinfo.get_pointed_object().is_udt()) or \
Helper.is_legal_type(func_data[0].type):
func_data[0].type = class_tinfo
func_data[0].name = "this"
func_tinfo.create_func(func_data)
func_tinfo.create_ptr(func_tinfo)
if func_tinfo.dstr() != self.tinfo.dstr():
self.tinfo = func_tinfo
self.tinfo_modified = True
for parent in self.parents:
parent.modified = True
else:
print "[Warning] function {0} probably have wrong type".format(self.name)
def open_function(self):
if self.address:
if idaapi.decompile(self.address):
idaapi.open_pseudocode(self.address, 0)
else:
idaapi.jumpto(self.address)
def commit(self):
if self.name_modified:
self.name_modified = False
if self.address:
idaapi.set_name(self.address, self.name)
if self.tinfo_modified:
self.tinfo_modified = False
if self.address:
idaapi.apply_tinfo2(self.address, self.tinfo.get_pointed_object(), idaapi.TINFO_DEFINITE)
def __eq__(self, other):
return self.address == other.address
def __repr__(self):
return self.name
class VirtualTable(object):
def __init__(self, ordinal, tinfo, class_):
self.ordinal = ordinal
self.tinfo = tinfo
self.class_ = [class_]
self.class_name = None
self.virtual_functions = []
self.name = self.tinfo.dstr()
self._modified = False
def update(self):
if self.modified:
vtable_tinfo = idaapi.tinfo_t()
udt_data = idaapi.udt_type_data_t()
vtable_tinfo.get_numbered_type(idaapi.cvar.idati, self.ordinal)
vtable_tinfo.get_udt_details(udt_data)
self.tinfo = vtable_tinfo
self.name = vtable_tinfo.dstr()
self.modified = False
if len(self.virtual_functions) == len(udt_data):
for current_function, other_function in zip(self.virtual_functions, udt_data):
current_function.update(other_function.name, other_function.type)
else:
print "[ERROR] Something have been modified in Local types. Please refresh this view"
def update_local_type(self):
if self.modified:
final_tinfo = idaapi.tinfo_t()
udt_data = idaapi.udt_type_data_t()
self.tinfo.get_udt_details(udt_data)
if len(udt_data) == len(self.virtual_functions):
for udt_member, virtual_function in zip(udt_data, self.virtual_functions):
udt_member.name = virtual_function.name
udt_member.type = virtual_function.tinfo
virtual_function.commit()
final_tinfo.create_udt(udt_data, idaapi.BTF_STRUCT)
final_tinfo.set_numbered_type(idaapi.cvar.idati, self.ordinal, idaapi.NTF_REPLACE, self.name)
self.modified = False
else:
print "[ERROR] Something have been modified in Local types. Please refresh this view"
def set_first_argument_type(self, class_name):
for function in self.virtual_functions:
function.set_first_argument_type(class_name)
@property
def modified(self):
return self._modified
@modified.setter
def modified(self, value):
self._modified = value
if value:
for class_ in self.class_:
class_.modified = True
@staticmethod
def create(tinfo, class_):
ordinal = idaapi.get_type_ordinal(idaapi.cvar.idati, tinfo.dstr())
result = all_virtual_tables.get(ordinal)
if result:
result.class_.append(class_)
else:
udt_data = idaapi.udt_type_data_t()
tinfo.get_udt_details(udt_data)
result = VirtualTable(ordinal, tinfo, class_)
virtual_functions = [VirtualMethod.create(func.type, func.name, result) for func in udt_data]
result.virtual_functions = virtual_functions
all_virtual_functions[ordinal] = result
return result
def get_class_tinfo(self):
if len(self.class_) == 1:
return self.class_.tinfo
def setData(self, column, value):
if column == 0:
if idaapi.isident(value) and self.name != value:
self.name = value
self.modified = True
return True
return False
def data(self, column):
if column == 0:
return self.name
@property
def color(self):
return QtGui.QColor("#d5f4e6")
@property
def tooltip(self):
pass
def font(self, column):
if self.modified:
return QtGui.QFont("Consolas", 12, italic=True)
return QtGui.QFont("Consolas", 12)
def flags(self, column):
if column == 0:
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEditable | QtCore.Qt.ItemIsEnabled
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled
@property
def children(self):
return self.virtual_functions
def __repr__(self):
return str(self.virtual_functions)
class Class(object):
def __init__(self, name, tinfo, ordinal):
self.name = name
self.class_name = name
self.ordinal = ordinal
self.parent = None
self.vtables = {}
self.modified = False
@staticmethod
def create_class(ordinal):
tinfo = idaapi.tinfo_t()
tinfo.get_numbered_type(idaapi.cvar.idati, ordinal)
vtables = {}
if tinfo.is_struct():
udt_data = idaapi.udt_type_data_t()
tinfo.get_udt_details(udt_data)
for field_udt in udt_data:
if field_udt.type.is_ptr():
possible_vtable = field_udt.type.get_pointed_object()
if possible_vtable.is_struct():
v_udt_data = idaapi.udt_type_data_t()
possible_vtable.get_udt_details(v_udt_data)
for possible_func_udt in v_udt_data:
if not possible_func_udt.type.is_funcptr():
break
else:
vtables[field_udt.offset / 8] = possible_vtable
if vtables:
class_ = Class(tinfo.dstr(), tinfo, ordinal)
for offset, vtable_tinfo in vtables.iteritems():
vtables[offset] = VirtualTable.create(vtable_tinfo, class_)
class_.vtables = vtables
return class_
def update_from_local_types(self):
try:
if self.modified:
class_ = self.create_class(self.ordinal)
if class_:
self.name = class_.name
self.modified = False
for offset, vtable in class_.vtables.iteritems():
self.vtables[offset].update()
else:
# TODO: drop class
raise IndexError
except IndexError:
print "[ERROR] Something have been modified in Local types. Please refresh this view"
def update_local_type(self):
if self.modified:
for vtable in self.vtables.values():
vtable.update_local_type()
udt_data = idaapi.udt_type_data_t()
tinfo = idaapi.tinfo_t()
self.tinfo.get_udt_details(udt_data)
tinfo.create_udt(udt_data, idaapi.BTF_STRUCT)
tinfo.set_numbered_type(idaapi.cvar.idati, self.ordinal, idaapi.NTF_REPLACE, self.name)
self.modified = False
def set_first_argument_type(self, class_name):
if 0 in self.vtables:
self.vtables[0].set_first_argument_type(class_name)
def has_function(self, regexp):
for vtable in self.vtables.values():
if filter(lambda func: regexp.indexIn(func.name) >= 0, vtable.virtual_functions):
return True
return False
def data(self, column):
if column == 0:
return self.name
def setData(self, column, value):
if column == 0:
if idaapi.isident(value) and self.name != value:
self.name = value
self.modified = True
return True
return False
def font(self, column):
if self.modified:
return QtGui.QFont("Consolas", 12, QtGui.QFont.Bold, italic=True)
return QtGui.QFont("Consolas", 12, QtGui.QFont.Bold)
@property
def color(self):
# return QtGui.QBrush(QtGui.QColor("#ffb3ff")):
return QtGui.QColor("#80ced6")
@property
def tooltip(self):
return None
def flags(self, column):
if column == 0:
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEditable | QtCore.Qt.ItemIsEnabled
return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled
@property
def children(self):
return self.vtables.values()
@property
def tinfo(self):
tinfo = idaapi.tinfo_t()
tinfo.get_numbered_type(idaapi.cvar.idati, self.ordinal)
return tinfo
def open_function(self):
pass
def __repr__(self):
return self.name + " ^_^ " + str(self.vtables)
class TreeItem:
def __init__(self, item, row, parent):
self.item = item
self.parent = parent
self.row = row
self.children = []
def __repr__(self):
return str(self.item)
class TreeModel(QtCore.QAbstractItemModel):
# TODO: Add higlighting if eip in function, consider setting breakpoints
refreshed = QtCore.Signal()
def __init__(self):
super(TreeModel, self).__init__()
self.tree_data = []
self.headers = ["Name", "Declaration", "Address"]
self.init()
# import pydevd
# pydevd.settrace("localhost", port=12345, stdoutToServer=True, stderrToServer=True)
def init(self):
idaapi.show_wait_box("Looking for classes...")
all_virtual_functions.clear()
all_virtual_tables.clear()
classes = []
for ordinal in xrange(1, idaapi.get_ordinal_qty(idaapi.cvar.idati)):
result = Class.create_class(ordinal)
if result:
classes.append(result)
for class_row, class_ in enumerate(classes):
class_item = TreeItem(class_, class_row, None)
for vtable_row, vtable in class_.vtables.iteritems():
vtable_item = TreeItem(vtable, vtable_row, class_item)
vtable_item.children = [TreeItem(function, 0, vtable_item) for function in vtable.virtual_functions]
class_item.children.append(vtable_item)
self.tree_data.append(class_item)
idaapi.hide_wait_box()
def flags(self, index):
if index.isValid():
return index.internalPointer().item.flags(index.column())
def index(self, row, column, parent=QtCore.QModelIndex()):
if parent.isValid():
node = parent.internalPointer()
return self.createIndex(row, column, node.children[row])
else:
return self.createIndex(row, column, self.tree_data[row])
def parent(self, index):
if index.isValid():
node = index.internalPointer()
if node.parent:
return self.createIndex(0, 0, node.parent)
return QtCore.QModelIndex()
def rowCount(self, index=QtCore.QModelIndex()):
if index.isValid():
node = index.internalPointer()
if node:
return len(node.children)
return len(self.tree_data)
def columnCount(self, index=QtCore.QModelIndex()):
return 3
def data(self, index, role=QtCore.Qt.DisplayRole):
if role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole:
node = index.internalPointer()
return node.item.data(index.column())
elif role == QtCore.Qt.FontRole:
return index.internalPointer().item.font(index.column())
elif role == QtCore.Qt.ToolTipRole:
return index.internalPointer().item.tooltip
elif role == QtCore.Qt.BackgroundRole:
return index.internalPointer().item.color
elif role == QtCore.Qt.ForegroundRole:
return QtGui.QBrush(QtGui.QColor("#191919"))
return None
def setData(self, index, value, role=QtCore.Qt.DisplayRole):
result = False
if role == QtCore.Qt.EditRole and value != "":
node = index.internalPointer()
result = node.item.setData(index.column(), str(value))
return result
def headerData(self, section, orientation, role):
if role == QtCore.Qt.DisplayRole and orientation == QtCore.Qt.Horizontal:
return self.headers[section]
def set_first_argument_type(self, indexes):
indexes = filter(lambda x: x.column() == 0, indexes)
class_name = indexes[0].internalPointer().item.class_name
if not class_name:
classes = [[x.item.name] for x in self.tree_data]
class_chooser = HexRaysPyTools.Forms.MyChoose(classes, "Select Class", [["Name", 25]])
idx = class_chooser.Show(True)
if idx != -1:
class_name = classes[idx][0]
if class_name:
for index in indexes:
index.internalPointer().item.set_first_argument_type(class_name)
def refresh(self):
self.tree_data = []
self.modelReset.emit()
self.init()
self.refreshed.emit()
def rollback(self):
for class_item in self.tree_data:
class_item.item.update_from_local_types()
self.dataChanged.emit(QtCore.QModelIndex(), QtCore.QModelIndex())
def commit(self):
for class_item in self.tree_data:
if class_item.item.modified:
class_item.item.update_local_type()
def open_function(self, index):
if index.column() == 2:
index.internalPointer().item.open_function()
class ProxyModel(QtGui.QSortFilterProxyModel):
def __init__(self):
super(ProxyModel, self).__init__()
self.filter_by_function = False
def set_regexp_filter(self, regexp):
if regexp and regexp[0] == '!':
self.filter_by_function = True
self.setFilterRegExp(regexp[1:])
else:
self.filter_by_function = False
self.setFilterRegExp(regexp)
def filterAcceptsRow(self, row, parent):
if not parent.isValid() and self.filterRegExp():
if self.filter_by_function:
return self.sourceModel().tree_data[row].item.has_function(self.filterRegExp())
return self.filterRegExp().indexIn(self.sourceModel().tree_data[row].item.class_name) >= 0
return True

View File

@@ -0,0 +1,58 @@
import idaapi
import idc
EA64 = idc.__EA64__
EA_SIZE = 8 if EA64 else 4
COT_ARITHMETIC = (idaapi.cot_num, idaapi.cot_fnum, idaapi.cot_add, idaapi.cot_fadd, idaapi.cot_sub, idaapi.cot_fsub,
idaapi.cot_mul, idaapi.cot_fmul, idaapi.cot_fdiv)
VOID_TINFO = None
PVOID_TINFO = idaapi.tinfo_t()
CONST_VOID_TINFO = None
CONST_PVOID_TINFO = idaapi.tinfo_t()
CHAR_TINFO = None
PCHAR_TINFO = idaapi.tinfo_t()
CONST_PCHAR_TINFO = idaapi.tinfo_t()
BYTE_TINFO = None
PBYTE_TINFO = None
WORD_TINFO = None
PWORD_TINFO = idaapi.tinfo_t()
X_WORD_TINFO = None # DWORD for x32 and QWORD for x64
PX_WORD_TINFO = None
DUMMY_FUNC = None
LEGAL_TYPES = []
def init():
""" All tinfo should be reinitialized between session. Otherwise they could have wrong type """
global VOID_TINFO, PVOID_TINFO, CONST_PVOID_TINFO, BYTE_TINFO, PBYTE_TINFO, LEGAL_TYPES, X_WORD_TINFO, \
PX_WORD_TINFO, DUMMY_FUNC, CONST_PCHAR_TINFO, CHAR_TINFO, PCHAR_TINFO, CONST_VOID_TINFO, \
WORD_TINFO, PWORD_TINFO
VOID_TINFO = idaapi.tinfo_t(idaapi.BT_VOID)
PVOID_TINFO.create_ptr(VOID_TINFO)
CONST_VOID_TINFO = idaapi.tinfo_t(idaapi.BT_VOID | idaapi.BTM_CONST)
CONST_PVOID_TINFO.create_ptr(idaapi.tinfo_t(idaapi.BT_VOID | idaapi.BTM_CONST))
CONST_PCHAR_TINFO.create_ptr(idaapi.tinfo_t(idaapi.BTF_CHAR | idaapi.BTM_CONST))
CHAR_TINFO = idaapi.tinfo_t(idaapi.BTF_CHAR)
PCHAR_TINFO.create_ptr(idaapi.tinfo_t(idaapi.BTF_CHAR))
BYTE_TINFO = idaapi.tinfo_t(idaapi.BTF_BYTE)
PBYTE_TINFO = idaapi.dummy_ptrtype(1, False)
X_WORD_TINFO = idaapi.get_unk_type(EA_SIZE)
PX_WORD_TINFO = idaapi.dummy_ptrtype(EA_SIZE, False)
WORD_TINFO = idaapi.tinfo_t(idaapi.BT_UNK_WORD)
PWORD_TINFO.create_ptr(idaapi.tinfo_t(idaapi.BT_UNK_WORD))
func_data = idaapi.func_type_data_t()
func_data.rettype = PVOID_TINFO
func_data.cc = idaapi.CM_CC_UNKNOWN
DUMMY_FUNC = idaapi.tinfo_t()
DUMMY_FUNC.create_func(func_data, idaapi.BT_FUNC)
LEGAL_TYPES = [PVOID_TINFO, PX_WORD_TINFO, PWORD_TINFO, PBYTE_TINFO, X_WORD_TINFO]

View File

@@ -0,0 +1,312 @@
import collections
import logging
import idaapi
import idc
import HexRaysPyTools.Core.Cache as Cache
import HexRaysPyTools.Core.Const as Const
import HexRaysPyTools.Settings as Settings
logger = logging.getLogger(__name__)
def is_imported_ea(ea):
if idc.get_segm_name(ea) == ".plt":
return True
return ea in Cache.imported_ea
def is_code_ea(ea):
flags = idaapi.getFlags(ea) # flags_t
return idaapi.isCode(flags)
def is_rw_ea(ea):
seg = idaapi.getseg(ea)
return seg.perm & idaapi.SEGPERM_WRITE and seg.perm & idaapi.SEGPERM_READ
def get_virtual_func_address(name, tinfo=None, offset=None):
"""
:param name: method name
:param tinfo: class tinfo
:param offset: virtual table offset
:return: address of the method
"""
address = idc.LocByName(name)
if address != idaapi.BADADDR:
return address
address = Cache.demangled_names.get(name, idaapi.BADADDR)
if address != idaapi.BADADDR:
return address + idaapi.get_imagebase()
if tinfo is None or offset is None:
return
offset *= 8
udt_member = idaapi.udt_member_t()
while tinfo.is_struct():
address = Cache.demangled_names.get(tinfo.dstr() + '::' + name, idaapi.BADADDR)
if address != idaapi.BADADDR:
return address + idaapi.get_imagebase()
udt_member.offset = offset
tinfo.find_udt_member(idaapi.STRMEM_OFFSET, udt_member)
tinfo = udt_member.type
offset = offset - udt_member.offset
def get_func_argument_info(function, expression):
"""
Function is cexpr with opname == 'cot_call', expression is any son. Returns index of argument and it's type
:param function: idaapi.cexpr_t
:param expression: idaapi.cexpr_t
:return: (int, idaapi.tinfo_t)
"""
for idx, argument in enumerate(function.a):
if expression == argument.cexpr:
func_tinfo = function.x.type
if idx < func_tinfo.get_nargs():
return idx, func_tinfo.get_nth_arg(idx)
return idx, None
print "[ERROR] Wrong usage of 'Helper.get_func_argument_info()'"
def set_func_argument(func_tinfo, index, arg_tinfo):
func_data = idaapi.func_type_data_t()
func_tinfo.get_func_details(func_data)
func_data[index].type = arg_tinfo
func_tinfo.create_func(func_data)
def set_funcptr_argument(funcptr_tinfo, index, arg_tinfo):
func_tinfo = funcptr_tinfo.get_pointed_object()
set_func_argument(func_tinfo, index, arg_tinfo)
funcptr_tinfo.create_ptr(func_tinfo)
def get_nice_pointed_object(tinfo):
"""
Returns nice pointer name (if exist) or None.
For example if tinfo is PKSPIN_LOCK which is typedef of unsigned int *, then if in local types exist KSPIN_LOCK with
type unsigned int, this function returns KSPIN_LOCK
"""
try:
name = tinfo.dstr()
if name[0] == 'P':
pointed_tinfo = idaapi.tinfo_t()
if pointed_tinfo.get_named_type(idaapi.cvar.idati, name[1:]):
if tinfo.get_pointed_object().equals_to(pointed_tinfo):
return pointed_tinfo
except TypeError:
pass
def get_fields_at_offset(tinfo, offset):
"""
Given tinfo and offset of the structure or union, returns list of all tinfo at that offset.
This function helps to find appropriate structures by type of the offset
"""
result = []
if offset == 0:
result.append(tinfo)
udt_data = idaapi.udt_type_data_t()
tinfo.get_udt_details(udt_data)
udt_member = idaapi.udt_member_t()
udt_member.offset = offset * 8
idx = tinfo.find_udt_member(idaapi.STRMEM_OFFSET, udt_member)
if idx != -1:
while idx < tinfo.get_udt_nmembers() and udt_data[idx].offset <= offset * 8:
udt_member = udt_data[idx]
if udt_member.offset == offset * 8:
if udt_member.type.is_ptr():
result.append(idaapi.get_unk_type(Const.EA_SIZE))
result.append(udt_member.type)
result.append(idaapi.dummy_ptrtype(Const.EA_SIZE, False))
elif not udt_member.type.is_udt():
result.append(udt_member.type)
if udt_member.type.is_array():
if (offset - udt_member.offset / 8) % udt_member.type.get_array_element().get_size() == 0:
result.append(udt_member.type.get_array_element())
elif udt_member.type.is_udt():
result.extend(get_fields_at_offset(udt_member.type, offset - udt_member.offset / 8))
idx += 1
return result
def is_legal_type(tinfo):
tinfo.clr_const()
if tinfo.is_ptr() and tinfo.get_pointed_object().is_forward_decl():
return tinfo.get_pointed_object().get_size() == idaapi.BADSIZE
return Settings.SCAN_ANY_TYPE or bool(filter(lambda x: x.equals_to(tinfo), Const.LEGAL_TYPES))
def search_duplicate_fields(udt_data):
# Returns list of lists with duplicate fields
default_dict = collections.defaultdict(list)
for idx, udt_member in enumerate(udt_data):
default_dict[udt_member.name].append(idx)
return [indices for indices in default_dict.values() if len(indices) > 1]
def get_member_name(tinfo, offset):
udt_member = idaapi.udt_member_t()
udt_member.offset = offset * 8
tinfo.find_udt_member(idaapi.STRMEM_OFFSET, udt_member)
return udt_member.name
def change_member_name(struct_name, offset, name):
return idc.set_member_name(idc.get_struc_id(struct_name), offset, name)
def import_structure(name, tinfo):
cdecl_typedef = idaapi.print_tinfo(None, 4, 5, idaapi.PRTYPE_MULTI | idaapi.PRTYPE_TYPE | idaapi.PRTYPE_SEMI,
tinfo, name, None)
if idc.parse_decl(cdecl_typedef, idaapi.PT_TYP) is None:
return 0
previous_ordinal = idaapi.get_type_ordinal(idaapi.cvar.idati, name)
if previous_ordinal:
idaapi.del_numbered_type(idaapi.cvar.idati, previous_ordinal)
ordinal = idaapi.idc_set_local_type(previous_ordinal, cdecl_typedef, idaapi.PT_TYP)
else:
ordinal = idaapi.idc_set_local_type(-1, cdecl_typedef, idaapi.PT_TYP)
return ordinal
def get_funcs_calling_address(ea):
""" Returns all addresses of functions which make call to a function at `ea`"""
xref_ea = idaapi.get_first_cref_to(ea)
xrefs = set()
while xref_ea != idaapi.BADADDR:
xref_func_ea = idc.GetFunctionAttr(xref_ea, idc.FUNCATTR_START)
if xref_func_ea != idaapi.BADADDR:
xrefs.add(xref_func_ea)
else:
print "[Warning] Function not found at 0x{0:08X}".format(xref_ea)
xref_ea = idaapi.get_next_cref_to(ea, xref_ea)
return xrefs
class FunctionTouchVisitor(idaapi.ctree_parentee_t):
def __init__(self, cfunc):
super(FunctionTouchVisitor, self).__init__()
self.functions = set()
self.cfunc = cfunc
def visit_expr(self, expression):
if expression.op == idaapi.cot_call:
self.functions.add(expression.x.obj_ea)
return 0
def touch_all(self):
diff = self.functions.difference(Cache.touched_functions)
for address in diff:
if is_imported_ea(address):
continue
try:
cfunc = idaapi.decompile(address)
if cfunc:
FunctionTouchVisitor(cfunc).process()
except idaapi.DecompilationFailure:
logger.warn("IDA failed to decompile function at {}".format(to_hex(address)))
Cache.touched_functions.add(address)
idaapi.decompile(self.cfunc.entry_ea)
def process(self):
if self.cfunc.entry_ea not in Cache.touched_functions:
Cache.touched_functions.add(self.cfunc.entry_ea)
self.apply_to(self.cfunc.body, None)
self.touch_all()
return True
return False
def to_hex(ea):
""" Formats address so it could be double clicked at console """
if Const.EA64:
return "0x{:016X}".format(ea)
return "0x{:08X}".format(ea)
def to_nice_str(ea):
""" Shows address as function name + offset """
func_start_ea = idc.get_func_attr(ea, idc.FUNCATTR_START)
func_name = idc.Name(func_start_ea)
offset = ea - func_start_ea
return "{}+0x{:X}".format(func_name, offset)
def save_long_str_to_idb(array_name, value):
""" Overwrites old array completely in process """
id = idc.get_array_id(array_name)
if id != -1:
idc.delete_array(id)
id = idc.create_array(array_name)
r = []
for idx in xrange(len(value) / 1024 + 1):
s = value[idx * 1024: (idx + 1) * 1024]
r.append(s)
idc.set_array_string(id, idx, s)
def load_long_str_from_idb(array_name):
id = idc.get_array_id(array_name)
if id == -1:
return None
max_idx = idc.get_last_index(idc.AR_STR, id)
result = [idc.get_array_element(idc.AR_STR, id, idx) for idx in xrange(max_idx + 1)]
return "".join(result)
# ======================================================================
# Functions that extends IDA Pro capabilities
# ======================================================================
def _find_asm_address(self, cexpr):
""" Returns most close virtual address corresponding to cexpr """
ea = cexpr.ea
if ea != idaapi.BADADDR:
return ea
for p in reversed(self.parents):
if p.ea != idaapi.BADADDR:
return p.ea
def my_cexpr_t(*args, **kwargs):
""" Replacement of bugged cexpr_t() function """
if len(args) == 0:
return idaapi.cexpr_t()
if len(args) != 1:
raise NotImplementedError
cexpr = idaapi.cexpr_t()
cexpr.thisown = False
if type(args[0]) == idaapi.cexpr_t:
cexpr.assign(args[0])
else:
op = args[0]
cexpr._set_op(op)
if 'x' in kwargs:
cexpr._set_x(kwargs['x'])
if 'y' in kwargs:
cexpr._set_y(kwargs['y'])
if 'z' in kwargs:
cexpr._set_z(kwargs['z'])
return cexpr
def extend_ida():
idaapi.ctree_parentee_t._find_asm_address = _find_asm_address

View File

@@ -0,0 +1,255 @@
import re
import logging
import idaapi
import idc
import Helper
logger = logging.getLogger(__name__)
def parse_lvar_comment(lvar):
if lvar.type().is_ptr():
m = re.search('```(.+)```', lvar.cmt)
if m:
structure_name, offset = m.group(1).split('+')
offset = int(offset)
parent_tinfo = idaapi.tinfo_t()
if parent_tinfo.get_named_type(idaapi.cvar.idati, structure_name) and parent_tinfo.get_size() > offset:
member_name = dict(find_deep_members(parent_tinfo, lvar.type().get_pointed_object())).get(offset, None)
if member_name:
return NegativeLocalInfo(lvar.type().get_pointed_object(), parent_tinfo, offset, member_name)
return None
def find_deep_members(parent_tinfo, target_tinfo):
udt_data = idaapi.udt_type_data_t()
parent_tinfo.get_udt_details(udt_data)
result = []
for udt_member in udt_data:
if udt_member.type.equals_to(target_tinfo):
result.append((udt_member.offset / 8, udt_member.name))
elif udt_member.type.is_udt():
for offset, name in find_deep_members(udt_member.type, target_tinfo):
final_name = udt_member.name + '.' + name if udt_member.name else name
result.append((udt_member.offset / 8 + offset, final_name))
return result
class NegativeLocalInfo:
def __init__(self, tinfo, parent_tinfo, offset, member_name):
self.tinfo = tinfo
self.size = tinfo.get_size() if tinfo.is_udt else 0
self.parent_tinfo = parent_tinfo
self.offset = offset
self.member_name = member_name
def __repr__(self):
return "Type - {0}, parent type - {1}, offset - {2}, member_name - {3}".format(
self.tinfo.dstr(),
self.parent_tinfo.dstr(),
self.offset,
self.member_name
)
class NegativeLocalCandidate:
def __init__(self, tinfo, offset):
"""
Tinfo - type of the structure tha local variable points to. So it's stripped from pointer. Offset - is first
found offset that points outside of the structure.
:param tinfo: idaapi.tinfo_t
:param offset: int
"""
self.tinfo = tinfo
self.offsets = [offset]
def __repr__(self):
return self.tinfo.dstr() + ' ' + str(self.offsets)
def is_structure_offset(self, tinfo, offset):
# Checks if structure tinfo contains a member at given offset
# TODO:array checking
udt_member = idaapi.udt_member_t()
udt_member.offset = offset * 8
if offset >= 0 and tinfo.find_udt_member(idaapi.STRMEM_OFFSET, udt_member) != -1:
if udt_member.type.is_udt():
return self.is_structure_offset(udt_member.type, offset - udt_member.offset / 8)
return udt_member.offset == offset * 8
return False
def find_containing_structures(self, type_library):
"""
Given the type library creates a list of structures from this library, that contains this structure and
satisfy offset conditions.
:param type_library: idaapi.til_t
:returns: ordinal, offset, member_name, containing structure name
"""
min_offset = min(self.offsets)
min_offset = min_offset if min_offset < 0 else 0
max_offset = max(self.offsets)
max_offset = max_offset if max_offset > 0 else self.tinfo.get_size()
# TODO: Check if all offsets are legal
# Least acceptable size of the containing structure
min_struct_size = max_offset - min_offset
result = []
parent_tinfo = idaapi.tinfo_t()
target_tinfo = idaapi.tinfo_t()
if not target_tinfo.get_named_type(type_library, self.tinfo.dstr()):
print "[Warning] Such type doesn't exist in '{0}' library".format(type_library.name)
return result
for ordinal in xrange(1, idaapi.get_ordinal_qty(type_library)):
parent_tinfo.create_typedef(type_library, ordinal)
if parent_tinfo.get_size() >= min_struct_size:
for offset, name in find_deep_members(parent_tinfo, target_tinfo):
# print "[DEBUG] Found {0} at {1} in {2}".format(name, offset, parent_tinfo.dstr())
if offset + min_offset >= 0 and offset + max_offset <= parent_tinfo.get_size():
result.append((ordinal, offset, name, parent_tinfo.dstr()))
return result
class ReplaceVisitor(idaapi.ctree_parentee_t):
def __init__(self, negative_lvars):
super(ReplaceVisitor, self).__init__()
self.negative_lvars = negative_lvars
self.pvoid_tinfo = idaapi.tinfo_t(idaapi.BT_VOID)
self.pvoid_tinfo.create_ptr(self.pvoid_tinfo)
def visit_expr(self, expression):
if expression.op == idaapi.cot_add and expression.x.op == idaapi.cot_var and expression.y.op == idaapi.cot_num:
index = expression.x.v.idx
if index in self.negative_lvars:
offset = expression.y.numval()
if offset >= self.negative_lvars[index].size:
self.create_containing_record(expression, index, offset)
elif expression.op == idaapi.cot_sub and expression.x.op == idaapi.cot_var and expression.y.op == idaapi.cot_num:
index = expression.x.v.idx
if index in self.negative_lvars:
offset = -expression.y.n.value(idaapi.tinfo_t(idaapi.BT_INT))
self.create_containing_record(expression, index, offset)
return 0
def create_containing_record(self, expression, index, offset):
negative_lvar = self.negative_lvars[index]
logger.debug("Creating CONTAINING_RECORD macro, offset: {}, negative offset: {}, TYPE: {}".format(
negative_lvar.offset,
offset,
negative_lvar.parent_tinfo.dstr()
))
arg_address = idaapi.carg_t()
if expression.op == idaapi.cot_var:
arg_address.assign(expression)
else:
arg_address.assign(expression.x)
arg_type = idaapi.carg_t()
cexpr_helper = idaapi.create_helper(True, self.pvoid_tinfo, negative_lvar.parent_tinfo.dstr())
arg_type.assign(cexpr_helper)
arg_field = idaapi.carg_t()
cexpr_helper = idaapi.create_helper(
True,
self.pvoid_tinfo,
negative_lvar.member_name
)
arg_field.assign(cexpr_helper)
return_tinfo = idaapi.tinfo_t(negative_lvar.parent_tinfo)
return_tinfo.create_ptr(return_tinfo)
new_cexpr_call = idaapi.call_helper(return_tinfo, None, "CONTAINING_RECORD")
new_cexpr_call.a.push_back(arg_address)
new_cexpr_call.a.push_back(arg_type)
new_cexpr_call.a.push_back(arg_field)
new_cexpr_call.thisown = False
parent = reversed(self.parents).next().cexpr
diff = negative_lvar.offset + offset
if diff:
number = idaapi.make_num(diff)
number.thisown = False
new_cexpr_add = Helper.my_cexpr_t(idaapi.cot_add, x=new_cexpr_call, y=number)
new_cexpr_add.type = return_tinfo
if parent.op == idaapi.cot_ptr:
tmp_tinfo = idaapi.tinfo_t()
tmp_tinfo.create_ptr(parent.type)
new_cexpr_cast = Helper.my_cexpr_t(idaapi.cot_cast, x=new_cexpr_add)
new_cexpr_cast.thisown = False
new_cexpr_cast.type = tmp_tinfo
expression.assign(new_cexpr_cast)
else:
expression.assign(new_cexpr_add)
else:
if parent.op == idaapi.cot_ptr:
tmp_tinfo = idaapi.tinfo_t()
tmp_tinfo.create_ptr(parent.type)
new_cexpr_cast = Helper.my_cexpr_t(idaapi.cot_cast, x=new_cexpr_call)
new_cexpr_cast.thisown = False
new_cexpr_cast.type = tmp_tinfo
expression.assign(new_cexpr_cast)
else:
expression.assign(new_cexpr_call)
class SearchVisitor(idaapi.ctree_parentee_t):
def __init__(self, cfunc):
super(SearchVisitor, self).__init__()
self.cfunc = cfunc
self.result = {}
def visit_expr(self, expression):
if expression.op == idaapi.cot_call and expression.x.op == idaapi.cot_helper and len(expression.a) == 3:
if expression.x.helper == "CONTAINING_RECORD":
if expression.a[0].op == idaapi.cot_var:
idx = expression.a[0].v.idx
if expression.a[1].op == idaapi.cot_helper and expression.a[2].op == idaapi.cot_helper:
parent_name = expression.a[1].helper
member_name = expression.a[2].helper
parent_tinfo = idaapi.tinfo_t()
if not parent_tinfo.get_named_type(idaapi.cvar.idati, parent_name):
return 0
udt_data = idaapi.udt_type_data_t()
parent_tinfo.get_udt_details(udt_data)
udt_member = filter(lambda x: x.name == member_name, udt_data)
if udt_member:
tinfo = udt_member[0].type
self.result[idx] = NegativeLocalInfo(
tinfo,
parent_tinfo,
udt_member[0].offset / 8,
member_name
)
return 1
return 0
class AnalyseVisitor(idaapi.ctree_parentee_t):
def __init__(self, candidates, potential_negatives):
super(AnalyseVisitor, self).__init__()
self.candidates = candidates
self.potential_negatives = potential_negatives
self.potential_negatives.clear()
def visit_expr(self, expression):
if expression.op == idaapi.cot_add and expression.y.op == idaapi.cot_num:
if expression.x.op == idaapi.cot_var and expression.x.v.idx in self.candidates:
idx = expression.x.v.idx
number = expression.y.numval()
if self.candidates[idx].get_size() <= number:
if idx in self.potential_negatives:
self.potential_negatives[idx].offsets.append(number)
else:
self.potential_negatives[idx] = NegativeLocalCandidate(self.candidates[idx], number)
elif expression.op == idaapi.cot_sub and expression.y.op == idaapi.cot_num:
if expression.x.op == idaapi.cot_var and expression.x.v.idx in self.candidates:
idx = expression.x.v.idx
number = -expression.y.numval()
if idx in self.potential_negatives:
self.potential_negatives[idx].offsets.append(number)
else:
self.potential_negatives[idx] = NegativeLocalCandidate(self.candidates[idx], number)
return 0

View File

@@ -0,0 +1,120 @@
import idaapi
import idc
def inverse_if_condition(cif):
# cexpr_t has become broken but fortunately still exist `assing` method which copies one expr into another
cit_if_condition = cif.expr
tmp_cexpr = idaapi.cexpr_t()
tmp_cexpr.assign(cit_if_condition)
new_if_condition = idaapi.lnot(tmp_cexpr)
cif.expr.swap(new_if_condition)
del cit_if_condition
def inverse_if(cif):
inverse_if_condition(cif)
idaapi.qswap(cif.ithen, cif.ielse)
class InversionInfo(object):
ARRAY_NAME = "$HexRaysPyTools:IfThenElse:"
def __init__(self, func_ea):
self.__name = InversionInfo.ARRAY_NAME + hex(int(func_ea))
self.__id = idc.GetArrayId(self.__name)
def get_inverted(self):
if self.__id != -1:
array = idc.GetArrayElement(idc.AR_STR, self.__id, 0)
return set(map(int, array.split()))
return set()
def switch_inverted(self, address):
if self.__id == -1:
self.__id = idc.CreateArray(self.__name)
idc.SetArrayString(self.__id, 0, str(address))
else:
inverted = self.get_inverted()
try:
inverted.remove(address)
if not inverted:
idc.DeleteArray(self.__id)
except KeyError:
inverted.add(address)
idc.SetArrayString(self.__id, 0, " ".join(map(str, inverted)))
class SpaghettiVisitor(idaapi.ctree_parentee_t):
def __init__(self):
super(SpaghettiVisitor, self).__init__()
def visit_insn(self, instruction):
if instruction.op != idaapi.cit_block:
return 0
while True:
cblock = instruction.cblock
size = cblock.size()
# Find block that has "If" and "return" as last 2 statements
if size < 2:
break
if cblock.at(size - 2).op != idaapi.cit_if:
break
cif = cblock.at(size - 2).cif
if cblock.back().op != idaapi.cit_return or cif.ielse:
break
cit_then = cif.ithen
# Skip if only one (not "if") statement in "then" branch
if cit_then.cblock.size() == 1 and cit_then.cblock.front().op != idaapi.cit_if:
return 0
inverse_if_condition(cif)
# Take return from list of statements and later put it back
cit_return = idaapi.cinsn_t()
cit_return.assign(instruction.cblock.back())
cit_return.thisown = False
instruction.cblock.pop_back()
# Fill main block with statements from "Then" branch
while cit_then.cblock:
instruction.cblock.push_back(cit_then.cblock.front())
cit_then.cblock.pop_front()
# Put back main return if there's no another return or "GOTO" already
if instruction.cblock.back().op not in (idaapi.cit_return, idaapi.cit_goto):
new_return = idaapi.cinsn_t()
new_return.thisown = False
new_return.assign(cit_return)
instruction.cblock.push_back(new_return)
# Put return into "Then" branch
cit_then.cblock.push_back(cit_return)
return 0
class SwapThenElseVisitor(idaapi.ctree_parentee_t):
def __init__(self, func_ea):
super(SwapThenElseVisitor, self).__init__()
self.__inversion_info = InversionInfo(func_ea)
self.__inverted = self.__inversion_info.get_inverted()
def visit_insn(self, insn):
if insn.op != idaapi.cit_if or insn.cif.ielse is None:
return 0
if insn.ea in self.__inverted:
inverse_if(insn.cif)
return 0
def apply_to(self, *args):
if self.__inverted:
super(SwapThenElseVisitor, self).apply_to(*args)

View File

@@ -0,0 +1,181 @@
import time
import logging
from collections import namedtuple
import json
import idaapi
import Helper
import HexRaysPyTools.Settings as Settings
logger = logging.getLogger(__name__)
XrefInfo = namedtuple('XrefInfo', ['func_ea', 'offset', 'line', 'type'])
def singleton(cls):
instances = {}
def get_instance():
if cls not in instances:
instances[cls] = cls()
return instances[cls]
return get_instance
@singleton
class XrefStorage(object):
ARRAY_NAME = "$HexRaysPyTools:XrefStorage"
def __init__(self):
self.storage = None
def open(self):
if not Settings.STORE_XREFS:
self.storage = {}
return
result = Helper.load_long_str_from_idb(self.ARRAY_NAME)
if result:
try:
self.storage = json.loads(result, object_hook=self.json_keys_to_str)
return
except ValueError:
logger.error("Failed to read previous info about Xrefs. Try Ctrl+F5 to cache data")
self.storage = {}
def close(self):
self.save()
self.storage = None
def save(self):
if not Settings.STORE_XREFS:
return
if self.storage:
Helper.save_long_str_to_idb(self.ARRAY_NAME, json.dumps(self.storage))
def update_structure_info(self, ordinal, function_address, data):
""" Accepts data in form dictionary {structure offset -> list(offsets within function with field appealing) """
if ordinal not in self.storage:
self.storage[ordinal] = {}
image_base = idaapi.get_imagebase()
function_offset = function_address - image_base
self.storage[ordinal][function_offset] = data
def get_structure_info(self, ordinal, struct_offset):
""" By given ordinal and offset within a structure returns dictionary {func_address -> list(offsets)} """
result = []
if ordinal not in self.storage:
return result
for func_offset, data in self.storage[ordinal].items():
if struct_offset in data:
func_ea = func_offset + idaapi.get_imagebase()
for xref_info in data[struct_offset]:
offset, line, usage_type = xref_info
result.append(XrefInfo(func_ea, offset, line, usage_type))
return result
@staticmethod
def json_keys_to_str(x):
if isinstance(x, dict):
return {int(k): v for k, v in x.items()}
return x
def __len__(self):
return len(str(self.storage))
class StructXrefVisitor(idaapi.ctree_parentee_t):
def __init__(self, cfunc):
super(StructXrefVisitor, self).__init__()
self.__cfunc = cfunc
self.__image_base = idaapi.get_imagebase()
self.__function_address = cfunc.entry_ea
self.__result = {}
self.__storage = XrefStorage()
def visit_expr(self, expression):
# Checks if expression is reference by pointer or by value
if expression.op == idaapi.cot_memptr:
struct_type = expression.x.type.get_pointed_object()
elif expression.op == idaapi.cot_memref:
struct_type = expression.x.type
else:
return 0
# Getting information about structure, field offset, address and one line corresponding to code
ordinal = struct_type.get_ordinal()
if ordinal == 0:
t = idaapi.tinfo_t()
struct_name = struct_type.dstr().split()[-1] # Get rid of `struct` prefix or something else
t.get_named_type(idaapi.cvar.idati, struct_name)
ordinal = t.get_ordinal()
field_offset = expression.m
ea = self.__find_ref_address(expression)
usage_type = self.__get_type(expression)
if ea == idaapi.BADADDR or not ordinal:
logger.warning("Failed to parse at address {0}, ordinal - {1}, type - {2}".format(
Helper.to_hex(ea), ordinal, struct_type.dstr()
))
one_line = self.__get_line()
occurrence_offset = ea - self.__function_address
xref_info = (occurrence_offset, one_line, usage_type)
# Saving results
if ordinal not in self.__result:
self.__result[ordinal] = {field_offset: [xref_info]}
elif field_offset not in self.__result[ordinal]:
self.__result[ordinal][field_offset] = [xref_info]
else:
self.__result[ordinal][field_offset].append(xref_info)
return 0
def process(self):
t = time.time()
self.apply_to(self.__cfunc.body, None)
for ordinal, data in self.__result.items():
self.__storage.update_structure_info(ordinal, self.__function_address, data)
storage_mb_size = len(self.__storage) * 1.0 / 1024 ** 2
logger.debug("Xref processing: %f seconds passed, storage size - %.2f MB ", (time.time() - t), storage_mb_size)
def __find_ref_address(self, cexpr):
""" Returns most close virtual address corresponding to cexpr """
ea = cexpr.ea
if ea != idaapi.BADADDR:
return ea
for p in reversed(self.parents):
if p.ea != idaapi.BADADDR:
return p.ea
def __get_type(self, cexpr):
""" Returns one of the following types: 'R' - read value, 'W' - write value, 'A' - function argument"""
child = cexpr
for p in reversed(self.parents):
assert p, "Failed to get type at " + Helper.to_hex(self.__function_address)
if p.cexpr.op == idaapi.cot_call:
return 'Arg'
if not p.is_expr():
return 'R'
if p.cexpr.op == idaapi.cot_asg:
if p.cexpr.x == child:
return 'W'
return 'R'
child = p.cexpr
def __get_line(self):
for p in reversed(self.parents):
if not p.is_expr():
return idaapi.tag_remove(p.print1(self.__cfunc))
AssertionError("Parent instruction is not found")

View File

@@ -0,0 +1,184 @@
import idaapi
import idc
class LocalType:
def __init__(self, name, members_ordinals, hint, is_selected=False, is_typedef=False, is_enum=False, is_union=False):
self.name = name
self.members_ordinals = members_ordinals
self.hint = hint
self.is_selected = is_selected
self.is_typedef = is_typedef
self.is_enum = is_enum
self.is_union = is_union
def __call__(self):
return self.name, self.members_ordinals
def __str__(self):
return "<{0}, {1}>".format(self.name, self.members_ordinals)
def __repr__(self):
return self.__str__()
@property
def name_and_color(self):
if self.is_selected:
return self.name, 0x0000FF
elif self.is_typedef:
return self.name, 0x99FFFF
elif self.is_enum:
return self.name, 0x33FF33
elif self.is_union:
return self.name, 0xCCCC00
return self.name, 0xffdd99
class StructureGraph:
# TODO:Enum types display
def __init__(self, ordinal_list=None):
self.ordinal_list = ordinal_list if ordinal_list else xrange(1, idc.GetMaxLocalType())
self.local_types = {}
self.edges = []
self.final_edges = []
self.visited_downward = []
self.visited_upward = []
self.downward_edges = {}
self.upward_edges = {}
self.initialize_nodes()
self.calculate_edges()
def change_selected(self, selected):
self.visited_downward = []
self.visited_upward = []
self.final_edges = []
for ordinal in self.ordinal_list:
self.local_types[ordinal].is_selected = False
self.ordinal_list = set(self.local_types).intersection(selected)
for ordinal in self.ordinal_list:
self.local_types[ordinal].is_selected = True
@staticmethod
def get_ordinal(tinfo):
while tinfo.is_ptr() or tinfo.is_array():
tinfo.remove_ptr_or_array()
if tinfo.is_udt():
return tinfo.get_ordinal()
elif tinfo.is_enum():
return tinfo.get_ordinal()
elif tinfo.is_typeref():
typeref_ordinal = tinfo.get_ordinal()
if typeref_ordinal:
typeref_tinfo = StructureGraph.get_tinfo_by_ordinal(typeref_ordinal)
if typeref_tinfo.is_typeref() or typeref_tinfo.is_udt() or typeref_tinfo.is_ptr():
return typeref_ordinal
else:
return 0
@staticmethod
def get_members_ordinals(tinfo):
ordinals = []
if tinfo.is_udt():
udt_data = idaapi.udt_type_data_t()
tinfo.get_udt_details(udt_data)
for udt_member in udt_data:
ordinal = StructureGraph.get_ordinal(udt_member.type)
if ordinal:
ordinals.append(ordinal)
return ordinals
@staticmethod
def get_tinfo_by_ordinal(ordinal):
local_typestring = idc.GetLocalTinfo(ordinal)
if local_typestring:
p_type, fields = local_typestring
local_tinfo = idaapi.tinfo_t()
local_tinfo.deserialize(idaapi.cvar.idati, p_type, fields)
return local_tinfo
return None
def initialize_nodes(self):
for ordinal in xrange(1, idc.GetMaxLocalType()):
# if ordinal == 15:
# import pydevd
# pydevd.settrace("localhost", port=12345, stdoutToServer=True, stderrToServer=True)
local_tinfo = StructureGraph.get_tinfo_by_ordinal(ordinal)
if not local_tinfo:
return
name = idc.GetLocalTypeName(ordinal)
if local_tinfo.is_typeref():
typeref_ordinal = local_tinfo.get_ordinal()
members_ordinals = []
if typeref_ordinal:
typeref_tinfo = StructureGraph.get_tinfo_by_ordinal(typeref_ordinal)
if typeref_tinfo.is_typeref() or typeref_tinfo.is_udt() or typeref_tinfo.is_ptr():
members_ordinals = [typeref_ordinal]
cdecl_typedef = idaapi.print_tinfo(None, 4, 5, 0x3, local_tinfo, None, None)
self.local_types[ordinal] = LocalType(name, members_ordinals, cdecl_typedef, is_typedef=True)
elif local_tinfo.is_udt():
# udt_data = idaapi.udt_type_data_t()
# local_tinfo.get_udt_details(udt_data)
members_ordinals = StructureGraph.get_members_ordinals(local_tinfo)
cdecl_typedef = idaapi.print_tinfo(None, 4, 5, 0x1, local_tinfo, None, None)
self.local_types[ordinal] = LocalType(name, members_ordinals, cdecl_typedef, is_union=local_tinfo.is_union())
elif local_tinfo.is_ptr():
typeref_ordinal = StructureGraph.get_ordinal(local_tinfo)
members_ordinals = [typeref_ordinal] if typeref_ordinal else []
cdecl_typedef = idaapi.print_tinfo(None, 4, 5, 0x2, local_tinfo, None, None)
self.local_types[ordinal] = LocalType(
name,
members_ordinals,
cdecl_typedef + ' *',
is_typedef=True
)
elif local_tinfo.is_enum():
cdecl_typedef = idaapi.print_tinfo(None, 4, 5, 0x21, local_tinfo, None, None)
self.local_types[ordinal] = LocalType(name, [], cdecl_typedef, is_enum=True)
self.ordinal_list = set(self.ordinal_list).intersection(self.local_types)
for ordinal in self.ordinal_list:
self.local_types[ordinal].is_selected = True
def calculate_edges(self):
for first in self.local_types.keys():
for second in self.local_types[first].members_ordinals:
self.edges.append((first, second))
self.downward_edges = {key: [] for key in self.local_types.keys()}
self.upward_edges = {key: [] for key in self.local_types.keys()}
for key, value in self.edges:
self.downward_edges[key].append(value)
self.upward_edges[value].append(key)
def generate_final_edges_down(self, node):
if node not in self.visited_downward:
self.visited_downward.append(node)
else:
return
for next_node in self.downward_edges[node]:
self.final_edges.append((node, next_node))
for next_node in self.downward_edges[node]:
self.generate_final_edges_down(next_node)
def generate_final_edges_up(self, node):
if node not in self.visited_upward:
self.visited_upward.append(node)
else:
return
for next_node in self.upward_edges[node]:
self.final_edges.append((next_node, node))
for next_node in self.upward_edges[node]:
self.generate_final_edges_up(next_node)
def get_nodes(self):
for ordinal in self.ordinal_list:
if ordinal in self.local_types:
self.generate_final_edges_down(ordinal)
self.generate_final_edges_up(ordinal)
return set([node for nodes in self.final_edges for node in nodes])
def get_edges(self):
return self.final_edges

View File

@@ -0,0 +1,864 @@
import bisect
import idc
import idaapi
import re
import itertools
# import PySide.QtCore as QtCore
# import PySide.QtGui as QtGui
import Cache
from HexRaysPyTools.Cute import *
import Const
import Helper
import VariableScanner
import HexRaysPyTools.Api as Api
from HexRaysPyTools.Forms import MyChoose
SCORE_TABLE = dict((v, k) for k, v in enumerate(
['unsigned __int8 *', 'unsigned __int8', '__int8 *', '__int8', '_BYTE', '_BYTE *', '_BYTE **', 'const char **',
'signed __int16', 'unsigned __int16', '__int16', 'signed __int16 *', 'unsigned __int16 *', '__int16 *',
'_WORD *', '_WORD **',
'signed int*', 'signed int', 'unsigned int *', 'unsigned int', 'int **', 'char **', 'int *', 'void **',
'int', '_DWORD *', 'char', '_DWORD', '_WORD', 'void *', 'char *']
))
def parse_vtable_name(address):
name = idaapi.get_short_name(address)
if idaapi.is_valid_typename(name):
if name[0:3] == 'off':
# off_XXXXXXXX case
return "Vtable" + name[3:], False
elif "table" in name:
return name, True
print "[Warning] Weird virtual table name -", name
return "Vtable_" + name
else:
# Attempt to make nice and valid name from demangled RTTI name
try:
name = re.sub("^const ", "", name)
sliced_names = name.split("::")
name, for_part = "_for_".join(sliced_names[:-1]), sliced_names[-1]
print name, for_part
templates = re.search("<(.*)>", name)
if templates:
templates = templates.group(1)
name = re.sub("<.*>", "", name)
templates = re.sub("[^a-zA-Z0-9_*]", "_", templates)
templates = re.sub("\*", "PTR", templates)
name += '_' + templates
for_part = re.search("\{for `(.*)'\}", for_part)
if for_part:
for_part = for_part.group(1)
name += '_' + for_part
return 'Vtable_' + name, True
except (AttributeError, IndexError):
print "[Warning] Unable to parse virtual table name - "
return "Vtable_{0:X}".format(address), False
class AbstractMember:
def __init__(self, offset, scanned_variable, origin):
"""
Offset is the very very base of the structure
Origin is from which offset of the base structure the variable have been scanned
scanned_variable - information about context in which this variable was scanned. This is necessary for final
applying type after packing or finalizing structure.
:param offset: int
:param scanned_variable: ScannedVariable
:param origin: int
"""
self.offset = offset
self.origin = origin
self.enabled = True
self.is_array = False
self.scanned_variables = {scanned_variable} if scanned_variable else set()
self.tinfo = None
def type_equals_to(self, tinfo):
return self.tinfo.equals_to(tinfo)
def switch_array_flag(self):
self.is_array ^= True
def activate(self):
pass
def set_enabled(self, enable):
self.enabled = enable
self.is_array = False
def has_collision(self, other):
if self.offset <= other.offset:
return self.offset + self.size > other.offset
return other.offset + other.size >= self.offset
@property
def score(self):
""" More score of the member - it better suits as candidate for this offset """
try:
return SCORE_TABLE[self.type_name]
except KeyError:
if self.tinfo and self.tinfo.is_funcptr():
return 0x1000 + len(self.tinfo.dstr())
return 0xFFFF
@property
def type_name(self):
return self.tinfo.dstr()
@property
def size(self):
size = self.tinfo.get_size()
return size if size != idaapi.BADSIZE else 1
@property
def font(self):
return None
def __repr__(self):
return hex(self.offset) + ' ' + self.type_name
def __eq__(self, other):
""" I'm aware that it's dirty but have no time to refactor whole file to nice one """
if self.offset == other.offset and self.type_name == other.type_name:
self.scanned_variables |= other.scanned_variables
return True
return False
__ne__ = lambda self, other: self.offset != other.offset or self.type_name != other.type_name
__lt__ = lambda self, other: self.offset < other.offset or \
(self.offset == other.offset and self.type_name < other.type_name)
__le__ = lambda self, other: self.offset <= other.offset
__gt__ = lambda self, other: self.offset > other.offset or \
(self.offset == other.offset and self.type_name < other.type_name)
__ge__ = lambda self, other: self.offset >= other.offset
class VirtualFunction:
def __init__(self, address, offset):
self.address = address
self.offset = offset
self.visited = False
def get_ptr_tinfo(self):
# print self.tinfo.dstr()
ptr_tinfo = idaapi.tinfo_t()
ptr_tinfo.create_ptr(self.tinfo)
return ptr_tinfo
def get_udt_member(self):
udt_member = idaapi.udt_member_t()
udt_member.type = self.get_ptr_tinfo()
udt_member.offset = self.offset
udt_member.name = self.name
udt_member.size = Const.EA_SIZE
return udt_member
def get_information(self):
return [Helper.to_hex(self.address), self.name, self.tinfo.dstr()]
@property
def name(self):
name = idaapi.get_short_name(self.address)
name = name.split('(')[0]
result = re.search(r"(\[thunk\]:)?([^`]*)(.*\{(\d+)}.*)?", name)
name, adjuster = result.group(2), result.group(4)
if adjuster:
name += "_adj_" + adjuster
name = name.translate(None, "`'").replace(':', '_').replace(' ', '_').replace(',', '_').replace('~', 'DESTR__')
name = name.replace("==", "__eq__")
name = name.replace("=", "__asg__")
name = re.sub(r'[<>]', '_t_', name)
return name
@property
def tinfo(self):
try:
decompiled_function = idaapi.decompile(self.address)
if decompiled_function:
return idaapi.tinfo_t(decompiled_function.type)
return Const.DUMMY_FUNC
except idaapi.DecompilationFailure:
pass
print "[ERROR] Failed to decompile function at 0x{0:08X}".format(self.address)
return Const.DUMMY_FUNC
def show_location(self):
idaapi.open_pseudocode(self.address, 1)
class ImportedVirtualFunction(VirtualFunction):
def __init__(self, address, offset):
VirtualFunction.__init__(self, address, offset)
@property
def tinfo(self):
print "[INFO] Ignoring import function at 0x{0:08X}".format(self.address)
tinfo = idaapi.tinfo_t()
if idaapi.guess_tinfo2(self.address, tinfo):
return tinfo
return Const.DUMMY_FUNC
def show_location(self):
idaapi.jumpto(self.address)
class VirtualTable(AbstractMember):
class VirtualTableChoose(MyChoose):
def __init__(self, items, temp_struct, virtual_table):
MyChoose.__init__(
self,
items,
"Select Virtual Function",
[["Address", 10], ["Name", 15], ["Declaration", 45]],
13
)
self.popup_names = ["Scan All", "-", "Scan", "-"]
self.__temp_struct = temp_struct
self.__virtual_table = virtual_table
def OnGetLineAttr(self, n):
return [0xd9d9d9, 0x0] if self.__virtual_table.virtual_functions[n].visited else [0xffffff, 0x0]
def OnGetIcon(self, n):
return 32 if self.__virtual_table.virtual_functions[n].visited else 160
def OnInsertLine(self):
""" Scan All Functions menu """
self.__virtual_table.scan_virtual_functions()
def OnEditLine(self, n):
""" Scan menu """
self.__virtual_table.scan_virtual_function(n)
def __init__(self, offset, address, scanned_variable=None, origin=0):
AbstractMember.__init__(self, offset + origin, scanned_variable, origin)
self.address = address
self.virtual_functions = []
self.name = "vtable" + ("_{0:X}".format(self.offset) if self.offset else '')
self.vtable_name, self.have_nice_name = parse_vtable_name(address)
self.populate()
def populate(self):
address = self.address
while True:
if Const.EA64:
func_address = idaapi.get_64bit(address)
else:
func_address = idaapi.get_32bit(address)
if Helper.is_code_ea(func_address):
self.virtual_functions.append(VirtualFunction(func_address, address - self.address))
elif Helper.is_imported_ea(func_address):
self.virtual_functions.append(ImportedVirtualFunction(func_address, address - self.address))
else:
break
address += Const.EA_SIZE
if idaapi.get_first_dref_to(address) != idaapi.BADADDR:
break
def create_tinfo(self):
# print "(Virtual table) at address: 0x{0:08X} name: {1}".format(self.address, self.name)
udt_data = idaapi.udt_type_data_t()
for function in self.virtual_functions:
udt_data.push_back(function.get_udt_member())
for duplicates in Helper.search_duplicate_fields(udt_data):
first_entry_idx = duplicates.pop(0)
print "[Warning] Found duplicate virtual functions", udt_data[first_entry_idx].name
for num, dup in enumerate(duplicates):
udt_data[dup].name = "duplicate_{0}_{1}".format(first_entry_idx, num + 1)
tinfo = idaapi.tinfo_t()
tinfo.create_ptr(Const.DUMMY_FUNC)
udt_data[dup].type = tinfo
final_tinfo = idaapi.tinfo_t()
if final_tinfo.create_udt(udt_data, idaapi.BTF_STRUCT):
# print "\n\t(Final structure)\n" + idaapi.print_tinfo('\t', 4, 5, idaapi.PRTYPE_MULTI | idaapi.PRTYPE_TYPE
# | idaapi.PRTYPE_SEMI, final_tinfo, self.name, None)
return final_tinfo
print "[ERROR] Virtual table creation failed"
def import_to_structures(self, ask=False):
"""
Imports virtual tables and returns tid_t of new structure
:return: idaapi.tid_t
"""
cdecl_typedef = idaapi.print_tinfo(None, 4, 5, idaapi.PRTYPE_MULTI | idaapi.PRTYPE_TYPE | idaapi.PRTYPE_SEMI,
self.create_tinfo(), self.vtable_name, None)
if ask:
cdecl_typedef = idaapi.asktext(0x10000, cdecl_typedef, "The following new type will be created")
if not cdecl_typedef:
return
previous_ordinal = idaapi.get_type_ordinal(idaapi.cvar.idati, self.vtable_name)
if previous_ordinal:
idaapi.del_numbered_type(idaapi.cvar.idati, previous_ordinal)
ordinal = idaapi.idc_set_local_type(previous_ordinal, cdecl_typedef, idaapi.PT_TYP)
else:
ordinal = idaapi.idc_set_local_type(-1, cdecl_typedef, idaapi.PT_TYP)
if ordinal:
print "[Info] Virtual table " + self.vtable_name + " added to Local Types"
return idaapi.import_type(idaapi.cvar.idati, -1, self.vtable_name)
else:
print "[Error] Failed to create virtual table " + self.vtable_name
print "*" * 100
print cdecl_typedef
print "*" * 100
def show_virtual_functions(self):
function_chooser = self.VirtualTableChoose(
[function.get_information() for function in self.virtual_functions], Cache.temporary_structure, self)
idx = function_chooser.Show(True)
if idx != -1:
virtual_function = self.virtual_functions[idx]
virtual_function.visited = True
virtual_function.show_location()
def scan_virtual_function(self, index):
if Helper.is_imported_ea(self.virtual_functions[index].address):
print "[INFO] Ignoring import function at 0x{0:08X}".format(self.address)
return
try:
function = idaapi.decompile(self.virtual_functions[index].address)
except idaapi.DecompilationFailure:
print "[ERROR] Failed to decompile function at 0x{0:08X}".format(self.address)
return
if Helper.FunctionTouchVisitor(function).process():
function = idaapi.decompile(self.virtual_functions[index].address)
if function.arguments and function.arguments[0].is_arg_var and Helper.is_legal_type(function.arguments[0].tif):
print "[Info] Scanning virtual function at 0x{0:08X}".format(function.entry_ea)
# TODO: Remove usage `temporary_structure' as global
obj = Api.VariableObject(function.get_lvars()[0], 0)
scanner = VariableScanner.NewDeepSearchVisitor(function, self.offset, obj, Cache.temporary_structure)
scanner.process()
else:
print "[Warning] Bad type of first argument in virtual function at 0x{0:08X}".format(function.entry_ea)
def scan_virtual_functions(self):
for idx in xrange(len(self.virtual_functions)):
self.scan_virtual_function(idx)
def get_udt_member(self, offset=0):
udt_member = idaapi.udt_member_t()
tid = self.import_to_structures()
if tid != idaapi.BADADDR:
udt_member.name = self.name
tmp_tinfo = idaapi.create_typedef(self.vtable_name)
tmp_tinfo.create_ptr(tmp_tinfo)
udt_member.type = tmp_tinfo
udt_member.offset = self.offset - offset
udt_member.size = Const.EA_SIZE
return udt_member
def type_equals_to(self, tinfo):
udt_data = idaapi.udt_type_data_t()
if tinfo.is_ptr() and tinfo.get_pointed_object().get_udt_details(udt_data):
if udt_data[0].type.is_funcptr():
return True
return False
def switch_array_flag(self):
pass
def activate(self):
self.show_virtual_functions()
@staticmethod
def check_address(address):
# Checks if given address contains virtual table. Returns True if more than 2 function pointers found
# Also if table's addresses point to code in executable section, than tries to make functions at that addresses
functions_count = 0
while True:
func_address = idaapi.get_64bit(address) if Const.EA64 else idaapi.get_32bit(address)
# print "[INFO] Address 0x{0:08X}".format(func_address)
if Helper.is_code_ea(func_address) or Helper.is_imported_ea(func_address):
functions_count += 1
address += Const.EA_SIZE
else:
segment = idaapi.getseg(func_address)
if segment and segment.perm & idaapi.SEGPERM_EXEC:
idc.MakeUnknown(func_address, 1, idaapi.DOUNK_SIMPLE)
if idc.MakeFunction(func_address):
functions_count += 1
address += Const.EA_SIZE
continue
break
idaapi.autoWait()
return functions_count
@property
def type_name(self):
return self.vtable_name + " *"
@property
def font(self):
return QtGui.QFont("Consolas", 10, QtGui.QFont.Bold)
@property
def size(self):
return Const.EA_SIZE
class Member(AbstractMember):
def __init__(self, offset, tinfo, scanned_variable, origin=0):
AbstractMember.__init__(self, offset + origin, scanned_variable, origin)
self.tinfo = tinfo
self.name = "field_{0:X}".format(self.offset)
def get_udt_member(self, array_size=0, offset=0):
udt_member = idaapi.udt_member_t()
udt_member.name = "field_{0:X}".format(self.offset - offset) if self.name[:6] == "field_" else self.name
udt_member.type = self.tinfo
if array_size:
tmp = idaapi.tinfo_t(self.tinfo)
tmp.create_array(self.tinfo, array_size)
udt_member.type = tmp
udt_member.offset = self.offset - offset
udt_member.size = self.size
return udt_member
def activate(self):
new_type_declaration = idaapi.askstr(0x100, self.type_name, "Enter type:")
if new_type_declaration is None:
return
result = idc.ParseType(new_type_declaration, 0)
if result is None:
return
_, tp, fld = result
tinfo = idaapi.tinfo_t()
tinfo.deserialize(idaapi.cvar.idati, tp, fld, None)
self.tinfo = tinfo
self.is_array = False
class VoidMember(Member):
def __init__(self, offset, scanned_variable, origin=0, char=False):
tinfo = Const.CHAR_TINFO if char else Const.BYTE_TINFO
Member.__init__(self, offset, tinfo, scanned_variable, origin)
self.is_array = True
def type_equals_to(self, tinfo):
return True
def switch_array_flag(self):
pass
def set_enabled(self, enable):
self.enabled = enable
@property
def font(self):
return QtGui.QFont("Consolas", 10, italic=True)
class TemporaryStructureModel(QtCore.QAbstractTableModel):
def __init__(self, *args):
"""
Keeps information about currently found fields in possible structure
main_offset - is the base from where variables scanned. Can be set to different value if some field is passed by
reverence
items - array of candidates to fields
"""
super(TemporaryStructureModel, self).__init__(*args)
self.main_offset = 0
self.headers = ["Offset", "Type", "Name"]
self.items = []
self.collisions = []
self.structure_name = "CHANGE_MY_NAME"
# OVERLOADED METHODS #
def rowCount(self, *args):
return len(self.items)
def columnCount(self, *args):
return len(self.headers)
def data(self, index, role):
row, col = index.row(), index.column()
item = self.items[row]
if role == QtCore.Qt.DisplayRole:
if col == 0:
return "0x{0:08X}".format(item.offset)
elif col == 1:
if item.is_array and item.size > 0:
array_size = self.calculate_array_size(row)
if array_size:
return item.type_name + "[{}]".format(array_size)
return item.type_name
elif col == 2:
return item.name
elif role == QtCore.Qt.ToolTipRole:
if col == 0:
return self.items[row].offset
elif col == 1:
return self.items[row].size * (self.calculate_array_size(row) if self.items[row].is_array else 1)
elif role == QtCore.Qt.EditRole:
if col == 2:
return self.items[row].name
elif role == QtCore.Qt.FontRole:
if col == 1:
return item.font
elif role == QtCore.Qt.BackgroundRole:
if not item.enabled:
return QtGui.QColor(QtCore.Qt.gray)
if item.offset == self.main_offset:
if col == 0:
return QtGui.QBrush(QtGui.QColor("#ff8080"))
if self.have_collision(row):
return QtGui.QBrush(QtGui.QColor("#ffff99"))
elif role == QtCore.Qt.ForegroundRole:
if self.have_collision(row):
return QtGui.QBrush(QtGui.QColor("#191919"))
def setData(self, index, value, role):
row, col = index.row(), index.column()
if role == QtCore.Qt.EditRole and idaapi.isident(str(value)):
self.items[row].name = str(value)
self.dataChanged.emit(index, index)
return True
return False
def headerData(self, section, orientation, role):
if role == QtCore.Qt.DisplayRole and orientation == QtCore.Qt.Horizontal:
return self.headers[section]
def flags(self, index):
if index.column() == 2:
return super(TemporaryStructureModel, self).flags(index) | QtGui.QAbstractItemView.DoubleClicked
return super(TemporaryStructureModel, self).flags(index)
# HELPER METHODS #
def pack(self, start=0, stop=None):
if self.collisions[start:stop].count(True):
print "[Warning] Collisions detected"
return
final_tinfo = idaapi.tinfo_t()
udt_data = idaapi.udt_type_data_t()
origin = self.items[start].offset if start else 0
offset = origin
for item in filter(lambda x: x.enabled, self.items[start:stop]): # Filter disabled members
gap_size = item.offset - offset
if gap_size:
udt_data.push_back(TemporaryStructureModel.get_padding_member(offset - origin, gap_size))
if item.is_array:
array_size = self.calculate_array_size(bisect.bisect_left(self.items, item))
if array_size:
udt_data.push_back(item.get_udt_member(array_size, offset=origin))
offset = item.offset + item.size * array_size
continue
udt_data.push_back(item.get_udt_member(offset=origin))
offset = item.offset + item.size
final_tinfo.create_udt(udt_data, idaapi.BTF_STRUCT)
cdecl = idaapi.print_tinfo(None, 4, 5, idaapi.PRTYPE_MULTI | idaapi.PRTYPE_TYPE | idaapi.PRTYPE_SEMI,
final_tinfo, self.structure_name, None)
cdecl = idaapi.asktext(0x10000, '#pragma pack(push, 1)\n' + cdecl, "The following new type will be created")
if cdecl:
structure_name = idaapi.idc_parse_decl(idaapi.cvar.idati, cdecl, idaapi.PT_TYP)[0]
previous_ordinal = idaapi.get_type_ordinal(idaapi.cvar.idati, structure_name)
if previous_ordinal:
reply = QtGui.QMessageBox.question(
None,
"HexRaysPyTools",
"Structure already exist. Do you want to overwrite it?",
QtGui.QMessageBox.Yes | QtGui.QMessageBox.No
)
if reply == QtGui.QMessageBox.Yes:
idaapi.del_numbered_type(idaapi.cvar.idati, previous_ordinal)
ordinal = idaapi.idc_set_local_type(previous_ordinal, cdecl, idaapi.PT_TYP)
else:
return
else:
ordinal = idaapi.idc_set_local_type(-1, cdecl, idaapi.PT_TYP)
if ordinal:
print "[Info] New type {0} was added to Local Types".format(structure_name)
tid = idaapi.import_type(idaapi.cvar.idati, -1, structure_name)
if tid:
tinfo = idaapi.create_typedef(structure_name)
ptr_tinfo = idaapi.tinfo_t()
ptr_tinfo.create_ptr(tinfo)
for scanned_var in self.get_unique_scanned_variables(origin):
scanned_var.apply_type(ptr_tinfo)
return tinfo
else:
print "[ERROR] Structure {0} probably already exist".format(structure_name)
def have_member(self, member):
if self.items:
idx = bisect.bisect_left(self.items, member)
if idx < self.rowCount():
return self.items[idx] == member
return False
def have_collision(self, row):
return self.collisions[row]
def refresh_collisions(self):
self.collisions = [False for _ in xrange(len(self.items))]
if (len(self.items)) > 1:
curr = 0
while curr < len(self.items):
if self.items[curr].enabled:
break
curr += 1
next = curr + 1
while next < len(self.items):
if self.items[next].enabled:
if self.items[curr].offset + self.items[curr].size > self.items[next].offset:
self.collisions[curr] = True
self.collisions[next] = True
if self.items[curr].offset + self.items[curr].size < self.items[next].offset + self.items[next].size:
curr = next
else:
curr = next
next += 1
def add_row(self, member):
if not self.have_member(member):
bisect.insort(self.items, member)
self.refresh_collisions()
self.modelReset.emit()
def get_unique_scanned_variables(self, origin=0):
scan_objects = itertools.chain.from_iterable(
[list(item.scanned_variables) for item in self.items if item.origin == origin])
return dict(((item.function_name, item.name), item) for item in scan_objects).values()
def get_next_enabled(self, row):
row += 1
while row < self.rowCount():
if self.items[row].enabled:
return row
row += 1
return None
def calculate_array_size(self, row):
next_row = self.get_next_enabled(row)
if next_row:
return (self.items[next_row].offset - self.items[row].offset) / self.items[row].size
return 0
def get_recognized_shape(self, start=0, stop=-1):
if not self.items:
return None
result = []
if stop != -1:
base = self.items[start].offset
enabled_items = filter(lambda x: x.enabled, self.items[start:stop])
else:
base = 0
enabled_items = filter(lambda x: x.enabled, self.items)
offsets = set(map(lambda x: x.offset, enabled_items))
if not enabled_items:
return
min_size = enabled_items[-1].offset + enabled_items[-1].size - base
tinfo = idaapi.tinfo_t()
for ordinal in xrange(1, idaapi.get_ordinal_qty(idaapi.cvar.idati)):
tinfo.get_numbered_type(idaapi.cvar.idati, ordinal)
if tinfo.is_udt() and tinfo.get_size() >= min_size:
is_found = False
for offset in offsets:
is_found = False
items = filter(lambda x: x.offset == offset, enabled_items)
potential_members = Helper.get_fields_at_offset(tinfo, offset - base)
for item in items:
for potential_member in potential_members:
if item.type_equals_to(potential_member):
is_found = True
break
if is_found:
break
if not is_found:
break
if is_found:
result.append((ordinal, idaapi.tinfo_t(tinfo)))
chooser = MyChoose(
[[str(x), "0x{0:08X}".format(y.get_size()), y.dstr()] for x, y in result],
"Select Structure",
[["Ordinal", 5], ["Size", 10], ["Structure name", 50]]
)
idx = chooser.Show(modal=True)
if idx != -1:
return result[idx][1]
return None
@staticmethod
def get_padding_member(offset, size):
udt_member = idaapi.udt_member_t()
if size == 1:
udt_member.name = "gap_{0:X}".format(offset)
udt_member.type = Const.BYTE_TINFO
udt_member.size = Const.BYTE_TINFO.get_size()
udt_member.offset = offset
return udt_member
array_data = idaapi.array_type_data_t()
array_data.base = 0
array_data.elem_type = Const.BYTE_TINFO
array_data.nelems = size
tmp_tinfo = idaapi.tinfo_t()
tmp_tinfo.create_array(array_data)
udt_member.name = "gap_{0:X}".format(offset)
udt_member.type = tmp_tinfo
udt_member.size = size
udt_member.offset = offset
return udt_member
# SLOTS #
def finalize(self):
if self.pack():
self.clear()
def disable_rows(self, indices):
for idx in indices:
if self.items[idx.row()].enabled:
self.items[idx.row()].set_enabled(False)
self.refresh_collisions()
self.modelReset.emit()
def enable_rows(self, indices):
for idx in indices:
if not self.items[idx.row()].enabled:
self.items[idx.row()].enabled = True
self.refresh_collisions()
self.modelReset.emit()
def set_origin(self, indices):
if indices:
self.main_offset = self.items[indices[0].row()].offset
self.modelReset.emit()
def make_array(self, indices):
if indices:
self.items[indices[0].row()].switch_array_flag()
self.dataChanged.emit(indices[0], indices[0])
def pack_substructure(self, indices):
if indices:
indices = sorted(indices)
self.dataChanged.emit(indices[0], indices[-1])
start, stop = indices[0].row(), indices[-1].row() + 1
tinfo = self.pack(start, stop)
if tinfo:
offset = self.items[start].offset
self.items = self.items[0:start] + self.items[stop:]
self.add_row(Member(offset, tinfo, None))
def unpack_substructure(self, indices):
if indices is None or len(indices) != 1:
return
item = self.items[indices[0].row()]
if item.tinfo is not None and item.tinfo.is_udt():
self.remove_items(indices)
offset = item.offset
udt_data = idaapi.udt_type_data_t()
if item.tinfo.get_udt_details(udt_data):
for udt_item in udt_data:
member = Member(offset + udt_item.offset / 8, udt_item.type, None)
member.name = udt_item.name
self.add_row(member)
def resolve_types(self):
current_item = None
current_item_score = 0
for item in self.items:
if not item.enabled:
continue
if current_item is None:
current_item = item
current_item_score = current_item.score
continue
item_score = item.score
if current_item.has_collision(item):
if item_score <= current_item_score:
item.set_enabled(False)
continue
elif item_score > current_item_score:
current_item.set_enabled(False)
current_item = item
current_item_score = item_score
self.refresh_collisions()
self.modelReset.emit()
def remove_items(self, indices):
rows = map(lambda x: x.row(), indices)
if rows:
self.items = [item for item in self.items if self.items.index(item) not in rows]
self.modelReset.emit()
def clear(self):
self.items = []
self.main_offset = 0
self.modelReset.emit()
def recognize_shape(self, indices):
min_idx = max_idx = None
if indices:
min_idx, max_idx = min(indices), max(indices, key=lambda x: (x.row(), x.column()))
if min_idx == max_idx:
tinfo = self.get_recognized_shape()
if tinfo:
tinfo.create_ptr(tinfo)
for scanned_var in self.get_unique_scanned_variables(origin=0):
scanned_var.apply_type(tinfo)
self.clear()
else:
# indices = sorted(indices)
start, stop = min_idx.row(), max_idx.row() + 1
base = self.items[start].offset
tinfo = self.get_recognized_shape(start, stop)
if tinfo:
ptr_tinfo = idaapi.tinfo_t()
ptr_tinfo.create_ptr(tinfo)
for scanned_var in self.get_unique_scanned_variables(base):
scanned_var.apply_type(ptr_tinfo)
self.items = filter(lambda x: x.offset < base or x.offset >= base + tinfo.get_size(), self.items)
self.add_row(Member(base, tinfo, None))
def activated(self, index):
# Double click on offset, opens window with variables
if index.column() == 0:
item = self.items[index.row()]
scanned_variables = list(item.scanned_variables)
variable_chooser = MyChoose(
map(lambda x: x.to_list(), scanned_variables),
"Select Variable",
[["Origin", 4], ["Function name", 25], ["Variable name", 25], ["Expression address", 10]]
)
row = variable_chooser.Show(modal=True)
if row != -1:
idaapi.open_pseudocode(scanned_variables[row].expression_address, 0)
# Double click on type. If type is virtual table than opens windows with virtual methods
elif index.column() == 1:
self.items[index.row()].activate()

View File

@@ -0,0 +1,339 @@
import logging
import idaapi
import idc
import Const
import Helper
import TemporaryStructure
import HexRaysPyTools.Api as Api
logger = logging.getLogger(__name__)
# If disabled then recursion will be triggered only for variable passed as first argument to function
SETTING_SCAN_ALL_ARGUMENTS = True
# Global set which is populated when deep scanning and cleared after completion
scanned_functions = set()
debug_scan_tree = []
class ScannedObject(object):
def __init__(self, name, expression_address, origin, applicable=True):
"""
:param name: Object name
:param expression_address: ea_t
:param origin: which offset had structure at scan moment
:param applicable: whether to apply type after creating structure
"""
self.name = name
self.expression_address = expression_address
self.func_ea = idc.get_func_attr(self.expression_address, idc.FUNCATTR_START)
self.origin = origin
self._applicable = applicable
@property
def function_name(self):
return idaapi.get_short_name(self.func_ea)
def apply_type(self, tinfo):
""" Finally apply Class'es tinfo to this variable """
raise NotImplemented
@staticmethod
def create(obj, expression_address, origin, applicable):
""" Creates suitable instance of ScannedObject depending on obj """
if obj.id == Api.SO_GLOBAL_OBJECT:
return ScannedGlobalObject(obj.ea, obj.name, expression_address, origin, applicable)
elif obj.id == Api.SO_LOCAL_VARIABLE:
return ScannedVariableObject(obj.lvar, obj.name, expression_address, origin, applicable)
elif obj.id in (Api.SO_STRUCT_REFERENCE, Api.SO_STRUCT_POINTER):
return ScannedStructureMemberObject(obj.struct_name, obj.offset, expression_address, origin, applicable)
else:
raise AssertionError
def to_list(self):
""" Creates list that is acceptable to MyChoose2 viewer """
return [
"0x{0:04X}".format(self.origin),
self.function_name,
self.name,
Helper.to_hex(self.expression_address)
]
def __eq__(self, other):
return self.func_ea == other.func_ea and self.name == other.name and \
self.expression_address == other.expression_address
def __hash__(self):
return hash((self.func_ea, self.name, self.expression_address))
def __repr__(self):
return "{} : {}".format(self.name, Helper.to_hex(self.expression_address))
class ScannedGlobalObject(ScannedObject):
def __init__(self, obj_ea, name, expression_address, origin, applicable=True):
super(ScannedGlobalObject, self).__init__(name, expression_address, origin, applicable)
self.__obj_ea = obj_ea
def apply_type(self, tinfo):
if self._applicable:
idaapi.set_tinfo2(self.__obj_ea, tinfo)
class ScannedVariableObject(ScannedObject):
def __init__(self, lvar, name, expression_address, origin, applicable=True):
super(ScannedVariableObject, self).__init__(name, expression_address, origin, applicable)
self.__lvar = idaapi.lvar_locator_t(lvar.location, lvar.defea)
def apply_type(self, tinfo):
if not self._applicable:
return
hx_view = idaapi.open_pseudocode(self.func_ea, -1)
if hx_view:
logger.debug("Applying tinfo to variable {0} in function {1}".format(self.name, self.function_name))
# Finding lvar of new window that have the same name that saved one and applying tinfo_t
lvar = filter(lambda x: x == self.__lvar, hx_view.cfunc.get_lvars())
if lvar:
logger.debug("Successful")
hx_view.set_lvar_type(lvar[0], tinfo)
else:
logger.warn("Failed to find previously scanned local variable {} from {}".format(
self.name, Helper.to_hex(self.expression_address)))
class ScannedStructureMemberObject(ScannedObject):
def __init__(self, struct_name, struct_offset, name, expression_address, origin, applicable=True):
super(ScannedStructureMemberObject, self).__init__(name, expression_address, origin, applicable)
self.__struct_name = struct_name
self.__struct_offset = struct_offset
def apply_type(self, tinfo):
if self._applicable:
logger.warn("Changing type of structure field is not yet implemented. Address - {}".format(
Helper.to_hex(self.expression_address)))
class SearchVisitor(Api.ObjectVisitor):
def __init__(self, cfunc, origin, obj, temporary_structure):
super(SearchVisitor, self).__init__(cfunc, obj, None, True)
self.__origin = origin
self.__temporary_structure = temporary_structure
def _manipulate(self, cexpr, obj):
super(SearchVisitor, self)._manipulate(cexpr, obj)
if obj.tinfo and not Helper.is_legal_type(obj.tinfo):
logger.warn("Variable obj.name has weird type at {}".format(Helper.to_hex(self._find_asm_address(cexpr))))
return
if cexpr.type.is_ptr():
member = self.__extract_member_from_pointer(cexpr, obj)
else:
member = self.__extract_member_from_xword(cexpr, obj)
if member:
logger.debug("\tCreating member with type {}, {}, offset - {}".format(
member.type_name, member.scanned_variables, member.offset))
self.__temporary_structure.add_row(member)
def _get_member(self, offset, cexpr, obj, tinfo=None, obj_ea=None):
if offset < 0:
logger.error("Considered to be imposible: offset - {}, obj - {}".format(
offset, Helper.to_hex(self._find_asm_address(cexpr))))
raise AssertionError
applicable = not self.crippled
cexpr_ea = self._find_asm_address(cexpr)
scan_obj = ScannedObject.create(obj, cexpr_ea, self.__origin, applicable)
if obj_ea:
if TemporaryStructure.VirtualTable.check_address(obj_ea):
return TemporaryStructure.VirtualTable(offset, obj_ea, scan_obj, self.__origin)
if Helper.is_code_ea(obj_ea):
cfunc = Api.decompile_function(obj_ea)
if cfunc:
tinfo = cfunc.type
tinfo.create_ptr(tinfo)
else:
tinfo = Const.DUMMY_FUNC
return TemporaryStructure.Member(offset, tinfo, scan_obj, self.__origin)
# logger.warn("Want to see this ea - {},".format(Helper.to_hex(cexpr_ea)))
if not tinfo or tinfo.equals_to(Const.VOID_TINFO) or tinfo.equals_to(Const.CONST_VOID_TINFO):
return TemporaryStructure.VoidMember(offset, scan_obj, self.__origin)
if tinfo.equals_to(Const.CHAR_TINFO):
return TemporaryStructure.VoidMember(offset, scan_obj, self.__origin, char=True)
if tinfo.equals_to(Const.CONST_PCHAR_TINFO):
tinfo = Const.PCHAR_TINFO
elif tinfo.equals_to(Const.CONST_PVOID_TINFO):
tinfo = Const.PVOID_TINFO
else:
tinfo.clr_const()
return TemporaryStructure.Member(offset, tinfo, scan_obj, self.__origin)
def _parse_call(self, call_cexpr, arg_cexpr, offset):
_, tinfo = Helper.get_func_argument_info(call_cexpr, arg_cexpr)
if tinfo:
return self.__deref_tinfo(tinfo)
# TODO: Find example with UTF-16 strings
return Const.CHAR_TINFO
def _parse_left_assignee(self, cexpr, offset):
pass
def __extract_member_from_pointer(self, cexpr, obj):
parents_type = map(lambda x: idaapi.get_ctype_name(x.cexpr.op), list(self.parents)[:0:-1])
parents = map(lambda x: x.cexpr, list(self.parents)[:0:-1])
logger.debug("Parsing expression {}. Parents - {}".format(obj.name, parents_type))
# Extracting offset and removing expression parents making this offset
if parents_type[0] in ('idx', 'add'):
# `obj[idx]' or `(TYPE *) + x'
if parents[0].y.op != idaapi.cot_num:
# There's no way to handle with dynamic offset
return
offset = parents[0].y.numval() * cexpr.type.get_ptrarr_objsize()
cexpr = self.parent_expr()
if parents_type[0] == 'add':
del parents_type[0]
del parents[0]
elif parents_type[0:2] == ['cast', 'add']:
# (TYPE *)obj + offset or (TYPE)obj + offset
if parents[1].y.op != idaapi.cot_num:
return
if parents[0].type.is_ptr():
size = parents[0].type.get_ptrarr_objsize()
else:
size = 1
offset = parents[1].theother(parents[0]).numval() * size
cexpr = parents[1]
del parents_type[0:2]
del parents[0:2]
else:
offset = 0
return self.__extract_member(cexpr, obj, offset, parents, parents_type)
def __extract_member_from_xword(self, cexpr, obj):
parents_type = map(lambda x: idaapi.get_ctype_name(x.cexpr.op), list(self.parents)[:0:-1])
parents = map(lambda x: x.cexpr, list(self.parents)[:0:-1])
logger.debug("Parsing expression {}. Parents - {}".format(obj.name, parents_type))
if parents_type[0] == 'add':
if parents[0].theother(cexpr).op != idaapi.cot_num:
return
offset = parents[0].theother(cexpr).numval()
cexpr = self.parent_expr()
del parents_type[0]
del parents[0]
else:
offset = 0
return self.__extract_member(cexpr, obj, offset, parents, parents_type)
def __extract_member(self, cexpr, obj, offset, parents, parents_type):
if parents_type[0] == 'cast':
default_tinfo = parents[0].type
cexpr = parents[0]
del parents_type[0]
del parents[0]
else:
default_tinfo = Const.PX_WORD_TINFO
if parents_type[0] in ('idx', 'ptr'):
if parents_type[1] == 'cast':
default_tinfo = parents[1].type
cexpr = parents[0]
del parents_type[0]
del parents[0]
else:
default_tinfo = self.__deref_tinfo(default_tinfo)
if parents_type[1] == 'asg':
if parents[1].x == parents[0]:
# *(TYPE *)(var + x) = ???
obj_ea = self.__extract_obj_ea(parents[1].y)
return self._get_member(offset, cexpr, obj, default_tinfo, obj_ea)
return self._get_member(offset, cexpr, obj, parents[1].x.type)
elif parents_type[1] == 'call':
if parents[1].x == parents[0]:
# ((type (__some_call *)(..., ..., ...)var[idx])(..., ..., ...)
# ((type (__some_call *)(..., ..., ...)*(TYPE *)(var + x))(..., ..., ...)
return self._get_member(offset, cexpr, obj, parents[0].type)
_, tinfo = Helper.get_func_argument_info(parents[1], parents[0])
if tinfo is None:
tinfo = Const.PCHAR_TINFO
return self._get_member(offset, cexpr, obj, tinfo)
return self._get_member(offset, cexpr, obj, default_tinfo)
elif parents_type[0] == 'call':
# call(..., (TYPE)(var + x), ...)
tinfo = self._parse_call(parents[0], cexpr, offset)
return self._get_member(offset, cexpr, obj, tinfo)
elif parents_type[0] == 'asg':
if parents[0].y == cexpr:
# other_obj = (TYPE) (var + offset)
self._parse_left_assignee(parents[1].x, offset)
return self._get_member(offset, cexpr, obj, self.__deref_tinfo(default_tinfo))
@staticmethod
def __extract_obj_ea(cexpr):
if cexpr.op == idaapi.cot_ref:
cexpr = cexpr.x
if cexpr.op == idaapi.cot_obj:
if cexpr.obj_ea != idaapi.BADADDR:
return cexpr.obj_ea
@staticmethod
def __deref_tinfo(tinfo):
if tinfo.is_ptr():
if tinfo.get_ptrarr_objsize() == 1:
if tinfo.equals_to(Const.PCHAR_TINFO) or tinfo.equals_to(Const.CONST_PCHAR_TINFO):
return Const.CHAR_TINFO
return None # Turns into VoidMember
return tinfo.get_pointed_object()
return tinfo
class NewShallowSearchVisitor(SearchVisitor, Api.ObjectDownwardsVisitor):
def __init__(self, cfunc, origin, obj, temporary_structure):
super(NewShallowSearchVisitor, self).__init__(cfunc, origin, obj, temporary_structure)
class NewDeepSearchVisitor(SearchVisitor, Api.RecursiveObjectDownwardsVisitor):
def __init__(self, cfunc, origin, obj, temporary_structure):
super(NewDeepSearchVisitor, self).__init__(cfunc, origin, obj, temporary_structure)
class DeepReturnVisitor(NewDeepSearchVisitor):
def __init__(self, cfunc, origin, obj, temporary_structure):
super(DeepReturnVisitor, self).__init__(cfunc, origin, obj, temporary_structure)
self.__callers_ea = Helper.get_funcs_calling_address(cfunc.entry_ea)
self.__call_obj = obj
def _start(self):
for ea in self.__callers_ea:
self._add_scan_tree_info(ea, -1)
assert self.__prepare_scanner()
def _finish(self):
if self.__prepare_scanner():
self._recursive_process()
def __prepare_scanner(self):
try:
cfunc = self.__iter_callers().next()
except StopIteration:
return False
self.prepare_new_scan(cfunc, -1, self.__call_obj)
return True
def __iter_callers(self):
for ea in self.__callers_ea:
cfunc = Api.decompile_function(ea)
if cfunc:
yield cfunc

View File

View File

@@ -0,0 +1,116 @@
"""
Cute - a crossQt compatibility module for IDAPython.
Feel free to copy this code into your own projects.
Latest version can be found at https://github.com/tmr232/Cute
The MIT License (MIT)
Copyright (c) 2015 Tamir Bahar
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
import os
import sys
import idaapi
# This nasty piece of code is here to force the loading of IDA's PySide.
# Without it, Python attempts to load PySide from the site-packages directory,
# and failing, as it does not play nicely with IDA.
old_path = sys.path[:]
try:
ida_python_path = os.path.dirname(idaapi.__file__)
sys.path.insert(0, ida_python_path)
if idaapi.IDA_SDK_VERSION >= 690:
from PyQt5 import QtGui, QtCore, QtWidgets
import sip
QtCore.Signal = QtCore.pyqtSignal
QtCore.Slot = QtCore.pyqtSlot
QtGui.QTreeView = QtWidgets.QTreeView
QtGui.QTableView = QtWidgets.QTableView
QtGui.QAction = QtWidgets.QAction
QtGui.QMenu = QtWidgets.QMenu
QtGui.QLabel = QtWidgets.QLabel
QtGui.QMessageBox = QtWidgets.QMessageBox
QtGui.QHeaderView = QtWidgets.QHeaderView
QtGui.QAbstractItemView = QtWidgets.QAbstractItemView
QtGui.QVBoxLayout = QtWidgets.QVBoxLayout
QtGui.QHBoxLayout = QtWidgets.QHBoxLayout
QtGui.QGridLayout = QtWidgets.QGridLayout
QtGui.QPushButton = QtWidgets.QPushButton
QtGui.QSpacerItem = QtWidgets.QSpacerItem
QtGui.QSizePolicy = QtWidgets.QSizePolicy
QtGui.QSortFilterProxyModel = QtCore.QSortFilterProxyModel
QtGui.QLineEdit = QtWidgets.QLineEdit
use_qt5 = True
else:
from PySide import QtGui, QtCore
from PySide import QtGui as QtWidgets
QtGui.QHeaderView.setSectionResizeMode = QtGui.QHeaderView.setResizeMode
use_qt5 = False
finally:
sys.path = old_path
def connect(sender, signal, callback):
"""Connect a signal.
Use this function only in cases where code should work with both Qt5 and Qt4, as it is an ugly hack.
Args:
sender: The Qt object emitting the signal
signal: A string, containing the signal signature (as in Qt4 and PySide)
callback: The function to be called upon receiving the signal
"""
if use_qt5:
return getattr(sender, signal.split('(', 1)[0]).connect(callback)
else:
return sender.connect(QtCore.SIGNAL(signal), callback)
def disconnect(sender, signal, callback):
"""Disconnect a signal.
Use this function only in cases where code should work with both Qt5 and Qt4, as it is an ugly hack.
Args:
sender: The Qt object emitting the signal
signal: A string, containing the signal signature (as in Qt4 and PySide)
callback: The function to be called upon receiving the signal
"""
if use_qt5:
return getattr(sender, signal.split('(', 1)[0]).disconnect(callback)
else:
return sender.disconnect(QtCore.SIGNAL(signal), callback)
def form_to_widget(tform):
"""Get the tform's widget.
IDA has two different form-to-widget functions, one for PyQt and one for PySide.
This function uses the relevant one based on the version of IDA you are using.
Args:
tform: The IDA TForm to get the widget from.
"""
class Ctx(object):
QtGui = QtGui
if use_qt5:
QtWidgets = QtWidgets
sip = sip
if use_qt5:
return idaapi.PluginForm.FormToPyQtWidget(tform, ctx=Ctx())
else:
return idaapi.PluginForm.FormToPySideWidget(tform, ctx=Ctx())

View File

@@ -0,0 +1,246 @@
import idaapi
# import PySide.QtGui as QtGui
# import PySide.QtCore as QtCore
from HexRaysPyTools.Cute import *
import Core.Classes
class MyChoose(idaapi.Choose2):
def __init__(self, items, title, cols, icon=-1):
idaapi.Choose2.__init__(self, title, cols, flags=idaapi.Choose2.CH_MODAL, icon=icon)
self.items = items
def OnClose(self):
pass
def OnGetLine(self, n):
return self.items[n]
def OnGetSize(self):
return len(self.items)
class StructureBuilder(idaapi.PluginForm):
def __init__(self, structure_model):
super(StructureBuilder, self).__init__()
self.structure_model = structure_model
self.parent = None
def OnCreate(self, form):
self.parent = form_to_widget(form)
self.init_ui()
def init_ui(self):
self.parent.setStyleSheet(
"QTableView {background-color: transparent; selection-background-color: #87bdd8;}"
"QHeaderView::section {background-color: transparent; border: 0.5px solid;}"
"QPushButton {width: 50px; height: 20px;}"
# "QPushButton::pressed {background-color: #ccccff}"
)
self.parent.resize(400, 600)
self.parent.setWindowTitle('Structure Builder')
btn_finalize = QtGui.QPushButton("&Finalize")
btn_disable = QtGui.QPushButton("&Disable")
btn_enable = QtGui.QPushButton("&Enable")
btn_origin = QtGui.QPushButton("&Origin")
btn_array = QtGui.QPushButton("&Array")
btn_pack = QtGui.QPushButton("&Pack")
btn_unpack = QtGui.QPushButton("&Unpack")
btn_remove = QtGui.QPushButton("&Remove")
btn_resolve = QtGui.QPushButton("Resolve")
btn_clear = QtGui.QPushButton("Clear") # Clear button doesn't have shortcut because it can fuck up all work
btn_recognize = QtGui.QPushButton("Recognize Shape")
btn_recognize.setStyleSheet("QPushButton {width: 100px; height: 20px;}")
btn_finalize.setShortcut("f")
btn_disable.setShortcut("d")
btn_enable.setShortcut("e")
btn_origin.setShortcut("o")
btn_array.setShortcut("a")
btn_pack.setShortcut("p")
btn_unpack.setShortcut("u")
btn_remove.setShortcut("r")
struct_view = QtGui.QTableView()
struct_view.setModel(self.structure_model)
# struct_view.setSelectionMode(QtGui.QAbstractItemView.ExtendedSelection)
struct_view.verticalHeader().setVisible(False)
struct_view.verticalHeader().setDefaultSectionSize(24)
struct_view.horizontalHeader().setStretchLastSection(True)
struct_view.horizontalHeader().setSectionResizeMode(QtGui.QHeaderView.ResizeToContents)
grid_box = QtGui.QGridLayout()
grid_box.setSpacing(0)
grid_box.addWidget(btn_finalize, 0, 0)
grid_box.addWidget(btn_enable, 0, 1)
grid_box.addWidget(btn_disable, 0, 2)
grid_box.addWidget(btn_origin, 0, 3)
grid_box.addItem(QtGui.QSpacerItem(20, 20, QtGui.QSizePolicy.Expanding), 0, 5)
grid_box.addWidget(btn_array, 1, 0)
grid_box.addWidget(btn_pack, 1, 1)
grid_box.addWidget(btn_unpack, 1, 2)
grid_box.addWidget(btn_remove, 1, 3)
grid_box.addWidget(btn_resolve, 0, 4)
grid_box.addItem(QtGui.QSpacerItem(20, 20, QtGui.QSizePolicy.Expanding), 1, 5)
grid_box.addWidget(btn_recognize, 0, 6)
grid_box.addWidget(btn_clear, 1, 6)
vertical_box = QtGui.QVBoxLayout()
vertical_box.addWidget(struct_view)
vertical_box.addLayout(grid_box)
self.parent.setLayout(vertical_box)
btn_finalize.clicked.connect(lambda: self.structure_model.finalize())
btn_disable.clicked.connect(lambda: self.structure_model.disable_rows(struct_view.selectedIndexes()))
btn_enable.clicked.connect(lambda: self.structure_model.enable_rows(struct_view.selectedIndexes()))
btn_origin.clicked.connect(lambda: self.structure_model.set_origin(struct_view.selectedIndexes()))
btn_array.clicked.connect(lambda: self.structure_model.make_array(struct_view.selectedIndexes()))
btn_pack.clicked.connect(lambda: self.structure_model.pack_substructure(struct_view.selectedIndexes()))
btn_unpack.clicked.connect(lambda: self.structure_model.unpack_substructure(struct_view.selectedIndexes()))
btn_remove.clicked.connect(lambda: self.structure_model.remove_items(struct_view.selectedIndexes()))
btn_resolve.clicked.connect(lambda: self.structure_model.resolve_types())
btn_clear.clicked.connect(lambda: self.structure_model.clear())
btn_recognize.clicked.connect(lambda: self.structure_model.recognize_shape(struct_view.selectedIndexes()))
struct_view.activated[QtCore.QModelIndex].connect(self.structure_model.activated)
self.structure_model.dataChanged.connect(struct_view.clearSelection)
def OnClose(self, form):
pass
def Show(self, caption=None, options=0):
return idaapi.PluginForm.Show(self, caption, options=options)
class StructureGraphViewer(idaapi.GraphViewer):
def __init__(self, title, graph):
idaapi.GraphViewer.__init__(self, title)
self.graph = graph
self.nodes_id = {}
def OnRefresh(self):
self.Clear()
self.nodes_id.clear()
for node in self.graph.get_nodes():
self.nodes_id[node] = self.AddNode(node)
for first, second in self.graph.get_edges():
self.AddEdge(self.nodes_id[first], self.nodes_id[second])
return True
def OnGetText(self, node_id):
return self.graph.local_types[self[node_id]].name_and_color
def OnHint(self, node_id):
""" Try-catch clause because IDA sometimes attempts to use old information to get hint """
try:
ordinal = self[node_id]
return self.graph.local_types[ordinal].hint
except KeyError:
return
def OnDblClick(self, node_id):
self.change_selected([self[node_id]])
def change_selected(self, ordinals):
self.graph.change_selected(ordinals)
self.Refresh()
self.Select(self.nodes_id[ordinals[0]])
class ClassViewer(idaapi.PluginForm):
def __init__(self):
super(ClassViewer, self).__init__()
self.parent = None
self.class_tree = QtGui.QTreeView()
self.line_edit_filter = QtGui.QLineEdit()
self.action_collapse = QtGui.QAction("Collapse all", self.class_tree)
self.action_expand = QtGui.QAction("Expand all", self.class_tree)
self.action_set_arg = QtGui.QAction("Set First Argument Type", self.class_tree)
self.action_rollback = QtGui.QAction("Rollback", self.class_tree)
self.action_refresh = QtGui.QAction("Refresh", self.class_tree)
self.action_commit = QtGui.QAction("Commit", self.class_tree)
self.menu = QtGui.QMenu(self.parent)
self.proxy_model = Core.Classes.ProxyModel()
def OnCreate(self, form):
# self.parent = self.FormToPySideWidget(form)
self.parent = form_to_widget(form)
self.init_ui()
def init_ui(self):
self.parent.setWindowTitle('Classes')
self.parent.setStyleSheet(
# "QTreeView::item:!has-children { background-color: #fefbd8; border: 0.5px solid lightgray ;}"
# "QTreeView::item:has-children { background-color: #80ced6; border-top: 1px solid black ;}"
# "QTreeView::item:selected { background-color: #618685; show-decoration-selected: 1;}"
"QTreeView {background-color: transparent; }"
"QHeaderView::section {background-color: transparent; border: 1px solid;}"
)
hbox_layout = QtGui.QHBoxLayout()
label_filter = QtGui.QLabel("&Filter:")
label_filter.setBuddy(self.line_edit_filter)
hbox_layout.addWidget(label_filter)
hbox_layout.addWidget(self.line_edit_filter)
class_model = Core.Classes.TreeModel()
self.proxy_model.setSourceModel(class_model)
self.proxy_model.setFilterCaseSensitivity(QtCore.Qt.CaseInsensitive)
self.class_tree.setModel(self.proxy_model)
self.class_tree.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
self.class_tree.expandAll()
self.class_tree.header().setStretchLastSection(True)
self.class_tree.header().setSectionResizeMode(QtGui.QHeaderView.ResizeToContents)
self.class_tree.setSelectionMode(QtGui.QAbstractItemView.ExtendedSelection)
self.action_collapse.triggered.connect(self.class_tree.collapseAll)
self.action_expand.triggered.connect(self.class_tree.expandAll)
self.action_set_arg.triggered.connect(
lambda: class_model.set_first_argument_type(
map(self.proxy_model.mapToSource, self.class_tree.selectedIndexes())
)
)
self.action_rollback.triggered.connect(lambda: class_model.rollback())
self.action_refresh.triggered.connect(lambda: class_model.refresh())
self.action_commit.triggered.connect(lambda: class_model.commit())
class_model.refreshed.connect(self.class_tree.expandAll)
self.menu.addAction(self.action_collapse)
self.menu.addAction(self.action_expand)
self.menu.addAction(self.action_refresh)
self.menu.addAction(self.action_set_arg)
self.menu.addAction(self.action_rollback)
self.menu.addAction(self.action_commit)
vertical_box = QtGui.QVBoxLayout()
vertical_box.addWidget(self.class_tree)
vertical_box.addLayout(hbox_layout)
self.parent.setLayout(vertical_box)
self.class_tree.activated[QtCore.QModelIndex].connect(
lambda x: class_model.open_function(self.proxy_model.mapToSource(x))
)
self.class_tree.customContextMenuRequested[QtCore.QPoint].connect(self.show_menu)
self.line_edit_filter.textChanged[str].connect(self.proxy_model.set_regexp_filter)
# proxy_model.rowsInserted[object].connect(lambda: self.class_tree.setExpanded(object, True))
def OnClose(self, form):
pass
def Show(self, caption=None, options=0):
return idaapi.PluginForm.Show(self, caption, options=options)
def show_menu(self, point):
self.action_set_arg.setEnabled(True)
indexes = map(
self.proxy_model.mapToSource,
filter(lambda x: x.column() == 0, self.class_tree.selectedIndexes())
)
if len(indexes) > 1:
if filter(lambda x: len(x.internalPointer().children) > 0, indexes):
self.action_set_arg.setEnabled(False)
self.menu.exec_(self.class_tree.mapToGlobal(point))

View File

@@ -0,0 +1,57 @@
import os
import ConfigParser
import idc
import logging
CONFIG_FILE_PATH = os.path.join(idc.GetIdaDirectory(), 'cfg', 'HexRaysPyTools.cfg')
DEBUG_MESSAGE_LEVEL = logging.INFO
# Whether propagate names (Propagate name feature) through all names or only defaults like v11, a3, this, field_4
PROPAGATE_THROUGH_ALL_NAMES = False
# Store Xref information in database. I don't know how much size it consumes yet
STORE_XREFS = True
# There're some types that can be pointers to structures like int, PVOID etc and by default plugin scans only them
# Full list can be found in `Const.LEGAL_TYPES`.
# But if set this option to True than variable of every type could be possible to scan
SCAN_ANY_TYPE = False
def add_default_settings(config):
updated = False
if not config.has_option("DEFAULT", "DEBUG_MESSAGE_LEVEL"):
print DEBUG_MESSAGE_LEVEL
config.set(None, 'DEBUG_MESSAGE_LEVEL', str(DEBUG_MESSAGE_LEVEL))
updated = True
if not config.has_option("DEFAULT", "PROPAGATE_THROUGH_ALL_NAMES"):
config.set(None, 'PROPAGATE_THROUGH_ALL_NAMES', str(PROPAGATE_THROUGH_ALL_NAMES))
updated = True
if not config.has_option("DEFAULT", "STORE_XREFS"):
config.set(None, 'STORE_XREFS', str(STORE_XREFS))
updated = True
if not config.has_option("DEFAULT", "SCAN_ANY_TYPE"):
config.set(None, 'SCAN_ANY_TYPE', str(SCAN_ANY_TYPE))
updated = True
if updated:
try:
with open(CONFIG_FILE_PATH, "w") as f:
config.write(f)
except IOError:
print "[ERROR] Failed to write or update config file at {}. Default settings will be used instead.\n" \
"Consider running IDA Pro under administrator once".format(CONFIG_FILE_PATH)
def load_settings():
global DEBUG_MESSAGE_LEVEL, PROPAGATE_THROUGH_ALL_NAMES, STORE_XREFS, SCAN_ANY_TYPE
config = ConfigParser.ConfigParser()
if os.path.isfile(CONFIG_FILE_PATH):
config.read(CONFIG_FILE_PATH)
add_default_settings(config)
DEBUG_MESSAGE_LEVEL = config.getint("DEFAULT", 'DEBUG_MESSAGE_LEVEL')
PROPAGATE_THROUGH_ALL_NAMES = config.getboolean("DEFAULT", 'PROPAGATE_THROUGH_ALL_NAMES')
STORE_XREFS = config.getboolean("DEFAULT", 'STORE_XREFS')
SCAN_ANY_TYPE = config.getboolean("DEFAULT", 'SCAN_ANY_TYPE')

View File

@@ -0,0 +1,7 @@
import logging
import Settings
Settings.load_settings()
logging.basicConfig(format='[%(levelname)s] %(message)s\t(%(module)s:%(funcName)s)')
logging.root.setLevel(Settings.DEBUG_MESSAGE_LEVEL)