1441 lines
51 KiB
Python
1441 lines
51 KiB
Python
import ctypes
|
|
import sys
|
|
import re
|
|
import logging
|
|
|
|
import idaapi
|
|
import idc
|
|
|
|
import HexRaysPyTools.Forms as Forms
|
|
import HexRaysPyTools.Core.Const as Const
|
|
import HexRaysPyTools.Core.Helper as Helper
|
|
import HexRaysPyTools.Api as Api
|
|
import Settings
|
|
from HexRaysPyTools.Core.StructureGraph import StructureGraph
|
|
from HexRaysPyTools.Core.TemporaryStructure import VirtualTable, TemporaryStructureModel
|
|
from HexRaysPyTools.Core.VariableScanner import NewShallowSearchVisitor, NewDeepSearchVisitor, DeepReturnVisitor
|
|
from HexRaysPyTools.Core.Helper import FunctionTouchVisitor
|
|
from HexRaysPyTools.Core.SpaghettiCode import *
|
|
from HexRaysPyTools.Core.StructXrefs import XrefStorage
|
|
|
|
RECAST_LOCAL_VARIABLE = 0
|
|
RECAST_GLOBAL_VARIABLE = 1
|
|
RECAST_ARGUMENT = 2
|
|
RECAST_RETURN = 3
|
|
RECAST_STRUCTURE = 4
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def register(action, *args):
|
|
idaapi.register_action(
|
|
idaapi.action_desc_t(
|
|
action.name,
|
|
action.description,
|
|
action(*args),
|
|
action.hotkey
|
|
)
|
|
)
|
|
|
|
|
|
def unregister(action):
|
|
idaapi.unregister_action(action.name)
|
|
|
|
|
|
class TypeLibrary:
|
|
|
|
class til_t(ctypes.Structure):
|
|
pass
|
|
|
|
til_t._fields_ = [
|
|
("name", ctypes.c_char_p),
|
|
("desc", ctypes.c_char_p),
|
|
("nbases", ctypes.c_int),
|
|
("base", ctypes.POINTER(ctypes.POINTER(til_t)))
|
|
]
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@staticmethod
|
|
def enable_library_ordinals(library_num):
|
|
idaname = "ida64" if Const.EA64 else "ida"
|
|
if sys.platform == "win32":
|
|
dll = ctypes.windll[idaname + ".wll"]
|
|
elif sys.platform == "linux2":
|
|
dll = ctypes.cdll["lib" + idaname + ".so"]
|
|
elif sys.platform == "darwin":
|
|
dll = ctypes.cdll["lib" + idaname + ".dylib"]
|
|
else:
|
|
print "[ERROR] Failed to enable ordinals"
|
|
return
|
|
|
|
idati = ctypes.POINTER(TypeLibrary.til_t).in_dll(dll, "idati")
|
|
dll.enable_numbered_types(idati.contents.base[library_num], True)
|
|
|
|
@staticmethod
|
|
def choose_til():
|
|
idati = idaapi.cvar.idati
|
|
list_type_library = [(idati, idati.name, idati.desc)]
|
|
for idx in xrange(idaapi.cvar.idati.nbases):
|
|
type_library = idaapi.cvar.idati.base(idx) # idaapi.til_t type
|
|
list_type_library.append((type_library, type_library.name, type_library.desc))
|
|
|
|
library_chooser = Forms.MyChoose(
|
|
list(map(lambda x: [x[1], x[2]], list_type_library)),
|
|
"Select Library",
|
|
[["Library", 10 | idaapi.Choose2.CHCOL_PLAIN], ["Description", 30 | idaapi.Choose2.CHCOL_PLAIN]],
|
|
69
|
|
)
|
|
library_num = library_chooser.Show(True)
|
|
if library_num != -1:
|
|
selected_library = list_type_library[library_num][0]
|
|
max_ordinal = idaapi.get_ordinal_qty(selected_library)
|
|
if max_ordinal == idaapi.BADORD:
|
|
TypeLibrary.enable_library_ordinals(library_num - 1)
|
|
max_ordinal = idaapi.get_ordinal_qty(selected_library)
|
|
print "[DEBUG] Maximal ordinal of lib {0} = {1}".format(selected_library.name, max_ordinal)
|
|
return selected_library, max_ordinal, library_num == 0
|
|
return None
|
|
|
|
@staticmethod
|
|
def import_type(library, name):
|
|
if library.name != idaapi.cvar.idati.name:
|
|
last_ordinal = idaapi.get_ordinal_qty(idaapi.cvar.idati)
|
|
type_id = idaapi.import_type(library, -1, name) # tid_t
|
|
if type_id != idaapi.BADORD:
|
|
return last_ordinal
|
|
return None
|
|
|
|
|
|
class RemoveArgument(idaapi.action_handler_t):
|
|
|
|
name = "my:RemoveArgument"
|
|
description = "Remove Argument"
|
|
hotkey = None
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
def activate(self, ctx):
|
|
vu = idaapi.get_widget_vdui(ctx.widget)
|
|
function_tinfo = idaapi.tinfo_t()
|
|
if not vu.cfunc.get_func_type(function_tinfo):
|
|
return
|
|
function_details = idaapi.func_type_data_t()
|
|
function_tinfo.get_func_details(function_details)
|
|
del_arg = vu.item.get_lvar() # lvar_t
|
|
|
|
function_details.erase(filter(lambda x: x.name == del_arg.name, function_details)[0])
|
|
|
|
function_tinfo.create_func(function_details)
|
|
idaapi.apply_tinfo2(vu.cfunc.entry_ea, function_tinfo, idaapi.TINFO_DEFINITE)
|
|
vu.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
return idaapi.AST_ENABLE_ALWAYS
|
|
|
|
|
|
class AddRemoveReturn(idaapi.action_handler_t):
|
|
|
|
name = "my:RemoveReturn"
|
|
description = "Add/Remove Return"
|
|
hotkey = None
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
def activate(self, ctx):
|
|
# ctx - action_activation_ctx_t
|
|
vu = idaapi.get_widget_vdui(ctx.widget)
|
|
function_tinfo = idaapi.tinfo_t()
|
|
if not vu.cfunc.get_func_type(function_tinfo):
|
|
return
|
|
function_details = idaapi.func_type_data_t()
|
|
function_tinfo.get_func_details(function_details)
|
|
if function_details.rettype.equals_to(Const.VOID_TINFO):
|
|
function_details.rettype = idaapi.tinfo_t(Const.PVOID_TINFO)
|
|
else:
|
|
function_details.rettype = idaapi.tinfo_t(idaapi.BT_VOID)
|
|
function_tinfo.create_func(function_details)
|
|
idaapi.apply_tinfo2(vu.cfunc.entry_ea, function_tinfo, idaapi.TINFO_DEFINITE)
|
|
vu.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
return idaapi.AST_ENABLE_ALWAYS
|
|
|
|
|
|
class ConvertToUsercall(idaapi.action_handler_t):
|
|
|
|
name = "my:ConvertToUsercall"
|
|
description = "Convert to __usercall"
|
|
hotkey = None
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
def activate(self, ctx):
|
|
# ctx - action_activation_ctx_t
|
|
vu = idaapi.get_widget_vdui(ctx.widget)
|
|
function_tinfo = idaapi.tinfo_t()
|
|
if not vu.cfunc.get_func_type(function_tinfo):
|
|
return
|
|
function_details = idaapi.func_type_data_t()
|
|
function_tinfo.get_func_details(function_details)
|
|
convention = idaapi.CM_CC_MASK & function_details.cc
|
|
if convention == idaapi.CM_CC_CDECL:
|
|
function_details.cc = idaapi.CM_CC_SPECIAL
|
|
elif convention in (idaapi.CM_CC_STDCALL, idaapi.CM_CC_FASTCALL, idaapi.CM_CC_PASCAL, idaapi.CM_CC_THISCALL):
|
|
function_details.cc = idaapi.CM_CC_SPECIALP
|
|
elif convention == idaapi.CM_CC_ELLIPSIS:
|
|
function_details.cc = idaapi.CM_CC_SPECIALE
|
|
else:
|
|
return
|
|
function_tinfo.create_func(function_details)
|
|
idaapi.apply_tinfo2(vu.cfunc.entry_ea, function_tinfo, idaapi.TINFO_DEFINITE)
|
|
vu.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
return idaapi.AST_ENABLE_ALWAYS
|
|
|
|
|
|
class GetStructureBySize(idaapi.action_handler_t):
|
|
# TODO: apply type automatically if expression like `var = new(size)`
|
|
|
|
name = "my:WhichStructHaveThisSize"
|
|
description = "Structures with this size"
|
|
hotkey = "W"
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def select_structure_by_size(size):
|
|
result = TypeLibrary.choose_til()
|
|
if result:
|
|
selected_library, max_ordinal, is_local_type = result
|
|
matched_types = []
|
|
tinfo = idaapi.tinfo_t()
|
|
for ordinal in xrange(1, max_ordinal):
|
|
tinfo.create_typedef(selected_library, ordinal)
|
|
if tinfo.get_size() == size:
|
|
name = tinfo.dstr()
|
|
description = idaapi.print_tinfo(None, 0, 0, idaapi.PRTYPE_DEF, tinfo, None, None)
|
|
matched_types.append([str(ordinal), name, description])
|
|
|
|
type_chooser = Forms.MyChoose(
|
|
matched_types,
|
|
"Select Type",
|
|
[["Ordinal", 5 | idaapi.Choose2.CHCOL_HEX], ["Type Name", 25], ["Declaration", 50]],
|
|
165
|
|
)
|
|
selected_type = type_chooser.Show(True)
|
|
if selected_type != -1:
|
|
if is_local_type:
|
|
return int(matched_types[selected_type][0])
|
|
return TypeLibrary.import_type(selected_library, matched_types[selected_type][1])
|
|
return None
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
if hx_view.item.citype != idaapi.VDI_EXPR or hx_view.item.e.op != idaapi.cot_num:
|
|
return
|
|
ea = ctx.cur_ea
|
|
c_number = hx_view.item.e.n
|
|
number_value = c_number._value
|
|
ordinal = GetStructureBySize.select_structure_by_size(number_value)
|
|
if ordinal:
|
|
number_format_old = c_number.nf
|
|
number_format_new = idaapi.number_format_t()
|
|
number_format_new.flags = idaapi.FF_1STRO | idaapi.FF_0STRO
|
|
operand_number = number_format_old.opnum
|
|
number_format_new.opnum = operand_number
|
|
number_format_new.props = number_format_old.props
|
|
number_format_new.type_name = idaapi.create_numbered_type_name(ordinal)
|
|
|
|
c_function = hx_view.cfunc
|
|
number_formats = c_function.numforms # idaapi.user_numforms_t
|
|
operand_locator = idaapi.operand_locator_t(ea, ord(operand_number) if operand_number else 0)
|
|
if operand_locator in number_formats:
|
|
del number_formats[operand_locator]
|
|
|
|
number_formats[operand_locator] = number_format_new
|
|
c_function.save_user_numforms()
|
|
hx_view.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class ShallowScanVariable(idaapi.action_handler_t):
|
|
|
|
name = "my:ShallowScanVariable"
|
|
description = "Scan Variable"
|
|
hotkey = "F"
|
|
|
|
def __init__(self, temporary_structure):
|
|
self.temporary_structure = temporary_structure
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(cfunc, ctree_item):
|
|
lvar = ctree_item.get_lvar()
|
|
if lvar is not None:
|
|
return Helper.is_legal_type(lvar.type())
|
|
|
|
if ctree_item.citype != idaapi.VDI_EXPR:
|
|
return False
|
|
|
|
obj = Api.ScanObject.create(cfunc, ctree_item.e)
|
|
return obj and Helper.is_legal_type(obj.tinfo)
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
cfunc = hx_view.cfunc
|
|
origin = self.temporary_structure.main_offset
|
|
|
|
if self.check(cfunc, hx_view.item):
|
|
obj = Api.ScanObject.create(cfunc, hx_view.item)
|
|
visitor = NewShallowSearchVisitor(cfunc, origin, obj, self.temporary_structure)
|
|
visitor.process()
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class DeepScanVariable(idaapi.action_handler_t):
|
|
|
|
name = "my:DeepScanVariable"
|
|
description = "Deep Scan Variable"
|
|
hotkey = "shift+F"
|
|
|
|
def __init__(self, temporary_structure):
|
|
self.temporary_structure = temporary_structure
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
cfunc = hx_view.cfunc
|
|
origin = self.temporary_structure.main_offset
|
|
|
|
if ShallowScanVariable.check(cfunc, hx_view.item):
|
|
obj = Api.ScanObject.create(cfunc, hx_view.item)
|
|
if FunctionTouchVisitor(cfunc).process():
|
|
hx_view.refresh_view(True)
|
|
visitor = NewDeepSearchVisitor(hx_view.cfunc, origin, obj, self.temporary_structure)
|
|
visitor.process()
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class DeepScanReturn(idaapi.action_handler_t):
|
|
name = "my:DeepScanReturn"
|
|
description = "Deep Scan Returned Variables"
|
|
hotkey = None
|
|
|
|
def __init__(self, temporary_structure):
|
|
self.temp_struct = temporary_structure
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(hx_view):
|
|
tinfo = idaapi.tinfo_t()
|
|
hx_view.cfunc.get_func_type(tinfo)
|
|
return not tinfo.get_rettype().equals_to(Const.VOID_TINFO)
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
func_ea = hx_view.cfunc.entry_ea
|
|
|
|
obj = Api.ReturnedObject(func_ea)
|
|
visitor = DeepReturnVisitor(hx_view.cfunc, self.temp_struct.main_offset, obj, self.temp_struct)
|
|
visitor.process()
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class DeepScanFunctions(idaapi.action_handler_t):
|
|
|
|
name = "my:DeepScanFunctions"
|
|
description = "Scan First Argument"
|
|
hotkey = None
|
|
|
|
def __init__(self, temporary_structure):
|
|
self.temporary_structure = temporary_structure
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
def activate(self, ctx):
|
|
for idx in ctx.chooser_selection:
|
|
func_ea = idaapi.getn_func(idx - 1).startEA
|
|
cfunc = Api.decompile_function(func_ea)
|
|
obj = Api.VariableObject(cfunc.get_lvars()[0], 0)
|
|
if cfunc:
|
|
NewDeepSearchVisitor(cfunc, 0, obj, self.temporary_structure).process()
|
|
|
|
def update(self, ctx):
|
|
if ctx.form_type == idaapi.BWN_FUNCS:
|
|
idaapi.attach_action_to_popup(ctx.widget, None, self.name)
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class RecognizeShape(idaapi.action_handler_t):
|
|
|
|
name = "my:RecognizeShape"
|
|
description = "Recognize Shape"
|
|
hotkey = None
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
cfunc = hx_view.cfunc
|
|
|
|
if not ShallowScanVariable.check(cfunc, hx_view.item):
|
|
return
|
|
|
|
obj = Api.ScanObject.create(cfunc, hx_view.item)
|
|
temp_struct = TemporaryStructureModel()
|
|
visitor = NewShallowSearchVisitor(cfunc, 0, obj, temp_struct)
|
|
visitor.process()
|
|
tinfo = temp_struct.get_recognized_shape()
|
|
if tinfo:
|
|
tinfo.create_ptr(tinfo)
|
|
if obj.id == Api.SO_LOCAL_VARIABLE:
|
|
hx_view.set_lvar_type(obj.lvar, tinfo)
|
|
elif obj.id == Api.SO_GLOBAL_OBJECT:
|
|
idaapi.apply_tinfo2(obj.obj_ea, tinfo, idaapi.TINFO_DEFINITE)
|
|
hx_view.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class CreateNewField(idaapi.action_handler_t):
|
|
name = "my:CreateNewField"
|
|
description = "Create New Field"
|
|
hotkey = None
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(cfunc, ctree_item):
|
|
if ctree_item.citype != idaapi.VDI_EXPR:
|
|
return
|
|
|
|
item = ctree_item.it.to_specific_type
|
|
if item.op != idaapi.cot_memptr:
|
|
return
|
|
|
|
parent = cfunc.body.find_parent_of(ctree_item.it).to_specific_type
|
|
if parent.op != idaapi.cot_idx or parent.y.op != idaapi.cot_num:
|
|
idx = 0
|
|
else:
|
|
idx = parent.y.numval()
|
|
|
|
struct_type = item.x.type.get_pointed_object()
|
|
udt_member = idaapi.udt_member_t()
|
|
udt_member.offset = item.m * 8
|
|
struct_type.find_udt_member(idaapi.STRMEM_OFFSET, udt_member)
|
|
if udt_member.name[0:3] != "gap":
|
|
return
|
|
|
|
return struct_type, udt_member.offset // 8, idx
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
result = self.check(hx_view.cfunc, hx_view.item)
|
|
if result is None:
|
|
return
|
|
|
|
struct_tinfo, offset, idx = result
|
|
ordinal = struct_tinfo.get_ordinal()
|
|
struct_name = struct_tinfo.dstr()
|
|
|
|
if (offset + idx) % 2:
|
|
default_field_type = "_BYTE"
|
|
elif (offset + idx) % 4:
|
|
default_field_type = "_WORD"
|
|
else:
|
|
default_field_type = "_DWORD"
|
|
|
|
declaration = idaapi.asktext(
|
|
0x10000, "{0} field_{1:X}".format(default_field_type, offset + idx), "Enter new structure member:"
|
|
)
|
|
if declaration is None:
|
|
return
|
|
|
|
result = self.parse_declaration(declaration)
|
|
if result is None:
|
|
logger.warn("Bad member declaration")
|
|
return
|
|
|
|
field_tinfo, field_name = result
|
|
field_size = field_tinfo.get_size()
|
|
udt_data = idaapi.udt_type_data_t()
|
|
udt_member = idaapi.udt_member_t()
|
|
|
|
struct_tinfo.get_udt_details(udt_data)
|
|
udt_member.offset = offset * 8
|
|
struct_tinfo.find_udt_member(idaapi.STRMEM_OFFSET, udt_member)
|
|
gap_size = udt_member.size // 8
|
|
|
|
gap_leftover = gap_size - idx - field_size
|
|
|
|
if gap_leftover < 0:
|
|
print "[ERROR] Too big size for the field. Type with maximum {0} bytes can be used".format(gap_size - idx)
|
|
return
|
|
|
|
iterator = udt_data.find(udt_member)
|
|
iterator = udt_data.erase(iterator)
|
|
|
|
if gap_leftover > 0:
|
|
udt_data.insert(iterator, TemporaryStructureModel.get_padding_member(offset + idx + field_size, gap_leftover))
|
|
|
|
udt_member = idaapi.udt_member_t()
|
|
udt_member.offset = offset * 8 + idx
|
|
udt_member.name = field_name
|
|
udt_member.type = field_tinfo
|
|
udt_member.size = field_size
|
|
|
|
iterator = udt_data.insert(iterator, udt_member)
|
|
|
|
if idx > 0:
|
|
udt_data.insert(iterator, TemporaryStructureModel.get_padding_member(offset, idx))
|
|
|
|
struct_tinfo.create_udt(udt_data, idaapi.BTF_STRUCT)
|
|
struct_tinfo.set_numbered_type(idaapi.cvar.idati, ordinal, idaapi.BTF_STRUCT, struct_name)
|
|
hx_view.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
@staticmethod
|
|
def parse_declaration(declaration):
|
|
m = re.search(r"^(\w+[ *]+)(\w+)(\[(\d+)\])?$", declaration)
|
|
if m is None:
|
|
return
|
|
|
|
type_name, field_name, _, arr_size = m.groups()
|
|
if field_name[0].isdigit():
|
|
print "[ERROR] Bad field name"
|
|
return
|
|
|
|
result = idc.ParseType(type_name, 0)
|
|
if result is None:
|
|
return
|
|
|
|
_, tp, fld = result
|
|
tinfo = idaapi.tinfo_t()
|
|
tinfo.deserialize(idaapi.cvar.idati, tp, fld, None)
|
|
if arr_size:
|
|
assert tinfo.create_array(tinfo, int(arr_size))
|
|
return tinfo, field_name
|
|
|
|
|
|
class ShowGraph(idaapi.action_handler_t):
|
|
name = "my:ShowGraph"
|
|
description = "Show graph"
|
|
hotkey = "G"
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
self.graph = None
|
|
self.graph_view = None
|
|
|
|
def activate(self, ctx):
|
|
"""
|
|
:param ctx: idaapi.action_activation_ctx_t
|
|
:return: None
|
|
"""
|
|
form = self.graph_view.GetTForm() if self.graph_view else None
|
|
if form:
|
|
self.graph_view.change_selected([sel + 1 for sel in ctx.chooser_selection])
|
|
self.graph_view.Show()
|
|
else:
|
|
self.graph = StructureGraph([sel + 1 for sel in ctx.chooser_selection])
|
|
self.graph_view = Forms.StructureGraphViewer("Structure Graph", self.graph)
|
|
self.graph_view.Show()
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_LOCTYPS:
|
|
idaapi.attach_action_to_popup(ctx.widget, None, self.name)
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class ShowClasses(idaapi.action_handler_t):
|
|
|
|
name = "my:ShowClasses"
|
|
description = "Classes"
|
|
hotkey = "Alt+F1"
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
def activate(self, ctx):
|
|
"""
|
|
:param ctx: idaapi.action_activation_ctx_t
|
|
:return: None
|
|
"""
|
|
tform = idaapi.find_tform('Classes')
|
|
if not tform:
|
|
class_viewer = Forms.ClassViewer()
|
|
class_viewer.Show()
|
|
else:
|
|
idaapi.switchto_tform(tform, True)
|
|
|
|
def update(self, ctx):
|
|
return idaapi.AST_ENABLE_ALWAYS
|
|
|
|
|
|
class CreateVtable(idaapi.action_handler_t):
|
|
|
|
name = "my:CreateVtable"
|
|
description = "Create Virtual Table"
|
|
hotkey = "V"
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
def activate(self, ctx):
|
|
ea = ctx.cur_ea
|
|
if ea != idaapi.BADADDR and VirtualTable.check_address(ea):
|
|
vtable = VirtualTable(0, ea)
|
|
vtable.import_to_structures(True)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_DISASM:
|
|
idaapi.attach_action_to_popup(ctx.widget, None, self.name)
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class SelectContainingStructure(idaapi.action_handler_t):
|
|
|
|
name = "my:SelectContainingStructure"
|
|
description = "Select Containing Structure"
|
|
hotkey = None
|
|
|
|
def __init__(self, potential_negatives):
|
|
idaapi.action_handler_t.__init__(self)
|
|
self.potential_negative = potential_negatives
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
result = TypeLibrary.choose_til()
|
|
if result:
|
|
selected_library, max_ordinal, is_local_types = result
|
|
lvar_idx = hx_view.item.e.v.idx
|
|
candidate = self.potential_negative[lvar_idx]
|
|
structures = candidate.find_containing_structures(selected_library)
|
|
items = map(lambda x: [str(x[0]), "0x{0:08X}".format(x[1]), x[2], x[3]], structures)
|
|
structure_chooser = Forms.MyChoose(
|
|
items,
|
|
"Select Containing Structure",
|
|
[["Ordinal", 5], ["Offset", 10], ["Member_name", 20], ["Structure Name", 20]],
|
|
165
|
|
)
|
|
selected_idx = structure_chooser.Show(modal=True)
|
|
if selected_idx != -1:
|
|
if not is_local_types:
|
|
TypeLibrary.import_type(selected_library, items[selected_idx][3])
|
|
lvar = hx_view.cfunc.get_lvars()[lvar_idx]
|
|
lvar_cmt = re.sub("```.*```", '', lvar.cmt)
|
|
hx_view.set_lvar_cmt(
|
|
lvar,
|
|
lvar_cmt + "```{0}+{1}```".format(
|
|
structures[selected_idx][3],
|
|
structures[selected_idx][1])
|
|
)
|
|
hx_view.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
return idaapi.AST_ENABLE_ALWAYS
|
|
|
|
|
|
class ResetContainingStructure(idaapi.action_handler_t):
|
|
|
|
name = "my:ResetContainingStructure"
|
|
description = "Reset Containing Structure"
|
|
hotkey = None
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(lvar):
|
|
return True if re.search("```.*```", lvar.cmt) else False
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
lvar = hx_view.cfunc.get_lvars()[hx_view.item.e.v.idx]
|
|
hx_view.set_lvar_cmt(lvar, re.sub("```.*```", '', lvar.cmt))
|
|
hx_view.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
return idaapi.AST_ENABLE_ALWAYS
|
|
|
|
|
|
class RecastItemLeft(idaapi.action_handler_t):
|
|
|
|
name = "my:RecastItemLeft"
|
|
description = "Recast Item"
|
|
hotkey = "Shift+L"
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(cfunc, ctree_item):
|
|
if ctree_item.citype != idaapi.VDI_EXPR:
|
|
return
|
|
|
|
expression = ctree_item.it.to_specific_type
|
|
child = None
|
|
|
|
# Look through parents until we found Return, Assignment or Call
|
|
while expression and expression.op not in (idaapi.cot_asg, idaapi.cit_return, idaapi.cot_call):
|
|
child = expression.to_specific_type
|
|
expression = cfunc.body.find_parent_of(expression)
|
|
if not expression:
|
|
return
|
|
|
|
expression = expression.to_specific_type
|
|
if expression.opname == 'asg':
|
|
|
|
if expression.x.opname not in ('var', 'obj', 'memptr', 'memref'):
|
|
return
|
|
|
|
right_expr = expression.y
|
|
right_tinfo = right_expr.x.type if right_expr.op == idaapi.cot_cast else right_expr.type
|
|
|
|
# Check if both left and right parts of expression are of the same types.
|
|
# If not then we can recast then.
|
|
if right_tinfo.dstr() == expression.x.type.dstr():
|
|
return
|
|
|
|
if expression.x.op == idaapi.cot_var:
|
|
# var = (TYPE ) ...;
|
|
variable = cfunc.get_lvars()[expression.x.v.idx]
|
|
idaapi.update_action_label(RecastItemLeft.name, 'Recast Variable "{0}"'.format(variable.name))
|
|
return RECAST_LOCAL_VARIABLE, right_tinfo, variable
|
|
elif expression.x.op == idaapi.cot_obj:
|
|
# g_var = (TYPE ) ...;
|
|
idaapi.update_action_label(RecastItemLeft.name, 'Recast Global')
|
|
return RECAST_GLOBAL_VARIABLE, right_tinfo, expression.x.obj_ea
|
|
elif expression.x.op == idaapi.cot_memptr:
|
|
# struct->member = (TYPE ) ...;
|
|
idaapi.update_action_label(RecastItemLeft.name, 'Recast Field')
|
|
return RECAST_STRUCTURE, expression.x.x.type.get_pointed_object().dstr(), expression.x.m, right_tinfo
|
|
elif expression.x.op == idaapi.cot_memref:
|
|
# struct.member = (TYPE ) ...;
|
|
idaapi.update_action_label(RecastItemLeft.name, 'Recast Field')
|
|
return RECAST_STRUCTURE, expression.x.x.type.dstr(), expression.x.m, right_tinfo
|
|
|
|
elif expression.op == idaapi.cit_return:
|
|
|
|
idaapi.update_action_label(RecastItemLeft.name, "Recast Return")
|
|
child = child or expression.creturn.expr
|
|
|
|
if child.op == idaapi.cot_cast:
|
|
# return (TYPE) ...;
|
|
return RECAST_RETURN, child.x.type, None
|
|
|
|
func_tinfo = idaapi.tinfo_t()
|
|
cfunc.get_func_type(func_tinfo)
|
|
rettype = func_tinfo.get_rettype()
|
|
|
|
if func_tinfo.get_rettype().dstr() != child.type.dstr():
|
|
# return ...;
|
|
# This's possible when returned type and value are both pointers to different types
|
|
return RECAST_RETURN, child.type, None
|
|
|
|
elif expression.op == idaapi.cot_call:
|
|
if expression.x.op == idaapi.cot_memptr:
|
|
if expression.x == child:
|
|
return
|
|
|
|
arg_index, arg_tinfo = Helper.get_func_argument_info(expression, child)
|
|
if child.op == idaapi.cot_cast:
|
|
# struct_ptr->func(..., (TYPE) var, ...);
|
|
new_arg_tinfo = child.x.type
|
|
else:
|
|
# struct_ptr->func(..., var, ...); When `var` and `arg` are different pointers
|
|
if arg_tinfo.equals_to(child.type):
|
|
return
|
|
new_arg_tinfo = child.type
|
|
|
|
struct_type = expression.x.x.type.get_pointed_object()
|
|
funcptr_tinfo = expression.x.type
|
|
Helper.set_funcptr_argument(funcptr_tinfo, arg_index, new_arg_tinfo)
|
|
return RECAST_STRUCTURE, struct_type.dstr(), expression.x.m, funcptr_tinfo
|
|
|
|
if child and child.op == idaapi.cot_cast:
|
|
if child.cexpr.x.op == idaapi.cot_memptr and expression.ea == idaapi.BADADDR:
|
|
idaapi.update_action_label(RecastItemLeft.name, 'Recast Virtual Function')
|
|
return RECAST_STRUCTURE, child.cexpr.x.x.type.get_pointed_object().dstr(), child.cexpr.x.m, child.type
|
|
|
|
if expression.x == child.cexpr:
|
|
return
|
|
arg_index, _ = Helper.get_func_argument_info(expression, child.cexpr)
|
|
idaapi.update_action_label(RecastItemLeft.name, "Recast Argument")
|
|
return (
|
|
RECAST_ARGUMENT,
|
|
arg_index,
|
|
expression.x.type.get_pointed_object(),
|
|
child.x.type,
|
|
expression.x.obj_ea
|
|
)
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
result = self.check(hx_view.cfunc, hx_view.item)
|
|
|
|
if not result:
|
|
return
|
|
|
|
if result[0] == RECAST_LOCAL_VARIABLE:
|
|
logger.debug("Recasting local variable. Type - %s", result[1].dstr())
|
|
tinfo, lvar = result[1:]
|
|
if hx_view.set_lvar_type(lvar, tinfo):
|
|
hx_view.refresh_view(True)
|
|
|
|
elif result[0] == RECAST_GLOBAL_VARIABLE:
|
|
logger.debug("Recasting global. Type - %s. Address - %s", result[1].dstr(), Helper.to_hex(result[2]))
|
|
tinfo, address = result[1:]
|
|
if idaapi.apply_tinfo2(address, tinfo, idaapi.TINFO_DEFINITE):
|
|
hx_view.refresh_view(True)
|
|
|
|
elif result[0] == RECAST_ARGUMENT:
|
|
arg_index, func_tinfo, arg_tinfo, address = result[1:]
|
|
if arg_tinfo.is_array():
|
|
arg_tinfo.convert_array_to_ptr()
|
|
|
|
func_data = idaapi.func_type_data_t()
|
|
func_tinfo.get_func_details(func_data)
|
|
func_data[arg_index].type = arg_tinfo
|
|
new_func_tinfo = idaapi.tinfo_t()
|
|
new_func_tinfo.create_func(func_data)
|
|
if idaapi.apply_tinfo2(address, new_func_tinfo, idaapi.TINFO_DEFINITE):
|
|
hx_view.refresh_view(True)
|
|
|
|
elif result[0] == RECAST_RETURN:
|
|
return_type, func_address = result[1:]
|
|
try:
|
|
cfunc = idaapi.decompile(func_address) if func_address else hx_view.cfunc
|
|
except idaapi.DecompilationFailure:
|
|
print "[ERROR] Ida failed to decompile function at 0x{0:08X}".format(func_address)
|
|
return
|
|
|
|
function_tinfo = idaapi.tinfo_t()
|
|
cfunc.get_func_type(function_tinfo)
|
|
func_data = idaapi.func_type_data_t()
|
|
function_tinfo.get_func_details(func_data)
|
|
func_data.rettype = return_type
|
|
function_tinfo.create_func(func_data)
|
|
if idaapi.apply_tinfo2(cfunc.entry_ea, function_tinfo, idaapi.TINFO_DEFINITE):
|
|
hx_view.refresh_view(True)
|
|
|
|
elif result[0] == RECAST_STRUCTURE:
|
|
structure_name, field_offset, new_type = result[1:]
|
|
tinfo = idaapi.tinfo_t()
|
|
tinfo.get_named_type(idaapi.cvar.idati, structure_name)
|
|
|
|
ordinal = idaapi.get_type_ordinal(idaapi.cvar.idati, structure_name)
|
|
|
|
if ordinal:
|
|
udt_member = idaapi.udt_member_t()
|
|
udt_member.offset = field_offset * 8
|
|
idx = tinfo.find_udt_member(idaapi.STRMEM_OFFSET, udt_member)
|
|
if udt_member.offset != field_offset * 8:
|
|
print "[Info] Can't handle with arrays yet"
|
|
elif udt_member.type.get_size() != new_type.get_size():
|
|
print "[Info] Can't recast different sizes yet"
|
|
else:
|
|
udt_data = idaapi.udt_type_data_t()
|
|
tinfo.get_udt_details(udt_data)
|
|
udt_data[idx].type = new_type
|
|
tinfo.create_udt(udt_data, idaapi.BTF_STRUCT)
|
|
tinfo.set_numbered_type(idaapi.cvar.idati, ordinal, idaapi.NTF_REPLACE, structure_name)
|
|
hx_view.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class RecastItemRight(RecastItemLeft):
|
|
|
|
name = "my:RecastItemRight"
|
|
description = "Recast Item"
|
|
hotkey = "Shift+R"
|
|
|
|
def __init__(self):
|
|
RecastItemLeft.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(cfunc, ctree_item):
|
|
if ctree_item.citype != idaapi.VDI_EXPR:
|
|
return
|
|
|
|
expression = ctree_item.it
|
|
|
|
result = RecastItemRight._check_potential_array(cfunc, expression)
|
|
if result:
|
|
return result
|
|
|
|
# Look through parents until we found Cast
|
|
while expression and expression.op != idaapi.cot_cast:
|
|
expression = expression.to_specific_type
|
|
expression = cfunc.body.find_parent_of(expression)
|
|
if not expression:
|
|
return
|
|
|
|
expression = expression.to_specific_type
|
|
|
|
# Find `(TYPE) something;` or `(TYPE *) &something;` and calculate appropriate type for recast
|
|
if expression.x.op == idaapi.cot_ref:
|
|
new_type = expression.type.get_pointed_object()
|
|
expression = expression.x
|
|
else:
|
|
new_type = expression.type
|
|
|
|
if expression.x.op == idaapi.cot_var:
|
|
# (TYPE) var;
|
|
variable = cfunc.get_lvars()[expression.x.v.idx]
|
|
idaapi.update_action_label(RecastItemRight.name, 'Recast Variable "{0}"'.format(variable.name))
|
|
return RECAST_LOCAL_VARIABLE, new_type, variable
|
|
|
|
elif expression.x.op == idaapi.cot_obj:
|
|
# (TYPE) g_var;
|
|
if Helper.is_code_ea(expression.x.obj_ea) and new_type.is_funcptr():
|
|
# (TYPE) sub_XXXXXX;
|
|
new_type = new_type.get_pointed_object()
|
|
|
|
idaapi.update_action_label(RecastItemRight.name, 'Recast Global')
|
|
return RECAST_GLOBAL_VARIABLE, new_type, expression.x.obj_ea
|
|
|
|
elif expression.x.op == idaapi.cot_call:
|
|
# (TYPE) call();
|
|
idaapi.update_action_label(RecastItemRight.name, "Recast Return")
|
|
return RECAST_RETURN, new_type, expression.x.x.obj_ea
|
|
|
|
elif expression.x.op == idaapi.cot_memptr:
|
|
# (TYPE) var->member;
|
|
idaapi.update_action_label(RecastItemRight.name, "Recast Field")
|
|
return RECAST_STRUCTURE, expression.x.x.type.get_pointed_object().dstr(), expression.x.m, new_type
|
|
|
|
@staticmethod
|
|
def _check_potential_array(cfunc, expr):
|
|
""" Checks `call(..., &buffer, ..., number)` and returns information for recasting """
|
|
if expr.op != idaapi.cot_var:
|
|
return
|
|
|
|
var_expr = expr.to_specific_type
|
|
parent = cfunc.body.find_parent_of(expr)
|
|
if parent.op != idaapi.cot_ref:
|
|
return
|
|
|
|
parent = cfunc.body.find_parent_of(parent)
|
|
if parent.op != idaapi.cot_call:
|
|
return
|
|
|
|
call_expr = parent.to_specific_type
|
|
for arg_expr in call_expr.a:
|
|
if arg_expr.op == idaapi.cot_num:
|
|
number = arg_expr.numval()
|
|
if number:
|
|
variable = cfunc.lvars[var_expr.v.idx]
|
|
char_array_tinfo = idaapi.tinfo_t()
|
|
char_array_tinfo.create_array(idaapi.tinfo_t(idaapi.BTF_CHAR), number)
|
|
idaapi.update_action_label(RecastItemRight.name, 'Recast Variable "{}" to "{}"'.format(
|
|
variable.name, char_array_tinfo.dstr()
|
|
))
|
|
return RECAST_LOCAL_VARIABLE, char_array_tinfo, variable
|
|
|
|
|
|
class RenameOther(idaapi.action_handler_t):
|
|
name = "my:RenameOther"
|
|
description = "Take other name"
|
|
hotkey = "Ctrl+N"
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(cfunc, ctree_item):
|
|
if ctree_item.citype != idaapi.VDI_EXPR:
|
|
return
|
|
|
|
expression = ctree_item.it.to_specific_type
|
|
if expression.op != idaapi.cot_var:
|
|
return
|
|
|
|
parent = cfunc.body.find_parent_of(expression).to_specific_type
|
|
if parent.op != idaapi.cot_asg:
|
|
return
|
|
|
|
other = parent.theother(expression)
|
|
if other.op != idaapi.cot_var:
|
|
return
|
|
|
|
this_lvar = ctree_item.get_lvar()
|
|
other_lvar = cfunc.get_lvars()[other.v.idx]
|
|
if (other_lvar.has_user_name or other_lvar.is_arg_var and re.search("a\d*$", other_lvar.name) is None) \
|
|
and this_lvar.name.lstrip('_') != other_lvar.name.lstrip('_'):
|
|
return '_' + other_lvar.name, this_lvar
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
result = self.check(hx_view.cfunc, hx_view.item)
|
|
|
|
if result:
|
|
name, lvar = result
|
|
hx_view.rename_lvar(lvar, name, True)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class RenameInside(idaapi.action_handler_t):
|
|
name = "my:RenameInto"
|
|
description = "Rename inside argument"
|
|
hotkey = "Shift+N"
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(cfunc, ctree_item):
|
|
if ctree_item.citype != idaapi.VDI_EXPR:
|
|
return False
|
|
|
|
expression = ctree_item.it.to_specific_type
|
|
if expression.op == idaapi.cot_var:
|
|
lvar = ctree_item.get_lvar()
|
|
# Check if it's either variable with user name or argument with not standard `aX` name
|
|
if lvar.has_user_name or lvar.is_arg_var and re.search("a\d*$", lvar.name) is None:
|
|
parent = cfunc.body.find_parent_of(expression).to_specific_type
|
|
if parent.op == idaapi.cot_call:
|
|
arg_index, _ = Helper.get_func_argument_info(parent, expression)
|
|
func_tinfo = parent.x.type.get_pointed_object()
|
|
func_data = idaapi.func_type_data_t()
|
|
func_tinfo.get_func_details(func_data)
|
|
if arg_index < func_tinfo.get_nargs() and lvar.name.lstrip('_') != func_data[arg_index].name:
|
|
return func_tinfo, parent.x.obj_ea, arg_index, lvar.name.lstrip('_')
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
result = self.check(hx_view.cfunc, hx_view.item)
|
|
|
|
if result:
|
|
func_tinfo, address, arg_index, name = result
|
|
|
|
func_data = idaapi.func_type_data_t()
|
|
func_tinfo.get_func_details(func_data)
|
|
func_data[arg_index].name = name
|
|
new_func_tinfo = idaapi.tinfo_t()
|
|
new_func_tinfo.create_func(func_data)
|
|
idaapi.apply_tinfo2(address, new_func_tinfo, idaapi.TINFO_DEFINITE)
|
|
hx_view.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class RenameOutside(idaapi.action_handler_t):
|
|
name = "my:RenameOutside"
|
|
description = "Take argument name"
|
|
hotkey = "Ctrl+Shift+N"
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(cfunc, ctree_item):
|
|
if ctree_item.citype != idaapi.VDI_EXPR:
|
|
return False
|
|
|
|
expression = ctree_item.it.to_specific_type
|
|
if expression.op == idaapi.cot_var:
|
|
lvar = ctree_item.get_lvar()
|
|
parent = cfunc.body.find_parent_of(expression).to_specific_type
|
|
|
|
if parent.op == idaapi.cot_call:
|
|
arg_index, _ = Helper.get_func_argument_info(parent, expression)
|
|
func_tinfo = parent.x.type.get_pointed_object()
|
|
if func_tinfo.get_nargs() < arg_index:
|
|
return
|
|
func_data = idaapi.func_type_data_t()
|
|
func_tinfo.get_func_details(func_data)
|
|
name = func_data[arg_index].name
|
|
if name and re.search("a\d*$", name) is None and name != 'this' and name != lvar.name:
|
|
return name, lvar
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
result = self.check(hx_view.cfunc, hx_view.item)
|
|
|
|
if result:
|
|
name, lvar = result
|
|
hx_view.rename_lvar(lvar, name, True)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class RenameUsingAssertVisitor(idaapi.ctree_parentee_t):
|
|
|
|
def __init__(self, cfunc, func_addr, arg_idx):
|
|
idaapi.ctree_parentee_t.__init__(self)
|
|
self.__cfunc = cfunc
|
|
self.__func_addr = func_addr
|
|
self.__arg_idx = arg_idx
|
|
self.__possible_names = set()
|
|
|
|
def visit_expr(self, expr):
|
|
if expr.op == idaapi.cot_call and expr.x.op == idaapi.cot_obj and expr.x.obj_ea == self.__func_addr:
|
|
arg_expr = expr.a[self.__arg_idx]
|
|
if arg_expr.op != idaapi.cot_obj:
|
|
logger.error("Argument is not string at {}".format(Helper.to_hex(self._find_asm_address(expr))))
|
|
return 1
|
|
self.__add_func_name(arg_expr)
|
|
return 0
|
|
|
|
def process(self):
|
|
self.apply_to(self.__cfunc.body, None)
|
|
if len(self.__possible_names) == 1:
|
|
self.__rename_func()
|
|
else:
|
|
logger.error("Function at {} has more than one candidate for renaming: {}".format(
|
|
Helper.to_hex(self.__cfunc.entry_ea), ", ".join(self.__possible_names)))
|
|
|
|
def __add_func_name(self, arg_expr):
|
|
new_name = idc.get_strlit_contents(arg_expr.obj_ea)
|
|
if not idaapi.is_valid_typename(new_name):
|
|
logger.warn("Argument has weird name `{}` at {}".format(
|
|
new_name, Helper.to_hex(self._find_asm_address(arg_expr))))
|
|
return
|
|
|
|
self.__possible_names.add(new_name)
|
|
|
|
def __rename_func(self):
|
|
idc.set_name(self.__cfunc.entry_ea, self.__possible_names.pop())
|
|
|
|
|
|
class RenameUsingAssert(idaapi.action_handler_t):
|
|
|
|
name = "my:RenameUsingAssert"
|
|
description = "Rename as assert argument"
|
|
hotkey = None
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(cfunc, ctree_item):
|
|
if ctree_item.citype != idaapi.VDI_EXPR:
|
|
return False
|
|
|
|
expression = ctree_item.it.to_specific_type
|
|
if expression.op != idaapi.cot_obj:
|
|
return False
|
|
|
|
parent = cfunc.body.find_parent_of(expression).to_specific_type
|
|
if parent.op != idaapi.cot_call or parent.x.op != idaapi.cot_obj:
|
|
return False
|
|
|
|
obj_ea = expression.obj_ea
|
|
if not Helper.is_code_ea(obj_ea) and idc.get_str_type(obj_ea) == idc.STRTYPE_C:
|
|
str_potential_name = idc.get_strlit_contents(obj_ea)
|
|
return idaapi.is_valid_typename(str_potential_name)
|
|
return False
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
cfunc = hx_view.cfunc
|
|
ctree_item = hx_view.item
|
|
if not self.check(cfunc, ctree_item):
|
|
return
|
|
|
|
expr_arg = ctree_item.it.to_specific_type
|
|
expr_call = cfunc.body.find_parent_of(expr_arg).to_specific_type
|
|
|
|
arg_idx, _ = Helper.get_func_argument_info(expr_call, expr_arg)
|
|
|
|
assert_ea = expr_call.x.obj_ea
|
|
all_callers = Helper.get_funcs_calling_address(assert_ea)
|
|
|
|
for caller_ea in all_callers:
|
|
try:
|
|
cfunc = idaapi.decompile(caller_ea)
|
|
if not cfunc:
|
|
raise idaapi.DecompilationFailure
|
|
|
|
RenameUsingAssertVisitor(cfunc, assert_ea, arg_idx).process()
|
|
|
|
except idaapi.DecompilationFailure:
|
|
logger.warn("IDA failed to decompile at {}".format(Helper.to_hex(caller_ea)))
|
|
|
|
hx_view.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class PropagateName(idaapi.action_handler_t):
|
|
name = "my:PropagateName"
|
|
description = "Propagate name"
|
|
hotkey = "P"
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def callback_start(self):
|
|
hx_view, _ = self._data
|
|
hx_view.switch_to(self._cfunc, False)
|
|
|
|
@staticmethod
|
|
def callback_manipulate(self, cexpr, obj):
|
|
if self.crippled:
|
|
logger.debug("Skipping crippled function at {}".format(Helper.to_hex(self._cfunc.entry_ea)))
|
|
return
|
|
|
|
if obj.id == Api.SO_GLOBAL_OBJECT:
|
|
old_name = idaapi.get_short_name(cexpr.obj_ea)
|
|
if Settings.PROPAGATE_THROUGH_ALL_NAMES or PropagateName._is_default_name(old_name):
|
|
_, name = self._data
|
|
new_name = PropagateName.rename(lambda x: idaapi.set_name(cexpr.obj_ea, x), name)
|
|
logger.debug("Renamed global variable from {} to {}".format(old_name, new_name))
|
|
elif obj.id == Api.SO_LOCAL_VARIABLE:
|
|
lvar = self._cfunc.get_lvars()[cexpr.v.idx]
|
|
old_name = lvar.name
|
|
if Settings.PROPAGATE_THROUGH_ALL_NAMES or PropagateName._is_default_name(old_name):
|
|
hx_view, name = self._data
|
|
new_name = PropagateName.rename(lambda x: hx_view.rename_lvar(lvar, x, True), name)
|
|
logger.debug("Renamed local variable from {} to {}".format(old_name, new_name))
|
|
elif obj.id in (Api.SO_STRUCT_POINTER, Api.SO_STRUCT_REFERENCE):
|
|
struct_tinfo = cexpr.x.type
|
|
offset = cexpr.m
|
|
struct_tinfo.remove_ptr_or_array()
|
|
old_name = Helper.get_member_name(struct_tinfo, offset)
|
|
if Settings.PROPAGATE_THROUGH_ALL_NAMES or PropagateName._is_default_name(old_name):
|
|
_, name = self._data
|
|
new_name = PropagateName.rename(lambda x: Helper.change_member_name(struct_tinfo.dstr(), offset, x), name)
|
|
logger.debug("Renamed struct member from {} to {}".format(old_name, new_name))
|
|
|
|
@staticmethod
|
|
def rename(rename_func, name):
|
|
while not rename_func(name):
|
|
name = "_" + name
|
|
return name
|
|
|
|
@staticmethod
|
|
def _is_default_name(string):
|
|
return re.match(r"[av]\d+$", string) is not None or \
|
|
re.match(r"this|[qd]?word|field_|off_", string) is not None
|
|
|
|
@staticmethod
|
|
def check(cfunc, ctree_item):
|
|
if ctree_item.citype != idaapi.VDI_EXPR:
|
|
return
|
|
|
|
obj = Api.ScanObject.create(cfunc, ctree_item)
|
|
if obj and not PropagateName._is_default_name(obj.name):
|
|
return obj
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
obj = self.check(hx_view.cfunc, hx_view.item)
|
|
if obj:
|
|
cfunc = hx_view.cfunc
|
|
visitor = Api.RecursiveObjectDownwardsVisitor(cfunc, obj, (hx_view, obj.name), True)
|
|
visitor.set_callbacks(
|
|
manipulate=PropagateName.callback_manipulate,
|
|
start_iteration=PropagateName.callback_start,
|
|
finish=lambda x: hx_view.switch_to(cfunc, True)
|
|
)
|
|
visitor.process()
|
|
hx_view.refresh_view(True)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class GuessAllocation(idaapi.action_handler_t):
|
|
name = "my:ActionApi"
|
|
description = "Guess allocation"
|
|
hotkey = None
|
|
|
|
class StructAllocChoose(Forms.MyChoose):
|
|
def __init__(self, items):
|
|
Forms.MyChoose.__init__(
|
|
self, items, "Possible structure allocations",
|
|
[["Function", 30], ["Variable", 10], ["Line", 50], ["Type", 10]]
|
|
)
|
|
|
|
def OnSelectLine(self, n):
|
|
idaapi.jumpto(self.items[n][0])
|
|
|
|
def OnGetLine(self, n):
|
|
func_ea, var, line, alloc_type = self.items[n]
|
|
return [Helper.to_nice_str(func_ea), var, line, alloc_type]
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(cfunc, ctree_item):
|
|
if ctree_item.citype != idaapi.VDI_EXPR:
|
|
return
|
|
return Api.ScanObject.create(cfunc, ctree_item)
|
|
|
|
@staticmethod
|
|
def callback_manipulate(self, cexpr, obj):
|
|
if obj.id == Api.SO_LOCAL_VARIABLE:
|
|
parent = self.parent_expr()
|
|
if parent.op == idaapi.cot_asg:
|
|
alloc_obj = Api.MemoryAllocationObject.create(self._cfunc, self.parent_expr().y)
|
|
if alloc_obj:
|
|
self._data.append([alloc_obj.ea, obj.name, self._get_line(), "HEAP"])
|
|
elif self.parent_expr().op == idaapi.cot_ref:
|
|
self._data.append([self._find_asm_address(cexpr), obj.name, self._get_line(), "STACK"])
|
|
elif obj.id == Api.SO_GLOBAL_OBJECT:
|
|
self._data.append([self._find_asm_address(cexpr), obj.name, self._get_line(), "GLOBAL"])
|
|
|
|
@staticmethod
|
|
def callback_finish(self):
|
|
chooser = GuessAllocation.StructAllocChoose(self._data)
|
|
chooser.Show(False)
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
item = hx_view.item
|
|
obj = GuessAllocation.check(hx_view.cfunc, item)
|
|
if obj:
|
|
visitor = Api.RecursiveObjectUpwardsVisitor(hx_view.cfunc, obj, data=[], skip_after_object=True)
|
|
visitor.set_callbacks(
|
|
manipulate=self.callback_manipulate,
|
|
finish=self.callback_finish
|
|
)
|
|
visitor.process()
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class SwapThenElse(idaapi.action_handler_t):
|
|
name = "my:SwapIfElse"
|
|
description = "Swap then/else"
|
|
hotkey = "Shift+S"
|
|
|
|
def __init__(self):
|
|
idaapi.action_handler_t.__init__(self)
|
|
|
|
@staticmethod
|
|
def check(cfunc, ctree_item):
|
|
if ctree_item.citype != idaapi.VDI_EXPR:
|
|
return False
|
|
|
|
insn = ctree_item.it.to_specific_type
|
|
|
|
if insn.op != idaapi.cit_if or insn.cif.ielse is None:
|
|
return False
|
|
|
|
return insn.op == idaapi.cit_if and insn.cif.ielse
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
if self.check(hx_view.cfunc, hx_view.item):
|
|
insn = hx_view.item.it.to_specific_type
|
|
inverse_if(insn.cif)
|
|
hx_view.refresh_ctext()
|
|
|
|
InversionInfo(hx_view.cfunc.entry_ea).switch_inverted(insn.ea)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|
|
|
|
|
|
class FindFieldXrefs(idaapi.action_handler_t):
|
|
name = "my:FindFieldXrefs"
|
|
description = "Field Xrefs"
|
|
hotkey = "Ctrl+X"
|
|
|
|
@staticmethod
|
|
def check(ctree_item):
|
|
return ctree_item.citype == idaapi.VDI_EXPR and \
|
|
ctree_item.it.to_specific_type.op in (idaapi.cot_memptr, idaapi.cot_memref)
|
|
|
|
def activate(self, ctx):
|
|
hx_view = idaapi.get_widget_vdui(ctx.widget)
|
|
item = hx_view.item
|
|
|
|
if not self.check(item):
|
|
return
|
|
|
|
data = []
|
|
offset = item.e.m
|
|
struct_type = idaapi.remove_pointer(item.e.x.type)
|
|
ordinal = struct_type.get_ordinal()
|
|
result = XrefStorage().get_structure_info(ordinal, offset)
|
|
for xref_info in result:
|
|
data.append([
|
|
idaapi.get_short_name(xref_info.func_ea) + "+" + hex(int(xref_info.offset)),
|
|
xref_info.type,
|
|
xref_info.line
|
|
])
|
|
|
|
field_name = Helper.get_member_name(struct_type, offset)
|
|
chooser = Forms.MyChoose(
|
|
data,
|
|
"Cross-references to {0}::{1}".format(struct_type.dstr(), field_name),
|
|
[["Function", 20 | idaapi.Choose2.CHCOL_PLAIN],
|
|
["Type", 2 | idaapi.Choose2.CHCOL_PLAIN],
|
|
["Line", 40 | idaapi.Choose2.CHCOL_PLAIN]]
|
|
)
|
|
idx = chooser.Show(True)
|
|
if idx == -1:
|
|
return
|
|
|
|
xref = result[idx]
|
|
idaapi.open_pseudocode(xref.func_ea + xref.offset, False)
|
|
|
|
def update(self, ctx):
|
|
if ctx.widget_type == idaapi.BWN_PSEUDOCODE:
|
|
return idaapi.AST_ENABLE_FOR_FORM
|
|
return idaapi.AST_DISABLE_FOR_FORM
|