├── HexRaysPyTools ├── __init__.py ├── core │ ├── __init__.py │ ├── const.py │ ├── cache.py │ ├── type_library.py │ ├── struct_xrefs.py │ ├── common.py │ ├── structure_graph.py │ ├── helper.py │ ├── variable_scanner.py │ └── classes.py ├── callbacks │ ├── __init__.py │ ├── callbacks.py │ ├── virtual_table_creation.py │ ├── struct_xref_representation.py │ ├── member_double_click.py │ ├── guess_allocation.py │ ├── form_requests.py │ ├── actions.py │ ├── structs_by_size.py │ ├── function_signature_modifiers.py │ ├── struct_xref_collector.py │ ├── new_field_creation.py │ ├── scanners.py │ ├── swap_if.py │ ├── recasts.py │ ├── renames.py │ └── negative_offsets.py ├── settings.py ├── forms.py └── api.py ├── Img ├── bad.JPG ├── good.JPG ├── builder.JPG ├── classes.JPG ├── fields_xref.JPG ├── structure_builder.JPG └── virtual_functions.JPG ├── HexRaysPyTools.py ├── .gitignore └── readme.md /HexRaysPyTools/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /HexRaysPyTools/core/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /Img/bad.JPG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/igogo-x86/HexRaysPyTools/HEAD/Img/bad.JPG -------------------------------------------------------------------------------- /Img/good.JPG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/igogo-x86/HexRaysPyTools/HEAD/Img/good.JPG -------------------------------------------------------------------------------- /Img/builder.JPG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/igogo-x86/HexRaysPyTools/HEAD/Img/builder.JPG -------------------------------------------------------------------------------- /Img/classes.JPG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/igogo-x86/HexRaysPyTools/HEAD/Img/classes.JPG -------------------------------------------------------------------------------- /Img/fields_xref.JPG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/igogo-x86/HexRaysPyTools/HEAD/Img/fields_xref.JPG -------------------------------------------------------------------------------- /Img/structure_builder.JPG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/igogo-x86/HexRaysPyTools/HEAD/Img/structure_builder.JPG -------------------------------------------------------------------------------- /Img/virtual_functions.JPG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/igogo-x86/HexRaysPyTools/HEAD/Img/virtual_functions.JPG -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/__init__.py: -------------------------------------------------------------------------------- 1 | from .actions import * 2 | from .callbacks import * 3 | from . import form_requests 4 | from . import function_signature_modifiers 5 | from . import guess_allocation 6 | from . import member_double_click 7 | from . import negative_offsets 8 | from . import new_field_creation 9 | from . import recasts 10 | from . import renames 11 | from . import scanners 12 | from . import structs_by_size 13 | from . import struct_xref_collector 14 | from . import struct_xref_representation 15 | from . import swap_if 16 | from . import virtual_table_creation 17 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/callbacks.py: -------------------------------------------------------------------------------- 1 | from collections import defaultdict 2 | import idaapi 3 | 4 | 5 | class HexRaysCallbackManager(object): 6 | def __init__(self): 7 | self.__hexrays_event_handlers = defaultdict(list) 8 | 9 | def initialize(self): 10 | idaapi.install_hexrays_callback(self.__handle) 11 | 12 | def finalize(self): 13 | idaapi.remove_hexrays_callback(self.__handle) 14 | 15 | def register(self, event, handler): 16 | self.__hexrays_event_handlers[event].append(handler) 17 | 18 | def __handle(self, event, *args): 19 | for handler in self.__hexrays_event_handlers[event]: 20 | handler.handle(event, *args) 21 | # IDA expects zero 22 | return 0 23 | 24 | 25 | hx_callback_manager = HexRaysCallbackManager() 26 | 27 | 28 | class HexRaysEventHandler(object): 29 | def __init__(self): 30 | super(HexRaysEventHandler, self).__init__() 31 | 32 | def handle(self, event, *args): 33 | raise NotImplementedError("This is an abstract class") 34 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/virtual_table_creation.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | 3 | from . import actions 4 | from HexRaysPyTools.core.temporary_structure import VirtualTable 5 | 6 | 7 | class CreateVtable(actions.Action): 8 | description = "Create Virtual Table" 9 | hotkey = "V" 10 | 11 | def __init__(self): 12 | super(CreateVtable, self).__init__() 13 | 14 | @staticmethod 15 | def check(ea): 16 | return ea != idaapi.BADADDR and VirtualTable.check_address(ea) 17 | 18 | def activate(self, ctx): 19 | ea = ctx.cur_ea 20 | if self.check(ea): 21 | vtable = VirtualTable(0, ea) 22 | vtable.import_to_structures(True) 23 | 24 | def update(self, ctx): 25 | if ctx.widget_type == idaapi.BWN_DISASM: 26 | if self.check(ctx.cur_ea): 27 | idaapi.attach_action_to_popup(ctx.widget, None, self.name) 28 | return idaapi.AST_ENABLE 29 | idaapi.detach_action_from_popup(ctx.widget, self.name) 30 | return idaapi.AST_DISABLE 31 | return idaapi.AST_DISABLE_FOR_WIDGET 32 | 33 | 34 | actions.action_manager.register(CreateVtable()) 35 | -------------------------------------------------------------------------------- /HexRaysPyTools.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | import idaapi 4 | 5 | import HexRaysPyTools.core.cache as cache 6 | import HexRaysPyTools.core.const as const 7 | import HexRaysPyTools.settings as settings 8 | from HexRaysPyTools.callbacks import hx_callback_manager, action_manager 9 | from HexRaysPyTools.core.struct_xrefs import XrefStorage 10 | from HexRaysPyTools.core.temporary_structure import TemporaryStructureModel 11 | 12 | 13 | class MyPlugin(idaapi.plugin_t): 14 | flags = 0 15 | comment = "Plugin for automatic classes reconstruction" 16 | help = "See https://github.com/igogo-x86/HexRaysPyTools/blob/master/readme.md" 17 | wanted_name = "HexRaysPyTools" 18 | wanted_hotkey = "" 19 | 20 | @staticmethod 21 | def init(): 22 | if not idaapi.init_hexrays_plugin(): 23 | logging.error("Failed to initialize Hex-Rays SDK") 24 | return idaapi.PLUGIN_SKIP 25 | 26 | action_manager.initialize() 27 | hx_callback_manager.initialize() 28 | cache.temporary_structure = TemporaryStructureModel() 29 | const.init() 30 | XrefStorage().open() 31 | return idaapi.PLUGIN_KEEP 32 | 33 | @staticmethod 34 | def run(*args): 35 | pass 36 | 37 | @staticmethod 38 | def term(): 39 | action_manager.finalize() 40 | hx_callback_manager.finalize() 41 | XrefStorage().close() 42 | idaapi.term_hexrays_plugin() 43 | 44 | 45 | def PLUGIN_ENTRY(): 46 | settings.load_settings() 47 | logging.basicConfig(format='[%(levelname)s] %(message)s\t(%(module)s:%(funcName)s)') 48 | logging.root.setLevel(settings.DEBUG_MESSAGE_LEVEL) 49 | idaapi.notify_when(idaapi.NW_OPENIDB, cache.initialize_cache) 50 | return MyPlugin() 51 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/struct_xref_representation.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | 3 | from . import actions 4 | import HexRaysPyTools.core.helper as helper 5 | import HexRaysPyTools.core.struct_xrefs as struct_xrefs 6 | import HexRaysPyTools.forms as forms 7 | 8 | 9 | class FindFieldXrefs(actions.HexRaysPopupAction): 10 | description = "Field Xrefs" 11 | hotkey = "Ctrl+X" 12 | 13 | def __init__(self): 14 | super(FindFieldXrefs, self).__init__() 15 | 16 | def check(self, hx_view): 17 | return hx_view.item.citype == idaapi.VDI_EXPR and \ 18 | hx_view.item.it.to_specific_type.op in (idaapi.cot_memptr, idaapi.cot_memref) 19 | 20 | def activate(self, ctx): 21 | hx_view = idaapi.get_widget_vdui(ctx.widget) 22 | if not self.check(hx_view): 23 | return 24 | 25 | data = [] 26 | offset = hx_view.item.e.m 27 | struct_type = idaapi.remove_pointer(hx_view.item.e.x.type) 28 | ordinal = helper.get_ordinal(struct_type) 29 | result = struct_xrefs.XrefStorage().get_structure_info(ordinal, offset) 30 | for xref_info in result: 31 | data.append([ 32 | idaapi.get_short_name(xref_info.func_ea) + "+" + hex(int(xref_info.offset)), 33 | xref_info.type, 34 | xref_info.line 35 | ]) 36 | 37 | field_name = helper.get_member_name(struct_type, offset) 38 | chooser = forms.MyChoose( 39 | data, 40 | "Cross-references to {0}::{1}".format(struct_type.dstr(), field_name), 41 | [["Function", 20 | idaapi.Choose.CHCOL_PLAIN], 42 | ["Type", 2 | idaapi.Choose.CHCOL_PLAIN], 43 | ["Line", 40 | idaapi.Choose.CHCOL_PLAIN]] 44 | ) 45 | idx = chooser.Show(True) 46 | if idx == -1: 47 | return 48 | 49 | xref = result[idx] 50 | idaapi.open_pseudocode(xref.func_ea + xref.offset, False) 51 | 52 | actions.action_manager.register(FindFieldXrefs()) 53 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/member_double_click.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | 3 | from . import callbacks 4 | import HexRaysPyTools.core.helper as helper 5 | 6 | 7 | class MemberDoubleClick(callbacks.HexRaysEventHandler): 8 | def __init__(self): 9 | super(MemberDoubleClick, self).__init__() 10 | 11 | def handle(self, event, *args): 12 | hx_view = args[0] 13 | item = hx_view.item 14 | if item.citype == idaapi.VDI_EXPR and item.e.op in (idaapi.cot_memptr, idaapi.cot_memref): 15 | # Look if we double clicked on expression that is member pointer. Then get tinfo_t of the structure. 16 | # After that remove pointer and get member name with the same offset 17 | if item.e.x.op == idaapi.cot_memref and item.e.x.x.op == idaapi.cot_memptr: 18 | vtable_tinfo = item.e.x.type.get_pointed_object() 19 | method_offset = item.e.m 20 | class_tinfo = item.e.x.x.x.type.get_pointed_object() 21 | vtable_offset = item.e.x.x.m 22 | elif item.e.x.op == idaapi.cot_memptr: 23 | vtable_tinfo = item.e.x.type 24 | if vtable_tinfo.is_ptr(): 25 | vtable_tinfo = vtable_tinfo.get_pointed_object() 26 | method_offset = item.e.m 27 | class_tinfo = item.e.x.x.type.get_pointed_object() 28 | vtable_offset = item.e.x.m 29 | else: 30 | func_offset = item.e.m 31 | struct_tinfo = item.e.x.type.get_pointed_object() 32 | func_ea = helper.choose_virtual_func_address(helper.get_member_name(struct_tinfo, func_offset)) 33 | if func_ea: 34 | idaapi.jumpto(func_ea) 35 | return 0 36 | 37 | func_name = helper.get_member_name(vtable_tinfo, method_offset) 38 | func_ea = helper.choose_virtual_func_address(func_name, class_tinfo, vtable_offset) 39 | if func_ea: 40 | idaapi.open_pseudocode(func_ea, 0) 41 | return 1 42 | 43 | callbacks.hx_callback_manager.register(idaapi.hxe_double_click, MemberDoubleClick()) 44 | -------------------------------------------------------------------------------- /HexRaysPyTools/core/const.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | 3 | EA64 = None 4 | EA_SIZE = None 5 | 6 | COT_ARITHMETIC = (idaapi.cot_num, idaapi.cot_fnum, idaapi.cot_add, idaapi.cot_fadd, idaapi.cot_sub, idaapi.cot_fsub, 7 | idaapi.cot_mul, idaapi.cot_fmul, idaapi.cot_fdiv) 8 | 9 | VOID_TINFO = None 10 | PVOID_TINFO = idaapi.tinfo_t() 11 | CONST_VOID_TINFO = None 12 | CONST_PVOID_TINFO = idaapi.tinfo_t() 13 | CHAR_TINFO = None 14 | PCHAR_TINFO = idaapi.tinfo_t() 15 | CONST_PCHAR_TINFO = idaapi.tinfo_t() 16 | BYTE_TINFO = None 17 | PBYTE_TINFO = None 18 | 19 | WORD_TINFO = None 20 | PWORD_TINFO = idaapi.tinfo_t() 21 | 22 | X_WORD_TINFO = None # DWORD for x32 and QWORD for x64 23 | PX_WORD_TINFO = None 24 | 25 | DUMMY_FUNC = None 26 | 27 | LEGAL_TYPES = [] 28 | 29 | 30 | def init(): 31 | """ All tinfo should be reinitialized between session. Otherwise they could have wrong type """ 32 | global VOID_TINFO, PVOID_TINFO, CONST_PVOID_TINFO, BYTE_TINFO, PBYTE_TINFO, LEGAL_TYPES, X_WORD_TINFO, \ 33 | PX_WORD_TINFO, DUMMY_FUNC, CONST_PCHAR_TINFO, CHAR_TINFO, PCHAR_TINFO, CONST_VOID_TINFO, \ 34 | WORD_TINFO, PWORD_TINFO, EA64, EA_SIZE 35 | 36 | EA64 = idaapi.get_inf_structure().is_64bit() 37 | EA_SIZE = 8 if EA64 else 4 38 | 39 | VOID_TINFO = idaapi.tinfo_t(idaapi.BT_VOID) 40 | PVOID_TINFO.create_ptr(VOID_TINFO) 41 | CONST_VOID_TINFO = idaapi.tinfo_t(idaapi.BT_VOID | idaapi.BTM_CONST) 42 | CONST_PVOID_TINFO.create_ptr(idaapi.tinfo_t(idaapi.BT_VOID | idaapi.BTM_CONST)) 43 | CONST_PCHAR_TINFO.create_ptr(idaapi.tinfo_t(idaapi.BTF_CHAR | idaapi.BTM_CONST)) 44 | CHAR_TINFO = idaapi.tinfo_t(idaapi.BTF_CHAR) 45 | PCHAR_TINFO.create_ptr(idaapi.tinfo_t(idaapi.BTF_CHAR)) 46 | BYTE_TINFO = idaapi.tinfo_t(idaapi.BTF_BYTE) 47 | PBYTE_TINFO = idaapi.dummy_ptrtype(1, False) 48 | X_WORD_TINFO = idaapi.get_unk_type(EA_SIZE) 49 | PX_WORD_TINFO = idaapi.dummy_ptrtype(EA_SIZE, False) 50 | 51 | WORD_TINFO = idaapi.tinfo_t(idaapi.BT_UNK_WORD) 52 | PWORD_TINFO.create_ptr(idaapi.tinfo_t(idaapi.BT_UNK_WORD)) 53 | 54 | func_data = idaapi.func_type_data_t() 55 | func_data.rettype = PVOID_TINFO 56 | func_data.cc = idaapi.CM_CC_UNKNOWN 57 | DUMMY_FUNC = idaapi.tinfo_t() 58 | DUMMY_FUNC.create_func(func_data, idaapi.BT_FUNC) 59 | 60 | LEGAL_TYPES = [PVOID_TINFO, PX_WORD_TINFO, PWORD_TINFO, PBYTE_TINFO, X_WORD_TINFO] 61 | -------------------------------------------------------------------------------- /HexRaysPyTools/core/cache.py: -------------------------------------------------------------------------------- 1 | import collections 2 | 3 | import idaapi 4 | import idautils 5 | import idc 6 | 7 | from . import common 8 | 9 | # All virtual addresses where imported by module function pointers are stored 10 | imported_ea = set() 11 | 12 | # Map from demangled and simplified to C-language compatible names of functions to their addresses 13 | demangled_names = collections.defaultdict(set) 14 | 15 | # Functions that went through "touching" decompilation. This is done before Deep Scanning and 16 | # enhance arguments parsing for subroutines called by scanned functions. 17 | touched_functions = set() 18 | 19 | # This is where all information about structure being reconstructed stored 20 | # TODO: Make some way to store several structures and switch between them. See issue #22 (3) 21 | temporary_structure = None # type: temporary_structure.TemporaryStructureModel 22 | 23 | 24 | def _init_imported_ea(): 25 | 26 | def imp_cb(ea, name, ord): 27 | imported_ea.add(ea - idaapi.get_imagebase()) 28 | # True -> Continue enumeration 29 | # False -> Stop enumeration 30 | return True 31 | 32 | print("[Info] Collecting information about imports") 33 | imported_ea.clear() 34 | nimps = idaapi.get_import_module_qty() 35 | 36 | for i in range(0, nimps): 37 | name = idaapi.get_import_module_name(i) 38 | if not name: 39 | print("[Warning] Failed to get import module name for #%d" % i) 40 | continue 41 | 42 | # print "Walking-> %s" % name 43 | idaapi.enum_import_names(i, imp_cb) 44 | print("[Info] Done...") 45 | 46 | 47 | def _init_demangled_names(): 48 | """ 49 | Creates dictionary of demangled names => set of address, that will be used further when user makes double click 50 | on methods in Decompiler output. 51 | """ 52 | demangled_names.clear() 53 | for address, name in idautils.Names(): 54 | short_name = idc.demangle_name(name, idc.INF_SHORT_DN) 55 | if short_name: 56 | short_name = common.demangled_name_to_c_str(short_name) 57 | demangled_names[short_name].add(address - idaapi.get_imagebase()) 58 | print("[DEBUG] Demangled names have been initialized") 59 | 60 | 61 | def _reset_touched_functions(*args): 62 | global touched_functions 63 | 64 | touched_functions = set() 65 | 66 | 67 | def initialize_cache(*args): 68 | global temporary_structure 69 | 70 | _init_demangled_names() 71 | _init_imported_ea() 72 | _reset_touched_functions() 73 | -------------------------------------------------------------------------------- /HexRaysPyTools/settings.py: -------------------------------------------------------------------------------- 1 | import os 2 | import logging 3 | import idc 4 | 5 | try: 6 | import configparser 7 | except ImportError: 8 | # for python 2 9 | import ConfigParser as configparser 10 | 11 | CONFIG_FILE_PATH = os.path.join(idc.idadir(), 'cfg', 'HexRaysPyTools.cfg') 12 | 13 | DEBUG_MESSAGE_LEVEL = logging.INFO 14 | # Whether propagate names (Propagate name feature) through all names or only defaults like v11, a3, this, field_4 15 | PROPAGATE_THROUGH_ALL_NAMES = False 16 | # Store Xref information in database. I don't know how much size it consumes yet 17 | STORE_XREFS = True 18 | # There're some types that can be pointers to structures like int, PVOID etc and by default plugin scans only them 19 | # Full list can be found in `Const.LEGAL_TYPES`. 20 | # But if set this option to True than variable of every type could be possible to scan 21 | SCAN_ANY_TYPE = False 22 | 23 | 24 | def add_default_settings(config): 25 | updated = False 26 | if not config.has_option("DEFAULT", "DEBUG_MESSAGE_LEVEL"): 27 | config.set(None, 'DEBUG_MESSAGE_LEVEL', str(DEBUG_MESSAGE_LEVEL)) 28 | updated = True 29 | if not config.has_option("DEFAULT", "PROPAGATE_THROUGH_ALL_NAMES"): 30 | config.set(None, 'PROPAGATE_THROUGH_ALL_NAMES', str(PROPAGATE_THROUGH_ALL_NAMES)) 31 | updated = True 32 | if not config.has_option("DEFAULT", "STORE_XREFS"): 33 | config.set(None, 'STORE_XREFS', str(STORE_XREFS)) 34 | updated = True 35 | if not config.has_option("DEFAULT", "SCAN_ANY_TYPE"): 36 | config.set(None, 'SCAN_ANY_TYPE', str(SCAN_ANY_TYPE)) 37 | updated = True 38 | 39 | if updated: 40 | try: 41 | with open(CONFIG_FILE_PATH, "w") as f: 42 | config.write(f) 43 | except IOError: 44 | print("[ERROR] Failed to write or update config file at {}. Default settings will be used instead.\n" \ 45 | "Consider running IDA Pro under administrator once".format(CONFIG_FILE_PATH)) 46 | 47 | 48 | def load_settings(): 49 | global DEBUG_MESSAGE_LEVEL, PROPAGATE_THROUGH_ALL_NAMES, STORE_XREFS, SCAN_ANY_TYPE 50 | 51 | config = configparser.ConfigParser() 52 | if os.path.isfile(CONFIG_FILE_PATH): 53 | config.read(CONFIG_FILE_PATH) 54 | 55 | add_default_settings(config) 56 | 57 | DEBUG_MESSAGE_LEVEL = config.getint("DEFAULT", 'DEBUG_MESSAGE_LEVEL') 58 | PROPAGATE_THROUGH_ALL_NAMES = config.getboolean("DEFAULT", 'PROPAGATE_THROUGH_ALL_NAMES') 59 | STORE_XREFS = config.getboolean("DEFAULT", 'STORE_XREFS') 60 | SCAN_ANY_TYPE = config.getboolean("DEFAULT", 'SCAN_ANY_TYPE') 61 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/guess_allocation.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | 3 | from . import actions 4 | import HexRaysPyTools.api as api 5 | import HexRaysPyTools.forms as forms 6 | import HexRaysPyTools.core.helper as helper 7 | 8 | 9 | class _StructAllocChoose(forms.MyChoose): 10 | def __init__(self, items): 11 | forms.MyChoose.__init__( 12 | self, items, "Possible structure allocations", 13 | [["Function", 30], ["Variable", 10], ["Line", 50], ["Type", 10]] 14 | ) 15 | 16 | def OnSelectLine(self, n): 17 | idaapi.jumpto(self.items[n][0]) 18 | 19 | def OnGetLine(self, n): 20 | func_ea, var, line, alloc_type = self.items[n] 21 | return [helper.to_nice_str(func_ea), var, line, alloc_type] 22 | 23 | 24 | class _GuessAllocationVisitor(api.RecursiveObjectUpwardsVisitor): 25 | def __init__(self, cfunc, obj): 26 | super(_GuessAllocationVisitor, self).__init__(cfunc, obj, skip_after_object=True) 27 | self._data = [] 28 | 29 | def _manipulate(self, cexpr, obj): 30 | if obj.id == api.SO_LOCAL_VARIABLE: 31 | parent = self.parent_expr() 32 | if parent.op == idaapi.cot_asg: 33 | alloc_obj = api.MemoryAllocationObject.create(self._cfunc, self.parent_expr().y) 34 | if alloc_obj: 35 | self._data.append([alloc_obj.ea, obj.name, self._get_line(), "HEAP"]) 36 | elif self.parent_expr().op == idaapi.cot_ref: 37 | self._data.append([helper.find_asm_address(cexpr, self.parents), obj.name, self._get_line(), "STACK"]) 38 | elif obj.id == api.SO_GLOBAL_OBJECT: 39 | self._data.append([helper.find_asm_address(cexpr, self.parents), obj.name, self._get_line(), "GLOBAL"]) 40 | 41 | def _finish(self): 42 | chooser = _StructAllocChoose(self._data) 43 | chooser.Show(False) 44 | 45 | 46 | class GuessAllocation(actions.HexRaysPopupAction): 47 | description = "Guess allocation" 48 | hotkey = None 49 | 50 | def __init__(self): 51 | super(GuessAllocation, self).__init__() 52 | 53 | def check(self, hx_view): 54 | if hx_view.item.citype != idaapi.VDI_EXPR: 55 | return False 56 | return api.ScanObject.create(hx_view.cfunc, hx_view.item) is not None 57 | 58 | def activate(self, ctx): 59 | hx_view = idaapi.get_widget_vdui(ctx.widget) 60 | obj = api.ScanObject.create(hx_view.cfunc, hx_view.item) 61 | if obj: 62 | visitor = _GuessAllocationVisitor(hx_view.cfunc, obj) 63 | visitor.process() 64 | 65 | actions.action_manager.register(GuessAllocation()) 66 | -------------------------------------------------------------------------------- /HexRaysPyTools/core/type_library.py: -------------------------------------------------------------------------------- 1 | import ctypes 2 | import sys 3 | 4 | import idaapi 5 | 6 | from . import const 7 | import HexRaysPyTools.forms as forms 8 | 9 | 10 | class til_t(ctypes.Structure): 11 | pass 12 | 13 | 14 | til_t._fields_ = [ 15 | ("name", ctypes.c_char_p), 16 | ("desc", ctypes.c_char_p), 17 | ("nbases", ctypes.c_int), 18 | ("base", ctypes.POINTER(ctypes.POINTER(til_t))) 19 | ] 20 | 21 | 22 | def _enable_library_ordinals(library_num): 23 | idaname = "ida64" if const.EA64 else "ida" 24 | if sys.platform == "win32": 25 | dll = ctypes.windll[idaname + ".dll"] 26 | elif sys.platform == "linux2": 27 | dll = ctypes.cdll["lib" + idaname + ".so"] 28 | elif sys.platform == "darwin": 29 | dll = ctypes.cdll["lib" + idaname + ".dylib"] 30 | else: 31 | print("[ERROR] Failed to enable ordinals") 32 | return 33 | 34 | dll.get_idati.restype = ctypes.POINTER(til_t) 35 | idati = dll.get_idati() 36 | dll.enable_numbered_types(idati.contents.base[library_num], True) 37 | 38 | 39 | def choose_til(): 40 | # type: () -> (idaapi.til_t, int, bool) 41 | """ Creates a list of loaded libraries, asks user to take one of them and returns it with 42 | information about max ordinal and whether it's local or imported library """ 43 | idati = idaapi.cvar.idati 44 | list_type_library = [(idati, idati.name, idati.desc)] 45 | for idx in range(idaapi.cvar.idati.nbases): 46 | type_library = idaapi.cvar.idati.base(idx) # type: idaapi.til_t 47 | list_type_library.append((type_library, type_library.name, type_library.desc)) 48 | 49 | library_chooser = forms.MyChoose( 50 | list([[x[1], x[2]] for x in list_type_library]), 51 | "Select Library", 52 | [["Library", 10 | idaapi.Choose.CHCOL_PLAIN], ["Description", 30 | idaapi.Choose.CHCOL_PLAIN]], 53 | 69 54 | ) 55 | library_num = library_chooser.Show(True) 56 | if library_num != -1: 57 | selected_library = list_type_library[library_num][0] # type: idaapi.til_t 58 | max_ordinal = idaapi.get_ordinal_qty(selected_library) 59 | if max_ordinal == idaapi.BADORD: 60 | _enable_library_ordinals(library_num - 1) 61 | max_ordinal = idaapi.get_ordinal_qty(selected_library) 62 | print("[DEBUG] Maximal ordinal of lib {0} = {1}".format(selected_library.name, max_ordinal)) 63 | return selected_library, max_ordinal, library_num == 0 64 | 65 | 66 | def import_type(library, name): 67 | if library.name != idaapi.cvar.idati.name: 68 | last_ordinal = idaapi.get_ordinal_qty(idaapi.cvar.idati) 69 | type_id = idaapi.import_type(library, -1, name) # tid_t 70 | if type_id != idaapi.BADORD: 71 | return last_ordinal 72 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/form_requests.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | 3 | from . import actions 4 | import HexRaysPyTools.core.cache as cache 5 | import HexRaysPyTools.core.classes as classes 6 | from HexRaysPyTools.core.structure_graph import StructureGraph 7 | from HexRaysPyTools.forms import StructureGraphViewer, ClassViewer, StructureBuilder 8 | 9 | 10 | class ShowGraph(actions.Action): 11 | description = "Show graph" 12 | hotkey = "G" 13 | 14 | def __init__(self): 15 | super(ShowGraph, self).__init__() 16 | self.graph = None 17 | self.graph_view = None 18 | 19 | def activate(self, ctx): 20 | widget = self.graph_view.GetWidget() if self.graph_view else None 21 | if widget: 22 | self.graph_view.change_selected([sel + 1 for sel in ctx.chooser_selection]) 23 | self.graph_view.Show() 24 | else: 25 | self.graph = StructureGraph([sel + 1 for sel in ctx.chooser_selection]) 26 | self.graph_view = StructureGraphViewer("Structure Graph", self.graph) 27 | self.graph_view.Show() 28 | 29 | def update(self, ctx): 30 | if ctx.widget_type == idaapi.BWN_LOCTYPS: 31 | idaapi.attach_action_to_popup(ctx.widget, None, self.name) 32 | return idaapi.AST_ENABLE_FOR_WIDGET 33 | return idaapi.AST_DISABLE_FOR_WIDGET 34 | 35 | 36 | actions.action_manager.register(ShowGraph()) 37 | 38 | 39 | class ShowClasses(actions.Action): 40 | description = "Classes" 41 | hotkey = "Alt+F1" 42 | 43 | def __init__(self): 44 | super(ShowClasses, self).__init__() 45 | 46 | def activate(self, ctx): 47 | tform = idaapi.find_widget('Classes') 48 | if not tform: 49 | class_viewer = ClassViewer(classes.ProxyModel(), classes.TreeModel()) 50 | class_viewer.Show() 51 | else: 52 | idaapi.activate_widget(tform, True) 53 | 54 | def update(self, ctx): 55 | return idaapi.AST_ENABLE_ALWAYS 56 | 57 | 58 | show_classes = ShowClasses() 59 | actions.action_manager.register(show_classes) 60 | idaapi.attach_action_to_menu('View/Open subviews/Local types', show_classes.name, idaapi.SETMENU_APP) 61 | 62 | 63 | class ShowStructureBuilder(actions.HexRaysPopupAction): 64 | description = "Show Structure Builder" 65 | hotkey = "Alt+F8" 66 | 67 | def __init__(self): 68 | super(ShowStructureBuilder, self).__init__() 69 | 70 | def check(self, hx_view): 71 | return True 72 | 73 | def activate(self, ctx): 74 | tform = idaapi.find_widget("Structure Builder") 75 | if tform: 76 | idaapi.activate_widget(tform, True) 77 | else: 78 | StructureBuilder(cache.temporary_structure).Show() 79 | 80 | def update(self, ctx): 81 | return idaapi.AST_ENABLE_ALWAYS 82 | 83 | 84 | actions.action_manager.register(ShowStructureBuilder()) 85 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/actions.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | 3 | from .callbacks import hx_callback_manager, HexRaysEventHandler 4 | 5 | 6 | class ActionManager(object): 7 | def __init__(self): 8 | self.__actions = [] 9 | 10 | def register(self, action): 11 | self.__actions.append(action) 12 | idaapi.register_action( 13 | idaapi.action_desc_t(action.name, action.description, action, action.hotkey) 14 | ) 15 | if isinstance(action, HexRaysPopupAction): 16 | hx_callback_manager.register(idaapi.hxe_populating_popup, HexRaysPopupRequestHandler(action)) 17 | 18 | def initialize(self): 19 | pass 20 | 21 | def finalize(self): 22 | for action in self.__actions: 23 | idaapi.unregister_action(action.name) 24 | 25 | 26 | action_manager = ActionManager() 27 | 28 | 29 | class Action(idaapi.action_handler_t): 30 | """ 31 | Convenience wrapper with name property allowing to be registered in IDA using ActionManager 32 | """ 33 | description = None 34 | hotkey = None 35 | 36 | def __init__(self): 37 | super(Action, self).__init__() 38 | 39 | @property 40 | def name(self): 41 | return "HexRaysPyTools:" + type(self).__name__ 42 | 43 | def activate(self, ctx): 44 | # type: (idaapi.action_activation_ctx_t) -> None 45 | raise NotImplementedError 46 | 47 | def update(self, ctx): 48 | # type: (idaapi.action_activation_ctx_t) -> None 49 | raise NotImplementedError 50 | 51 | 52 | class HexRaysPopupAction(Action): 53 | """ 54 | Wrapper around Action. Represents Action which can be added to menu after right-clicking in Decompile window. 55 | Has `check` method that should tell whether Action should be added to popup menu when different items 56 | are right-clicked. 57 | Children of this class can also be fired by hot-key without right-clicking if one provided in `hotkey` 58 | static member. 59 | """ 60 | 61 | def __init__(self): 62 | super(HexRaysPopupAction, self).__init__() 63 | 64 | def activate(self, ctx): 65 | # type: (idaapi.action_activation_ctx_t) -> None 66 | raise NotImplementedError 67 | 68 | def check(self, hx_view): 69 | # type: (idaapi.vdui_t) -> bool 70 | raise NotImplementedError 71 | 72 | def update(self, ctx): 73 | if ctx.widget_type == idaapi.BWN_PSEUDOCODE: 74 | return idaapi.AST_ENABLE_FOR_WIDGET 75 | return idaapi.AST_DISABLE_FOR_WIDGET 76 | 77 | 78 | class HexRaysPopupRequestHandler(HexRaysEventHandler): 79 | """ 80 | This is wrapper around HexRaysPopupAction which allows to dynamically decide whether to add Action to popup 81 | menu or not. 82 | Register this in CallbackManager. 83 | """ 84 | def __init__(self, action): 85 | super(HexRaysPopupRequestHandler, self).__init__() 86 | self.__action = action 87 | 88 | def handle(self, event, *args): 89 | form, popup, hx_view = args 90 | if self.__action.check(hx_view): 91 | idaapi.attach_action_to_popup(form, popup, self.__action.name, None) 92 | return 0 93 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/structs_by_size.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | 3 | from . import actions 4 | import HexRaysPyTools.forms as forms 5 | import HexRaysPyTools.core.type_library as type_library 6 | 7 | 8 | def _choose_structure_by_size(size): 9 | result = type_library.choose_til() 10 | if result: 11 | selected_library, max_ordinal, is_local_type = result 12 | matched_types = [] 13 | tinfo = idaapi.tinfo_t() 14 | for ordinal in range(1, max_ordinal): 15 | tinfo.create_typedef(selected_library, ordinal) 16 | if tinfo.get_size() == size: 17 | name = tinfo.dstr() 18 | description = idaapi.print_tinfo(None, 0, 0, idaapi.PRTYPE_DEF, tinfo, None, None) 19 | matched_types.append([str(ordinal), name, description]) 20 | 21 | type_chooser = forms.MyChoose( 22 | matched_types, 23 | "Select Type", 24 | [["Ordinal", 5 | idaapi.Choose.CHCOL_HEX], ["Type Name", 25], ["Declaration", 50]], 25 | 165 26 | ) 27 | selected_type = type_chooser.Show(True) 28 | if selected_type != -1: 29 | if is_local_type: 30 | return int(matched_types[selected_type][0]) 31 | return type_library.import_type(selected_library, matched_types[selected_type][1]) 32 | return None 33 | 34 | 35 | class GetStructureBySize(actions.HexRaysPopupAction): 36 | # TODO: apply type automatically if expression like `var = new(size)` 37 | description = "Structures with this size" 38 | 39 | def __init__(self): 40 | super(GetStructureBySize, self).__init__() 41 | 42 | def check(self, hx_view): 43 | return hx_view.item.citype == idaapi.VDI_EXPR and hx_view.item.e.op == idaapi.cot_num 44 | 45 | def activate(self, ctx): 46 | hx_view = idaapi.get_widget_vdui(ctx.widget) 47 | if not self.check(hx_view): 48 | return 49 | ea = ctx.cur_ea 50 | c_number = hx_view.item.e 51 | number_value = c_number.numval() 52 | ordinal = _choose_structure_by_size(number_value) 53 | if ordinal: 54 | number_format_old = c_number.n.nf 55 | number_format_new = idaapi.number_format_t() 56 | number_format_new.flags = idaapi.FF_1STRO | idaapi.FF_0STRO 57 | operand_number = number_format_old.opnum 58 | number_format_new.opnum = operand_number 59 | number_format_new.props = number_format_old.props 60 | number_format_new.type_name = idaapi.get_numbered_type_name(idaapi.cvar.idati, ordinal) 61 | 62 | c_function = hx_view.cfunc 63 | number_formats = c_function.numforms # type: idaapi.user_numforms_t 64 | # print "(number) flags: {0:#010X}, type_name: {1}, opnum: {2}".format( 65 | # number_format.flags, 66 | # number_format.type_name, 67 | # number_format.opnum 68 | # ) 69 | operand_locator = idaapi.operand_locator_t(ea, ord(operand_number) if operand_number else 0) 70 | if operand_locator in number_formats: 71 | del number_formats[operand_locator] 72 | 73 | number_formats[operand_locator] = number_format_new 74 | c_function.save_user_numforms() 75 | hx_view.refresh_view(True) 76 | 77 | actions.action_manager.register(GetStructureBySize()) 78 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | ## Ignore Visual Studio temporary files, build results, and 2 | ## files generated by popular Visual Studio add-ons. 3 | 4 | # User-specific files 5 | *.suo 6 | *.user 7 | *.sln.docstates 8 | *.coati* 9 | *.pyc 10 | 11 | # Build results 12 | [Dd]ebug/ 13 | [Dd]ebugPublic/ 14 | [Rr]elease/ 15 | [Rr]eleases/ 16 | .idea 17 | x64/ 18 | x86/ 19 | build/ 20 | bld/ 21 | [Bb]in/ 22 | [Oo]bj/ 23 | 24 | # Roslyn cache directories 25 | *.ide/ 26 | 27 | # MSTest test Results 28 | [Tt]est[Rr]esult*/ 29 | [Bb]uild[Ll]og.* 30 | 31 | #NUNIT 32 | *.VisualState.xml 33 | TestResult.xml 34 | 35 | # Build Results of an ATL Project 36 | [Dd]ebugPS/ 37 | [Rr]eleasePS/ 38 | dlldata.c 39 | 40 | *_i.c 41 | *_p.c 42 | *_i.h 43 | *.ilk 44 | *.meta 45 | *.obj 46 | *.pch 47 | *.pdb 48 | *.pgc 49 | *.pgd 50 | *.rsp 51 | *.sbr 52 | *.tlb 53 | *.tli 54 | *.tlh 55 | *.tmp 56 | *.tmp_proj 57 | *.log 58 | *.vspscc 59 | *.vssscc 60 | .builds 61 | *.pidb 62 | *.svclog 63 | *.scc 64 | 65 | # Chutzpah Test files 66 | _Chutzpah* 67 | 68 | # Visual C++ cache files 69 | ipch/ 70 | *.aps 71 | *.ncb 72 | *.opensdf 73 | *.sdf 74 | *.cachefile 75 | 76 | # Visual Studio profiler 77 | *.psess 78 | *.vsp 79 | *.vspx 80 | 81 | # TFS 2012 Local Workspace 82 | $tf/ 83 | 84 | # Guidance Automation Toolkit 85 | *.gpState 86 | 87 | # ReSharper is a .NET coding add-in 88 | _ReSharper*/ 89 | *.[Rr]e[Ss]harper 90 | *.DotSettings.user 91 | 92 | # JustCode is a .NET coding addin-in 93 | .JustCode 94 | 95 | # TeamCity is a build add-in 96 | _TeamCity* 97 | 98 | # DotCover is a Code Coverage Tool 99 | *.dotCover 100 | 101 | # NCrunch 102 | _NCrunch_* 103 | .*crunch*.local.xml 104 | 105 | # MightyMoose 106 | *.mm.* 107 | AutoTest.Net/ 108 | 109 | # Web workbench (sass) 110 | .sass-cache/ 111 | 112 | # Installshield output folder 113 | [Ee]xpress/ 114 | 115 | # DocProject is a documentation generator add-in 116 | DocProject/buildhelp/ 117 | DocProject/Help/*.HxT 118 | DocProject/Help/*.HxC 119 | DocProject/Help/*.hhc 120 | DocProject/Help/*.hhk 121 | DocProject/Help/*.hhp 122 | DocProject/Help/Html2 123 | DocProject/Help/html 124 | 125 | # Click-Once directory 126 | publish/ 127 | 128 | # Publish Web Output 129 | *.[Pp]ublish.xml 130 | *.azurePubxml 131 | # TODO: Comment the next line if you want to checkin your web deploy settings 132 | # but database connection strings (with potential passwords) will be unencrypted 133 | *.pubxml 134 | *.publishproj 135 | 136 | # NuGet Packages 137 | *.nupkg 138 | # The packages folder can be ignored because of Package Restore 139 | **/packages/* 140 | # except build/, which is used as an MSBuild target. 141 | !**/packages/build/ 142 | # If using the old MSBuild-Integrated Package Restore, uncomment this: 143 | #!**/packages/repositories.config 144 | 145 | # Windows Azure Build Output 146 | csx/ 147 | *.build.csdef 148 | 149 | # Windows Store app package directory 150 | AppPackages/ 151 | 152 | # Others 153 | sql/ 154 | *.Cache 155 | ClientBin/ 156 | [Ss]tyle[Cc]op.* 157 | ~$* 158 | *~ 159 | *.dbmdl 160 | *.dbproj.schemaview 161 | *.pfx 162 | *.publishsettings 163 | node_modules/ 164 | 165 | # RIA/Silverlight projects 166 | Generated_Code/ 167 | 168 | # Backup & report files from converting an old project file 169 | # to a newer Visual Studio version. Backup files are not needed, 170 | # because we have git ;-) 171 | _UpgradeReport_Files/ 172 | Backup*/ 173 | UpgradeLog*.XML 174 | UpgradeLog*.htm 175 | 176 | # SQL Server files 177 | *.mdf 178 | *.ldf 179 | 180 | # Business Intelligence projects 181 | *.rdl.data 182 | *.bim.layout 183 | *.bim_*.settings 184 | 185 | # Microsoft Fakes 186 | FakesAssemblies/ 187 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/function_signature_modifiers.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | from . import actions 3 | import HexRaysPyTools.core.const as const 4 | 5 | 6 | class ConvertToUsercall(actions.HexRaysPopupAction): 7 | description = "Convert to __usercall" 8 | 9 | def __init__(self): 10 | super(ConvertToUsercall, self).__init__() 11 | 12 | def check(self, hx_view): 13 | return hx_view.item.citype == idaapi.VDI_FUNC 14 | 15 | def activate(self, ctx): 16 | vu = idaapi.get_widget_vdui(ctx.widget) 17 | function_tinfo = idaapi.tinfo_t() 18 | if not vu.cfunc.get_func_type(function_tinfo): 19 | return 20 | function_details = idaapi.func_type_data_t() 21 | function_tinfo.get_func_details(function_details) 22 | convention = idaapi.CM_CC_MASK & function_details.cc 23 | if convention == idaapi.CM_CC_CDECL: 24 | function_details.cc = idaapi.CM_CC_SPECIAL 25 | elif convention in (idaapi.CM_CC_STDCALL, idaapi.CM_CC_FASTCALL, idaapi.CM_CC_PASCAL, idaapi.CM_CC_THISCALL): 26 | function_details.cc = idaapi.CM_CC_SPECIALP 27 | elif convention == idaapi.CM_CC_ELLIPSIS: 28 | function_details.cc = idaapi.CM_CC_SPECIALE 29 | else: 30 | return 31 | function_tinfo.create_func(function_details) 32 | idaapi.apply_tinfo(vu.cfunc.entry_ea, function_tinfo, idaapi.TINFO_DEFINITE) 33 | vu.refresh_view(True) 34 | 35 | 36 | class AddRemoveReturn(actions.HexRaysPopupAction): 37 | description = "Add/Remove Return" 38 | 39 | def __init__(self): 40 | super(AddRemoveReturn, self).__init__() 41 | 42 | def check(self, hx_view): 43 | return hx_view.item.citype == idaapi.VDI_FUNC 44 | 45 | def activate(self, ctx): 46 | vu = idaapi.get_widget_vdui(ctx.widget) 47 | function_tinfo = idaapi.tinfo_t() 48 | if not vu.cfunc.get_func_type(function_tinfo): 49 | return 50 | function_details = idaapi.func_type_data_t() 51 | function_tinfo.get_func_details(function_details) 52 | if function_details.rettype.equals_to(const.VOID_TINFO): 53 | function_details.rettype = idaapi.tinfo_t(const.PVOID_TINFO) 54 | else: 55 | function_details.rettype = idaapi.tinfo_t(idaapi.BT_VOID) 56 | function_tinfo.create_func(function_details) 57 | idaapi.apply_tinfo(vu.cfunc.entry_ea, function_tinfo, idaapi.TINFO_DEFINITE) 58 | vu.refresh_view(True) 59 | 60 | 61 | class RemoveArgument(actions.HexRaysPopupAction): 62 | description = "Remove Argument" 63 | 64 | def __init__(self): 65 | super(RemoveArgument, self).__init__() 66 | 67 | def check(self, hx_view): 68 | if hx_view.item.citype != idaapi.VDI_LVAR: 69 | return False 70 | local_variable = hx_view.item.get_lvar() # type:idaapi.lvar_t 71 | return local_variable.is_arg_var 72 | 73 | def activate(self, ctx): 74 | vu = idaapi.get_widget_vdui(ctx.widget) 75 | function_tinfo = idaapi.tinfo_t() 76 | if not vu.cfunc.get_func_type(function_tinfo): 77 | return 78 | function_details = idaapi.func_type_data_t() 79 | function_tinfo.get_func_details(function_details) 80 | del_arg = vu.item.get_lvar() 81 | 82 | function_details.erase([x for x in function_details if x.name == del_arg.name][0]) 83 | 84 | function_tinfo.create_func(function_details) 85 | idaapi.apply_tinfo(vu.cfunc.entry_ea, function_tinfo, idaapi.TINFO_DEFINITE) 86 | vu.refresh_view(True) 87 | 88 | 89 | actions.action_manager.register(ConvertToUsercall()) 90 | actions.action_manager.register(AddRemoveReturn()) 91 | actions.action_manager.register(RemoveArgument()) 92 | -------------------------------------------------------------------------------- /HexRaysPyTools/core/struct_xrefs.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple, defaultdict 2 | import json 3 | import logging 4 | 5 | import idaapi 6 | from . import helper 7 | import HexRaysPyTools.settings as settings 8 | 9 | logger = logging.getLogger(__name__) 10 | 11 | XrefInfo = namedtuple('XrefInfo', ['func_ea', 'offset', 'line', 'type']) 12 | 13 | 14 | def singleton(cls): 15 | instances = {} 16 | 17 | def get_instance(): 18 | if cls not in instances: 19 | instances[cls] = cls() 20 | return instances[cls] 21 | return get_instance 22 | 23 | 24 | @singleton 25 | class XrefStorage(object): 26 | ARRAY_NAME = "$HexRaysPyTools:XrefStorage" 27 | 28 | def __init__(self): 29 | """ 30 | storage - {ordinal: {func_offset: (code_offset, line, usage_type)}} 31 | __delete_items_helper - {func_offset: set(ordinals)} 32 | """ 33 | self.storage = None 34 | self.__delete_items_helper = defaultdict(set) 35 | 36 | def open(self): 37 | if not settings.STORE_XREFS: 38 | self.storage = {} 39 | return 40 | 41 | result = helper.load_long_str_from_idb(self.ARRAY_NAME) 42 | if result: 43 | try: 44 | self.storage = json.loads(result, object_hook=self.json_keys_to_str) 45 | self.__init_delete_helper() 46 | return 47 | except ValueError: 48 | logger.error("Failed to read previous info about Xrefs. Try Ctrl+F5 to cache data") 49 | self.storage = {} 50 | 51 | def close(self): 52 | self.save() 53 | self.storage = None 54 | self.__delete_items_helper = defaultdict(set) 55 | 56 | def save(self): 57 | if not settings.STORE_XREFS: 58 | return 59 | 60 | if self.storage: 61 | helper.save_long_str_to_idb(self.ARRAY_NAME, json.dumps(self.storage)) 62 | 63 | def update(self, function_offset, data): 64 | """ data - {ordinal : (code_offset, line, usage_type)} """ 65 | for ordinal, info in list(data.items()): 66 | self.__update_ordinal_info(ordinal, function_offset, info) 67 | 68 | deleted_ordinals = self.__delete_items_helper[function_offset].difference(list(data.keys())) 69 | for ordinal in deleted_ordinals: 70 | self.__remove_ordinal_info(ordinal, function_offset) 71 | 72 | def get_structure_info(self, ordinal, struct_offset): 73 | """ By given ordinal and offset within a structure returns dictionary {func_address -> list(offsets)} """ 74 | result = [] 75 | 76 | if ordinal not in self.storage: 77 | return result 78 | 79 | for func_offset, info in list(self.storage[ordinal].items()): 80 | if struct_offset in info: 81 | func_ea = func_offset + idaapi.get_imagebase() 82 | for xref_info in info[struct_offset]: 83 | offset, line, usage_type = xref_info 84 | result.append(XrefInfo(func_ea, offset, line, usage_type)) 85 | return result 86 | 87 | @staticmethod 88 | def json_keys_to_str(x): 89 | if isinstance(x, dict): 90 | return {int(k): v for k, v in list(x.items())} 91 | return x 92 | 93 | def __len__(self): 94 | return len(str(self.storage)) 95 | 96 | def __init_delete_helper(self): 97 | for ordinal, data in list(self.storage.items()): 98 | for func_offset in data: 99 | self.__delete_items_helper[func_offset].add(ordinal) 100 | 101 | def __remove_ordinal_info(self, ordinal, function_offset): 102 | del self.storage[ordinal][function_offset] 103 | self.__delete_items_helper[function_offset].remove(ordinal) 104 | 105 | def __update_ordinal_info(self, ordinal, function_offset, info): 106 | if ordinal not in self.storage: 107 | self.storage[ordinal] = {} 108 | self.storage[ordinal][function_offset] = info 109 | self.__delete_items_helper[function_offset].add(ordinal) 110 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/struct_xref_collector.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import time 3 | 4 | import idaapi 5 | 6 | from . import callbacks 7 | import HexRaysPyTools.core.struct_xrefs as struct_xrefs 8 | import HexRaysPyTools.core.helper as helper 9 | 10 | logger = logging.getLogger(__name__) 11 | 12 | 13 | class StructXrefCollectorVisitor(idaapi.ctree_parentee_t): 14 | def __init__(self, cfunc, storage): 15 | super(StructXrefCollectorVisitor, self).__init__() 16 | self.__cfunc = cfunc 17 | self.__function_address = cfunc.entry_ea 18 | self.__result = {} 19 | self.__storage = storage 20 | 21 | def visit_expr(self, expression): 22 | # Checks if expression is reference by pointer or by value 23 | if expression.op == idaapi.cot_memptr: 24 | struct_type = expression.x.type.get_pointed_object() 25 | elif expression.op == idaapi.cot_memref: 26 | struct_type = expression.x.type 27 | else: 28 | return 0 29 | 30 | # Getting information about structure, field offset, address and one line corresponding to code 31 | ordinal = helper.get_ordinal(struct_type) 32 | field_offset = expression.m 33 | ea = self.__find_ref_address(expression) 34 | usage_type = self.__get_type(expression) 35 | 36 | if ea == idaapi.BADADDR or not ordinal: 37 | logger.warning("Failed to parse at address {0}, ordinal - {1}, type - {2}".format( 38 | helper.to_hex(ea), ordinal, struct_type.dstr() 39 | )) 40 | 41 | one_line = self.__get_line() 42 | 43 | occurrence_offset = ea - self.__function_address 44 | xref_info = (occurrence_offset, one_line, usage_type) 45 | 46 | # Saving results 47 | if ordinal not in self.__result: 48 | self.__result[ordinal] = {field_offset: [xref_info]} 49 | elif field_offset not in self.__result[ordinal]: 50 | self.__result[ordinal][field_offset] = [xref_info] 51 | else: 52 | self.__result[ordinal][field_offset].append(xref_info) 53 | return 0 54 | 55 | def process(self): 56 | t = time.time() 57 | self.apply_to(self.__cfunc.body, None) 58 | self.__storage.update(self.__function_address - idaapi.get_imagebase(), self.__result) 59 | 60 | storage_mb_size = len(self.__storage) * 1.0 // 1024 ** 2 61 | logger.debug("Xref processing: %f seconds passed, storage size - %.2f MB ", (time.time() - t), storage_mb_size) 62 | 63 | def __find_ref_address(self, cexpr): 64 | """ Returns most close virtual address corresponding to cexpr """ 65 | 66 | ea = cexpr.ea 67 | if ea != idaapi.BADADDR: 68 | return ea 69 | 70 | for p in reversed(self.parents): 71 | if p.ea != idaapi.BADADDR: 72 | return p.ea 73 | 74 | def __get_type(self, cexpr): 75 | """ Returns one of the following types: 'R' - read value, 'W' - write value, 'A' - function argument""" 76 | child = cexpr 77 | for p in reversed(self.parents): 78 | assert p, "Failed to get type at " + helper.to_hex(self.__function_address) 79 | 80 | if p.cexpr.op == idaapi.cot_call: 81 | return 'Arg' 82 | if not p.is_expr(): 83 | return 'R' 84 | if p.cexpr.op == idaapi.cot_asg: 85 | if p.cexpr.x == child: 86 | return 'W' 87 | return 'R' 88 | child = p.cexpr 89 | 90 | def __get_line(self): 91 | for p in reversed(self.parents): 92 | if not p.is_expr(): 93 | return idaapi.tag_remove(p.print1(self.__cfunc)) 94 | AssertionError("Parent instruction is not found") 95 | 96 | 97 | class StructXrefCollector(callbacks.HexRaysEventHandler): 98 | def __init__(self): 99 | super(StructXrefCollector, self).__init__() 100 | 101 | def handle(self, event, *args): 102 | cfunc, level_of_maturity = args 103 | if level_of_maturity == idaapi.CMAT_FINAL: 104 | StructXrefCollectorVisitor(cfunc, struct_xrefs.XrefStorage()).process() 105 | 106 | 107 | callbacks.hx_callback_manager.register(idaapi.hxe_maturity, StructXrefCollector()) 108 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/new_field_creation.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import re 3 | 4 | import idaapi 5 | import idc 6 | 7 | from . import actions 8 | import HexRaysPyTools.core.helper as helper 9 | import HexRaysPyTools.core.const as const 10 | 11 | logger = logging.getLogger(__name__) 12 | 13 | 14 | def _is_gap_field(cexpr): 15 | if cexpr.op not in (idaapi.cot_memptr, idaapi.cot_memref): 16 | return False 17 | struct_type = cexpr.x.type 18 | struct_type.remove_ptr_or_array() 19 | return helper.get_member_name(struct_type, cexpr.m)[0:3] == "gap" 20 | 21 | 22 | class CreateNewField(actions.HexRaysPopupAction): 23 | description = "Create New Field" 24 | hotkey = "Ctrl+F" 25 | 26 | def __init__(self): 27 | super(CreateNewField, self).__init__() 28 | 29 | def check(self, hx_view): 30 | if hx_view.item.citype != idaapi.VDI_EXPR: 31 | return False 32 | 33 | return _is_gap_field(hx_view.item.it.to_specific_type) 34 | 35 | def activate(self, ctx): 36 | hx_view = idaapi.get_widget_vdui(ctx.widget) 37 | if not self.check(hx_view): 38 | return 39 | 40 | item = hx_view.item.it.to_specific_type 41 | parent = hx_view.cfunc.body.find_parent_of(item).to_specific_type 42 | if parent.op != idaapi.cot_idx or parent.y.op != idaapi.cot_num: 43 | idx = 0 44 | else: 45 | idx = parent.y.numval() 46 | 47 | struct_tinfo = item.x.type 48 | struct_tinfo.remove_ptr_or_array() 49 | 50 | offset = item.m 51 | ordinal = struct_tinfo.get_ordinal() 52 | struct_name = struct_tinfo.dstr() 53 | 54 | if (offset + idx) % 2: 55 | default_field_type = "_BYTE" 56 | elif (offset + idx) % 4: 57 | default_field_type = "_WORD" 58 | elif (offset + idx) % 8: 59 | default_field_type = "_DWORD" 60 | else: 61 | default_field_type = "_QWORD" if const.EA64 else "_DWORD" 62 | 63 | declaration = idaapi.ask_text( 64 | 0x10000, "{0} field_{1:X}".format(default_field_type, offset + idx), "Enter new structure member:" 65 | ) 66 | if declaration is None: 67 | return 68 | 69 | result = self.parse_declaration(declaration) 70 | if result is None: 71 | logger.warn("Bad member declaration") 72 | return 73 | 74 | field_tinfo, field_name = result 75 | field_size = field_tinfo.get_size() 76 | udt_data = idaapi.udt_type_data_t() 77 | udt_member = idaapi.udt_member_t() 78 | 79 | struct_tinfo.get_udt_details(udt_data) 80 | udt_member.offset = offset * 8 81 | struct_tinfo.find_udt_member(udt_member, idaapi.STRMEM_OFFSET) 82 | gap_size = udt_member.size // 8 83 | 84 | gap_leftover = gap_size - idx - field_size 85 | 86 | if gap_leftover < 0: 87 | logger.error("Too big size for the field. Type with maximum {0} bytes can be used".format(gap_size - idx)) 88 | return 89 | 90 | iterator = udt_data.find(udt_member) 91 | iterator = udt_data.erase(iterator) 92 | 93 | if gap_leftover > 0: 94 | udt_data.insert(iterator, helper.create_padding_udt_member(offset + idx + field_size, gap_leftover)) 95 | 96 | udt_member = idaapi.udt_member_t() 97 | udt_member.offset = offset * 8 + idx 98 | udt_member.name = field_name 99 | udt_member.type = field_tinfo 100 | udt_member.size = field_size 101 | 102 | iterator = udt_data.insert(iterator, udt_member) 103 | 104 | if idx > 0: 105 | udt_data.insert(iterator, helper.create_padding_udt_member(offset, idx)) 106 | 107 | struct_tinfo.create_udt(udt_data, idaapi.BTF_STRUCT) 108 | struct_tinfo.set_numbered_type(idaapi.cvar.idati, ordinal, idaapi.BTF_STRUCT, struct_name) 109 | hx_view.refresh_view(True) 110 | 111 | @staticmethod 112 | def parse_declaration(declaration): 113 | m = re.search(r"^(\w+[ *]+)(\w+)(\[(\d+)\])?$", declaration) 114 | if m is None: 115 | logger.error("Member declaration should be like `TYPE_NAME NAME[SIZE]` (Array is optional)") 116 | return 117 | 118 | type_name, field_name, _, arr_size = m.groups() 119 | if field_name[0].isdigit(): 120 | logger.error("Bad field name") 121 | return 122 | 123 | result = idc.parse_decl(type_name, 0) 124 | if result is None: 125 | logger.error("Failed to parse member type. It should be like `TYPE_NAME NAME[SIZE]` (Array is optional)") 126 | return 127 | 128 | _, tp, fld = result 129 | tinfo = idaapi.tinfo_t() 130 | tinfo.deserialize(idaapi.cvar.idati, tp, fld, None) 131 | if arr_size: 132 | assert tinfo.create_array(tinfo, int(arr_size)) 133 | return tinfo, field_name 134 | 135 | actions.action_manager.register(CreateNewField()) 136 | 137 | # TODO: All this stuff can be done automatically when we either use ctrl+F5 or regular F5 138 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/scanners.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | from . import actions 3 | import HexRaysPyTools.api as api 4 | import HexRaysPyTools.core.cache as cache 5 | import HexRaysPyTools.core.helper as helper 6 | from ..core.variable_scanner import NewShallowSearchVisitor, NewDeepSearchVisitor, DeepReturnVisitor 7 | from ..core.temporary_structure import TemporaryStructureModel 8 | 9 | 10 | class Scanner(actions.HexRaysPopupAction): 11 | """ 12 | Abstract class containing common check of whether object can be scanned or not. 13 | Concrete class implement actual scan process in activate method 14 | """ 15 | def __init__(self): 16 | super(Scanner, self).__init__() 17 | 18 | def _can_be_scanned(self, cfunc, ctree_item): 19 | obj = api.ScanObject.create(cfunc, ctree_item) 20 | return obj and helper.is_legal_type(obj.tinfo) 21 | 22 | def check(self, hx_view): 23 | cfunc, ctree_item = hx_view.cfunc, hx_view.item 24 | return self._can_be_scanned(cfunc, ctree_item) 25 | 26 | 27 | class ShallowScanVariable(Scanner): 28 | description = "Scan Variable" 29 | hotkey = "F" 30 | 31 | def __init__(self): 32 | super(ShallowScanVariable, self).__init__() 33 | 34 | def activate(self, ctx): 35 | hx_view = idaapi.get_widget_vdui(ctx.widget) 36 | cfunc = hx_view.cfunc 37 | origin = cache.temporary_structure.main_offset 38 | 39 | if self._can_be_scanned(cfunc, hx_view.item): 40 | obj = api.ScanObject.create(cfunc, hx_view.item) 41 | visitor = NewShallowSearchVisitor(cfunc, origin, obj, cache.temporary_structure) 42 | visitor.process() 43 | 44 | 45 | class DeepScanVariable(Scanner): 46 | description = "Deep Scan Variable" 47 | hotkey = "shift+F" 48 | 49 | def __init__(self): 50 | super(DeepScanVariable, self).__init__() 51 | 52 | def activate(self, ctx): 53 | hx_view = idaapi.get_widget_vdui(ctx.widget) 54 | cfunc = hx_view.cfunc 55 | origin = cache.temporary_structure.main_offset 56 | 57 | if self._can_be_scanned(cfunc, hx_view.item): 58 | obj = api.ScanObject.create(cfunc, hx_view.item) 59 | if helper.FunctionTouchVisitor(cfunc).process(): 60 | hx_view.refresh_view(True) 61 | visitor = NewDeepSearchVisitor(hx_view.cfunc, origin, obj, cache.temporary_structure) 62 | visitor.process() 63 | 64 | 65 | class RecognizeShape(Scanner): 66 | description = "Recognize Shape" 67 | 68 | def __init__(self): 69 | super(RecognizeShape, self).__init__() 70 | 71 | def activate(self, ctx): 72 | hx_view = idaapi.get_widget_vdui(ctx.widget) 73 | cfunc = hx_view.cfunc 74 | 75 | if not self._can_be_scanned(cfunc, hx_view.item): 76 | return 77 | 78 | obj = api.ScanObject.create(cfunc, hx_view.item) 79 | tmp_struct = TemporaryStructureModel() 80 | visitor = NewShallowSearchVisitor(cfunc, 0, obj, tmp_struct) 81 | visitor.process() 82 | tinfo = tmp_struct.get_recognized_shape() 83 | if tinfo: 84 | tinfo.create_ptr(tinfo) 85 | if obj.id == api.SO_LOCAL_VARIABLE: 86 | hx_view.set_lvar_type(obj.lvar, tinfo) 87 | elif obj.id == api.SO_GLOBAL_OBJECT: 88 | idaapi.apply_tinfo(obj.obj_ea, tinfo, idaapi.TINFO_DEFINITE) 89 | hx_view.refresh_view(True) 90 | 91 | 92 | class DeepScanReturn(Scanner): 93 | description = "Deep Scan Returned Variables" 94 | 95 | def __init__(self): 96 | super(DeepScanReturn, self).__init__() 97 | 98 | def check(self, hx_view): 99 | cfunc, ctree_item = hx_view.cfunc, hx_view.item 100 | if ctree_item.citype != idaapi.VDI_FUNC: 101 | return False 102 | tinfo = idaapi.tinfo_t() 103 | hx_view.cfunc.get_func_type(tinfo) 104 | return helper.is_legal_type(tinfo.get_rettype()) 105 | 106 | def activate(self, ctx): 107 | hx_view = idaapi.get_widget_vdui(ctx.widget) 108 | func_ea = hx_view.cfunc.entry_ea 109 | obj = api.ReturnedObject(func_ea) 110 | origin = cache.temporary_structure.main_offset 111 | visitor = DeepReturnVisitor(hx_view.cfunc, origin, obj, cache.temporary_structure) 112 | visitor.process() 113 | 114 | 115 | class DeepScanFunctions(actions.Action): 116 | description = "Scan First Argument" 117 | 118 | def __init__(self): 119 | super(DeepScanFunctions, self).__init__() 120 | 121 | def activate(self, ctx): 122 | for idx in ctx.chooser_selection: 123 | func_ea = idaapi.getn_func(idx - 1).start_ea 124 | cfunc = helper.decompile_function(func_ea) 125 | obj = api.VariableObject(cfunc.get_lvars()[0], 0) 126 | if cfunc: 127 | NewDeepSearchVisitor(cfunc, 0, obj, cache.temporary_structure).process() 128 | 129 | def update(self, ctx): 130 | if ctx.form_type == idaapi.BWN_FUNCS: 131 | idaapi.attach_action_to_popup(ctx.widget, None, self.name) 132 | return idaapi.AST_ENABLE_FOR_WIDGET 133 | return idaapi.AST_DISABLE_FOR_WIDGET 134 | 135 | 136 | actions.action_manager.register(ShallowScanVariable()) 137 | actions.action_manager.register(DeepScanVariable()) 138 | actions.action_manager.register(RecognizeShape()) 139 | actions.action_manager.register(DeepScanReturn()) 140 | actions.action_manager.register(DeepScanFunctions()) 141 | -------------------------------------------------------------------------------- /HexRaysPyTools/core/common.py: -------------------------------------------------------------------------------- 1 | import re 2 | 3 | 4 | BAD_C_NAME_PATTERN = re.compile('[^a-zA-Z_0-9:]') 5 | 6 | 7 | def demangled_name_to_c_str(name): 8 | """ 9 | Removes or replaces characters from demangled symbol so that it was possible to create legal C structure from it 10 | """ 11 | if not BAD_C_NAME_PATTERN.findall(name): 12 | return name 13 | 14 | # FIXME: This is very ugly way to find and replace illegal characters 15 | idx = name.find("::operator") 16 | if idx >= 0: 17 | idx += len("::operator") 18 | if idx == len(name) or name[idx].isalpha(): 19 | # `operator` is part of name of some name and not a keyword 20 | pass 21 | elif name[idx:idx + 2] == "==": 22 | name = name.replace("operator==", "operator_EQ_") 23 | elif name[idx:idx + 2] == "!=": 24 | name = name.replace("operator!=", "operator_NEQ_") 25 | elif name[idx] == "=": 26 | name = name.replace("operator=", "operator_ASSIGN_") 27 | elif name[idx:idx + 2] == "+=": 28 | name = name.replace("operator+=", "operator_PLUS_ASSIGN_") 29 | elif name[idx:idx + 2] == "-=": 30 | name = name.replace("operator-=", "operator_MINUS_ASSIGN_") 31 | elif name[idx:idx + 2] == "*=": 32 | name = name.replace("operator*=", "operator_MUL_ASSIGN_") 33 | elif name[idx:idx + 2] == "/=": 34 | name = name.replace("operator/=", "operator_DIV_ASSIGN_") 35 | elif name[idx:idx + 2] == "%=": 36 | name = name.replace("operator%=", "operator_MODULO_DIV_ASSIGN_") 37 | elif name[idx:idx + 2] == "|=": 38 | name = name.replace("operator|=", "operator_OR_ASSIGN_") 39 | elif name[idx:idx + 2] == "&=": 40 | name = name.replace("operator&=", "operator_AND_ASSIGN_") 41 | elif name[idx:idx + 2] == "^=": 42 | name = name.replace("operator^=", "operator_XOR_ASSIGN_") 43 | elif name[idx:idx + 3] == "<<=": 44 | name = name.replace("operator<<=", "operator_LEFT_SHIFT_ASSIGN_") 45 | elif name[idx:idx + 3] == ">>=": 46 | name = name.replace("operator>>=", "operator_RIGHT_SHIFT_ASSIGN_") 47 | elif name[idx:idx + 2] == "++": 48 | name = name.replace("operator++", "operator_INC_") 49 | elif name[idx:idx + 2] == "--": 50 | name = name.replace("operator--", "operator_PTR_") 51 | elif name[idx:idx + 2] == "->": 52 | name = name.replace("operator->", "operator_REF_") 53 | elif name[idx:idx + 2] == "[]": 54 | name = name.replace("operator[]", "operator_IDX_") 55 | elif name[idx] == "*": 56 | name = name.replace("operator*", "operator_STAR_") 57 | elif name[idx:idx + 2] == "&&": 58 | name = name.replace("operator&&", "operator_LAND_") 59 | elif name[idx:idx + 2] == "||": 60 | name = name.replace("operator||", "operator_LOR_") 61 | elif name[idx] == "!": 62 | name = name.replace("operator!", "operator_LNOT_") 63 | elif name[idx] == "&": 64 | name = name.replace("operator&", "operator_AND_") 65 | elif name[idx] == "|": 66 | name = name.replace("operator|", "operator_OR_") 67 | elif name[idx] == "^": 68 | name = name.replace("operator^", "operator_XOR_") 69 | elif name[idx:idx + 2] == "<<": 70 | name = name.replace("operator<<", "operator_LEFT_SHIFT_") 71 | elif name[idx:idx + 2] == ">>": 72 | name = name.replace("operator>", "operator_GREATER_") 73 | elif name[idx:idx + 2] == "<=": 74 | name = name.replace("operator<=", "operator_LESS_EQUAL_") 75 | elif name[idx:idx + 2] == ">=": 76 | name = name.replace("operator>>", "operator_RIGHT_SHIFT_") 77 | elif name[idx] == "<": 78 | name = name.replace("operator<", "operator_LESS_") 79 | elif name[idx] == ">": 80 | name = name.replace("operator>=", "operator_GREATER_EQUAL_") 81 | elif name[idx] == "+": 82 | name = name.replace("operator+", "operator_ADD_") 83 | elif name[idx] == "-": 84 | name = name.replace("operator-", "operator_SUB_") 85 | elif name[idx] == "/": 86 | name = name.replace("operator/", "operator_DIV_") 87 | elif name[idx] == "%": 88 | name = name.replace("operator%", "operator_MODULO_DIV_") 89 | elif name[idx:idx + 2] == "()": 90 | name = name.replace("operator()", "operator_CALL_") 91 | elif name[idx: idx + 6] == " new[]": 92 | name = name.replace("operator new[]", "operator_NEW_ARRAY_") 93 | elif name[idx: idx + 9] == " delete[]": 94 | name = name.replace("operator delete[]", "operator_DELETE_ARRAY_") 95 | elif name[idx: idx + 4] == " new": 96 | name = name.replace("operator new", "operator_NEW_") 97 | elif name[idx: idx + 7] == " delete": 98 | name = name.replace("operator delete", "operator_DELETE_") 99 | elif name[idx:idx + 2] == "\"\" ": 100 | name = name.replace("operator\"\" ", "operator_LITERAL_") 101 | elif name[idx] == "~": 102 | name = name.replace("operator~", "operator_NOT_") 103 | elif name[idx] == ' ': 104 | pass 105 | else: 106 | raise AssertionError("Replacement of demangled string by c-string for keyword `operatorXXX` is not yet" 107 | "implemented ({}). You can do it by yourself or create an issue".format(name)) 108 | 109 | name = name.replace("public:", "") 110 | name = name.replace("protected:", "") 111 | name = name.replace("private:", "") 112 | name = name.replace("~", "DESTRUCTOR_") 113 | name = name.replace("*", "_PTR") 114 | name = name.replace("<", "_t_") 115 | name = name.replace(">", "_t_") 116 | name = "_".join(filter(len, BAD_C_NAME_PATTERN.split(name))) 117 | return name 118 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/swap_if.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | import idc 3 | 4 | from . import actions 5 | from . import callbacks 6 | 7 | 8 | def inverse_if_condition(cif): 9 | # cexpr_t has become broken but fortunately still exist `assing` method which copies one expr into another 10 | cit_if_condition = cif.expr 11 | tmp_cexpr = idaapi.cexpr_t() 12 | tmp_cexpr.assign(cit_if_condition) 13 | new_if_condition = idaapi.lnot(tmp_cexpr) 14 | cif.expr.swap(new_if_condition) 15 | del cit_if_condition 16 | 17 | 18 | def inverse_if(cif): 19 | inverse_if_condition(cif) 20 | idaapi.qswap(cif.ithen, cif.ielse) 21 | 22 | 23 | _ARRAY_STORAGE_PREFIX = "$HexRaysPyTools:IfThenElse:" 24 | 25 | 26 | def has_inverted(func_ea): 27 | # Find if function has any swapped THEN-ELSE branches 28 | internal_name = _ARRAY_STORAGE_PREFIX + hex(int(func_ea - idaapi.get_imagebase())) 29 | internal_id = idc.get_array_id(internal_name) 30 | return internal_id != -1 31 | 32 | 33 | def get_inverted(func_ea): 34 | # Returns set of relative virtual addresses which are tied to IF and swapped 35 | internal_name = _ARRAY_STORAGE_PREFIX + hex(int(func_ea - idaapi.get_imagebase())) 36 | internal_id = idc.get_array_id(internal_name) 37 | array = idc.get_array_element(idc.AR_STR, internal_id, 0) 38 | return set(map(int, array.split())) 39 | 40 | 41 | def invert(func_ea, if_ea): 42 | # Store information about swaps (affected through actions) 43 | iv_rva = if_ea - idaapi.get_imagebase() 44 | func_rva = func_ea - idaapi.get_imagebase() 45 | internal_name = _ARRAY_STORAGE_PREFIX + hex(int(func_rva)) 46 | internal_id = idc.get_array_id(internal_name) 47 | if internal_id == -1: 48 | internal_id = idc.create_array(internal_name) 49 | idc.set_array_string(internal_id, 0, str(iv_rva)) 50 | else: 51 | inverted = get_inverted(func_ea) 52 | try: 53 | inverted.remove(iv_rva) 54 | if not inverted: 55 | idc.delete_array(internal_id) 56 | 57 | except KeyError: 58 | inverted.add(iv_rva) 59 | 60 | idc.set_array_string(internal_id, 0, " ".join(map(str, inverted))) 61 | 62 | 63 | class SwapThenElse(actions.HexRaysPopupAction): 64 | description = "Swap then/else" 65 | hotkey = "Shift+S" 66 | 67 | def __init__(self): 68 | super(SwapThenElse, self).__init__() 69 | 70 | def check(self, hx_view): 71 | # Checks if we clicked on IF and this if has both THEN and ELSE branches 72 | if hx_view.item.citype != idaapi.VDI_EXPR: 73 | return False 74 | insn = hx_view.item.it.to_specific_type 75 | if insn.op != idaapi.cit_if or insn.cif.ielse is None: 76 | return False 77 | return insn.op == idaapi.cit_if and insn.cif.ielse 78 | 79 | def activate(self, ctx): 80 | hx_view = idaapi.get_widget_vdui(ctx.widget) 81 | if self.check(hx_view): 82 | insn = hx_view.item.it.to_specific_type 83 | inverse_if(insn.cif) 84 | hx_view.refresh_ctext() 85 | 86 | invert(hx_view.cfunc.entry_ea, insn.ea) 87 | 88 | def update(self, ctx): 89 | if ctx.widget_type == idaapi.BWN_PSEUDOCODE: 90 | return idaapi.AST_ENABLE_FOR_WIDGET 91 | return idaapi.AST_DISABLE_FOR_WIDGET 92 | 93 | 94 | actions.action_manager.register(SwapThenElse()) 95 | 96 | 97 | class SwapThenElseVisitor(idaapi.ctree_parentee_t): 98 | def __init__(self, inverted): 99 | super(SwapThenElseVisitor, self).__init__() 100 | self.__inverted = inverted 101 | 102 | def visit_insn(self, insn): 103 | if insn.op != idaapi.cit_if or insn.cif.ielse is None: 104 | return 0 105 | 106 | if insn.ea in self.__inverted: 107 | inverse_if(insn.cif) 108 | 109 | return 0 110 | 111 | def apply_to(self, *args): 112 | if self.__inverted: 113 | super(SwapThenElseVisitor, self).apply_to(*args) 114 | 115 | 116 | class SpaghettiVisitor(idaapi.ctree_parentee_t): 117 | def __init__(self): 118 | super(SpaghettiVisitor, self).__init__() 119 | 120 | def visit_insn(self, instruction): 121 | if instruction.op != idaapi.cit_block: 122 | return 0 123 | 124 | while True: 125 | cblock = instruction.cblock 126 | size = cblock.size() 127 | # Find block that has "If" and "return" as last 2 statements 128 | if size < 2: 129 | break 130 | 131 | if cblock.at(size - 2).op != idaapi.cit_if: 132 | break 133 | 134 | cif = cblock.at(size - 2).cif 135 | if cblock.back().op != idaapi.cit_return or cif.ielse: 136 | break 137 | 138 | cit_then = cif.ithen 139 | 140 | # Skip if only one (not "if") statement in "then" branch 141 | if cit_then.cblock.size() == 1 and cit_then.cblock.front().op != idaapi.cit_if: 142 | return 0 143 | 144 | inverse_if_condition(cif) 145 | 146 | # Take return from list of statements and later put it back 147 | cit_return = idaapi.cinsn_t() 148 | cit_return.assign(instruction.cblock.back()) 149 | cit_return.thisown = False 150 | instruction.cblock.pop_back() 151 | 152 | # Fill main block with statements from "Then" branch 153 | while cit_then.cblock: 154 | instruction.cblock.push_back(cit_then.cblock.front()) 155 | cit_then.cblock.pop_front() 156 | 157 | # Put back main return if there's no another return or "GOTO" already 158 | if instruction.cblock.back().op not in (idaapi.cit_return, idaapi.cit_goto): 159 | new_return = idaapi.cinsn_t() 160 | new_return.thisown = False 161 | new_return.assign(cit_return) 162 | instruction.cblock.push_back(new_return) 163 | 164 | # Put return into "Then" branch 165 | cit_then.cblock.push_back(cit_return) 166 | return 0 167 | 168 | 169 | class SilentIfSwapper(callbacks.HexRaysEventHandler): 170 | 171 | def __init__(self): 172 | super(SilentIfSwapper, self).__init__() 173 | 174 | def handle(self, event, *args): 175 | cfunc, level_of_maturity = args 176 | if level_of_maturity == idaapi.CMAT_TRANS1 and has_inverted(cfunc.entry_ea): 177 | # Make RVA from VA of IF instructions that should be inverted 178 | inverted = [n + idaapi.get_imagebase() for n in get_inverted(cfunc.entry_ea)] 179 | visitor = SwapThenElseVisitor(inverted) 180 | visitor.apply_to(cfunc.body, None) 181 | elif level_of_maturity == idaapi.CMAT_TRANS2: 182 | visitor = SpaghettiVisitor() 183 | visitor.apply_to(cfunc.body, None) 184 | 185 | 186 | callbacks.hx_callback_manager.register(idaapi.hxe_maturity, SilentIfSwapper()) 187 | -------------------------------------------------------------------------------- /HexRaysPyTools/core/structure_graph.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | import idaapi 4 | import idc 5 | 6 | logger = logging.getLogger(__name__) 7 | 8 | 9 | class LocalType: 10 | def __init__(self, name, members_ordinals, hint, is_selected=False, is_typedef=False, is_enum=False, is_union=False): 11 | self.name = name 12 | self.members_ordinals = members_ordinals 13 | self.hint = hint 14 | self.is_selected = is_selected 15 | self.is_typedef = is_typedef 16 | self.is_enum = is_enum 17 | self.is_union = is_union 18 | 19 | def __call__(self): 20 | return self.name, self.members_ordinals 21 | 22 | def __str__(self): 23 | return "<{0}, {1}>".format(self.name, self.members_ordinals) 24 | 25 | def __repr__(self): 26 | return self.__str__() 27 | 28 | @property 29 | def name_and_color(self): 30 | if self.is_selected: 31 | return self.name, 0x0000FF 32 | elif self.is_typedef: 33 | return self.name, 0x99FFFF 34 | elif self.is_enum: 35 | return self.name, 0x33FF33 36 | elif self.is_union: 37 | return self.name, 0xCCCC00 38 | return self.name, 0xffdd99 39 | 40 | 41 | class StructureGraph: 42 | # TODO:Enum types display 43 | def __init__(self, ordinal_list=None): 44 | self.ordinal_list = ordinal_list if ordinal_list else range(1, idc.get_ordinal_qty()) 45 | self.local_types = {} 46 | self.edges = [] 47 | self.final_edges = [] 48 | self.visited_downward = [] 49 | self.visited_upward = [] 50 | self.downward_edges = {} 51 | self.upward_edges = {} 52 | self.initialize_nodes() 53 | self.calculate_edges() 54 | 55 | def change_selected(self, selected): 56 | self.visited_downward = [] 57 | self.visited_upward = [] 58 | self.final_edges = [] 59 | for ordinal in self.ordinal_list: 60 | self.local_types[ordinal].is_selected = False 61 | self.ordinal_list = set(self.local_types).intersection(selected) 62 | for ordinal in self.ordinal_list: 63 | self.local_types[ordinal].is_selected = True 64 | 65 | @staticmethod 66 | def get_ordinal(tinfo): 67 | while tinfo.is_ptr() or tinfo.is_array(): 68 | tinfo.remove_ptr_or_array() 69 | if tinfo.is_udt(): 70 | return tinfo.get_ordinal() 71 | elif tinfo.is_enum(): 72 | return tinfo.get_ordinal() 73 | elif tinfo.is_typeref(): 74 | typeref_ordinal = tinfo.get_ordinal() 75 | if typeref_ordinal: 76 | typeref_tinfo = StructureGraph.get_tinfo_by_ordinal(typeref_ordinal) 77 | if typeref_tinfo is None: 78 | logger.warn("You have dependencies of deleted %s type", tinfo.dstr()) 79 | return 0 80 | 81 | if typeref_tinfo.is_typeref() or typeref_tinfo.is_udt() or typeref_tinfo.is_ptr(): 82 | return typeref_ordinal 83 | return 0 84 | 85 | @staticmethod 86 | def get_members_ordinals(tinfo): 87 | ordinals = [] 88 | if tinfo.is_udt(): 89 | udt_data = idaapi.udt_type_data_t() 90 | tinfo.get_udt_details(udt_data) 91 | for udt_member in udt_data: 92 | ordinal = StructureGraph.get_ordinal(udt_member.type) 93 | if ordinal: 94 | ordinals.append(ordinal) 95 | return ordinals 96 | 97 | @staticmethod 98 | def get_tinfo_by_ordinal(ordinal): 99 | local_typestring = idc.get_local_tinfo(ordinal) 100 | if local_typestring: 101 | p_type, fields = local_typestring 102 | local_tinfo = idaapi.tinfo_t() 103 | local_tinfo.deserialize(idaapi.cvar.idati, p_type, fields) 104 | return local_tinfo 105 | return None 106 | 107 | def initialize_nodes(self): 108 | for ordinal in range(1, idc.get_ordinal_qty()): 109 | # if ordinal == 15: 110 | # import pydevd 111 | # pydevd.settrace("localhost", port=12345, stdoutToServer=True, stderrToServer=True) 112 | 113 | local_tinfo = StructureGraph.get_tinfo_by_ordinal(ordinal) 114 | if not local_tinfo: 115 | continue 116 | name = idc.get_numbered_type_name(ordinal) 117 | 118 | if local_tinfo.is_typeref(): 119 | typeref_ordinal = local_tinfo.get_ordinal() 120 | members_ordinals = [] 121 | if typeref_ordinal: 122 | typeref_tinfo = StructureGraph.get_tinfo_by_ordinal(typeref_ordinal) 123 | if typeref_tinfo.is_typeref() or typeref_tinfo.is_udt() or typeref_tinfo.is_ptr(): 124 | members_ordinals = [typeref_ordinal] 125 | cdecl_typedef = idaapi.print_tinfo(None, 4, 5, 0x3, local_tinfo, None, None) 126 | self.local_types[ordinal] = LocalType(name, members_ordinals, cdecl_typedef, is_typedef=True) 127 | elif local_tinfo.is_udt(): 128 | # udt_data = idaapi.udt_type_data_t() 129 | # local_tinfo.get_udt_details(udt_data) 130 | members_ordinals = StructureGraph.get_members_ordinals(local_tinfo) 131 | cdecl_typedef = idaapi.print_tinfo(None, 4, 5, 0x1, local_tinfo, None, None) 132 | self.local_types[ordinal] = LocalType(name, members_ordinals, cdecl_typedef, is_union=local_tinfo.is_union()) 133 | elif local_tinfo.is_ptr(): 134 | typeref_ordinal = StructureGraph.get_ordinal(local_tinfo) 135 | members_ordinals = [typeref_ordinal] if typeref_ordinal else [] 136 | cdecl_typedef = idaapi.print_tinfo(None, 4, 5, 0x2, local_tinfo, None, None) 137 | self.local_types[ordinal] = LocalType( 138 | name, 139 | members_ordinals, 140 | cdecl_typedef + ' *', 141 | is_typedef=True 142 | ) 143 | elif local_tinfo.is_enum(): 144 | cdecl_typedef = idaapi.print_tinfo(None, 4, 5, 0x21, local_tinfo, None, None) 145 | self.local_types[ordinal] = LocalType(name, [], cdecl_typedef, is_enum=True) 146 | 147 | self.ordinal_list = set(self.ordinal_list).intersection(self.local_types) 148 | for ordinal in self.ordinal_list: 149 | self.local_types[ordinal].is_selected = True 150 | 151 | def calculate_edges(self): 152 | for first in list(self.local_types.keys()): 153 | for second in self.local_types[first].members_ordinals: 154 | self.edges.append((first, second)) 155 | 156 | self.downward_edges = {key: [] for key in list(self.local_types.keys())} 157 | self.upward_edges = {key: [] for key in list(self.local_types.keys())} 158 | 159 | for key, value in self.edges: 160 | self.downward_edges[key].append(value) 161 | self.upward_edges[value].append(key) 162 | 163 | def generate_final_edges_down(self, node): 164 | if node not in self.visited_downward: 165 | self.visited_downward.append(node) 166 | else: 167 | return 168 | for next_node in self.downward_edges[node]: 169 | self.final_edges.append((node, next_node)) 170 | for next_node in self.downward_edges[node]: 171 | self.generate_final_edges_down(next_node) 172 | 173 | def generate_final_edges_up(self, node): 174 | if node not in self.visited_upward: 175 | self.visited_upward.append(node) 176 | else: 177 | return 178 | for next_node in self.upward_edges[node]: 179 | self.final_edges.append((next_node, node)) 180 | for next_node in self.upward_edges[node]: 181 | self.generate_final_edges_up(next_node) 182 | 183 | def get_nodes(self): 184 | for ordinal in self.ordinal_list: 185 | if ordinal in self.local_types: 186 | self.generate_final_edges_down(ordinal) 187 | self.generate_final_edges_up(ordinal) 188 | return set([node for nodes in self.final_edges for node in nodes]) 189 | 190 | def get_edges(self): 191 | return self.final_edges 192 | -------------------------------------------------------------------------------- /HexRaysPyTools/forms.py: -------------------------------------------------------------------------------- 1 | from PyQt5 import QtCore, QtWidgets 2 | 3 | import idaapi 4 | 5 | 6 | class MyChoose(idaapi.Choose): 7 | def __init__(self, items, title, cols, icon=-1): 8 | idaapi.Choose.__init__(self, title, cols, flags=idaapi.Choose.CH_MODAL, icon=icon) 9 | self.items = items 10 | 11 | def OnClose(self): 12 | pass 13 | 14 | def OnGetLine(self, n): 15 | return self.items[n] 16 | 17 | def OnGetSize(self): 18 | return len(self.items) 19 | 20 | 21 | class StructureBuilder(idaapi.PluginForm): 22 | def __init__(self, structure_model): 23 | super(StructureBuilder, self).__init__() 24 | self.structure_model = structure_model 25 | self.parent = None 26 | 27 | def OnCreate(self, form): 28 | self.parent = idaapi.PluginForm.FormToPyQtWidget(form) 29 | self.init_ui() 30 | 31 | def init_ui(self): 32 | self.parent.setStyleSheet( 33 | "QTableView {background-color: transparent; selection-background-color: #87bdd8;}" 34 | "QHeaderView::section {background-color: transparent; border: 0.5px solid;}" 35 | "QPushButton {width: 50px; height: 20px;}" 36 | # "QPushButton::pressed {background-color: #ccccff}" 37 | ) 38 | self.parent.resize(400, 600) 39 | self.parent.setWindowTitle('Structure Builder') 40 | 41 | btn_finalize = QtWidgets.QPushButton("&Finalize") 42 | btn_disable = QtWidgets.QPushButton("&Disable") 43 | btn_enable = QtWidgets.QPushButton("&Enable") 44 | btn_origin = QtWidgets.QPushButton("&Origin") 45 | btn_array = QtWidgets.QPushButton("&Array") 46 | btn_pack = QtWidgets.QPushButton("&Pack") 47 | btn_unpack = QtWidgets.QPushButton("&Unpack") 48 | btn_remove = QtWidgets.QPushButton("&Remove") 49 | btn_resolve = QtWidgets.QPushButton("Resolve") 50 | btn_clear = QtWidgets.QPushButton("Clear") # Clear button doesn't have shortcut because it can fuck up all work 51 | btn_recognize = QtWidgets.QPushButton("Recognize Shape") 52 | btn_recognize.setStyleSheet("QPushButton {width: 100px; height: 20px;}") 53 | 54 | btn_finalize.setShortcut("f") 55 | btn_disable.setShortcut("d") 56 | btn_enable.setShortcut("e") 57 | btn_origin.setShortcut("o") 58 | btn_array.setShortcut("a") 59 | btn_pack.setShortcut("p") 60 | btn_unpack.setShortcut("u") 61 | btn_remove.setShortcut("r") 62 | 63 | struct_view = QtWidgets.QTableView() 64 | struct_view.setModel(self.structure_model) 65 | # struct_view.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection) 66 | 67 | struct_view.verticalHeader().setVisible(False) 68 | struct_view.verticalHeader().setDefaultSectionSize(24) 69 | struct_view.horizontalHeader().setStretchLastSection(True) 70 | struct_view.horizontalHeader().setSectionResizeMode(QtWidgets.QHeaderView.ResizeToContents) 71 | 72 | grid_box = QtWidgets.QGridLayout() 73 | grid_box.setSpacing(0) 74 | grid_box.addWidget(btn_finalize, 0, 0) 75 | grid_box.addWidget(btn_enable, 0, 1) 76 | grid_box.addWidget(btn_disable, 0, 2) 77 | grid_box.addWidget(btn_origin, 0, 3) 78 | grid_box.addItem(QtWidgets.QSpacerItem(20, 20, QtWidgets.QSizePolicy.Expanding), 0, 5) 79 | grid_box.addWidget(btn_array, 1, 0) 80 | grid_box.addWidget(btn_pack, 1, 1) 81 | grid_box.addWidget(btn_unpack, 1, 2) 82 | grid_box.addWidget(btn_remove, 1, 3) 83 | grid_box.addWidget(btn_resolve, 0, 4) 84 | grid_box.addItem(QtWidgets.QSpacerItem(20, 20, QtWidgets.QSizePolicy.Expanding), 1, 5) 85 | grid_box.addWidget(btn_recognize, 0, 6) 86 | grid_box.addWidget(btn_clear, 1, 6) 87 | 88 | vertical_box = QtWidgets.QVBoxLayout() 89 | vertical_box.addWidget(struct_view) 90 | vertical_box.addLayout(grid_box) 91 | self.parent.setLayout(vertical_box) 92 | 93 | btn_finalize.clicked.connect(lambda: self.structure_model.finalize()) 94 | btn_disable.clicked.connect(lambda: self.structure_model.disable_rows(struct_view.selectedIndexes())) 95 | btn_enable.clicked.connect(lambda: self.structure_model.enable_rows(struct_view.selectedIndexes())) 96 | btn_origin.clicked.connect(lambda: self.structure_model.set_origin(struct_view.selectedIndexes())) 97 | btn_array.clicked.connect(lambda: self.structure_model.make_array(struct_view.selectedIndexes())) 98 | btn_pack.clicked.connect(lambda: self.structure_model.pack_substructure(struct_view.selectedIndexes())) 99 | btn_unpack.clicked.connect(lambda: self.structure_model.unpack_substructure(struct_view.selectedIndexes())) 100 | btn_remove.clicked.connect(lambda: self.structure_model.remove_items(struct_view.selectedIndexes())) 101 | btn_resolve.clicked.connect(lambda: self.structure_model.resolve_types()) 102 | btn_clear.clicked.connect(lambda: self.structure_model.clear()) 103 | btn_recognize.clicked.connect(lambda: self.structure_model.recognize_shape(struct_view.selectedIndexes())) 104 | struct_view.activated[QtCore.QModelIndex].connect(self.structure_model.activated) 105 | self.structure_model.dataChanged.connect(struct_view.clearSelection) 106 | 107 | def OnClose(self, form): 108 | pass 109 | 110 | def Show(self, caption=None, options=0): 111 | return idaapi.PluginForm.Show(self, caption, options=options) 112 | 113 | 114 | class StructureGraphViewer(idaapi.GraphViewer): 115 | def __init__(self, title, graph): 116 | idaapi.GraphViewer.__init__(self, title) 117 | self.graph = graph 118 | self.nodes_id = {} 119 | 120 | def OnRefresh(self): 121 | self.Clear() 122 | self.nodes_id.clear() 123 | for node in self.graph.get_nodes(): 124 | self.nodes_id[node] = self.AddNode(node) 125 | for first, second in self.graph.get_edges(): 126 | self.AddEdge(self.nodes_id[first], self.nodes_id[second]) 127 | return True 128 | 129 | def OnGetText(self, node_id): 130 | return self.graph.local_types[self[node_id]].name_and_color 131 | 132 | def OnHint(self, node_id): 133 | """ Try-catch clause because IDA sometimes attempts to use old information to get hint """ 134 | try: 135 | ordinal = self[node_id] 136 | return self.graph.local_types[ordinal].hint 137 | except KeyError: 138 | return 139 | 140 | def OnDblClick(self, node_id): 141 | self.change_selected([self[node_id]]) 142 | 143 | def change_selected(self, ordinals): 144 | self.graph.change_selected(ordinals) 145 | self.Refresh() 146 | self.Select(self.nodes_id[ordinals[0]]) 147 | 148 | 149 | class ClassViewer(idaapi.PluginForm): 150 | def __init__(self, proxy_model, class_model): 151 | super(ClassViewer, self).__init__() 152 | self.parent = None 153 | self.class_tree = QtWidgets.QTreeView() 154 | self.line_edit_filter = QtWidgets.QLineEdit() 155 | 156 | self.action_collapse = QtWidgets.QAction("Collapse all", self.class_tree) 157 | self.action_expand = QtWidgets.QAction("Expand all", self.class_tree) 158 | self.action_set_arg = QtWidgets.QAction("Set First Argument Type", self.class_tree) 159 | self.action_rollback = QtWidgets.QAction("Rollback", self.class_tree) 160 | self.action_refresh = QtWidgets.QAction("Refresh", self.class_tree) 161 | self.action_commit = QtWidgets.QAction("Commit", self.class_tree) 162 | 163 | self.menu = QtWidgets.QMenu(self.parent) 164 | 165 | self.proxy_model = proxy_model 166 | self.class_model = class_model 167 | 168 | def OnCreate(self, form): 169 | # self.parent = self.FormToPySideWidget(form) 170 | self.parent = idaapi.PluginForm.FormToPyQtWidget(form) 171 | self.init_ui() 172 | 173 | def init_ui(self): 174 | self.parent.setWindowTitle('Classes') 175 | self.parent.setStyleSheet( 176 | # "QTreeView::item:!has-children { background-color: #fefbd8; border: 0.5px solid lightgray ;}" 177 | # "QTreeView::item:has-children { background-color: #80ced6; border-top: 1px solid black ;}" 178 | # "QTreeView::item:selected { background-color: #618685; show-decoration-selected: 1;}" 179 | "QTreeView {background-color: transparent; }" 180 | "QHeaderView::section {background-color: transparent; border: 1px solid;}" 181 | ) 182 | 183 | hbox_layout = QtWidgets.QHBoxLayout() 184 | label_filter = QtWidgets.QLabel("&Filter:") 185 | label_filter.setBuddy(self.line_edit_filter) 186 | hbox_layout.addWidget(label_filter) 187 | hbox_layout.addWidget(self.line_edit_filter) 188 | 189 | self.proxy_model.setSourceModel(self.class_model) 190 | self.proxy_model.setFilterCaseSensitivity(QtCore.Qt.CaseInsensitive) 191 | self.class_tree.setModel(self.proxy_model) 192 | self.class_tree.setContextMenuPolicy(QtCore.Qt.CustomContextMenu) 193 | self.class_tree.expandAll() 194 | self.class_tree.header().setStretchLastSection(True) 195 | self.class_tree.header().setSectionResizeMode(QtWidgets.QHeaderView.ResizeToContents) 196 | self.class_tree.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection) 197 | 198 | self.action_collapse.triggered.connect(self.class_tree.collapseAll) 199 | self.action_expand.triggered.connect(self.class_tree.expandAll) 200 | self.action_set_arg.triggered.connect( 201 | lambda: self.class_model.set_first_argument_type( 202 | list(map(self.proxy_model.mapToSource, self.class_tree.selectedIndexes())) 203 | ) 204 | ) 205 | self.action_rollback.triggered.connect(lambda: self.class_model.rollback()) 206 | self.action_refresh.triggered.connect(lambda: self.class_model.refresh()) 207 | self.action_commit.triggered.connect(lambda: self.class_model.commit()) 208 | self.class_model.refreshed.connect(self.class_tree.expandAll) 209 | 210 | self.menu.addAction(self.action_collapse) 211 | self.menu.addAction(self.action_expand) 212 | self.menu.addAction(self.action_refresh) 213 | self.menu.addAction(self.action_set_arg) 214 | self.menu.addAction(self.action_rollback) 215 | self.menu.addAction(self.action_commit) 216 | 217 | vertical_box = QtWidgets.QVBoxLayout() 218 | vertical_box.addWidget(self.class_tree) 219 | vertical_box.addLayout(hbox_layout) 220 | self.parent.setLayout(vertical_box) 221 | 222 | self.class_tree.activated[QtCore.QModelIndex].connect( 223 | lambda x: self.class_model.open_function(self.proxy_model.mapToSource(x)) 224 | ) 225 | self.class_tree.customContextMenuRequested[QtCore.QPoint].connect(self.show_menu) 226 | self.line_edit_filter.textChanged[str].connect(self.proxy_model.set_regexp_filter) 227 | # proxy_model.rowsInserted[object].connect(lambda: self.class_tree.setExpanded(object, True)) 228 | 229 | def OnClose(self, form): 230 | pass 231 | 232 | def Show(self, caption=None, options=0): 233 | return idaapi.PluginForm.Show(self, caption, options=options) 234 | 235 | def show_menu(self, point): 236 | self.action_set_arg.setEnabled(True) 237 | indexes = list(map( 238 | self.proxy_model.mapToSource, 239 | [x for x in self.class_tree.selectedIndexes() if x.column() == 0] 240 | )) 241 | if len(indexes) > 1: 242 | if [x for x in indexes if len(x.internalPointer().children) > 0]: 243 | self.action_set_arg.setEnabled(False) 244 | self.menu.exec_(self.class_tree.mapToGlobal(point)) 245 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/recasts.py: -------------------------------------------------------------------------------- 1 | from collections import namedtuple 2 | import idaapi 3 | from . import actions 4 | import HexRaysPyTools.core.helper as helper 5 | 6 | 7 | RecastLocalVariable = namedtuple('RecastLocalVariable', ['recast_tinfo', 'local_variable']) 8 | RecastGlobalVariable = namedtuple('RecastGlobalVariable', ['recast_tinfo', 'global_variable_ea']) 9 | RecastArgument = namedtuple('RecastArgument', ['recast_tinfo', 'arg_idx', 'func_ea', 'func_tinfo']) 10 | RecastReturn = namedtuple('RecastReturn', ['recast_tinfo', 'func_ea']) 11 | RecastStructure = namedtuple('RecastStructure', ['recast_tinfo', 'structure_name', 'field_offset']) 12 | 13 | 14 | class RecastItemLeft(actions.HexRaysPopupAction): 15 | 16 | description = "Recast Item" 17 | hotkey = "Shift+L" 18 | 19 | def __init__(self): 20 | super(RecastItemLeft, self).__init__() 21 | 22 | def extract_recast_info(self, cfunc, ctree_item): 23 | # type: (idaapi.cfunc_t, idaapi.ctree_item_t) -> namedtuple 24 | # Returns one of the Recast... namedtuple or None if nothing was found 25 | 26 | if ctree_item.citype != idaapi.VDI_EXPR: 27 | return 28 | 29 | expression = ctree_item.it.to_specific_type 30 | child = None 31 | 32 | # Look through parents until we found Return, Assignment or Call 33 | while expression and expression.op not in (idaapi.cot_asg, idaapi.cit_return, idaapi.cot_call): 34 | child = expression.to_specific_type 35 | expression = cfunc.body.find_parent_of(expression) 36 | if not expression: 37 | return 38 | 39 | expression = expression.to_specific_type 40 | if expression.op == idaapi.cot_asg: 41 | 42 | if expression.x.opname not in ('var', 'obj', 'memptr', 'memref'): 43 | return 44 | 45 | right_expr = expression.y 46 | right_tinfo = right_expr.x.type if right_expr.op == idaapi.cot_cast else right_expr.type 47 | 48 | # Check if both left and right parts of expression are of the same types. 49 | # If not then we can recast then. 50 | if right_tinfo.dstr() == expression.x.type.dstr(): 51 | return 52 | 53 | if expression.x.op == idaapi.cot_var: 54 | # var = (TYPE ) ...; 55 | variable = cfunc.get_lvars()[expression.x.v.idx] 56 | return RecastLocalVariable(right_tinfo, variable) 57 | 58 | elif expression.x.op == idaapi.cot_obj: 59 | # g_var = (TYPE ) ...; 60 | return RecastGlobalVariable(right_tinfo, expression.x.obj_ea) 61 | 62 | elif expression.x.op == idaapi.cot_memptr: 63 | # struct->member = (TYPE ) ...; 64 | struct_name = expression.x.x.type.get_pointed_object().dstr() 65 | struct_offset = expression.x.m 66 | return RecastStructure(right_tinfo, struct_name, struct_offset) 67 | 68 | elif expression.x.op == idaapi.cot_memref: 69 | # struct.member = (TYPE ) ...; 70 | struct_name = expression.x.x.type.dstr() 71 | struct_offset = expression.x.m 72 | return RecastStructure(right_tinfo, struct_name, struct_offset) 73 | 74 | elif expression.op == idaapi.cit_return: 75 | child = child or expression.creturn.expr 76 | if child.op == idaapi.cot_cast: 77 | # return (TYPE) ...; 78 | return RecastReturn(child.x.type, cfunc.entry_ea) 79 | 80 | func_tinfo = idaapi.tinfo_t() 81 | cfunc.get_func_type(func_tinfo) 82 | rettype = func_tinfo.get_rettype() 83 | if rettype.dstr() != child.type.dstr(): 84 | # return ...; 85 | # This's possible when returned type and value are both pointers to different types 86 | return RecastReturn(child.type, cfunc.entry_ea) 87 | 88 | elif expression.op == idaapi.cot_call: 89 | if expression.x == child: 90 | return 91 | func_ea = expression.x.obj_ea 92 | arg_index, param_tinfo = helper.get_func_argument_info(expression, child) 93 | if expression.x.op == idaapi.cot_memptr: 94 | if child.op == idaapi.cot_cast: 95 | # struct_ptr->func(..., (TYPE) var, ...); 96 | arg_tinfo = child.x.type 97 | else: 98 | # struct_ptr->func(..., var, ...); When `var` and `arg` are different pointers 99 | if param_tinfo.equals_to(child.type): 100 | return 101 | arg_tinfo = child.type 102 | 103 | struct_tinfo = expression.x.x.type.get_pointed_object() 104 | funcptr_tinfo = expression.x.type 105 | helper.set_funcptr_argument(funcptr_tinfo, arg_index, arg_tinfo) 106 | return RecastStructure(funcptr_tinfo, struct_tinfo.dstr(), expression.x.m) 107 | 108 | if child.op == idaapi.cot_ref: 109 | if child.x.op == idaapi.cot_memref and child.x.m == 0: 110 | # func(..., &struct.field_0, ...) 111 | arg_tinfo = idaapi.tinfo_t() 112 | arg_tinfo.create_ptr(child.x.x.type) 113 | elif child.x.op == idaapi.cot_memptr and child.x.m == 0: 114 | # func(..., &struct->field_0, ...) 115 | arg_tinfo = child.x.x.type 116 | else: 117 | # func(..., &var, ...) 118 | arg_tinfo = child.type 119 | elif child.op == idaapi.cot_cast: 120 | arg_tinfo = child.x.type 121 | else: 122 | arg_tinfo = child.type 123 | 124 | func_tinfo = expression.x.type.get_pointed_object() 125 | return RecastArgument(arg_tinfo, arg_index, func_ea, func_tinfo) 126 | 127 | def set_label(self, label): 128 | idaapi.update_action_label(self.name, label) 129 | 130 | def check(self, hx_view): 131 | cfunc, ctree_item = hx_view.cfunc, hx_view.item 132 | 133 | ri = self.extract_recast_info(cfunc, ctree_item) 134 | if not ri: 135 | return False 136 | 137 | if isinstance(ri, RecastLocalVariable): 138 | self.set_label('Recast Variable "{0}" to {1}'.format(ri.local_variable.name, ri.recast_tinfo.dstr())) 139 | elif isinstance(ri, RecastGlobalVariable): 140 | gvar_name = idaapi.get_name(ri.global_variable_ea) 141 | self.set_label('Recast Global Variable "{0}" to {1}'.format(gvar_name, ri.recast_tinfo.dstr())) 142 | elif isinstance(ri, RecastArgument): 143 | self.set_label("Recast Argument") 144 | elif isinstance(ri, RecastStructure): 145 | self.set_label("Recast Field of {0} structure".format(ri.structure_name)) 146 | elif isinstance(ri, RecastReturn): 147 | self.set_label("Recast Return to ".format(ri.recast_tinfo.dstr())) 148 | else: 149 | raise NotImplementedError 150 | return True 151 | 152 | def activate(self, ctx): 153 | hx_view = idaapi.get_widget_vdui(ctx.widget) 154 | ri = self.extract_recast_info(hx_view.cfunc, hx_view.item) 155 | if not ri: 156 | return 0 157 | 158 | if isinstance(ri, RecastLocalVariable): 159 | hx_view.set_lvar_type(ri.local_variable, ri.recast_tinfo) 160 | 161 | elif isinstance(ri, RecastGlobalVariable): 162 | idaapi.apply_tinfo(ri.global_variable_ea, ri.recast_tinfo, idaapi.TINFO_DEFINITE) 163 | 164 | elif isinstance(ri, RecastArgument): 165 | if ri.recast_tinfo.is_array(): 166 | ri.recast_tinfo.convert_array_to_ptr() 167 | helper.set_func_argument(ri.func_tinfo, ri.arg_idx, ri.recast_tinfo) 168 | idaapi.apply_tinfo(ri.func_ea, ri.func_tinfo, idaapi.TINFO_DEFINITE) 169 | 170 | elif isinstance(ri, RecastReturn): 171 | cfunc = helper.decompile_function(ri.func_ea) 172 | if not cfunc: 173 | return 0 174 | 175 | func_tinfo = idaapi.tinfo_t() 176 | cfunc.get_func_type(func_tinfo) 177 | helper.set_func_return(func_tinfo, ri.recast_tinfo) 178 | idaapi.apply_tinfo(cfunc.entry_ea, func_tinfo, idaapi.TINFO_DEFINITE) 179 | 180 | elif isinstance(ri, RecastStructure): 181 | tinfo = idaapi.tinfo_t() 182 | tinfo.get_named_type(idaapi.cvar.idati, ri.structure_name) 183 | ordinal = idaapi.get_type_ordinal(idaapi.cvar.idati, ri.structure_name) 184 | if ordinal == 0: 185 | return 0 186 | 187 | udt_member = idaapi.udt_member_t() 188 | udt_member.offset = ri.field_offset * 8 189 | idx = tinfo.find_udt_member(udt_member, idaapi.STRMEM_OFFSET) 190 | if udt_member.offset != ri.field_offset * 8: 191 | print("[Info] Can't handle with arrays yet") 192 | elif udt_member.type.get_size() != ri.recast_tinfo.get_size(): 193 | print("[Info] Can't recast different sizes yet") 194 | else: 195 | udt_data = idaapi.udt_type_data_t() 196 | tinfo.get_udt_details(udt_data) 197 | udt_data[idx].type = ri.recast_tinfo 198 | tinfo.create_udt(udt_data, idaapi.BTF_STRUCT) 199 | tinfo.set_numbered_type(idaapi.cvar.idati, ordinal, idaapi.NTF_REPLACE, ri.structure_name) 200 | else: 201 | raise NotImplementedError 202 | 203 | hx_view.refresh_view(True) 204 | return 0 205 | 206 | 207 | class RecastItemRight(RecastItemLeft): 208 | 209 | name = "my:RecastItemRight" 210 | description = "Recast Item" 211 | hotkey = "Shift+R" 212 | 213 | def __init__(self): 214 | super(RecastItemRight, self).__init__() 215 | 216 | def extract_recast_info(self, cfunc, ctree_item): 217 | if ctree_item.citype != idaapi.VDI_EXPR: 218 | return 219 | 220 | expression = ctree_item.it 221 | result = RecastItemRight._check_potential_array(cfunc, expression) 222 | if result: 223 | return result 224 | 225 | # Look through parents until we found Cast 226 | while expression and expression.op != idaapi.cot_cast: 227 | expression = expression.to_specific_type 228 | expression = cfunc.body.find_parent_of(expression) 229 | if not expression: 230 | return 231 | 232 | expression = expression.to_specific_type 233 | 234 | # Find `(TYPE) something;` or `(TYPE *) &something;` and calculate appropriate type for recast 235 | if expression.x.op == idaapi.cot_ref: 236 | tinfo = expression.type.get_pointed_object() 237 | expression = expression.x 238 | else: 239 | tinfo = expression.type 240 | 241 | if expression.x.op == idaapi.cot_var: 242 | # (TYPE) var; 243 | variable = cfunc.get_lvars()[expression.x.v.idx] 244 | return RecastLocalVariable(tinfo, variable) 245 | 246 | elif expression.x.op == idaapi.cot_obj: 247 | # (TYPE) g_var; 248 | if helper.is_code_ea(expression.x.obj_ea) and tinfo.is_funcptr(): 249 | # (TYPE) sub_XXXXXX; 250 | tinfo = tinfo.get_pointed_object() 251 | gvar_ea = expression.x.obj_ea 252 | return RecastGlobalVariable(tinfo, gvar_ea) 253 | 254 | elif expression.x.op == idaapi.cot_call: 255 | # (TYPE) call(); 256 | idaapi.update_action_label(RecastItemRight.name, "Recast Return") 257 | func_ea = expression.x.x.obj_ea 258 | return RecastReturn(tinfo, func_ea) 259 | 260 | elif expression.x.op == idaapi.cot_memptr: 261 | # (TYPE) var->member; 262 | idaapi.update_action_label(RecastItemRight.name, "Recast Field") 263 | struct_name = expression.x.x.type.get_pointed_object().dstr() 264 | struct_offset = expression.x.m 265 | return RecastStructure(tinfo, struct_name, struct_offset) 266 | 267 | @staticmethod 268 | def _check_potential_array(cfunc, expr): 269 | """ Checks `call(..., &buffer, ..., number)` and returns information for recasting """ 270 | if expr.op != idaapi.cot_var: 271 | return 272 | 273 | var_expr = expr.to_specific_type 274 | parent = cfunc.body.find_parent_of(expr) 275 | if parent.op != idaapi.cot_ref: 276 | return 277 | 278 | parent = cfunc.body.find_parent_of(parent) 279 | if parent.op != idaapi.cot_call: 280 | return 281 | 282 | call_expr = parent.to_specific_type 283 | for arg_expr in call_expr.a: 284 | if arg_expr.op == idaapi.cot_num: 285 | number = arg_expr.numval() 286 | if number: 287 | variable = cfunc.lvars[var_expr.v.idx] 288 | char_array_tinfo = idaapi.tinfo_t() 289 | char_array_tinfo.create_array(idaapi.tinfo_t(idaapi.BTF_CHAR), number) 290 | idaapi.update_action_label(RecastItemRight.name, 'Recast Variable "{}" to "{}"'.format( 291 | variable.name, char_array_tinfo.dstr() 292 | )) 293 | return RecastLocalVariable(char_array_tinfo, variable) 294 | 295 | 296 | actions.action_manager.register(RecastItemLeft()) 297 | actions.action_manager.register(RecastItemRight()) 298 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | Plugin for IDA Pro 2 | 3 | **Table of Contents** 4 | 5 | * [About](#user-content-about) 6 | * [Installation](#user-content-installation) 7 | * [Configuration](#user-content-configuration) 8 | * [Features](#user-content-features) 9 | * [Structure reconstruction](#user-content-structure) 10 | * [Decompiler output manipulation](#user-content-manipulation) 11 | * [Classes](#user-content-classes) 12 | * [Structure Graph](#user-content-graph) 13 | * [API](#user-content-api) 14 | * [Presentations](#user-content-presentations) 15 | 16 | About 17 | ===== 18 | 19 | The plugin assists in the creation of classes/structures and detection of virtual tables. It also facilitates transforming decompiler output faster and allows to do some stuff which is otherwise impossible. 20 | 21 | **Note**: The plugin supports IDA Pro 7.x with Python 2/3. 22 | 23 | Installation 24 | ============ 25 | 26 | Just copy `HexRaysPyTools.py` file and `HexRaysPyTools` directory to Ida plugins directory. 27 | 28 | Configuration 29 | ============ 30 | 31 | Can be found at `IDADIR\cfg\HexRaysPyTools.cfg` 32 | 33 | * `debug_message_level`. Set 10 if you have a bug and want to show the log along with the information about how it was encountered in the issue. 34 | * `propagate_through_all_names`. Set `True` if you want to rename not only the default variables for the [Propagate Name](#Propagate) feature. 35 | * `store_xrefs`. Specifies whether to store the cross-references collected during the decompilation phase inside the database. (Default - True) 36 | * `scan_any_type`. Set `True` if you want to apply scanning to any variable type. By default, it is possible to scan only basic types like `DWORD`, `QWORD`, `void *` e t.c. and pointers to non-defined structure declarations. 37 | 38 | Features 39 | ======== 40 | 41 | **[Recently added][feature_history]** 42 | 43 | Structure reconstruction 44 | ------------------------ 45 | 46 | The reconstruction process usually comprises the following steps: 47 | 48 | 1) Open structure builder. 49 | 2) Find a local variable that points to the structure you would like to reconstruct. 50 | 3) Apply "Scan variable". It will collect the information about the fields that were accessed in the boundaries of one function. As an option, you can apply "Deep Scan variable", which will do the same thing but will also recursively visit other functions that has the same variable as its argument. 51 | 4) After applying steps 2 and 3 enough times, resolve conflicts in the structure builder and finalize structure creation. All the scanned variables will get a new type. Also, cross-references will be remembered and usable anytime. 52 | 53 | Now, a few more details. 54 | 55 | ### Structure Builder (Alt + F8) 56 | 57 | The place where all the collected information about the scanned variables can be viewed and modified. Ways of collecting information: 58 | * Right Click on a variable -> Scan Variable. Recognizes fields usage within the current function. 59 | * Right Click on a variable -> Deep Scan Variable. First, recursively touches functions to make Ida recognize proper arguments (it happens only once for each function during a session). Then, it recursively applies the scanner to variables and functions, which get the structure pointer as their argument. 60 | * Right Click on a function -> Deep Scan Returned Value. If you have the singleton pattern or the constructor is called in many places, it is possible to scan all the places, where a pointer to an object was recieved or an object was created. 61 | * API [TODO] 62 | 63 | ![img][builder] 64 | 65 | Structure builder stores collected information and enables interaction: 66 | 67 | * Types with the __BOLD__ font are virtual tables. A double click opens the list with all virtual functions, which helps to visit them. The visited functions are marked with a cross and color: 68 | 69 | ![img][virtual_functions] 70 | 71 | * Types with the _ITALIC_ font have been found as passed argument. It can help in finding substructures. [TODO] 72 | * Double click on field `Name` or `Type` to edit. 73 | * Double click on `Offset` opens a window with all the places, where this field has been extracted. Click the "Ok" button to open a selected place in the pseudocode window: 74 | 75 | ![img][scanned_variables] 76 | 77 | Buttons serve the following purpose: 78 | 79 | __Finalize__ - opens a window with an editable C-like declaration and assigns new types to all scanned variables. 80 | 81 | __Disable__, __Enable__ - are used for collision resolution. 82 | 83 | __Origin__ - switches the base offset which is used to produce new fields to structure (this value will be added to every offset of a newly-scanned variable, default = 0). 84 | 85 | __Array__ - renders a selected field as an array the size of which is automatically calculated. 86 | 87 | __Pack__ - creates and substitutes a substructure for selected items (collisions for these items should be resolved). 88 | 89 | __Unpack__ - dismembers a selected structure and adds all its fields to the builder. 90 | 91 | __Remove__ - removes the information about selected fields. 92 | 93 | __Clear__ - clears all. 94 | 95 | __Recognize Shape__ - looks for appropriates structure for selected fields. 96 | 97 | __Resolve Conflicts (new)__ - attempts to disable less meaningful fields in favor of more useful ones. (`char` > `_BYTE`, `SOCKET` > `_DWORD` etc). Doesn't help to find arrays. 98 | 99 | ### Structure Cross-references (Ctrl + X) 100 | 101 | With HexRaysPyTools, every time the F5 button is pressed and code is decompiled, the information about addressing to fields is stored inside cache. It can be retrieved with the "Field Xrefs" menu. So, it is better to apply reconstructed types to as many locations as possible to have more information about the way structures are used. 102 | 103 | Note: IDA 7.4 has now an official implementation of this feature, available through Shift-X hotkey. 104 | 105 | ### Guessing Allocation 106 | 107 | **Warning!! Very raw feature.** The idea is to help find where a variable came from so as to run Deep Scan Process at the very top level and not to skip large amounts of code. 108 | 109 | ### Structures with given size 110 | 111 | Usage: 112 | 113 | 1. In Pseudocode viewer, right click on a number -> "Structures with this size". (hotkey "W") 114 | 2. Select a library to be looked for structures. 115 | 3. Select a structure. The Number will become `sizeof(Structure Name)`, and type will be imported to Local Types. 116 | 117 | ### Recognition of structures by shapes 118 | 119 | Helps find a suitable structure by the information gleaned from pseudocode after variable scanning. 120 | 121 | Usage: 122 | 123 | * _Method 1_ 124 | 1. Right click on a variable with -> Select "Recognize Shape". 125 | 2. Select Type Library. 126 | 3. Select structure. 127 | 4. Type of the variable will be changed automatically. 128 | * _Method 2_ 129 | 1. Clear Structure Builder if it's currently used. 130 | 2. Right click on the variables that are supposed to be the same -> "Scan Variable". 131 | 3. Edit types (will be implemented later), disable or remove uninteresting fields, and click the "Recognize Shape" button. 132 | 4. You can select several fields and try to recognize their shapes. If found and selected, they will be replaced with a new structure. 133 | 5. After final structure selection, types of all scanned variables will be changed automatically. 134 | 135 | ## Disassembler code manipulations 136 | 137 | ### Containing structures 138 | 139 | Helps find containing structure and makes code prettier by replacing pointers with [CONTAINING_RECORD][1] macro 140 | 141 | __Before:__ 142 | 143 | ![img][bad_structures] 144 | 145 | __After:__ 146 | 147 | ![img][good_structures] 148 | 149 | Usage: 150 | 151 | If a variable is a structure pointer and there's an access to outside of the boundaries of that structure, then: 152 | 153 | 1. Right click -> Select Containing Structure. 154 | 2. Select Type Library. 155 | 3. Select appropriate Structure and Offset. 156 | 4. If the result does not satisfy the requirements, then Right Click -> Reset Containing Structure and go back to step 1. 157 | 158 | ### Function signature manipulation 159 | 160 | 1. Right click first line -> "Remove Return" converts return type to void (or from void to _DWORD). 161 | 2. Right click on argument -> "Remove Argument" disposes of this argument. 162 | 3. Right click on convention -> "Convert to __usercall" switches to __usercall or __userpurge (same as __usercall but the callee cleans the stack). 163 | 164 | ### Recasting (Shift+R, Shift+L), Renaming (Shift+N, Ctrl+Shift+N) 165 | 166 | Every time you have two sides in an expression, where each side may be a local or global variable, argument or return value of the function signature, it is possible to right-click or press the hotkey to give both sides of the expression similar types. Below, there is the table of possible conversions: 167 | 168 | | Original | Shift+L | Shift+R 169 | | --- | --- | --- | 170 | | var = (TYPE) expr | var type -> TYPE | | 171 | | exp = (TYPE) var | | var type -> TYPE | 172 | | function(..., (TYPE) var, ...) | functions' argument -> TYPE | var type -> TYPE | 173 | | (TYPE) function(...) | | functions' return type -> TYPE | 174 | | return (TYPE) var | functions' return type -> TYPE | var type -> TYPE | 175 | | struct.field = (TYPE) var | type(field) -> TYPE | | 176 | | pstruct->field = (TYPE) var | type(field) -> TYPE | | 177 | 178 | When you have an expression like `function(..., some_good_name, ...)`, you can rename function parameter. 179 | 180 | When you have an expression like `function(..., v12, ...)`, and function has an appropriate parameter name, you can quickly apply this name to the variable. 181 | 182 | Also possible to rename `vXX = v_named` into `_v_named = v_named` and vice versa. 183 | 184 | And there's a feature for massive renaming functions using assert statements. If you find a function that looks like an assert, right-click the string argument with the function name and select "Rename as assert argument". All the functions where a call to assert statement has happened will be renamed (provided that there is no conflicts, either way, you'll see the warning in the output window) 185 | 186 | ### Name Propagation (P) 187 | 188 | This feature does the same recursive traversal over functions as the Deep Scan Variable does. But this time, all elements that have a connection with the selected one receive its name. It’s possible to rename it or use names of both local and global variables, as well as structure members. By default, the plugin propagates names only over default names like `v1`, `a2`. See **Configuration** in order to change that. 189 | 190 | ### Untangling 'if' statements 191 | 192 | * Clicking `if` manually allows to switch `then` and `else` branches 193 | * Automatically applies the following transformations: 194 | 195 | Before: 196 | 197 | ```c 198 | ... 199 | if (condition) { 200 | statement_1; 201 | statement_2; 202 | ... 203 | return another_value; 204 | } 205 | return value; 206 | ``` 207 | 208 | After: 209 | ```c 210 | ... 211 | if (opposite_condition) { 212 | return value; 213 | } 214 | statement_1; 215 | statement_2; 216 | ... 217 | return another_value; // if 'then' branch has no return, than `return value;` 218 | ``` 219 | 220 | Classes 221 | ------- 222 | 223 | Also, it can be found at _View->Open Subview->Classes_. Helps to manage classes (structures with virtual tables). 224 | 225 | ![img][classes] 226 | 227 | ##### !! Better to rename all functions before debugging, because Ida can mess up default names, and the information in virtual tables will be inconsistent. 228 | 229 | Class, virtual tables, and functions names are editable. Also a function's declaration can be edited. After editting, the altered items change font to _italic_. Right click opens the following menu options: 230 | 231 | * Expand All / Collapse All. 232 | * Refresh - clear all and rescan local types for information again. 233 | * Rollback - undo changes. 234 | * Commit - apply changes. Functions will be renamed and recasted both in virtual tables in Local Types and disassembly code. 235 | * Set First Argument type - allows selecting the first argument for a function among all classes. If right click was used on class name, then its type will be automatically applied to the virtual table at offset 0. 236 | 237 | You can also filter classes using Regexp either by class_name or by existence of specific functions. Simply input an expression in line edit for filtering by class_name or prepend it with "!" to filter by function name. 238 | 239 | Structure Graph 240 | --------------- 241 | 242 | Shows relationship between structures: 243 | 244 | ![img][structure_graph] 245 | 246 | Also: dark green node is union, light green - enum. 247 | 248 | Usage: 249 | 250 | 1. Open Local Types. 251 | 2. Select structures and right click -> "Show Graph" (Hotkey "G"). 252 | 3. Plugin creates a graph of all structures that have relationship with selected items. 253 | 4. Double clicking on a node recalculates the graph for it. 254 | 5. Every node has a hint message that shows C-like typedef. 255 | 256 | API 257 | --- 258 | 259 | **Under construction** 260 | 261 | Presentations 262 | ============= 263 | 264 | * [ZeroNights 2016](https://2016.zeronights.ru/wp-content/uploads/2016/12/zeronights_2016_Kirillov.pptx) 265 | * [Insomni'hack 2018](https://www.youtube.com/watch?v=pnPuwBtW2_4) 266 | * [Infosec in the City 2018](https://www.infosec-city.com/sg18-1-hex-rays) ([Slides](https://1drv.ms/p/s!AocQazyOQ8prgxNCpajrkwURQnPd)) 267 | 268 | [0]: https://sourceforge.net/projects/classinformer/ 269 | [1]: https://msdn.microsoft.com/en-us/library/windows/hardware/ff542043%28v=vs.85%29.aspx?f=255&MSPPError=-2147217396 270 | [structure_graph]: Img/structure_builder.JPG 271 | [bad_structures]: Img/bad.JPG 272 | [good_structures]: Img/good.JPG 273 | [builder]: Img/builder.JPG 274 | [virtual_functions]: Img/virtual_functions.JPG 275 | [scanned_variables]: Img/fields_xref.JPG 276 | [classes]: Img/classes.JPG 277 | [feature_history]: https://github.com/igogo-x86/HexRaysPyTools/wiki/History 278 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/renames.py: -------------------------------------------------------------------------------- 1 | import re 2 | import logging 3 | 4 | import idaapi 5 | import idc 6 | 7 | from . import actions 8 | import HexRaysPyTools.api as api 9 | import HexRaysPyTools.core.helper as helper 10 | import HexRaysPyTools.settings as settings 11 | 12 | 13 | logger = logging.getLogger(__name__) 14 | 15 | 16 | def _should_be_renamed(old_name, new_name): 17 | # type: (str, str) -> bool 18 | """ Checks if there's a point to rename a variable or argument """ 19 | 20 | # There's no point to rename into default name 21 | if _is_default_name(new_name): 22 | return False 23 | 24 | # Strip prefixes and check if names are the same 25 | return old_name.lstrip('_') != new_name.lstrip('_') 26 | 27 | 28 | def _is_default_name(string): 29 | return re.match(r"[av]\d+$", string) is not None or \ 30 | re.match(r"[qd]?word|field_|off_", string) is not None 31 | 32 | 33 | class RenameOther(actions.HexRaysPopupAction): 34 | description = "Take other name" 35 | hotkey = "Ctrl+N" 36 | 37 | def __init__(self): 38 | super(RenameOther, self).__init__() 39 | 40 | def check(self, hx_view): 41 | return self.__extract_rename_info(hx_view.cfunc, hx_view.item) is not None 42 | 43 | def activate(self, ctx): 44 | hx_view = idaapi.get_widget_vdui(ctx.widget) 45 | result = self.__extract_rename_info(hx_view.cfunc, hx_view.item) 46 | 47 | if result: 48 | lvar, name = result 49 | while not hx_view.rename_lvar(lvar, name, True): 50 | name = '_' + name 51 | 52 | @staticmethod 53 | def __extract_rename_info(cfunc, ctree_item): 54 | # type: (idaapi.cfunc_t, idaapi.ctree_item_t) -> (idaapi.lvar_t, str) 55 | 56 | if ctree_item.citype != idaapi.VDI_EXPR: 57 | return 58 | 59 | expression = ctree_item.it.to_specific_type 60 | if expression.op != idaapi.cot_var: 61 | return 62 | 63 | parent = cfunc.body.find_parent_of(expression).to_specific_type 64 | if parent.op != idaapi.cot_asg: 65 | return 66 | 67 | other = parent.theother(expression) 68 | if other.op != idaapi.cot_var: 69 | return 70 | 71 | this_lvar = ctree_item.get_lvar() 72 | other_lvar = cfunc.get_lvars()[other.v.idx] 73 | 74 | if _should_be_renamed(this_lvar.name, other_lvar.name): 75 | return this_lvar, other_lvar.name.lstrip('_') 76 | 77 | 78 | class RenameInside(actions.HexRaysPopupAction): 79 | description = "Rename inside argument" 80 | hotkey = "Shift+N" 81 | 82 | def __init__(self): 83 | super(RenameInside, self).__init__() 84 | 85 | def check(self, hx_view): 86 | return self.__extract_rename_info(hx_view.cfunc, hx_view.item) is not None 87 | 88 | def activate(self, ctx): 89 | hx_view = idaapi.get_widget_vdui(ctx.widget) 90 | result = self.__extract_rename_info(hx_view.cfunc, hx_view.item) 91 | 92 | if result: 93 | func_tinfo, address, arg_index, name = result 94 | helper.set_func_arg_name(func_tinfo, arg_index, name) 95 | idaapi.apply_tinfo(address, func_tinfo, idaapi.TINFO_DEFINITE) 96 | hx_view.refresh_view(True) 97 | 98 | @staticmethod 99 | def __extract_rename_info(cfunc, ctree_item): 100 | # type: (idaapi.cfunc_t, idaapi.ctree_item_t) -> (idaapi.tinfo_t, long, int, str) 101 | 102 | if ctree_item.citype != idaapi.VDI_EXPR: 103 | return False 104 | 105 | expression = ctree_item.it.to_specific_type 106 | if expression.op != idaapi.cot_var: 107 | return 108 | 109 | parent = cfunc.body.find_parent_of(expression).to_specific_type 110 | if parent.op != idaapi.cot_call or parent.x.obj_ea == idaapi.BADADDR: 111 | return 112 | 113 | lvar = ctree_item.get_lvar() 114 | arg_index, _ = helper.get_func_argument_info(parent, expression) 115 | func_tinfo = parent.x.type.get_pointed_object() 116 | arg_name = helper.get_func_arg_name(func_tinfo, arg_index) 117 | if _should_be_renamed(arg_name, lvar.name): 118 | return func_tinfo, parent.x.obj_ea, arg_index, lvar.name.lstrip('_') 119 | 120 | 121 | class RenameOutside(actions.HexRaysPopupAction): 122 | description = "Take argument name" 123 | hotkey = "Ctrl+Shift+N" 124 | 125 | def __init__(self): 126 | super(RenameOutside, self).__init__() 127 | 128 | def check(self, hx_view): 129 | return self.__extract_rename_info(hx_view.cfunc, hx_view.item) is not None 130 | 131 | def activate(self, ctx): 132 | hx_view = idaapi.get_widget_vdui(ctx.widget) 133 | result = self.__extract_rename_info(hx_view.cfunc, hx_view.item) 134 | 135 | if result: 136 | lvar, name = result 137 | while not hx_view.rename_lvar(lvar, name, True): 138 | name = '_' + name 139 | 140 | @staticmethod 141 | def __extract_rename_info(cfunc, ctree_item): 142 | # type: (idaapi.cfunc_t, idaapi.ctree_item_t) -> (idaapi.lvar_t, str) 143 | 144 | if ctree_item.citype != idaapi.VDI_EXPR: 145 | return 146 | 147 | expression = ctree_item.it.to_specific_type 148 | if expression.op != idaapi.cot_var: 149 | return 150 | 151 | parent = cfunc.body.find_parent_of(expression).to_specific_type 152 | if parent.op != idaapi.cot_call or parent.x.obj_ea == idaapi.BADADDR: 153 | return 154 | 155 | lvar = ctree_item.get_lvar() 156 | arg_index, _ = helper.get_func_argument_info(parent, expression) 157 | func_tinfo = parent.x.type.get_pointed_object() 158 | arg_name = helper.get_func_arg_name(func_tinfo, arg_index) 159 | if arg_name and _should_be_renamed(lvar.name, arg_name): 160 | return lvar, arg_name.lstrip("_") 161 | 162 | 163 | class _RenameUsingAssertVisitor(idaapi.ctree_parentee_t): 164 | 165 | def __init__(self, cfunc, func_addr, arg_idx): 166 | idaapi.ctree_parentee_t.__init__(self) 167 | self.__cfunc = cfunc 168 | self.__func_addr = func_addr 169 | self.__arg_idx = arg_idx 170 | self.__possible_names = set() 171 | 172 | def visit_expr(self, expr): 173 | if expr.op == idaapi.cot_call and expr.x.op == idaapi.cot_obj and expr.x.obj_ea == self.__func_addr: 174 | arg_expr = expr.a[self.__arg_idx] 175 | if arg_expr.op != idaapi.cot_obj: 176 | cexpr_ea = helper.find_asm_address(expr, self.parents) 177 | logger.error("Argument is a not string at {}".format(helper.to_hex(cexpr_ea))) 178 | return 1 179 | self.__add_func_name(arg_expr) 180 | return 0 181 | 182 | def process(self): 183 | self.apply_to(self.__cfunc.body, None) 184 | if len(self.__possible_names) == 1: 185 | # Only one potential name was found, rename function using it 186 | new_name = self.__possible_names.pop() 187 | logging.info("Renaming function at {} to `{}`".format(helper.to_hex(self.__cfunc.entry_ea), new_name)) 188 | idc.set_name(self.__cfunc.entry_ea, new_name) 189 | elif len(self.__possible_names) > 1: 190 | logger.error("Function at {} has more than one candidate for renaming: {}".format( 191 | helper.to_hex(self.__cfunc.entry_ea), ", ".join(self.__possible_names))) 192 | 193 | def __add_func_name(self, arg_expr): 194 | new_name = idc.get_strlit_contents(arg_expr.obj_ea) 195 | if type(new_name) is not str: 196 | # convert bytes to str (python 3) 197 | new_name = new_name.decode('ascii') 198 | if not idaapi.is_valid_typename(new_name): 199 | logger.warn("Argument has a weird name `{}` at {}".format( 200 | new_name, helper.to_hex(helper.find_asm_address(arg_expr, self.parents)))) 201 | return 202 | 203 | self.__possible_names.add(new_name) 204 | 205 | 206 | class RenameUsingAssert(actions.HexRaysPopupAction): 207 | description = "Rename as assert argument" 208 | hotkey = None 209 | 210 | def __init__(self): 211 | super(RenameUsingAssert, self).__init__() 212 | 213 | @staticmethod 214 | def __can_be_part_of_assert(cfunc, ctree_item): 215 | # type: (idaapi.cfunc_t, idaapi.ctree_item_t) -> bool 216 | """ 217 | Returns true if expression we clicked is an argument passed to a function 218 | and this argument is a string that can be a valid function name 219 | """ 220 | 221 | if ctree_item.citype != idaapi.VDI_EXPR: 222 | return False 223 | 224 | expression = ctree_item.it.to_specific_type 225 | if expression.op != idaapi.cot_obj: 226 | return False 227 | 228 | parent = cfunc.body.find_parent_of(expression).to_specific_type 229 | if parent.op != idaapi.cot_call or parent.x.op != idaapi.cot_obj: 230 | return False 231 | 232 | obj_ea = expression.obj_ea 233 | if not helper.is_code_ea(obj_ea) and idc.get_str_type(obj_ea) == idc.STRTYPE_C: 234 | str_potential_name = idc.get_strlit_contents(obj_ea) 235 | if type(str_potential_name) is not str: 236 | # convert bytes to str (python 3) 237 | str_potential_name = str_potential_name.decode('ascii') 238 | return idaapi.is_valid_typename(str_potential_name) 239 | return False 240 | 241 | def check(self, hx_view): 242 | return self.__can_be_part_of_assert(hx_view.cfunc, hx_view.item) 243 | 244 | def activate(self, ctx): 245 | hx_view = idaapi.get_widget_vdui(ctx.widget) 246 | if not self.__can_be_part_of_assert(hx_view.cfunc, hx_view.item): 247 | return 248 | 249 | # So we clicked on function an func argument that is a string. Now we extract 250 | # argument index and address of assert function 251 | expr_arg = hx_view.item.it.to_specific_type 252 | expr_call = hx_view.cfunc.body.find_parent_of(expr_arg).to_specific_type 253 | arg_idx, _ = helper.get_func_argument_info(expr_call, expr_arg) 254 | assert_func_ea = expr_call.x.obj_ea 255 | 256 | # Iterate through all places where assert function and rename using helper class 257 | all_callers = helper.get_funcs_calling_address(assert_func_ea) 258 | for caller_ea in all_callers: 259 | cfunc = helper.decompile_function(caller_ea) 260 | if cfunc: 261 | _RenameUsingAssertVisitor(cfunc, assert_func_ea, arg_idx).process() 262 | 263 | hx_view.refresh_view(True) 264 | 265 | 266 | class _NamePropagator(api.RecursiveObjectDownwardsVisitor): 267 | def __init__(self, hx_view, cfunc, obj): 268 | super(_NamePropagator, self).__init__(cfunc, obj, skip_until_object=True) 269 | self.__hx_view = hx_view 270 | self.__propagated_name = obj.name 271 | 272 | def _start_iteration(self): 273 | self.__hx_view.switch_to(self._cfunc, False) 274 | 275 | def _manipulate(self, cexpr, obj): 276 | if self.crippled: 277 | logger.debug("Skipping crippled function at {}".format(helper.to_hex(self._cfunc.entry_ea))) 278 | return 279 | 280 | if obj.id == api.SO_GLOBAL_OBJECT: 281 | old_name = idaapi.get_short_name(cexpr.obj_ea) 282 | if settings.PROPAGATE_THROUGH_ALL_NAMES or _is_default_name(old_name): 283 | new_name = self.__rename_with_prefix( 284 | lambda x: idaapi.set_name(cexpr.obj_ea, x), 285 | self.__propagated_name) 286 | logger.debug("Renamed global variable from {} to {}".format(old_name, new_name)) 287 | elif obj.id == api.SO_LOCAL_VARIABLE: 288 | lvar = self._cfunc.get_lvars()[cexpr.v.idx] 289 | old_name = lvar.name 290 | if settings.PROPAGATE_THROUGH_ALL_NAMES or _is_default_name(old_name): 291 | new_name = self.__rename_with_prefix( 292 | lambda x: self.__hx_view.rename_lvar(lvar, x, True), 293 | self.__propagated_name) 294 | logger.debug("Renamed local variable from {} to {}".format(old_name, new_name)) 295 | elif obj.id in (api.SO_STRUCT_POINTER, api.SO_STRUCT_REFERENCE): 296 | struct_tinfo = cexpr.x.type 297 | offset = cexpr.m 298 | struct_tinfo.remove_ptr_or_array() 299 | old_name = helper.get_member_name(struct_tinfo, offset) 300 | if settings.PROPAGATE_THROUGH_ALL_NAMES or _is_default_name(old_name): 301 | new_name = self.__rename_with_prefix( 302 | lambda x: helper.change_member_name(struct_tinfo.dstr(), offset, x), 303 | self.__propagated_name) 304 | logger.debug("Renamed struct member from {} to {}".format(old_name, new_name)) 305 | 306 | def _finish(self): 307 | self.__hx_view.switch_to(self._cfunc, True) 308 | 309 | @staticmethod 310 | def __rename_with_prefix(rename_func, name): 311 | while not rename_func(name): 312 | name = "_" + name 313 | return name 314 | 315 | 316 | class PropagateName(actions.HexRaysPopupAction): 317 | description = "Propagate name" 318 | hotkey = "P" 319 | 320 | def __init__(self): 321 | super(PropagateName, self).__init__() 322 | 323 | @staticmethod 324 | def __extract_propagate_info(cfunc, ctree_item): 325 | if ctree_item.citype != idaapi.VDI_EXPR: 326 | return 327 | 328 | obj = api.ScanObject.create(cfunc, ctree_item) 329 | if obj and not _is_default_name(obj.name): 330 | return obj 331 | 332 | def check(self, hx_view): 333 | return self.__extract_propagate_info(hx_view.cfunc, hx_view.item) is not None 334 | 335 | def activate(self, ctx): 336 | hx_view = idaapi.get_widget_vdui(ctx.widget) 337 | obj = self.__extract_propagate_info(hx_view.cfunc, hx_view.item) 338 | if obj: 339 | cfunc = hx_view.cfunc 340 | visitor = _NamePropagator(hx_view, cfunc, obj) 341 | visitor.process() 342 | hx_view.refresh_view(True) 343 | 344 | 345 | actions.action_manager.register(RenameOther()) 346 | actions.action_manager.register(RenameInside()) 347 | actions.action_manager.register(RenameOutside()) 348 | actions.action_manager.register(RenameUsingAssert()) 349 | actions.action_manager.register(PropagateName()) 350 | -------------------------------------------------------------------------------- /HexRaysPyTools/core/helper.py: -------------------------------------------------------------------------------- 1 | import collections 2 | import logging 3 | 4 | import idaapi 5 | import idc 6 | 7 | import HexRaysPyTools.core.cache as cache 8 | import HexRaysPyTools.core.const as const 9 | import HexRaysPyTools.settings as settings 10 | import HexRaysPyTools.forms as forms 11 | 12 | 13 | logger = logging.getLogger(__name__) 14 | 15 | 16 | def is_imported_ea(ea): 17 | if idc.get_segm_name(ea) == ".plt": 18 | return True 19 | return ea + idaapi.get_imagebase() in cache.imported_ea 20 | 21 | 22 | def is_code_ea(ea): 23 | if idaapi.cvar.inf.procname == "ARM": 24 | # In case of ARM code in THUMB mode we sometimes get pointers with thumb bit set 25 | flags = idaapi.get_full_flags(ea & -2) # flags_t 26 | else: 27 | flags = idaapi.get_full_flags(ea) 28 | return idaapi.is_code(flags) 29 | 30 | 31 | def is_rw_ea(ea): 32 | seg = idaapi.getseg(ea) 33 | return seg.perm & idaapi.SEGPERM_WRITE and seg.perm & idaapi.SEGPERM_READ 34 | 35 | 36 | def get_ptr(ea): 37 | """ Reads ptr at specified address. """ 38 | if const.EA64: 39 | return idaapi.get_64bit(ea) 40 | ptr = idaapi.get_32bit(ea) 41 | if idaapi.cvar.inf.procname == "ARM": 42 | ptr &= -2 # Clear thumb bit 43 | return ptr 44 | 45 | 46 | def get_ordinal(tinfo): 47 | """ Returns non-zero ordinal of tinfo if it exist in database """ 48 | ordinal = tinfo.get_ordinal() 49 | if ordinal == 0: 50 | t = idaapi.tinfo_t() 51 | struct_name = tinfo.dstr().split()[-1] # Get rid of `struct` prefix or something else 52 | t.get_named_type(idaapi.cvar.idati, struct_name) 53 | ordinal = t.get_ordinal() 54 | return ordinal 55 | 56 | 57 | def get_virtual_func_addresses(name, tinfo=None, offset=None): 58 | """ 59 | Returns set of possible addresses of virtual function by its name. 60 | If there're symbols in binary and name is the name of an overloaded function, then returns list of all address of 61 | this overloaded function. 62 | TODO: After implementing inheritance return set of methods of all child classes 63 | 64 | :param name: method name, can be mangled 65 | :param tinfo: class tinfo to which this method belong 66 | :param offset: virtual table offset 67 | :return: list of possible addresses 68 | """ 69 | 70 | address = idc.get_name_ea_simple(name) 71 | 72 | if address != idaapi.BADADDR: 73 | return [address] 74 | 75 | raw_addresses = cache.demangled_names.get(name) 76 | if raw_addresses: 77 | addresses = [ea + idaapi.get_imagebase() for ea in raw_addresses] 78 | return addresses 79 | 80 | if tinfo is None or offset is None: 81 | return [] 82 | 83 | offset *= 8 84 | udt_member = idaapi.udt_member_t() 85 | while tinfo.is_struct(): 86 | address = cache.demangled_names.get(tinfo.dstr() + '::' + name, idaapi.BADADDR) 87 | if address != idaapi.BADADDR: 88 | return [address + idaapi.get_imagebase()] 89 | udt_member.offset = offset 90 | tinfo.find_udt_member(udt_member, idaapi.STRMEM_OFFSET) 91 | tinfo = udt_member.type 92 | offset = offset - udt_member.offset 93 | 94 | 95 | def choose_virtual_func_address(name, tinfo=None, offset=None): 96 | addresses = get_virtual_func_addresses(name, tinfo, offset) 97 | if not addresses: 98 | return 99 | 100 | if len(addresses) == 1: 101 | return addresses[0] 102 | 103 | chooser = forms.MyChoose( 104 | [[to_hex(ea), idc.demangle_name(idc.get_name(ea), idc.INF_LONG_DN)] for ea in addresses], 105 | "Select Function", 106 | [["Address", 10], ["Full name", 50]] 107 | ) 108 | idx = chooser.Show(modal=True) 109 | if idx != -1: 110 | return addresses[idx] 111 | 112 | 113 | def get_func_argument_info(function, expression): 114 | """ 115 | Function is cexpr with opname == 'cot_call', expression is any son. Returns index of argument and it's type 116 | 117 | :param function: idaapi.cexpr_t 118 | :param expression: idaapi.cexpr_t 119 | :return: (int, idaapi.tinfo_t) 120 | """ 121 | for idx, argument in enumerate(function.a): 122 | if expression == argument.cexpr: 123 | func_tinfo = function.x.type 124 | if idx < func_tinfo.get_nargs(): 125 | return idx, func_tinfo.get_nth_arg(idx) 126 | return idx, None 127 | print("[ERROR] Wrong usage of 'Helper.get_func_argument_info()'") 128 | 129 | 130 | def set_func_argument(func_tinfo, index, arg_tinfo): 131 | func_data = idaapi.func_type_data_t() 132 | func_tinfo.get_func_details(func_data) 133 | func_data[index].type = arg_tinfo 134 | func_tinfo.create_func(func_data) 135 | 136 | 137 | def get_func_arg_name(func_tinfo, arg_idx): 138 | # type: (idaapi.tinfo_t, int) -> str 139 | 140 | func_data = idaapi.func_type_data_t() 141 | func_tinfo.get_func_details(func_data) 142 | if arg_idx < func_tinfo.get_nargs(): 143 | return func_data[arg_idx].name 144 | 145 | 146 | def set_func_arg_name(func_tinfo, arg_idx, name): 147 | # type: (idaapi.tinfo_t, int, str) -> None 148 | 149 | func_data = idaapi.func_type_data_t() 150 | func_tinfo.get_func_details(func_data) 151 | func_data[arg_idx].name = name 152 | func_tinfo.create_func(func_data) 153 | 154 | 155 | def set_funcptr_argument(funcptr_tinfo, index, arg_tinfo): 156 | func_tinfo = funcptr_tinfo.get_pointed_object() 157 | set_func_argument(func_tinfo, index, arg_tinfo) 158 | funcptr_tinfo.create_ptr(func_tinfo) 159 | 160 | 161 | def set_func_return(func_tinfo, return_tinfo): 162 | func_data = idaapi.func_type_data_t() 163 | func_tinfo.get_func_details(func_data) 164 | func_data.rettype = return_tinfo 165 | func_tinfo.create_func(func_data) 166 | 167 | 168 | def get_nice_pointed_object(tinfo): 169 | """ 170 | Returns nice pointer name (if exist) or None. 171 | For example if tinfo is PKSPIN_LOCK which is typedef of unsigned int *, then if in local types exist KSPIN_LOCK with 172 | type unsigned int, this function returns KSPIN_LOCK 173 | """ 174 | try: 175 | name = tinfo.dstr() 176 | if name[0] == 'P': 177 | pointed_tinfo = idaapi.tinfo_t() 178 | if pointed_tinfo.get_named_type(idaapi.cvar.idati, name[1:]): 179 | if tinfo.get_pointed_object().equals_to(pointed_tinfo): 180 | return pointed_tinfo 181 | except TypeError: 182 | pass 183 | 184 | 185 | def get_fields_at_offset(tinfo, offset): 186 | """ 187 | Given tinfo and offset of the structure or union, returns list of all tinfo at that offset. 188 | This function helps to find appropriate structures by type of the offset 189 | """ 190 | result = [] 191 | if offset == 0: 192 | result.append(tinfo) 193 | udt_data = idaapi.udt_type_data_t() 194 | tinfo.get_udt_details(udt_data) 195 | udt_member = idaapi.udt_member_t() 196 | udt_member.offset = offset * 8 197 | idx = tinfo.find_udt_member(udt_member, idaapi.STRMEM_OFFSET) 198 | if idx != -1: 199 | while idx < tinfo.get_udt_nmembers() and udt_data[idx].offset <= offset * 8: 200 | udt_member = udt_data[idx] 201 | if udt_member.offset == offset * 8: 202 | if udt_member.type.is_ptr(): 203 | result.append(idaapi.get_unk_type(const.EA_SIZE)) 204 | result.append(udt_member.type) 205 | result.append(idaapi.dummy_ptrtype(const.EA_SIZE, False)) 206 | elif not udt_member.type.is_udt(): 207 | result.append(udt_member.type) 208 | if udt_member.type.is_array(): 209 | if (offset - udt_member.offset // 8) % udt_member.type.get_array_element().get_size() == 0: 210 | result.append(udt_member.type.get_array_element()) 211 | elif udt_member.type.is_udt(): 212 | result.extend(get_fields_at_offset(udt_member.type, offset - udt_member.offset // 8)) 213 | idx += 1 214 | return result 215 | 216 | 217 | def is_legal_type(tinfo): 218 | tinfo.clr_const() 219 | if tinfo.is_ptr() and tinfo.get_pointed_object().is_forward_decl(): 220 | return tinfo.get_pointed_object().get_size() == idaapi.BADSIZE 221 | return settings.SCAN_ANY_TYPE or bool([x for x in const.LEGAL_TYPES if x.equals_to(tinfo)]) 222 | 223 | 224 | def search_duplicate_fields(udt_data): 225 | # Returns list of lists with duplicate fields 226 | 227 | default_dict = collections.defaultdict(list) 228 | for idx, udt_member in enumerate(udt_data): 229 | default_dict[udt_member.name].append(idx) 230 | return [indices for indices in list(default_dict.values()) if len(indices) > 1] 231 | 232 | 233 | def get_member_name(tinfo, offset): 234 | udt_member = idaapi.udt_member_t() 235 | udt_member.offset = offset * 8 236 | tinfo.find_udt_member(udt_member, idaapi.STRMEM_OFFSET) 237 | return udt_member.name 238 | 239 | 240 | def change_member_name(struct_name, offset, name): 241 | return idc.set_member_name(idc.get_struc_id(struct_name), offset, name) 242 | 243 | 244 | def import_structure(name, tinfo): 245 | cdecl_typedef = idaapi.print_tinfo(None, 4, 5, idaapi.PRTYPE_MULTI | idaapi.PRTYPE_TYPE | idaapi.PRTYPE_SEMI, 246 | tinfo, name, None) 247 | if idc.parse_decl(cdecl_typedef, idaapi.PT_TYP) is None: 248 | return 0 249 | 250 | previous_ordinal = idaapi.get_type_ordinal(idaapi.cvar.idati, name) 251 | if previous_ordinal: 252 | idaapi.del_numbered_type(idaapi.cvar.idati, previous_ordinal) 253 | ordinal = idaapi.idc_set_local_type(previous_ordinal, cdecl_typedef, idaapi.PT_TYP) 254 | else: 255 | ordinal = idaapi.idc_set_local_type(-1, cdecl_typedef, idaapi.PT_TYP) 256 | return ordinal 257 | 258 | 259 | def get_funcs_calling_address(ea): 260 | """ Returns all addresses of functions which make call to a function at `ea`""" 261 | xref_ea = idaapi.get_first_cref_to(ea) 262 | xrefs = set() 263 | while xref_ea != idaapi.BADADDR: 264 | xref_func_ea = idc.get_func_attr(xref_ea, idc.FUNCATTR_START) 265 | if xref_func_ea != idaapi.BADADDR: 266 | xrefs.add(xref_func_ea) 267 | else: 268 | print("[Warning] Function not found at 0x{0:08X}".format(xref_ea)) 269 | xref_ea = idaapi.get_next_cref_to(ea, xref_ea) 270 | return xrefs 271 | 272 | 273 | class FunctionTouchVisitor(idaapi.ctree_parentee_t): 274 | def __init__(self, cfunc): 275 | super(FunctionTouchVisitor, self).__init__() 276 | self.functions = set() 277 | self.cfunc = cfunc 278 | 279 | def visit_expr(self, expression): 280 | if expression.op == idaapi.cot_call: 281 | self.functions.add(expression.x.obj_ea) 282 | return 0 283 | 284 | def touch_all(self): 285 | diff = self.functions.difference(cache.touched_functions) 286 | for address in diff: 287 | if is_imported_ea(address): 288 | continue 289 | try: 290 | cfunc = idaapi.decompile(address) 291 | if cfunc: 292 | FunctionTouchVisitor(cfunc).process() 293 | except idaapi.DecompilationFailure: 294 | logger.warn("IDA failed to decompile function at {}".format(to_hex(address))) 295 | cache.touched_functions.add(address) 296 | idaapi.decompile(self.cfunc.entry_ea) 297 | 298 | def process(self): 299 | if self.cfunc.entry_ea not in cache.touched_functions: 300 | cache.touched_functions.add(self.cfunc.entry_ea) 301 | self.apply_to(self.cfunc.body, None) 302 | self.touch_all() 303 | return True 304 | return False 305 | 306 | 307 | def to_hex(ea): 308 | """ Formats address so it could be double clicked at console """ 309 | if const.EA64: 310 | return "0x{:016X}".format(ea) 311 | return "0x{:08X}".format(ea) 312 | 313 | 314 | def to_nice_str(ea): 315 | """ Shows address as function name + offset """ 316 | func_start_ea = idc.get_func_attr(ea, idc.FUNCATTR_START) 317 | func_name = idc.get_name(func_start_ea) 318 | offset = ea - func_start_ea 319 | return "{}+0x{:X}".format(func_name, offset) 320 | 321 | 322 | def save_long_str_to_idb(array_name, value): 323 | """ Overwrites old array completely in process """ 324 | id = idc.get_array_id(array_name) 325 | if id != -1: 326 | idc.delete_array(id) 327 | id = idc.create_array(array_name) 328 | r = [] 329 | for idx in range(len(value) // 1024 + 1): 330 | s = value[idx * 1024: (idx + 1) * 1024] 331 | r.append(s) 332 | idc.set_array_string(id, idx, s) 333 | 334 | 335 | def load_long_str_from_idb(array_name): 336 | id = idc.get_array_id(array_name) 337 | if id == -1: 338 | return None 339 | max_idx = idc.get_last_index(idc.AR_STR, id) 340 | result = [idc.get_array_element(idc.AR_STR, id, idx) for idx in range(max_idx + 1)] 341 | return b"".join(result).decode("utf-8") 342 | 343 | def create_padding_udt_member(offset, size): 344 | # type: (long, long) -> idaapi.udt_member_t 345 | """ Creates internal IDA structure with name gap_XXX and appropriate size and offset """ 346 | 347 | udt_member = idaapi.udt_member_t() 348 | udt_member.name = "gap_{0:X}".format(offset) 349 | udt_member.offset = offset 350 | udt_member.size = size 351 | 352 | if size == 1: 353 | udt_member.type = const.BYTE_TINFO 354 | else: 355 | array_data = idaapi.array_type_data_t() 356 | array_data.base = 0 357 | array_data.elem_type = const.BYTE_TINFO 358 | array_data.nelems = size 359 | tmp_tinfo = idaapi.tinfo_t() 360 | tmp_tinfo.create_array(array_data) 361 | udt_member.type = tmp_tinfo 362 | return udt_member 363 | 364 | 365 | def decompile_function(address): 366 | try: 367 | cfunc = idaapi.decompile(address) 368 | if cfunc: 369 | return cfunc 370 | except idaapi.DecompilationFailure: 371 | pass 372 | logger.warn("IDA failed to decompile function at 0x{address:08X}".format(address=address)) 373 | 374 | 375 | def find_asm_address(cexpr, parents): 376 | """ Returns most close virtual address corresponding to cexpr """ 377 | 378 | ea = cexpr.ea 379 | if ea != idaapi.BADADDR: 380 | return ea 381 | 382 | for p in reversed(parents): 383 | if p.ea != idaapi.BADADDR: 384 | return p.ea 385 | 386 | 387 | def my_cexpr_t(*args, **kwargs): 388 | """ Replacement of bugged cexpr_t() function """ 389 | 390 | if len(args) == 0: 391 | return idaapi.cexpr_t() 392 | 393 | if len(args) != 1: 394 | raise NotImplementedError 395 | 396 | cexpr = idaapi.cexpr_t() 397 | cexpr.thisown = False 398 | if type(args[0]) == idaapi.cexpr_t: 399 | cexpr.assign(args[0]) 400 | else: 401 | op = args[0] 402 | cexpr._set_op(op) 403 | 404 | if 'x' in kwargs: 405 | cexpr._set_x(kwargs['x']) 406 | if 'y' in kwargs: 407 | cexpr._set_y(kwargs['y']) 408 | if 'z' in kwargs: 409 | cexpr._set_z(kwargs['z']) 410 | return cexpr 411 | -------------------------------------------------------------------------------- /HexRaysPyTools/core/variable_scanner.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import idaapi 3 | import idc 4 | from . import const 5 | from . import helper 6 | from . import temporary_structure 7 | import HexRaysPyTools.api as api 8 | 9 | logger = logging.getLogger(__name__) 10 | 11 | # If disabled then recursion will be triggered only for variable passed as first argument to function 12 | SETTING_SCAN_ALL_ARGUMENTS = True 13 | 14 | # Global set which is populated when deep scanning and cleared after completion 15 | scanned_functions = set() 16 | debug_scan_tree = [] 17 | 18 | 19 | class ScannedObject(object): 20 | def __init__(self, name, expression_address, origin, applicable=True): 21 | """ 22 | :param name: Object name 23 | :param expression_address: ea_t 24 | :param origin: which offset had structure at scan moment 25 | :param applicable: whether to apply type after creating structure 26 | """ 27 | self.name = name 28 | self.expression_address = expression_address 29 | self.func_ea = idc.get_func_attr(self.expression_address, idc.FUNCATTR_START) 30 | self.origin = origin 31 | self._applicable = applicable 32 | 33 | @property 34 | def function_name(self): 35 | return idaapi.get_short_name(self.func_ea) 36 | 37 | def apply_type(self, tinfo): 38 | """ Finally apply Class'es tinfo to this variable """ 39 | raise NotImplementedError 40 | 41 | @staticmethod 42 | def create(obj, expression_address, origin, applicable): 43 | """ Creates suitable instance of ScannedObject depending on obj """ 44 | if obj.id == api.SO_GLOBAL_OBJECT: 45 | return ScannedGlobalObject(obj.ea, obj.name, expression_address, origin, applicable) 46 | elif obj.id == api.SO_LOCAL_VARIABLE: 47 | return ScannedVariableObject(obj.lvar, obj.name, expression_address, origin, applicable) 48 | elif obj.id in (api.SO_STRUCT_REFERENCE, api.SO_STRUCT_POINTER): 49 | return ScannedStructureMemberObject(obj.struct_name, obj.offset, expression_address, origin, applicable) 50 | else: 51 | raise AssertionError 52 | 53 | def to_list(self): 54 | """ Creates list that is acceptable to MyChoose2 viewer """ 55 | return [ 56 | "0x{0:04X}".format(self.origin), 57 | self.function_name, 58 | self.name, 59 | helper.to_hex(self.expression_address) 60 | ] 61 | 62 | def __eq__(self, other): 63 | return self.func_ea == other.func_ea and self.name == other.name and \ 64 | self.expression_address == other.expression_address 65 | 66 | def __hash__(self): 67 | return hash((self.func_ea, self.name, self.expression_address)) 68 | 69 | def __repr__(self): 70 | return "{} : {}".format(self.name, helper.to_hex(self.expression_address)) 71 | 72 | 73 | class ScannedGlobalObject(ScannedObject): 74 | def __init__(self, obj_ea, name, expression_address, origin, applicable=True): 75 | super(ScannedGlobalObject, self).__init__(name, expression_address, origin, applicable) 76 | self.__obj_ea = obj_ea 77 | 78 | def apply_type(self, tinfo): 79 | if self._applicable: 80 | idaapi.set_tinfo(self.__obj_ea, tinfo) 81 | 82 | 83 | class ScannedVariableObject(ScannedObject): 84 | def __init__(self, lvar, name, expression_address, origin, applicable=True): 85 | super(ScannedVariableObject, self).__init__(name, expression_address, origin, applicable) 86 | self.__lvar = idaapi.lvar_locator_t(lvar.location, lvar.defea) 87 | 88 | def apply_type(self, tinfo): 89 | if not self._applicable: 90 | return 91 | 92 | hx_view = idaapi.open_pseudocode(self.func_ea, -1) 93 | if hx_view: 94 | logger.debug("Applying tinfo to variable {0} in function {1}".format(self.name, self.function_name)) 95 | # Finding lvar of new window that have the same name that saved one and applying tinfo_t 96 | lvar = [x for x in hx_view.cfunc.get_lvars() if x == self.__lvar] 97 | if lvar: 98 | logger.debug("Successful") 99 | hx_view.set_lvar_type(lvar[0], tinfo) 100 | else: 101 | logger.warn("Failed to find previously scanned local variable {} from {}".format( 102 | self.name, helper.to_hex(self.expression_address))) 103 | 104 | 105 | class ScannedStructureMemberObject(ScannedObject): 106 | def __init__(self, struct_name, struct_offset, name, expression_address, origin, applicable=True): 107 | super(ScannedStructureMemberObject, self).__init__(name, expression_address, origin, applicable) 108 | self.__struct_name = struct_name 109 | self.__struct_offset = struct_offset 110 | 111 | def apply_type(self, tinfo): 112 | if self._applicable: 113 | logger.warn("Changing type of structure field is not yet implemented. Address - {}".format( 114 | helper.to_hex(self.expression_address))) 115 | 116 | 117 | class SearchVisitor(api.ObjectVisitor): 118 | def __init__(self, cfunc, origin, obj, temporary_structure): 119 | super(SearchVisitor, self).__init__(cfunc, obj, None, True) 120 | self.__origin = origin 121 | self.__temporary_structure = temporary_structure 122 | 123 | def _manipulate(self, cexpr, obj): 124 | super(SearchVisitor, self)._manipulate(cexpr, obj) 125 | 126 | if obj.tinfo and not helper.is_legal_type(obj.tinfo): 127 | cexpr_ea = helper.find_asm_address(cexpr, self.parents) 128 | logger.warn("Variable obj.name has weird type at {}".format(helper.to_hex(cexpr_ea))) 129 | return 130 | if cexpr.type.is_ptr(): 131 | member = self.__extract_member_from_pointer(cexpr, obj) 132 | else: 133 | member = self.__extract_member_from_xword(cexpr, obj) 134 | if member: 135 | logger.debug("\tCreating member with type {}, {}, offset - {}".format( 136 | member.type_name, member.scanned_variables, member.offset)) 137 | self.__temporary_structure.add_row(member) 138 | 139 | def _get_member(self, offset, cexpr, obj, tinfo=None, obj_ea=None): 140 | cexpr_ea = helper.find_asm_address(cexpr, self.parents) 141 | if offset < 0: 142 | logger.error("Considered to be impossible: offset - {}, obj - {}".format( 143 | offset, helper.to_hex(cexpr_ea))) 144 | raise AssertionError 145 | 146 | applicable = not self.crippled 147 | scan_obj = ScannedObject.create(obj, cexpr_ea, self.__origin, applicable) 148 | if obj_ea: 149 | if temporary_structure.VirtualTable.check_address(obj_ea): 150 | return temporary_structure.VirtualTable(offset, obj_ea, scan_obj, self.__origin) 151 | if helper.is_code_ea(obj_ea): 152 | cfunc = helper.decompile_function(obj_ea) 153 | if cfunc: 154 | tinfo = cfunc.type 155 | tinfo.create_ptr(tinfo) 156 | else: 157 | tinfo = const.DUMMY_FUNC 158 | return temporary_structure.Member(offset, tinfo, scan_obj, self.__origin) 159 | # logger.warn("Want to see this ea - {},".format(Helper.to_hex(cexpr_ea))) 160 | 161 | if not tinfo or tinfo.equals_to(const.VOID_TINFO) or tinfo.equals_to(const.CONST_VOID_TINFO): 162 | return temporary_structure.VoidMember(offset, scan_obj, self.__origin) 163 | 164 | if tinfo.equals_to(const.CONST_PCHAR_TINFO): 165 | tinfo = const.PCHAR_TINFO 166 | elif tinfo.equals_to(const.CONST_PVOID_TINFO): 167 | tinfo = const.PVOID_TINFO 168 | else: 169 | tinfo.clr_const() 170 | return temporary_structure.Member(offset, tinfo, scan_obj, self.__origin) 171 | 172 | def _parse_call(self, call_cexpr, arg_cexpr, offset): 173 | _, tinfo = helper.get_func_argument_info(call_cexpr, arg_cexpr) 174 | if tinfo: 175 | return self.__deref_tinfo(tinfo) 176 | # TODO: Find example with UTF-16 strings 177 | return const.CHAR_TINFO 178 | 179 | def _parse_left_assignee(self, cexpr, offset): 180 | pass 181 | 182 | def __extract_member_from_pointer(self, cexpr, obj): 183 | parents_type = [idaapi.get_ctype_name(x.cexpr.op) for x in list(self.parents)[:0:-1]] 184 | parents = [x.cexpr for x in list(self.parents)[:0:-1]] 185 | 186 | logger.debug("Parsing expression {}. Parents - {}".format(obj.name, parents_type)) 187 | 188 | # Extracting offset and removing expression parents making this offset 189 | if parents_type[0] in ('idx', 'add'): 190 | # `obj[idx]' or `(TYPE *) + x' 191 | if parents[0].y.op != idaapi.cot_num: 192 | # There's no way to handle with dynamic offset 193 | return 194 | offset = parents[0].y.numval() * cexpr.type.get_ptrarr_objsize() 195 | cexpr = self.parent_expr() 196 | if parents_type[0] == 'add': 197 | del parents_type[0] 198 | del parents[0] 199 | elif parents_type[0:2] == ['cast', 'add']: 200 | # (TYPE *)obj + offset or (TYPE)obj + offset 201 | if parents[1].y.op != idaapi.cot_num: 202 | return 203 | if parents[0].type.is_ptr(): 204 | size = parents[0].type.get_ptrarr_objsize() 205 | else: 206 | size = 1 207 | offset = parents[1].theother(parents[0]).numval() * size 208 | cexpr = parents[1] 209 | del parents_type[0:2] 210 | del parents[0:2] 211 | else: 212 | offset = 0 213 | 214 | return self.__extract_member(cexpr, obj, offset, parents, parents_type) 215 | 216 | def __extract_member_from_xword(self, cexpr, obj): 217 | parents_type = [idaapi.get_ctype_name(x.cexpr.op) for x in list(self.parents)[:0:-1]] 218 | parents = [x.cexpr for x in list(self.parents)[:0:-1]] 219 | 220 | logger.debug("Parsing expression {}. Parents - {}".format(obj.name, parents_type)) 221 | 222 | if parents_type[0] == 'add': 223 | if parents[0].theother(cexpr).op != idaapi.cot_num: 224 | return 225 | offset = parents[0].theother(cexpr).numval() 226 | cexpr = self.parent_expr() 227 | del parents_type[0] 228 | del parents[0] 229 | else: 230 | offset = 0 231 | 232 | return self.__extract_member(cexpr, obj, offset, parents, parents_type) 233 | 234 | def __extract_member(self, cexpr, obj, offset, parents, parents_type): 235 | if parents_type[0] == 'cast': 236 | default_tinfo = parents[0].type 237 | cexpr = parents[0] 238 | del parents_type[0] 239 | del parents[0] 240 | else: 241 | default_tinfo = const.PX_WORD_TINFO 242 | 243 | if parents_type[0] in ('idx', 'ptr'): 244 | if parents_type[1] == 'cast': 245 | default_tinfo = parents[1].type 246 | cexpr = parents[0] 247 | del parents_type[0] 248 | del parents[0] 249 | else: 250 | default_tinfo = self.__deref_tinfo(default_tinfo) 251 | 252 | if parents_type[1] == 'asg': 253 | if parents[1].x == parents[0]: 254 | # *(TYPE *)(var + x) = ??? 255 | obj_ea = self.__extract_obj_ea(parents[1].y) 256 | return self._get_member(offset, cexpr, obj, parents[1].y.type, obj_ea) 257 | return self._get_member(offset, cexpr, obj, parents[1].x.type) 258 | elif parents_type[1] == 'call': 259 | if parents[1].x == parents[0]: 260 | # ((type (__some_call *)(..., ..., ...)var[idx])(..., ..., ...) 261 | # ((type (__some_call *)(..., ..., ...)*(TYPE *)(var + x))(..., ..., ...) 262 | return self._get_member(offset, cexpr, obj, parents[0].type) 263 | _, tinfo = helper.get_func_argument_info(parents[1], parents[0]) 264 | if tinfo is None: 265 | tinfo = const.PCHAR_TINFO 266 | return self._get_member(offset, cexpr, obj, tinfo) 267 | return self._get_member(offset, cexpr, obj, default_tinfo) 268 | 269 | elif parents_type[0] == 'call': 270 | # call(..., (TYPE)(var + x), ...) 271 | tinfo = self._parse_call(parents[0], cexpr, offset) 272 | return self._get_member(offset, cexpr, obj, tinfo) 273 | 274 | elif parents_type[0] == 'asg': 275 | if parents[0].y == cexpr: 276 | # other_obj = (TYPE) (var + offset) 277 | self._parse_left_assignee(parents[1].x, offset) 278 | return self._get_member(offset, cexpr, obj, self.__deref_tinfo(default_tinfo)) 279 | 280 | @staticmethod 281 | def __extract_obj_ea(cexpr): 282 | if cexpr.op == idaapi.cot_ref: 283 | cexpr = cexpr.x 284 | if cexpr.op == idaapi.cot_obj: 285 | if cexpr.obj_ea != idaapi.BADADDR: 286 | return cexpr.obj_ea 287 | 288 | @staticmethod 289 | def __deref_tinfo(tinfo): 290 | if tinfo.is_ptr(): 291 | if tinfo.get_ptrarr_objsize() == 1: 292 | if tinfo.equals_to(const.PCHAR_TINFO) or tinfo.equals_to(const.CONST_PCHAR_TINFO): 293 | return const.CHAR_TINFO 294 | return None # Turns into VoidMember 295 | return tinfo.get_pointed_object() 296 | return tinfo 297 | 298 | 299 | class NewShallowSearchVisitor(SearchVisitor, api.ObjectDownwardsVisitor): 300 | def __init__(self, cfunc, origin, obj, temporary_structure): 301 | super(NewShallowSearchVisitor, self).__init__(cfunc, origin, obj, temporary_structure) 302 | 303 | 304 | class NewDeepSearchVisitor(SearchVisitor, api.RecursiveObjectDownwardsVisitor): 305 | def __init__(self, cfunc, origin, obj, temporary_structure): 306 | super(NewDeepSearchVisitor, self).__init__(cfunc, origin, obj, temporary_structure) 307 | 308 | 309 | class DeepReturnVisitor(NewDeepSearchVisitor): 310 | def __init__(self, cfunc, origin, obj, temporary_structure): 311 | super(DeepReturnVisitor, self).__init__(cfunc, origin, obj, temporary_structure) 312 | self.__callers_ea = helper.get_funcs_calling_address(cfunc.entry_ea) 313 | self.__call_obj = obj 314 | 315 | def _start(self): 316 | for ea in self.__callers_ea: 317 | self._add_scan_tree_info(ea, -1) 318 | assert self.__prepare_scanner() 319 | 320 | def _finish(self): 321 | if self.__prepare_scanner(): 322 | self._recursive_process() 323 | 324 | def __prepare_scanner(self): 325 | try: 326 | cfunc = next(self.__iter_callers()) 327 | except StopIteration: 328 | return False 329 | 330 | self.prepare_new_scan(cfunc, -1, self.__call_obj) 331 | return True 332 | 333 | def __iter_callers(self): 334 | for ea in self.__callers_ea: 335 | cfunc = helper.decompile_function(ea) 336 | if cfunc: 337 | yield cfunc 338 | -------------------------------------------------------------------------------- /HexRaysPyTools/callbacks/negative_offsets.py: -------------------------------------------------------------------------------- 1 | import re 2 | import logging 3 | import idaapi 4 | 5 | from . import actions 6 | from . import callbacks 7 | import HexRaysPyTools.core.helper as helper 8 | import HexRaysPyTools.core.type_library as type_library 9 | import HexRaysPyTools.forms as forms 10 | 11 | logger = logging.getLogger(__name__) 12 | potential_negatives = {} 13 | 14 | 15 | def _has_magic_comment(lvar): 16 | # type: (idaapi.lvar_t) -> bool 17 | # FIXME: Use internal IDA storage for CONTAINING_RECORD macro 18 | return bool(re.search("```.*```", lvar.cmt)) 19 | 20 | 21 | def _parse_magic_comment(lvar): 22 | if lvar.type().is_ptr(): 23 | m = re.search('```(.+)```', lvar.cmt) 24 | if m: 25 | structure_name, offset = m.group(1).split('+') 26 | offset = int(offset) 27 | parent_tinfo = idaapi.tinfo_t() 28 | if parent_tinfo.get_named_type(idaapi.cvar.idati, structure_name) and parent_tinfo.get_size() > offset: 29 | member_name = dict(find_deep_members(parent_tinfo, lvar.type().get_pointed_object())).get(offset, None) 30 | if member_name: 31 | return NegativeLocalInfo(lvar.type().get_pointed_object(), parent_tinfo, offset, member_name) 32 | return None 33 | 34 | 35 | def find_deep_members(parent_tinfo, target_tinfo): 36 | udt_data = idaapi.udt_type_data_t() 37 | parent_tinfo.get_udt_details(udt_data) 38 | result = [] 39 | for udt_member in udt_data: 40 | if udt_member.type.equals_to(target_tinfo): 41 | result.append((udt_member.offset // 8, udt_member.name)) 42 | elif udt_member.type.is_udt(): 43 | for offset, name in find_deep_members(udt_member.type, target_tinfo): 44 | final_name = udt_member.name + '.' + name if udt_member.name else name 45 | result.append((udt_member.offset // 8 + offset, final_name)) 46 | return result 47 | 48 | 49 | class NegativeLocalInfo: 50 | def __init__(self, tinfo, parent_tinfo, offset, member_name): 51 | self.tinfo = tinfo 52 | self.size = tinfo.get_size() if tinfo.is_udt else 0 53 | self.parent_tinfo = parent_tinfo 54 | self.offset = offset 55 | self.member_name = member_name 56 | 57 | def __repr__(self): 58 | return "Type - {0}, parent type - {1}, offset - {2}, member_name - {3}".format( 59 | self.tinfo.dstr(), 60 | self.parent_tinfo.dstr(), 61 | self.offset, 62 | self.member_name 63 | ) 64 | 65 | 66 | class NegativeLocalCandidate: 67 | def __init__(self, tinfo, offset): 68 | """ 69 | Tinfo - type of the structure tha local variable points to. So it's stripped from pointer. Offset - is first 70 | found offset that points outside of the structure. 71 | :param tinfo: idaapi.tinfo_t 72 | :param offset: int 73 | """ 74 | self.tinfo = tinfo 75 | self.offsets = [offset] 76 | 77 | def __repr__(self): 78 | return self.tinfo.dstr() + ' ' + str(self.offsets) 79 | 80 | def is_structure_offset(self, tinfo, offset): 81 | # Checks if structure tinfo contains a member at given offset 82 | # TODO:array checking 83 | udt_member = idaapi.udt_member_t() 84 | udt_member.offset = offset * 8 85 | if offset >= 0 and tinfo.find_udt_member(udt_member, idaapi.STRMEM_OFFSET) != -1: 86 | if udt_member.type.is_udt(): 87 | return self.is_structure_offset(udt_member.type, offset - udt_member.offset // 8) 88 | return udt_member.offset == offset * 8 89 | return False 90 | 91 | def find_containing_structures(self, type_library): 92 | """ 93 | Given the type library creates a list of structures from this library, that contains this structure and 94 | satisfy offset conditions. 95 | :param type_library: idaapi.til_t 96 | :returns: ordinal, offset, member_name, containing structure name 97 | """ 98 | 99 | min_offset = min(self.offsets) 100 | min_offset = min_offset if min_offset < 0 else 0 101 | max_offset = max(self.offsets) 102 | max_offset = max_offset if max_offset > 0 else self.tinfo.get_size() 103 | # TODO: Check if all offsets are legal 104 | 105 | # Least acceptable size of the containing structure 106 | min_struct_size = max_offset - min_offset 107 | result = [] 108 | parent_tinfo = idaapi.tinfo_t() 109 | target_tinfo = idaapi.tinfo_t() 110 | if not target_tinfo.get_named_type(type_library, self.tinfo.dstr()): 111 | print("[Warning] Such type doesn't exist in '{0}' library".format(type_library.name)) 112 | return result 113 | for ordinal in range(1, idaapi.get_ordinal_qty(type_library)): 114 | parent_tinfo.create_typedef(type_library, ordinal) 115 | if parent_tinfo.get_size() >= min_struct_size: 116 | for offset, name in find_deep_members(parent_tinfo, target_tinfo): 117 | # print "[DEBUG] Found {0} at {1} in {2}".format(name, offset, parent_tinfo.dstr()) 118 | if offset + min_offset >= 0 and offset + max_offset <= parent_tinfo.get_size(): 119 | result.append((ordinal, offset, name, parent_tinfo.dstr())) 120 | return result 121 | 122 | 123 | class ReplaceVisitor(idaapi.ctree_parentee_t): 124 | 125 | def __init__(self, negative_lvars): 126 | super(ReplaceVisitor, self).__init__() 127 | self.negative_lvars = negative_lvars 128 | self.pvoid_tinfo = idaapi.tinfo_t(idaapi.BT_VOID) 129 | self.pvoid_tinfo.create_ptr(self.pvoid_tinfo) 130 | 131 | def visit_expr(self, expression): 132 | if expression.op == idaapi.cot_add and expression.x.op == idaapi.cot_var and expression.y.op == idaapi.cot_num: 133 | index = expression.x.v.idx 134 | if index in self.negative_lvars: 135 | offset = expression.y.numval() 136 | if offset >= self.negative_lvars[index].size: 137 | self.create_containing_record(expression, index, offset) 138 | elif expression.op == idaapi.cot_sub and expression.x.op == idaapi.cot_var and expression.y.op == idaapi.cot_num: 139 | index = expression.x.v.idx 140 | if index in self.negative_lvars: 141 | offset = -expression.y.n.value(idaapi.tinfo_t(idaapi.BT_INT)) 142 | self.create_containing_record(expression, index, offset) 143 | return 0 144 | 145 | def create_containing_record(self, expression, index, offset): 146 | negative_lvar = self.negative_lvars[index] 147 | logger.debug("Creating CONTAINING_RECORD macro, offset: {}, negative offset: {}, TYPE: {}".format( 148 | negative_lvar.offset, 149 | offset, 150 | negative_lvar.parent_tinfo.dstr() 151 | )) 152 | 153 | arg_address = idaapi.carg_t() 154 | if expression.op == idaapi.cot_var: 155 | arg_address.assign(expression) 156 | else: 157 | arg_address.assign(expression.x) 158 | 159 | arg_type = idaapi.carg_t() 160 | cexpr_helper = idaapi.create_helper(True, self.pvoid_tinfo, negative_lvar.parent_tinfo.dstr()) 161 | arg_type.assign(cexpr_helper) 162 | 163 | arg_field = idaapi.carg_t() 164 | cexpr_helper = idaapi.create_helper( 165 | True, 166 | self.pvoid_tinfo, 167 | negative_lvar.member_name 168 | ) 169 | arg_field.assign(cexpr_helper) 170 | return_tinfo = idaapi.tinfo_t(negative_lvar.parent_tinfo) 171 | return_tinfo.create_ptr(return_tinfo) 172 | new_cexpr_call = idaapi.call_helper(return_tinfo, None, "CONTAINING_RECORD") 173 | new_cexpr_call.a.push_back(arg_address) 174 | new_cexpr_call.a.push_back(arg_type) 175 | new_cexpr_call.a.push_back(arg_field) 176 | new_cexpr_call.thisown = False 177 | 178 | parent = reversed(self.parents).next().cexpr 179 | 180 | diff = negative_lvar.offset + offset 181 | if diff: 182 | number = idaapi.make_num(diff) 183 | number.thisown = False 184 | new_cexpr_add = helper.my_cexpr_t(idaapi.cot_add, x=new_cexpr_call, y=number) 185 | new_cexpr_add.type = return_tinfo 186 | 187 | if parent.op == idaapi.cot_ptr: 188 | tmp_tinfo = idaapi.tinfo_t() 189 | tmp_tinfo.create_ptr(parent.type) 190 | new_cexpr_cast = helper.my_cexpr_t(idaapi.cot_cast, x=new_cexpr_add) 191 | new_cexpr_cast.thisown = False 192 | new_cexpr_cast.type = tmp_tinfo 193 | expression.assign(new_cexpr_cast) 194 | else: 195 | expression.assign(new_cexpr_add) 196 | else: 197 | if parent.op == idaapi.cot_ptr: 198 | tmp_tinfo = idaapi.tinfo_t() 199 | tmp_tinfo.create_ptr(parent.type) 200 | new_cexpr_cast = helper.my_cexpr_t(idaapi.cot_cast, x=new_cexpr_call) 201 | new_cexpr_cast.thisown = False 202 | new_cexpr_cast.type = tmp_tinfo 203 | expression.assign(new_cexpr_cast) 204 | else: 205 | expression.assign(new_cexpr_call) 206 | 207 | 208 | class SearchVisitor(idaapi.ctree_parentee_t): 209 | def __init__(self, cfunc): 210 | super(SearchVisitor, self).__init__() 211 | self.cfunc = cfunc 212 | self.result = {} 213 | 214 | def visit_expr(self, expression): 215 | if expression.op == idaapi.cot_call and expression.x.op == idaapi.cot_helper and len(expression.a) == 3: 216 | if expression.x.helper == "CONTAINING_RECORD": 217 | if expression.a[0].op == idaapi.cot_var: 218 | idx = expression.a[0].v.idx 219 | if expression.a[1].op == idaapi.cot_helper and expression.a[2].op == idaapi.cot_helper: 220 | parent_name = expression.a[1].helper 221 | member_name = expression.a[2].helper 222 | parent_tinfo = idaapi.tinfo_t() 223 | if not parent_tinfo.get_named_type(idaapi.cvar.idati, parent_name): 224 | return 0 225 | udt_data = idaapi.udt_type_data_t() 226 | parent_tinfo.get_udt_details(udt_data) 227 | udt_member = [x for x in udt_data if x.name == member_name] 228 | if udt_member: 229 | tinfo = udt_member[0].type 230 | self.result[idx] = NegativeLocalInfo( 231 | tinfo, 232 | parent_tinfo, 233 | udt_member[0].offset // 8, 234 | member_name 235 | ) 236 | return 1 237 | return 0 238 | 239 | 240 | class AnalyseVisitor(idaapi.ctree_parentee_t): 241 | def __init__(self, candidates): 242 | global potential_negatives 243 | super(AnalyseVisitor, self).__init__() 244 | self.candidates = candidates 245 | self.potential_negatives = potential_negatives 246 | self.potential_negatives.clear() 247 | 248 | def visit_expr(self, expression): 249 | if expression.op == idaapi.cot_add and expression.y.op == idaapi.cot_num: 250 | if expression.x.op == idaapi.cot_var and expression.x.v.idx in self.candidates: 251 | idx = expression.x.v.idx 252 | number = expression.y.numval() 253 | if self.candidates[idx].get_size() <= number: 254 | if idx in self.potential_negatives: 255 | self.potential_negatives[idx].offsets.append(number) 256 | else: 257 | self.potential_negatives[idx] = NegativeLocalCandidate(self.candidates[idx], number) 258 | elif expression.op == idaapi.cot_sub and expression.y.op == idaapi.cot_num: 259 | if expression.x.op == idaapi.cot_var and expression.x.v.idx in self.candidates: 260 | idx = expression.x.v.idx 261 | number = -expression.y.numval() 262 | if idx in self.potential_negatives: 263 | self.potential_negatives[idx].offsets.append(number) 264 | else: 265 | self.potential_negatives[idx] = NegativeLocalCandidate(self.candidates[idx], number) 266 | 267 | return 0 268 | 269 | 270 | class PotentialNegativeCollector(callbacks.HexRaysEventHandler): 271 | def __init__(self): 272 | super(PotentialNegativeCollector, self).__init__() 273 | 274 | def handle(self, event, *args): 275 | global potential_negatives 276 | 277 | cfunc, level_of_maturity = args 278 | if level_of_maturity == idaapi.CMAT_BUILT: 279 | # First search for CONTAINING_RECORD made by Ida 280 | visitor = SearchVisitor(cfunc) 281 | visitor.apply_to(cfunc.body, None) 282 | negative_lvars = visitor.result 283 | 284 | # Second get saved information from comments 285 | lvars = cfunc.get_lvars() 286 | for idx in range(len(lvars)): 287 | result = _parse_magic_comment(lvars[idx]) 288 | if result and result.tinfo.equals_to(lvars[idx].type().get_pointed_object()): 289 | negative_lvars[idx] = result 290 | 291 | # Third analyze local variables that are a structure pointers and have references going beyond 292 | # structure boundaries. This variables will be considered as potential pointers to substructure 293 | # and will get a special menu on right click 294 | 295 | # First collect all structure pointers 296 | structure_pointer_variables = {} 297 | for idx in set(range(len(lvars))) - set(negative_lvars.keys()): 298 | if lvars[idx].type().is_ptr(): 299 | pointed_tinfo = lvars[idx].type().get_pointed_object() 300 | if pointed_tinfo.is_udt(): 301 | structure_pointer_variables[idx] = pointed_tinfo 302 | 303 | # Then use them in order to find all potential negative offset situations 304 | if structure_pointer_variables: 305 | visitor = AnalyseVisitor(structure_pointer_variables) 306 | visitor.apply_to(cfunc.body, None) 307 | 308 | # If negative offsets were found, then we replace them with CONTAINING_RECORD macro 309 | if negative_lvars: 310 | visitor = ReplaceVisitor(negative_lvars) 311 | visitor.apply_to(cfunc.body, None) 312 | 313 | 314 | callbacks.hx_callback_manager.register(idaapi.hxe_maturity, PotentialNegativeCollector()) 315 | 316 | 317 | class ResetContainingStructure(actions.HexRaysPopupAction): 318 | description = "Reset Containing Structure" 319 | 320 | def __init__(self): 321 | super(ResetContainingStructure, self).__init__() 322 | 323 | def check(self, hx_view): 324 | ctree_item = hx_view.item 325 | if ctree_item.citype != idaapi.VDI_EXPR or ctree_item.e.op != idaapi.cot_var: 326 | return False 327 | return _has_magic_comment(hx_view.cfunc.get_lvars()[ctree_item.e.v.idx]) 328 | 329 | def activate(self, ctx): 330 | hx_view = idaapi.get_widget_vdui(ctx.widget) 331 | lvar = hx_view.cfunc.get_lvars()[hx_view.item.e.v.idx] 332 | hx_view.set_lvar_cmt(lvar, re.sub("```.*```", '', lvar.cmt)) 333 | hx_view.refresh_view(True) 334 | 335 | 336 | actions.action_manager.register(ResetContainingStructure()) 337 | 338 | 339 | class SelectContainingStructure(actions.HexRaysPopupAction): 340 | description = "Select Containing Structure" 341 | 342 | def __init__(self): 343 | super(SelectContainingStructure, self).__init__() 344 | 345 | def check(self, hx_view): 346 | ctree_item = hx_view.item 347 | if ctree_item.citype != idaapi.VDI_EXPR or ctree_item.e.op != idaapi.cot_var: 348 | return False 349 | return ctree_item.e.v.idx in potential_negatives 350 | 351 | def activate(self, ctx): 352 | global potential_negatives 353 | 354 | hx_view = idaapi.get_widget_vdui(ctx.widget) 355 | result = type_library.choose_til() 356 | if not result: 357 | return 358 | 359 | selected_library, max_ordinal, is_local_types = result 360 | lvar_idx = hx_view.item.e.v.idx 361 | candidate = potential_negatives[lvar_idx] 362 | structures = candidate.find_containing_structures(selected_library) 363 | items = [[str(x[0]), "0x{0:08X}".format(x[1]), x[2], x[3]] for x in structures] 364 | structure_chooser = forms.MyChoose( 365 | items, 366 | "Select Containing Structure", 367 | [["Ordinal", 5], ["Offset", 10], ["Member_name", 20], ["Structure Name", 20]], 368 | 165 369 | ) 370 | selected_idx = structure_chooser.Show(modal=True) 371 | if selected_idx != -1: 372 | if not is_local_types: 373 | type_library.import_type(selected_library, items[selected_idx][3]) 374 | lvar = hx_view.cfunc.get_lvars()[lvar_idx] 375 | lvar_cmt = re.sub("```.*```", '', lvar.cmt) 376 | hx_view.set_lvar_cmt( 377 | lvar, 378 | lvar_cmt + "```{0}+{1}```".format( 379 | structures[selected_idx][3], 380 | structures[selected_idx][1]) 381 | ) 382 | hx_view.refresh_view(True) 383 | 384 | 385 | actions.action_manager.register(SelectContainingStructure()) 386 | -------------------------------------------------------------------------------- /HexRaysPyTools/api.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import idaapi 3 | import idc 4 | from .core.helper import to_hex 5 | from .core import helper 6 | 7 | logger = logging.getLogger(__name__) 8 | 9 | 10 | SETTING_START_FROM_CURRENT_EXPR = True 11 | 12 | 13 | class ScanObject(object): 14 | def __init__(self): 15 | self.ea = idaapi.BADADDR 16 | self.name = None 17 | self.tinfo = None 18 | self.id = 0 19 | 20 | @staticmethod 21 | def create(cfunc, arg): 22 | # Creates object suitable for scaning either from cexpr_t or ctree_item_t 23 | if isinstance(arg, idaapi.ctree_item_t): 24 | lvar = arg.get_lvar() 25 | if lvar: 26 | index = list(cfunc.get_lvars()).index(lvar) 27 | result = VariableObject(lvar, index) 28 | if arg.e: 29 | result.ea = ScanObject.get_expression_address(cfunc, arg.e) 30 | return result 31 | if arg.citype != idaapi.VDI_EXPR: 32 | return None 33 | cexpr = arg.e 34 | else: 35 | cexpr = arg 36 | 37 | if cexpr.op == idaapi.cot_var: 38 | lvar = cfunc.get_lvars()[cexpr.v.idx] 39 | result = VariableObject(lvar, cexpr.v.idx) 40 | result.ea = ScanObject.get_expression_address(cfunc, cexpr) 41 | return result 42 | elif cexpr.op == idaapi.cot_memptr: 43 | t = cexpr.x.type.get_pointed_object() 44 | result = StructPtrObject(t.dstr(), cexpr.m) 45 | result.name = helper.get_member_name(t, cexpr.m) 46 | elif cexpr.op == idaapi.cot_memref: 47 | t = cexpr.x.type 48 | result = StructRefObject(t.dstr(), cexpr.m) 49 | result.name = helper.get_member_name(t, cexpr.m) 50 | elif cexpr.op == idaapi.cot_obj: 51 | result = GlobalVariableObject(cexpr.obj_ea) 52 | result.name = idaapi.get_short_name(cexpr.obj_ea) 53 | else: 54 | return 55 | result.tinfo = cexpr.type 56 | result.ea = ScanObject.get_expression_address(cfunc, cexpr) 57 | return result 58 | 59 | @staticmethod 60 | def get_expression_address(cfunc, cexpr): 61 | expr = cexpr 62 | 63 | while expr and expr.ea == idaapi.BADADDR: 64 | expr = expr.to_specific_type 65 | expr = cfunc.body.find_parent_of(expr) 66 | 67 | assert expr is not None 68 | return expr.ea 69 | 70 | def __hash__(self): 71 | return hash((self.id, self.name)) 72 | 73 | def __eq__(self, other): 74 | return self.id == other.id and self.name == other.name 75 | 76 | def __repr__(self): 77 | return self.name 78 | 79 | 80 | SO_LOCAL_VARIABLE = 1 # cexpr.op == idaapi.cot_var 81 | SO_STRUCT_POINTER = 2 # cexpr.op == idaapi.cot_memptr 82 | SO_STRUCT_REFERENCE = 3 # cexpr.op == idaapi.cot_memref 83 | SO_GLOBAL_OBJECT = 4 # cexpr.op == idaapi.cot_obj 84 | SO_CALL_ARGUMENT = 5 # cexpr.op == idaapi.cot_call 85 | SO_MEMORY_ALLOCATOR = 6 86 | SO_RETURNED_OBJECT = 7 87 | 88 | 89 | class VariableObject(ScanObject): 90 | # Represents `var` expression 91 | def __init__(self, lvar, index): 92 | super(VariableObject, self).__init__() 93 | self.lvar = lvar 94 | self.tinfo = lvar.type() 95 | self.name = lvar.name 96 | self.index = index 97 | self.id = SO_LOCAL_VARIABLE 98 | 99 | def is_target(self, cexpr): 100 | return cexpr.op == idaapi.cot_var and cexpr.v.idx == self.index 101 | 102 | 103 | class StructPtrObject(ScanObject): 104 | # Represents `x->m` expression 105 | def __init__(self, struct_name, offset): 106 | super(StructPtrObject, self).__init__() 107 | self.struct_name = struct_name 108 | self.offset = offset 109 | self.id = SO_STRUCT_POINTER 110 | 111 | def is_target(self, cexpr): 112 | return cexpr.op == idaapi.cot_memptr and cexpr.m == self.offset and \ 113 | cexpr.x.type.get_pointed_object().dstr() == self.struct_name 114 | 115 | 116 | class StructRefObject(ScanObject): 117 | # Represents `x.m` expression 118 | def __init__(self, struct_name, offset): 119 | super(StructRefObject, self).__init__() 120 | self.struct_name = struct_name 121 | self.offset = offset 122 | self.id = SO_STRUCT_REFERENCE 123 | 124 | def is_target(self, cexpr): 125 | return cexpr.op == idaapi.cot_memref and cexpr.m == self.offset and cexpr.x.type.dstr() == self.struct_name 126 | 127 | 128 | class GlobalVariableObject(ScanObject): 129 | # Represents global object 130 | def __init__(self, object_address): 131 | super(GlobalVariableObject, self).__init__() 132 | self.obj_ea = object_address 133 | self.id = SO_GLOBAL_OBJECT 134 | 135 | def is_target(self, cexpr): 136 | return cexpr.op == idaapi.cot_obj and self.obj_ea == cexpr.obj_ea 137 | 138 | 139 | class CallArgObject(ScanObject): 140 | # Represents call of a function and argument index 141 | def __init__(self, func_address, arg_idx): 142 | super(CallArgObject, self).__init__() 143 | self.func_ea = func_address 144 | self.arg_idx = arg_idx 145 | self.id = SO_CALL_ARGUMENT 146 | 147 | def is_target(self, cexpr): 148 | return cexpr.op == idaapi.cot_call and cexpr.x.obj_ea == self.func_ea 149 | 150 | def create_scan_obj(self, cfunc, cexpr): 151 | e = cexpr.a[self.arg_idx] 152 | while e.op in (idaapi.cot_cast, idaapi.cot_ref, idaapi.cot_add, idaapi.cot_sub, idaapi.cot_idx): 153 | e = e.x 154 | return ScanObject.create(cfunc, e) 155 | 156 | @staticmethod 157 | def create(cfunc, arg_idx): 158 | result = CallArgObject(cfunc.entry_ea, arg_idx) 159 | result.name = cfunc.get_lvars()[arg_idx].name 160 | result.tinfo = cfunc.type 161 | return result 162 | 163 | def __repr__(self): 164 | return "{}" 165 | 166 | 167 | class ReturnedObject(ScanObject): 168 | # Represents value returned by function 169 | def __init__(self, func_address): 170 | super(ReturnedObject, self).__init__() 171 | self.__func_ea = func_address 172 | self.id = SO_RETURNED_OBJECT 173 | 174 | def is_target(self, cexpr): 175 | return cexpr.op == idaapi.cot_call and cexpr.x.obj_ea == self.__func_ea 176 | 177 | 178 | class MemoryAllocationObject(ScanObject): 179 | # Represents `operator new()' or `malloc' 180 | def __init__(self, name, size): 181 | super(MemoryAllocationObject, self).__init__() 182 | self.name = name 183 | self.size = size 184 | self.id = SO_MEMORY_ALLOCATOR 185 | 186 | @staticmethod 187 | def create(cfunc, cexpr): 188 | if cexpr.op == idaapi.cot_call: 189 | e = cexpr 190 | elif cexpr.op == idaapi.cot_cast and cexpr.x.op == idaapi.cot_call: 191 | e = cexpr.x 192 | else: 193 | return 194 | 195 | func_name = idaapi.get_short_name(e.x.obj_ea) 196 | if "malloc" in func_name or "operator new" in func_name: 197 | carg = e.a[0] 198 | if carg.op == idaapi.cot_num: 199 | size = carg.numval() 200 | else: 201 | size = 0 202 | result = MemoryAllocationObject(func_name, size) 203 | result.ea = ScanObject.get_expression_address(cfunc, e) 204 | return result 205 | 206 | 207 | ASSIGNMENT_RIGHT = 1 208 | ASSIGNMENT_LEFT = 2 209 | 210 | 211 | class ObjectVisitor(idaapi.ctree_parentee_t): 212 | def __init__(self, cfunc, obj, data, skip_until_object): 213 | super(ObjectVisitor, self).__init__() 214 | self._cfunc = cfunc 215 | self._objects = [obj] 216 | self._init_obj = obj 217 | self._data = data 218 | self._start_ea = obj.ea 219 | self._skip = skip_until_object if self._start_ea != idaapi.BADADDR else False 220 | self.crippled = False 221 | 222 | def process(self): 223 | self.apply_to(self._cfunc.body, None) 224 | 225 | def set_callbacks(self, manipulate=None): 226 | if manipulate: 227 | self.__manipulate = manipulate.__get__(self, ObjectDownwardsVisitor) 228 | 229 | def _get_line(self): 230 | for p in reversed(self.parents): 231 | if not p.is_expr(): 232 | return idaapi.tag_remove(p.print1(self._cfunc.__ref__())) 233 | AssertionError("Parent instruction is not found") 234 | 235 | def _manipulate(self, cexpr, obj): 236 | """ 237 | Method called for every object having assignment relationship with starter object. This method should be 238 | reimplemented in order to do something useful 239 | 240 | :param cexpr: idaapi_cexpr_t 241 | :param id: one of the SO_* constants 242 | :return: None 243 | """ 244 | self.__manipulate(cexpr, obj) 245 | 246 | def __manipulate(self, cexpr, obj): 247 | logger.debug("Expression {} at {} Id - {}".format( 248 | cexpr.opname, 249 | to_hex(helper.find_asm_address(cexpr, self.parents)), 250 | obj.id)) 251 | 252 | 253 | class ObjectDownwardsVisitor(ObjectVisitor): 254 | def __init__(self, cfunc, obj, data=None, skip_until_object=False): 255 | super(ObjectDownwardsVisitor, self).__init__(cfunc, obj, data, skip_until_object) 256 | self.cv_flags |= idaapi.CV_POST 257 | 258 | def visit_expr(self, cexpr): 259 | if self._skip: 260 | if self._is_initial_object(cexpr): 261 | self._skip = False 262 | else: 263 | return 0 264 | 265 | if cexpr.op != idaapi.cot_asg: 266 | return 0 267 | 268 | x_cexpr = cexpr.x 269 | if cexpr.y.op == idaapi.cot_cast: 270 | y_cexpr = cexpr.y.x 271 | else: 272 | y_cexpr = cexpr.y 273 | 274 | for obj in self._objects: 275 | if obj.is_target(x_cexpr): 276 | if self.__is_object_overwritten(x_cexpr, obj, y_cexpr): 277 | logger.info("Removed object {} from scanning at {}".format( 278 | obj, to_hex(helper.find_asm_address(x_cexpr, self.parents)))) 279 | self._objects.remove(obj) 280 | return 0 281 | elif obj.is_target(y_cexpr): 282 | new_obj = ScanObject.create(self._cfunc, x_cexpr) 283 | if new_obj: 284 | self._objects.append(new_obj) 285 | return 0 286 | return 0 287 | 288 | def leave_expr(self, cexpr): 289 | if self._skip: 290 | return 0 291 | 292 | for obj in self._objects: 293 | if obj.is_target(cexpr) and obj.id != SO_RETURNED_OBJECT: 294 | self._manipulate(cexpr, obj) 295 | return 0 296 | return 0 297 | 298 | def _is_initial_object(self, cexpr): 299 | if cexpr.op == idaapi.cot_asg: 300 | cexpr = cexpr.y 301 | if cexpr.op == idaapi.cot_cast: 302 | cexpr = cexpr.x 303 | return self._init_obj.is_target(cexpr) and helper.find_asm_address(cexpr, self.parents) == self._start_ea 304 | 305 | def __is_object_overwritten(self, x_cexpr, obj, y_cexpr): 306 | if len(self._objects) < 2: 307 | return False 308 | 309 | if y_cexpr.op == idaapi.cot_cast: 310 | e = y_cexpr.x 311 | else: 312 | e = y_cexpr 313 | 314 | if e.op != idaapi.cot_call or len(e.a) == 0: 315 | return True 316 | 317 | for obj in self._objects: 318 | if obj.is_target(e. a[0]): 319 | return False 320 | return True 321 | 322 | 323 | class ObjectUpwardsVisitor(ObjectVisitor): 324 | STAGE_PREPARE = 1 325 | STAGE_PARSING = 2 326 | 327 | def __init__(self, cfunc, obj, data=None, skip_after_object=False): 328 | super(ObjectUpwardsVisitor, self).__init__(cfunc, obj, data, skip_after_object) 329 | self._stage = self.STAGE_PREPARE 330 | self._tree = {} 331 | self._call_obj = obj if obj.id == SO_CALL_ARGUMENT else None 332 | 333 | def visit_expr(self, cexpr): 334 | if self._stage == self.STAGE_PARSING: 335 | return 0 336 | 337 | if self._call_obj and self._call_obj.is_target(cexpr): 338 | obj = self._call_obj.create_scan_obj(self._cfunc, cexpr) 339 | if obj: 340 | self._objects.append(obj) 341 | return 0 342 | 343 | if cexpr.op != idaapi.cot_asg: 344 | return 0 345 | 346 | x_cexpr = cexpr.x 347 | if cexpr.y.op == idaapi.cot_cast: 348 | y_cexpr = cexpr.y.x 349 | else: 350 | y_cexpr = cexpr.y 351 | 352 | obj_left = ScanObject.create(self._cfunc, x_cexpr) 353 | obj_right = ScanObject.create(self._cfunc, y_cexpr) 354 | if obj_left and obj_right: 355 | self.__add_object_assignment(obj_left, obj_right) 356 | 357 | if self._skip and self._is_initial_object(cexpr): 358 | return 1 359 | return 0 360 | 361 | def leave_expr(self, cexpr): 362 | if self._stage == self.STAGE_PREPARE: 363 | return 0 364 | 365 | if self._skip and self._is_initial_object(cexpr): 366 | self._manipulate(cexpr, self._init_obj) 367 | return 1 368 | 369 | for obj in self._objects: 370 | if obj.is_target(cexpr): 371 | self._manipulate(cexpr, obj) 372 | return 0 373 | return 0 374 | 375 | def process(self): 376 | self._stage = self.STAGE_PREPARE 377 | self.cv_flags &= ~idaapi.CV_POST 378 | super(ObjectUpwardsVisitor, self).process() 379 | self._stage = self.STAGE_PARSING 380 | self.cv_flags |= idaapi.CV_POST 381 | self.__prepare() 382 | super(ObjectUpwardsVisitor, self).process() 383 | 384 | def _is_initial_object(self, cexpr): 385 | return self._init_obj.is_target(cexpr) and helper.find_asm_address(cexpr, self.parents) == self._start_ea 386 | 387 | def __add_object_assignment(self, from_obj, to_obj): 388 | if from_obj in self._tree: 389 | self._tree[from_obj].add(to_obj) 390 | else: 391 | self._tree[from_obj] = {to_obj} 392 | 393 | def __prepare(self): 394 | result = set() 395 | todo = set(self._objects) 396 | while todo: 397 | obj = todo.pop() 398 | result.add(obj) 399 | if obj.id == SO_CALL_ARGUMENT or obj not in self._tree: 400 | continue 401 | o = self._tree[obj] 402 | todo |= o - result 403 | result |= o 404 | self._objects = list(result) 405 | self._tree.clear() 406 | 407 | 408 | class RecursiveObjectVisitor(ObjectVisitor): 409 | def __init__(self, cfunc, obj, data=None, skip_until_object=False, visited=None): 410 | super(RecursiveObjectVisitor, self).__init__(cfunc, obj, data, skip_until_object) 411 | self._visited = visited if visited else set() 412 | self._new_for_visit = set() 413 | self.crippled = False 414 | self._arg_idx = -1 415 | self._debug_scan_tree = {} 416 | self.__debug_scan_tree_root = idc.get_name(self._cfunc.entry_ea) 417 | self.__debug_message = [] 418 | 419 | def visit_expr(self, cexpr): 420 | return super(RecursiveObjectVisitor, self).visit_expr(cexpr) 421 | 422 | def set_callbacks(self, manipulate=None, start=None, start_iteration=None, finish=None, finish_iteration=None): 423 | super(RecursiveObjectVisitor, self).set_callbacks(manipulate) 424 | if start: 425 | self._start = start.__get__(self, RecursiveObjectDownwardsVisitor) 426 | if start_iteration: 427 | self._start_iteration = start_iteration.__get__(self, RecursiveObjectDownwardsVisitor) 428 | if finish: 429 | self._finish = finish.__get__(self, RecursiveObjectDownwardsVisitor) 430 | if finish_iteration: 431 | self._finish_iteration = finish_iteration.__get__(self, RecursiveObjectDownwardsVisitor) 432 | 433 | def prepare_new_scan(self, cfunc, arg_idx, obj, skip=False): 434 | self._cfunc = cfunc 435 | self._arg_idx = arg_idx 436 | self._objects = [obj] 437 | self._init_obj = obj 438 | self._skip = False 439 | self.crippled = self.__is_func_crippled() 440 | 441 | def process(self): 442 | self._start() 443 | self._recursive_process() 444 | self._finish() 445 | self.dump_scan_tree() 446 | 447 | def dump_scan_tree(self): 448 | self.__prepare_debug_message() 449 | logger.info("{}\n---------------".format("\n".join(self.__debug_message))) 450 | 451 | def __prepare_debug_message(self, key=None, level=1): 452 | if key is None: 453 | key = (self.__debug_scan_tree_root, -1) 454 | self.__debug_message.append("--- Scan Tree---\n{}".format(self.__debug_scan_tree_root)) 455 | if key in self._debug_scan_tree: 456 | for func_name, arg_idx in self._debug_scan_tree[key]: 457 | prefix = " | " * (level - 1) + " |_ " 458 | self.__debug_message.append("{}{} (idx: {})".format(prefix, func_name, arg_idx)) 459 | self.__prepare_debug_message((func_name, arg_idx), level + 1) 460 | 461 | def _recursive_process(self): 462 | self._start_iteration() 463 | super(RecursiveObjectVisitor, self).process() 464 | self._finish_iteration() 465 | 466 | def _manipulate(self, cexpr, obj): 467 | self._check_call(cexpr) 468 | super(RecursiveObjectVisitor, self)._manipulate(cexpr, obj) 469 | 470 | def _check_call(self, cexpr): 471 | raise NotImplemented 472 | 473 | def _add_visit(self, func_ea, arg_idx): 474 | if (func_ea, arg_idx) not in self._visited: 475 | self._visited.add((func_ea, arg_idx)) 476 | self._new_for_visit.add((func_ea, arg_idx)) 477 | return True 478 | return False 479 | 480 | def _add_scan_tree_info(self, func_ea, arg_idx): 481 | head_node = (idc.get_name(self._cfunc.entry_ea), self._arg_idx) 482 | tail_node = (idc.get_name(func_ea), arg_idx) 483 | if head_node in self._debug_scan_tree: 484 | self._debug_scan_tree[head_node].add(tail_node) 485 | else: 486 | self._debug_scan_tree[head_node] = {tail_node} 487 | 488 | def _start(self): 489 | """ Called at the beginning of visiting """ 490 | pass 491 | 492 | def _start_iteration(self): 493 | """ Called every time new function visiting started """ 494 | pass 495 | 496 | def _finish(self): 497 | """ Called after all visiting happened """ 498 | pass 499 | 500 | def _finish_iteration(self): 501 | """ Called every time new function visiting finished """ 502 | pass 503 | 504 | def __is_func_crippled(self): 505 | # Check if function is just call to another function 506 | b = self._cfunc.body.cblock 507 | if b.size() == 1: 508 | e = b.at(0) 509 | return e.op == idaapi.cit_return or (e.op == idaapi.cit_expr and e.cexpr.op == idaapi.cot_call) 510 | return False 511 | 512 | 513 | class RecursiveObjectDownwardsVisitor(RecursiveObjectVisitor, ObjectDownwardsVisitor): 514 | def __init__(self, cfunc, obj, data=None, skip_until_object=False, visited=None): 515 | super(RecursiveObjectDownwardsVisitor, self).__init__(cfunc, obj, data, skip_until_object, visited) 516 | 517 | def _check_call(self, cexpr): 518 | parent = self.parent_expr() 519 | grandparent = self.parents.at(self.parents.size() - 2) 520 | if parent.op == idaapi.cot_call: 521 | call_cexpr = parent 522 | arg_cexpr = cexpr 523 | elif parent.op == idaapi.cot_cast and grandparent.op == idaapi.cot_call: 524 | call_cexpr = grandparent.cexpr 525 | arg_cexpr = parent 526 | else: 527 | return 528 | idx, _ = helper.get_func_argument_info(call_cexpr, arg_cexpr) 529 | func_ea = call_cexpr.x.obj_ea 530 | if func_ea == idaapi.BADADDR: 531 | return 532 | if self._add_visit(func_ea, idx): 533 | self._add_scan_tree_info(func_ea, idx) 534 | 535 | def _recursive_process(self): 536 | super(RecursiveObjectDownwardsVisitor, self)._recursive_process() 537 | 538 | while self._new_for_visit: 539 | func_ea, arg_idx = self._new_for_visit.pop() 540 | if helper.is_imported_ea(func_ea): 541 | continue 542 | cfunc = helper.decompile_function(func_ea) 543 | if cfunc: 544 | assert arg_idx < len(cfunc.get_lvars()), "Wrong argument at func {}".format(to_hex(func_ea)) 545 | obj = VariableObject(cfunc.get_lvars()[arg_idx], arg_idx) 546 | self.prepare_new_scan(cfunc, arg_idx, obj) 547 | self._recursive_process() 548 | 549 | 550 | class RecursiveObjectUpwardsVisitor(RecursiveObjectVisitor, ObjectUpwardsVisitor): 551 | def __init__(self, cfunc, obj, data=None, skip_after_object=False, visited=None): 552 | super(RecursiveObjectUpwardsVisitor, self).__init__(cfunc, obj, data, skip_after_object, visited) 553 | 554 | def prepare_new_scan(self, cfunc, arg_idx, obj, skip=False): 555 | super(RecursiveObjectUpwardsVisitor, self).prepare_new_scan(cfunc, arg_idx, obj, skip) 556 | self._call_obj = obj if obj.id == SO_CALL_ARGUMENT else None 557 | 558 | def _check_call(self, cexpr): 559 | if cexpr.op == idaapi.cot_var and self._cfunc.get_lvars()[cexpr.v.idx].is_arg_var: 560 | func_ea = self._cfunc.entry_ea 561 | arg_idx = cexpr.v.idx 562 | if self._add_visit(func_ea, arg_idx): 563 | for callee_ea in helper.get_funcs_calling_address(func_ea): 564 | self._add_scan_tree_info(callee_ea, arg_idx) 565 | 566 | def _recursive_process(self): 567 | super(RecursiveObjectUpwardsVisitor, self)._recursive_process() 568 | 569 | while self._new_for_visit: 570 | new_visit = list(self._new_for_visit) 571 | self._new_for_visit.clear() 572 | for func_ea, arg_idx in new_visit: 573 | funcs = helper.get_funcs_calling_address(func_ea) 574 | obj = CallArgObject.create(idaapi.decompile(func_ea), arg_idx) 575 | for callee_ea in funcs: 576 | cfunc = helper.decompile_function(callee_ea) 577 | if cfunc: 578 | self.prepare_new_scan(cfunc, arg_idx, obj, False) 579 | super(RecursiveObjectUpwardsVisitor, self)._recursive_process() 580 | -------------------------------------------------------------------------------- /HexRaysPyTools/core/classes.py: -------------------------------------------------------------------------------- 1 | from PyQt5 import QtCore, QtGui 2 | 3 | import idaapi 4 | 5 | import HexRaysPyTools.forms 6 | from . import helper 7 | 8 | 9 | all_virtual_functions = {} # name -> VirtualMethod 10 | all_virtual_tables = {} # ordinal -> VirtualTable 11 | 12 | 13 | class VirtualMethod(object): 14 | def __init__(self, tinfo, name, parent): 15 | self.tinfo = tinfo 16 | self.tinfo_modified = False 17 | self.name = name 18 | self.class_name = None 19 | self.name_modified = False 20 | self.parents = [parent] 21 | image_base = idaapi.get_imagebase() 22 | self.ra_addresses = [ea - image_base for ea in helper.get_virtual_func_addresses(name)] 23 | 24 | self.rowcount = 0 25 | self.children = [] 26 | 27 | @staticmethod 28 | def create(tinfo, name, parent): 29 | result = all_virtual_functions.get(name) 30 | if result: 31 | result.parents.append(parent) 32 | return result 33 | result = VirtualMethod(tinfo, name, parent) 34 | all_virtual_functions[name] = result 35 | return result 36 | 37 | def update(self, name, tinfo): 38 | self.name = name 39 | self.tinfo = tinfo 40 | self.name_modified = False 41 | self.tinfo_modified = False 42 | 43 | def data(self, column): 44 | if column == 0: 45 | return self.name 46 | elif column == 1: 47 | return self.tinfo.get_pointed_object().dstr() 48 | elif column == 2: 49 | addresses = self.addresses 50 | if len(addresses) > 1: 51 | return "LIST" 52 | elif len(addresses) == 1: 53 | return helper.to_hex(addresses[0]) 54 | 55 | def setData(self, column, value): 56 | if column == 0: 57 | if idaapi.is_ident(value) and self.name != value: 58 | self.name = value 59 | self.name_modified = True 60 | for parent in self.parents: 61 | parent.modified = True 62 | return True 63 | elif column == 1: 64 | tinfo = idaapi.tinfo_t() 65 | split = value.split('(') 66 | if len(split) == 2: 67 | value = split[0] + ' ' + self.name + '(' + split[1] + ';' 68 | if idaapi.parse_decl(tinfo, idaapi.cvar.idati, value, idaapi.PT_TYP) is not None: 69 | if tinfo.is_func(): 70 | tinfo.create_ptr(tinfo) 71 | if tinfo.dstr() != self.tinfo.dstr(): 72 | self.tinfo = tinfo 73 | self.tinfo_modified = True 74 | for parent in self.parents: 75 | parent.modified = True 76 | return True 77 | return False 78 | 79 | def font(self, column): 80 | if column == 0 and self.name_modified: 81 | return QtGui.QFont("Consolas", 10, italic=True) 82 | elif column == 1 and self.tinfo_modified: 83 | return QtGui.QFont("Consolas", 10, italic=True) 84 | return QtGui.QFont("Consolas", 10, 0) 85 | 86 | def flags(self, column): 87 | if column != 2: 88 | if len(self.addresses) == 1: 89 | # Virtual function has only one address. Allow to modify its signature and name 90 | return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsEditable 91 | return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled 92 | 93 | 94 | @property 95 | def color(self): 96 | # return QtGui.QBrush(QtGui.QColor("#ffffb3")) 97 | return QtGui.QColor("#fefbd8") 98 | 99 | @property 100 | def tooltip(self): 101 | return None 102 | 103 | @property 104 | def addresses(self): 105 | image_base = idaapi.get_imagebase() 106 | return [ra + image_base for ra in self.ra_addresses] 107 | 108 | def set_first_argument_type(self, name): 109 | func_data = idaapi.func_type_data_t() 110 | func_tinfo = self.tinfo.get_pointed_object() 111 | class_tinfo = idaapi.tinfo_t() 112 | if func_tinfo.get_func_details(func_data) and func_tinfo.get_nargs() and \ 113 | class_tinfo.get_named_type(idaapi.cvar.idati, name): 114 | class_tinfo.create_ptr(class_tinfo) 115 | first_arg_tinfo = func_data[0].type 116 | if (first_arg_tinfo.is_ptr() and first_arg_tinfo.get_pointed_object().is_udt()) or \ 117 | helper.is_legal_type(func_data[0].type): 118 | func_data[0].type = class_tinfo 119 | func_data[0].name = "this" 120 | func_tinfo.create_func(func_data) 121 | func_tinfo.create_ptr(func_tinfo) 122 | if func_tinfo.dstr() != self.tinfo.dstr(): 123 | self.tinfo = func_tinfo 124 | self.tinfo_modified = True 125 | for parent in self.parents: 126 | parent.modified = True 127 | else: 128 | print("[Warning] function {0} probably have wrong type".format(self.name)) 129 | 130 | def open_function(self): 131 | addresses = self.addresses 132 | if len(addresses) > 1: 133 | address = helper.choose_virtual_func_address(self.name) 134 | if not address: 135 | return 136 | elif len(addresses) == 1: 137 | address = addresses[0] 138 | else: 139 | return 140 | 141 | if helper.decompile_function(address): 142 | idaapi.open_pseudocode(address, 0) 143 | else: 144 | idaapi.jumpto(address) 145 | 146 | def commit(self): 147 | addresses = self.addresses 148 | if self.name_modified: 149 | self.name_modified = False 150 | if len(addresses) == 1: 151 | idaapi.set_name(addresses[0], self.name) 152 | if self.tinfo_modified: 153 | self.tinfo_modified = False 154 | if len(addresses) == 1: 155 | idaapi.apply_tinfo(addresses[0], self.tinfo.get_pointed_object(), idaapi.TINFO_DEFINITE) 156 | 157 | def __eq__(self, other): 158 | return self.addresses == other.addresses 159 | 160 | def __repr__(self): 161 | return self.name 162 | 163 | 164 | class VirtualTable(object): 165 | def __init__(self, ordinal, tinfo, class_): 166 | self.ordinal = ordinal 167 | self.tinfo = tinfo 168 | self.class_ = [class_] 169 | self.class_name = None 170 | self.virtual_functions = [] 171 | self.name = self.tinfo.dstr() 172 | self._modified = False 173 | 174 | def update(self): 175 | if self.modified: 176 | vtable_tinfo = idaapi.tinfo_t() 177 | udt_data = idaapi.udt_type_data_t() 178 | vtable_tinfo.get_numbered_type(idaapi.cvar.idati, self.ordinal) 179 | vtable_tinfo.get_udt_details(udt_data) 180 | self.tinfo = vtable_tinfo 181 | self.name = vtable_tinfo.dstr() 182 | self.modified = False 183 | if len(self.virtual_functions) == len(udt_data): 184 | for current_function, other_function in zip(self.virtual_functions, udt_data): 185 | current_function.update(other_function.name, other_function.type) 186 | else: 187 | print("[ERROR] Something have been modified in Local types. Please refresh this view") 188 | 189 | def update_local_type(self): 190 | if self.modified: 191 | final_tinfo = idaapi.tinfo_t() 192 | udt_data = idaapi.udt_type_data_t() 193 | self.tinfo.get_udt_details(udt_data) 194 | if len(udt_data) == len(self.virtual_functions): 195 | for udt_member, virtual_function in zip(udt_data, self.virtual_functions): 196 | udt_member.name = virtual_function.name 197 | udt_member.type = virtual_function.tinfo 198 | virtual_function.commit() 199 | final_tinfo.create_udt(udt_data, idaapi.BTF_STRUCT) 200 | final_tinfo.set_numbered_type(idaapi.cvar.idati, self.ordinal, idaapi.NTF_REPLACE, self.name) 201 | self.modified = False 202 | else: 203 | print("[ERROR] Something have been modified in Local types. Please refresh this view") 204 | 205 | def set_first_argument_type(self, class_name): 206 | for function in self.virtual_functions: 207 | function.set_first_argument_type(class_name) 208 | 209 | @property 210 | def modified(self): 211 | return self._modified 212 | 213 | @modified.setter 214 | def modified(self, value): 215 | self._modified = value 216 | if value: 217 | for class_ in self.class_: 218 | class_.modified = True 219 | 220 | @staticmethod 221 | def create(tinfo, class_): 222 | ordinal = idaapi.get_type_ordinal(idaapi.cvar.idati, tinfo.dstr()) 223 | if ordinal == 0: 224 | if idaapi.import_type(idaapi.cvar.idati, -1, tinfo.dstr(), 0) == idaapi.BADNODE: 225 | raise ImportError("unable to import type to idb ({})".format(tinfo.dstr())) 226 | ordinal = idaapi.get_type_ordinal(idaapi.cvar.idati, tinfo.dstr()) 227 | 228 | result = all_virtual_tables.get(ordinal) 229 | if result: 230 | result.class_.append(class_) 231 | else: 232 | udt_data = idaapi.udt_type_data_t() 233 | tinfo.get_udt_details(udt_data) 234 | result = VirtualTable(ordinal, tinfo, class_) 235 | virtual_functions = [VirtualMethod.create(func.type, func.name, result) for func in udt_data] 236 | result.virtual_functions = virtual_functions 237 | all_virtual_functions[ordinal] = result 238 | return result 239 | 240 | def get_class_tinfo(self): 241 | if len(self.class_) == 1: 242 | return self.class_.tinfo 243 | 244 | def setData(self, column, value): 245 | if column == 0: 246 | if idaapi.is_ident(value) and self.name != value: 247 | self.name = value 248 | self.modified = True 249 | return True 250 | return False 251 | 252 | def data(self, column): 253 | if column == 0: 254 | return self.name 255 | 256 | @property 257 | def color(self): 258 | return QtGui.QColor("#d5f4e6") 259 | 260 | @property 261 | def tooltip(self): 262 | pass 263 | 264 | def font(self, column): 265 | if self.modified: 266 | return QtGui.QFont("Consolas", 12, italic=True) 267 | return QtGui.QFont("Consolas", 12) 268 | 269 | def flags(self, column): 270 | if column == 0: 271 | return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEditable | QtCore.Qt.ItemIsEnabled 272 | return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled 273 | 274 | @property 275 | def children(self): 276 | return self.virtual_functions 277 | 278 | def __repr__(self): 279 | return str(self.virtual_functions) 280 | 281 | 282 | class Class(object): 283 | def __init__(self, name, tinfo, ordinal): 284 | self.name = name 285 | self.class_name = name 286 | self.ordinal = ordinal 287 | self.parent = None 288 | self.vtables = {} 289 | self.modified = False 290 | 291 | @staticmethod 292 | def create_class(ordinal): 293 | tinfo = idaapi.tinfo_t() 294 | tinfo.get_numbered_type(idaapi.cvar.idati, ordinal) 295 | vtables = {} 296 | if tinfo.is_struct(): 297 | udt_data = idaapi.udt_type_data_t() 298 | tinfo.get_udt_details(udt_data) 299 | for field_udt in udt_data: 300 | if field_udt.type.is_ptr(): 301 | possible_vtable = field_udt.type.get_pointed_object() 302 | if possible_vtable.is_struct(): 303 | v_udt_data = idaapi.udt_type_data_t() 304 | possible_vtable.get_udt_details(v_udt_data) 305 | for possible_func_udt in v_udt_data: 306 | if not possible_func_udt.type.is_funcptr(): 307 | break 308 | else: 309 | vtables[field_udt.offset // 8] = possible_vtable 310 | if vtables: 311 | class_ = Class(tinfo.dstr(), tinfo, ordinal) 312 | for offset, vtable_tinfo in vtables.items(): 313 | vtables[offset] = VirtualTable.create(vtable_tinfo, class_) 314 | class_.vtables = vtables 315 | return class_ 316 | 317 | def update_from_local_types(self): 318 | try: 319 | if self.modified: 320 | class_ = self.create_class(self.ordinal) 321 | if class_: 322 | self.name = class_.name 323 | self.modified = False 324 | for offset, vtable in class_.vtables.items(): 325 | self.vtables[offset].update() 326 | else: 327 | # TODO: drop class 328 | raise IndexError 329 | except IndexError: 330 | print("[ERROR] Something have been modified in Local types. Please refresh this view") 331 | 332 | def update_local_type(self): 333 | if self.modified: 334 | for vtable in list(self.vtables.values()): 335 | vtable.update_local_type() 336 | udt_data = idaapi.udt_type_data_t() 337 | tinfo = idaapi.tinfo_t() 338 | self.tinfo.get_udt_details(udt_data) 339 | tinfo.create_udt(udt_data, idaapi.BTF_STRUCT) 340 | tinfo.set_numbered_type(idaapi.cvar.idati, self.ordinal, idaapi.NTF_REPLACE, self.name) 341 | self.modified = False 342 | 343 | def set_first_argument_type(self, class_name): 344 | if 0 in self.vtables: 345 | self.vtables[0].set_first_argument_type(class_name) 346 | 347 | def has_function(self, regexp): 348 | for vtable in list(self.vtables.values()): 349 | if [func for func in vtable.virtual_functions if regexp.indexIn(func.name) >= 0]: 350 | return True 351 | return False 352 | 353 | def data(self, column): 354 | if column == 0: 355 | return self.name 356 | 357 | def setData(self, column, value): 358 | if column == 0: 359 | if idaapi.is_ident(value) and self.name != value: 360 | self.name = value 361 | self.modified = True 362 | return True 363 | return False 364 | 365 | def font(self, column): 366 | if self.modified: 367 | return QtGui.QFont("Consolas", 12, QtGui.QFont.Bold, italic=True) 368 | return QtGui.QFont("Consolas", 12, QtGui.QFont.Bold) 369 | 370 | @property 371 | def color(self): 372 | # return QtGui.QBrush(QtGui.QColor("#ffb3ff")): 373 | return QtGui.QColor("#80ced6") 374 | 375 | @property 376 | def tooltip(self): 377 | return None 378 | 379 | def flags(self, column): 380 | if column == 0: 381 | return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEditable | QtCore.Qt.ItemIsEnabled 382 | return QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled 383 | 384 | @property 385 | def children(self): 386 | return list(self.vtables.values()) 387 | 388 | @property 389 | def tinfo(self): 390 | tinfo = idaapi.tinfo_t() 391 | tinfo.get_numbered_type(idaapi.cvar.idati, self.ordinal) 392 | return tinfo 393 | 394 | def open_function(self): 395 | pass 396 | 397 | def __repr__(self): 398 | return self.name + " ^_^ " + str(self.vtables) 399 | 400 | 401 | class TreeItem(object): 402 | def __init__(self, item, parent=None): 403 | self.item = item 404 | self.parent = parent 405 | self.children = [] 406 | 407 | def __repr__(self): 408 | return str(self.item) 409 | 410 | def appendChild(self, item): 411 | self.children.append(item) 412 | 413 | def child(self, row): 414 | return self.children[row] 415 | 416 | def childCount(self): 417 | return len(self.children) 418 | 419 | def columnCount(self): 420 | return len(self.item) 421 | 422 | def data(self, column): 423 | try: 424 | return self.item[column] 425 | except IndexError: 426 | return None 427 | 428 | def row(self): 429 | if self.parent: 430 | return self.parent.children.index(self) 431 | 432 | return 0 433 | 434 | 435 | class TreeModel(QtCore.QAbstractItemModel): 436 | # TODO: Add higlighting if eip in function, consider setting breakpoints 437 | 438 | refreshed = QtCore.pyqtSignal() 439 | 440 | def __init__(self, parent=None): 441 | super(TreeModel, self).__init__(parent) 442 | 443 | self.rootItem = TreeItem(("Name", "Declaration", "Address")) 444 | self.setupModelData(self.rootItem) 445 | 446 | def setupModelData(self, root): 447 | idaapi.show_wait_box("Looking for classes...") 448 | 449 | all_virtual_functions.clear() 450 | all_virtual_tables.clear() 451 | 452 | classes = [] 453 | for ordinal in range(1, idaapi.get_ordinal_qty(idaapi.cvar.idati)): 454 | result = Class.create_class(ordinal) 455 | if result: 456 | classes.append(result) 457 | 458 | for class_ in classes: 459 | class_item = TreeItem(class_, root) 460 | for vtable in class_.vtables.values(): 461 | vtable_item = TreeItem(vtable, class_item) 462 | vtable_item.children = [TreeItem(function, vtable_item) for function in vtable.virtual_functions] 463 | class_item.appendChild(vtable_item) 464 | root.appendChild(class_item) 465 | 466 | idaapi.hide_wait_box() 467 | 468 | def flags(self, index): 469 | if index.isValid(): 470 | return index.internalPointer().item.flags(index.column()) 471 | else: 472 | return QtCore.Qt.NoItemFlags 473 | 474 | def index(self, row, column, parent): 475 | if not self.hasIndex(row, column, parent): 476 | return QtCore.QModelIndex() 477 | 478 | if not parent.isValid(): 479 | parentItem = self.rootItem 480 | else: 481 | parentItem = parent.internalPointer() 482 | 483 | childItem = parentItem.child(row) 484 | if childItem: 485 | return self.createIndex(row, column, childItem) 486 | else: 487 | return QtCore.QModelIndex() 488 | 489 | def parent(self, index): 490 | if not index.isValid(): 491 | return QtCore.QModelIndex() 492 | 493 | childItem = index.internalPointer() 494 | parentItem = childItem.parent 495 | 496 | if parentItem == self.rootItem: 497 | return QtCore.QModelIndex() 498 | 499 | return self.createIndex(parentItem.row(), 0, parentItem) 500 | 501 | def rowCount(self, parent): 502 | if parent.column() > 0: 503 | return 0 504 | 505 | if not parent.isValid(): 506 | parentItem = self.rootItem 507 | else: 508 | parentItem = parent.internalPointer() 509 | 510 | return parentItem.childCount() 511 | 512 | def columnCount(self, parent): 513 | return self.rootItem.columnCount() 514 | 515 | def data(self, index, role=QtCore.Qt.DisplayRole): 516 | if not index.isValid(): 517 | return None 518 | 519 | if role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole: 520 | node = index.internalPointer() 521 | return node.item.data(index.column()) 522 | elif role == QtCore.Qt.FontRole: 523 | return index.internalPointer().item.font(index.column()) 524 | elif role == QtCore.Qt.ToolTipRole: 525 | return index.internalPointer().item.tooltip 526 | elif role == QtCore.Qt.BackgroundRole: 527 | return index.internalPointer().item.color 528 | elif role == QtCore.Qt.ForegroundRole: 529 | return QtGui.QBrush(QtGui.QColor("#191919")) 530 | 531 | return None 532 | 533 | def setData(self, index, value, role=QtCore.Qt.EditRole): 534 | result = False 535 | value = str(value) 536 | if role == QtCore.Qt.EditRole and value: 537 | node = index.internalPointer() 538 | result = node.item.setData(index.column(), value) 539 | 540 | return result 541 | 542 | def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole): 543 | if role == QtCore.Qt.DisplayRole and orientation == QtCore.Qt.Horizontal: 544 | return self.rootItem.data(section) 545 | 546 | return None 547 | 548 | def set_first_argument_type(self, indexes): 549 | indexes = [x for x in indexes if x.column() == 0] 550 | class_name = indexes[0].internalPointer().item.class_name 551 | if not class_name: 552 | classes = [[x.item.name] for x in self.rootItem.children] 553 | class_chooser = HexRaysPyTools.forms.MyChoose(classes, "Select Class", [["Name", 25]]) 554 | idx = class_chooser.Show(True) 555 | if idx != -1: 556 | class_name = classes[idx][0] 557 | if class_name: 558 | for index in indexes: 559 | index.internalPointer().item.set_first_argument_type(class_name) 560 | 561 | def refresh(self): 562 | self.beginResetModel() 563 | self.rootItem.children = [] 564 | self.endResetModel() 565 | 566 | self.layoutAboutToBeChanged.emit() 567 | self.setupModelData(self.rootItem) 568 | self.layoutChanged.emit() 569 | 570 | self.refreshed.emit() 571 | 572 | def rollback(self): 573 | self.layoutAboutToBeChanged.emit() 574 | for class_item in self.rootItem.children: 575 | class_item.item.update_from_local_types() 576 | self.layoutChanged.emit() 577 | 578 | def commit(self): 579 | for class_item in self.rootItem.children: 580 | if class_item.item.modified: 581 | class_item.item.update_local_type() 582 | 583 | def open_function(self, index): 584 | item = index.internalPointer().item 585 | if isinstance(item, VirtualMethod): 586 | index.internalPointer().item.open_function() 587 | 588 | 589 | class ProxyModel(QtCore.QSortFilterProxyModel): 590 | def __init__(self): 591 | super(ProxyModel, self).__init__() 592 | self.filter_by_function = False 593 | 594 | def set_regexp_filter(self, regexp): 595 | if regexp and regexp[0] == '!': 596 | self.filter_by_function = True 597 | self.setFilterRegExp(regexp[1:]) 598 | else: 599 | self.filter_by_function = False 600 | self.setFilterRegExp(regexp) 601 | 602 | def filterAcceptsRow(self, row, parent): 603 | filter_regexp = self.filterRegExp() 604 | if filter_regexp: 605 | index = self.sourceModel().index(row, 0, parent) 606 | item = index.internalPointer().item 607 | 608 | if self.filter_by_function and isinstance(item, Class): 609 | return item.has_function(filter_regexp) 610 | else: 611 | return filter_regexp.indexIn(item.class_name) >= 0 612 | 613 | return True 614 | --------------------------------------------------------------------------------