├── .gitignore ├── CustomTypeAssign.py ├── CustomTypeAssign ├── __init__.py ├── actions.py └── callbacks.py └── README.md /.gitignore: -------------------------------------------------------------------------------- 1 | /.idea/ 2 | -------------------------------------------------------------------------------- /CustomTypeAssign.py: -------------------------------------------------------------------------------- 1 | import ida_idaapi 2 | import ida_hexrays 3 | import logging 4 | import ida_kernwin 5 | import ida_struct 6 | import ida_typeinf 7 | from collections import namedtuple 8 | import re 9 | 10 | DEBUG = False 11 | if DEBUG: 12 | import pydevd_pycharm 13 | 14 | from CustomTypeAssign.actions import HexRaysPopupAction, action_manager 15 | from CustomTypeAssign.callbacks import hx_callback_manager 16 | 17 | RecastLocalVariable = namedtuple('RecastLocalVariable', ['curr_ti', 'local_variable']) 18 | RecastGlobalVariable = namedtuple('RecastGlobalVariable', ['curr_ti', 'global_variable_ea']) 19 | RecastArgument = namedtuple('RecastArgument', ['curr_ti', 'arg_idx', 'func_ea', 'func_tinfo']) 20 | RecastReturn = namedtuple('RecastReturn', ['curr_ti', 'func_ea']) 21 | RecastStructure = namedtuple('RecastStructure', ['curr_ti', 'structure_name', 'field_offset']) 22 | 23 | prefixes = ['unsigned', 'signed'] 24 | 25 | def parse_decl(tif,inputtype,flags=0,til=None): 26 | if len(inputtype) != 0 and inputtype[-1] != ';': 27 | inputtype = inputtype + ';' 28 | 29 | if ida_typeinf.parse_decl(tif,til,inputtype,flags) is not None: 30 | return True 31 | return False 32 | 33 | def parse_input_type(decl): 34 | pointer_count = 0 35 | type_name = decl 36 | tail = None 37 | tif = ida_typeinf.tinfo_t() 38 | if ' ' in decl: 39 | match = re.search(r"(.*<.*>)(.*)",decl) 40 | if match: 41 | tail = match.groups()[1].strip() 42 | type_name = match.groups()[0].strip() 43 | elif decl.count(' ') == 1: 44 | type_name, tail = decl.split("") 45 | else: 46 | ida_kernwin.warning("During type applying error occured.\nIn type declaration '%s' too many spaces symbols."%decl) 47 | type_name = None 48 | 49 | if tail: 50 | pointer_count += tail.count('*') 51 | if type_name: 52 | while type_name[-1] == '*': 53 | pointer_count += 1 54 | type_name = type_name[:-1] 55 | if tif.get_named_type(None, type_name): 56 | for i in range(pointer_count): 57 | tif.create_ptr(tif) 58 | return tif 59 | 60 | 61 | class CustomTypeAssignHandler(HexRaysPopupAction): 62 | 63 | hotkey = "shift-y" 64 | description = "Set custom type..." 65 | 66 | def activate(self, ctx): # type: (ida_kernwin.action_activation_ctx_t) -> None 67 | if DEBUG: 68 | pydevd_pycharm.settrace('localhost', port=3333, stdoutToServer=True, stderrToServer=True, suspend=False) 69 | vdui = ida_hexrays.get_widget_vdui(ctx.widget) 70 | vdui.get_current_item(ida_hexrays.USE_KEYBOARD) 71 | target = vdui.item 72 | if target: 73 | cfunc = vdui.cfunc 74 | ri = self.get_target_type(cfunc,target) 75 | if ri: 76 | type_decl = ida_kernwin.ask_str("%s"%ri.curr_ti.dstr(),10,"Please enter the type declaration") 77 | if type_decl and type_decl != ri.curr_ti.dstr(): 78 | tif = ida_typeinf.tinfo_t() 79 | if not parse_decl(tif,type_decl,ida_typeinf.PT_SIL): 80 | type_decl = type_decl.strip() 81 | tif = parse_input_type(type_decl) 82 | if not tif.empty(): 83 | if isinstance(ri, RecastLocalVariable): 84 | vdui.set_lvar_type(ri.local_variable, tif) 85 | elif isinstance(ri, RecastGlobalVariable): 86 | ida_typeinf.apply_tinfo(ri.global_variable_ea, tif, ida_typeinf.TINFO_DEFINITE) 87 | elif isinstance(ri,RecastStructure): 88 | sid = ida_struct.get_struc_id(ri.structure_name) 89 | if sid != ida_idaapi.BADADDR: 90 | sptr = ida_struct.get_struc(sid) 91 | mptr = ida_struct.get_member(sptr, ri.field_offset) 92 | if mptr: 93 | rc = ida_struct.set_member_tinfo(sptr, mptr, ri.field_offset, tif, 94 | ida_struct.SET_MEMTI_MAY_DESTROY) 95 | if rc != 1: 96 | print("set_member_tinfo rc = %d" % rc) 97 | vdui.refresh_view(True) 98 | else: 99 | ida_kernwin.warning("During type applying error occured.\nInput type declaration: '%s'." % type_decl) 100 | 101 | def check(self, hx_view): # type: (ida_hexrays.vdui_t) -> bool 102 | if DEBUG: 103 | pydevd_pycharm.settrace('localhost', port=3333, stdoutToServer=True, stderrToServer=True, suspend=False) 104 | hx_view.get_current_item(ida_hexrays.USE_KEYBOARD) 105 | target = hx_view.item 106 | cfunc = hx_view.cfunc 107 | if self.get_target_type(cfunc,target) is not None: 108 | return True 109 | return False 110 | 111 | def get_target_type(self,cfunc,ctree_item): 112 | # type: (ida_hexrays.cfunc_t,ida_hexrays.ctree_item_t) -> namedtuple 113 | if ctree_item.citype == ida_hexrays.VDI_EXPR and ctree_item.it.is_expr(): 114 | expression = ctree_item.e 115 | if expression.opname not in ('var', 'obj', 'memptr', 'memref'): 116 | return None 117 | 118 | if expression.op == ida_hexrays.cot_var: 119 | variable = cfunc.get_lvars()[expression.v.idx] 120 | curr_ti = variable.tif 121 | return RecastLocalVariable(curr_ti,variable) 122 | 123 | elif expression.op == ida_hexrays.cot_obj: 124 | return RecastGlobalVariable(expression.type,expression.obj_ea) 125 | 126 | elif expression.op == ida_hexrays.cot_memptr: 127 | struct_name = expression.x.type.get_pointed_object().dstr() 128 | struct_offset = expression.m 129 | return RecastStructure(expression.type, struct_name, struct_offset) 130 | 131 | elif expression.op == ida_hexrays.cot_memref: 132 | struct_name = expression.x.type.dstr() 133 | struct_offset = expression.m 134 | return RecastStructure(expression.type, struct_name, struct_offset) 135 | 136 | elif ctree_item.citype == ida_hexrays.VDI_LVAR: 137 | variable = ctree_item.l 138 | curr_ti = variable.tif 139 | return RecastLocalVariable(curr_ti, variable) 140 | return None 141 | 142 | 143 | def register_actions(): 144 | pass 145 | 146 | 147 | class CustomTypeAssign(ida_idaapi.plugin_t): 148 | flags = 0 149 | comment = "Type assign for HexRays without \"bad\" characters filtration" 150 | help = "Type assign for HexRays without \"bad\" characters filtration" 151 | wanted_name = "MyPlugins:CustomTypeAssign" 152 | wanted_hotkey = "" 153 | 154 | @staticmethod 155 | def init(): 156 | if DEBUG: 157 | pydevd_pycharm.settrace('localhost', port=3333, stdoutToServer=True, stderrToServer=True, suspend=False) 158 | if not ida_hexrays.init_hexrays_plugin(): 159 | logging.error("Failed to initialize Hex-Rays SDK") 160 | return ida_idaapi.PLUGIN_SKIP 161 | 162 | action_manager.register(CustomTypeAssignHandler()) 163 | action_manager.initialize() 164 | hx_callback_manager.initialize() 165 | return ida_idaapi.PLUGIN_KEEP 166 | 167 | 168 | @staticmethod 169 | def term(): 170 | hx_callback_manager.finalize() 171 | action_manager.finalize() 172 | 173 | @staticmethod 174 | def run(): 175 | pass 176 | 177 | def PLUGIN_ENTRY(): 178 | logging.basicConfig(format='[%(levelname)s] %(message)s\t(%(module)s:%(funcName)s)') 179 | logging.root.setLevel(0) 180 | return CustomTypeAssign() -------------------------------------------------------------------------------- /CustomTypeAssign/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rigmar/CustomTypeAssign/79d25720d84baf70a499f243e7ffc967bc15d066/CustomTypeAssign/__init__.py -------------------------------------------------------------------------------- /CustomTypeAssign/actions.py: -------------------------------------------------------------------------------- 1 | import idaapi 2 | 3 | from CustomTypeAssign.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 | if isinstance(action, HexRaysPopupAction): 13 | hx_callback_manager.register(idaapi.hxe_populating_popup, HexRaysPopupRequestHandler(action)) 14 | 15 | def initialize(self): 16 | for action in self.__actions: 17 | idaapi.register_action( 18 | idaapi.action_desc_t(action.name, action.description, action, action.hotkey) 19 | ) 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 "CustomTypeAssign:" + 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 | -------------------------------------------------------------------------------- /CustomTypeAssign/callbacks.py: -------------------------------------------------------------------------------- 1 | from collections import defaultdict 2 | import ida_hexrays 3 | 4 | 5 | class HexRaysCallbackManager(object): 6 | def __init__(self): 7 | self.__hexrays_event_handlers = defaultdict(list) 8 | 9 | def initialize(self): 10 | ida_hexrays.install_hexrays_callback(self.__handle) 11 | 12 | def finalize(self): 13 | ida_hexrays.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 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # CustomTypeAssign 2 | IDA pro plugin for assign type declarations with 'bad' characters. Like c++ stl templates Less/Greater brackets.\ 3 | Adds 'Set custom type...' action to right-click menu in HexRays pseudocode windows. Or can be obtained through 'shift-y' hotkey. 4 | 5 | 6 | ## Works with: 7 | - Local variables 8 | - Global variables 9 | - Struct members 10 | - Function arguments 11 | 12 | ## Not works for now with: 13 | - Function declaration like 'void __thiscall CMap_int_int_int_int___CMap_int_int_int_int_(CMap *this, int nBlockSize)' for example. 14 | --------------------------------------------------------------------------------