# ----------------------------------------------------------------------- # This is an example illustrating how to use custom data types in Python # (c) Hex-Rays # from idaapi import data_type_t, data_format_t, NW_OPENIDB, NW_CLOSEIDB, NW_TERMIDA, NW_REMOVE, COLSTR import struct import ctypes import platform # # ----------------------------------------------------------------------- class pascal_data_type(data_type_t): def __init__(self): data_type_t.__init__(self, name="py_pascal_string", value_size = 2, menu_name = "Pascal string", asm_keyword = "pstr") def calc_item_size(self, ea, maxsize): # Custom data types may be used in structure definitions. If this case # ea is a member id. Check for this situation and return 1 if _idaapi.is_member_id(ea): return 1 # get the length byte n = _idaapi.get_byte(ea) # string too big? if n > maxsize: return 0 # ok, accept the string return n + 1 class pascal_data_format(data_format_t): FORMAT_NAME = "py_pascal_string_pstr" def __init__(self): data_format_t.__init__(self, name=pascal_data_format.FORMAT_NAME) def printf(self, value, current_ea, operand_num, dtid): # Take the length byte n = ord(value[0]) o = ['"'] for ch in value[1:]: b = ord(ch) if b < 0x20 or b > 128: o.append(r'\x%02x' % ord(ch)) else: o.append(ch) o.append('"') return "".join(o) # ----------------------------------------------------------------------- class simplevm_data_type(data_type_t): ASM_KEYWORD = "svm_emit" def __init__(self): data_type_t.__init__(self, name="py_simple_vm", value_size = 1, menu_name = "SimpleVM", asm_keyword = simplevm_data_type.ASM_KEYWORD) def calc_item_size(self, ea, maxsize): if _idaapi.is_member_id(ea): return 1 # get the opcode and see if it has an imm n = 5 if (_idaapi.get_byte(ea) & 3) == 0 else 1 # string too big? if n > maxsize: return 0 # ok, accept return n class simplevm_data_format(data_format_t): def __init__(self): data_format_t.__init__(self, name="py_simple_vm_format", menu_name = "SimpleVM") # Some tables for the disassembler INST = {1: 'add', 2: 'mul', 3: 'sub', 4: 'xor', 5: 'mov'} REGS = {1: 'r1', 2: 'r2', 3: 'r3'} def disasm(self, inst): """A simple local disassembler. In reality one can use a full-blown disassembler to render the text""" opbyte = ord(inst[0]) op = opbyte >> 4 if not (1<=op<=5): return None r1 = (opbyte & 0xf) >> 2 r2 = opbyte & 3 sz = 0 if r2 == 0: if len(inst) != 5: return None imm = struct.unpack_from('L', inst, 1)[0] sz = 5 else: imm = None sz = 1 text = "%s %s, %s" % ( COLSTR(simplevm_data_format.INST[op], idaapi.SCOLOR_INSN), COLSTR(simplevm_data_format.REGS[r1], idaapi.SCOLOR_REG), COLSTR("0x%08X" % imm, idaapi.SCOLOR_NUMBER) if imm is not None else COLSTR(simplevm_data_format.REGS[r2], idaapi.SCOLOR_REG)) return (sz, text) def printf(self, value, current_ea, operand_num, dtid): r = self.disasm(value) if not r: return None if dtid == 0: return "%s(%s)" % (simplevm_data_type.ASM_KEYWORD, r[1]) return r[1] # ----------------------------------------------------------------------- # This format will display DWORD values as MAKE_DWORD(0xHI, 0xLO) class makedword_data_format(data_format_t): def __init__(self): data_format_t.__init__(self, name="py_makedword", value_size = 4, menu_name = "Make DWORD") def printf(self, value, current_ea, operand_num, dtid): if len(value) != 4: return None w1 = struct.unpack_from("H", value, 0)[0] w2 = struct.unpack_from("H", value, 2)[0] return "MAKE_DWORD(0x%04X, 0x%04X)" % (w2, w1) # ----------------------------------------------------------------------- # This format will try to load a resource string given a number # So instead of displaying: # push 66h # call message_box_from_rsrc_string # It can be rendered as; # push RSRC("The message") # call message_box_from_rsrc_string # # The get_rsrc_string() is not optimal since it loads/unloads the # DLL each time for a new string. It can be improved in many ways. class rsrc_string_format(data_format_t): def __init__(self): data_format_t.__init__(self, name="py_w32rsrcstring", value_size = 1, menu_name = "Resource string") self.cache_node = idaapi.netnode("$ py_w32rsrcstring", 0, 1) def get_rsrc_string(self, fn, id): """ Simple method that loads the input file as a DLL with LOAD_LIBRARY_AS_DATAFILE flag. It then tries to LoadString() """ k32 = ctypes.windll.kernel32 u32 = ctypes.windll.user32 hinst = k32.LoadLibraryExA(fn, 0, 0x2) if hinst == 0: return "" buf = ctypes.create_string_buffer(1024) r = u32.LoadStringA(hinst, id, buf, 1024-1) k32.FreeLibrary(hinst) return buf.value if r else "" def printf(self, value, current_ea, operand_num, dtid): # Is it already cached? val = self.cache_node.supval(current_ea) # Not cached? if val == None: # Retrieve it num = idaapi.struct_unpack(value) val = self.get_rsrc_string(idaapi.get_input_file_path(), num) # Cache it self.cache_node.supset(current_ea, val) # Failed to retrieve? if val == "" or val == "\x00": return None # Return the format return "RSRC_STR(\"%s\")" % COLSTR(val, idaapi.SCOLOR_IMPNAME) # ----------------------------------------------------------------------- # Table of formats and types to be registered/unregistered # If a tuple has one element then it is the format to be registered with dtid=0 # If the tuple has more than one element, the tuple[0] is the data type and tuple[1:] are the data formats new_formats = [ (pascal_data_type(), pascal_data_format()), (simplevm_data_type(), simplevm_data_format()), (makedword_data_format(),), (simplevm_data_format(),) ] if platform.system() == 'Windows': new_formats.append((rsrc_string_format(),)) # # ----------------------------------------------------------------------- def nw_handler(code, old=0): # delete notifications if code == NW_OPENIDB: idaapi.register_data_types_and_formats(new_formats) elif code == NW_CLOSEIDB: idaapi.unregister_data_types_and_formats(new_formats) elif code == NW_TERMIDA: idaapi.notify_when(NW_TERMIDA | NW_OPENIDB | NW_CLOSEIDB | NW_REMOVE, nw_handler) # ----------------------------------------------------------------------- # Check if already installed if idaapi.find_custom_data_type(pascal_data_format.FORMAT_NAME) == -1: if not idaapi.register_data_types_and_formats(new_formats): print "Failed to register types!" else: idaapi.notify_when(NW_TERMIDA | NW_OPENIDB | NW_CLOSEIDB, nw_handler) print "Formats installed!" else: print "Formats already installed!"