def convert_reg(reg_name, width): """Convert given register name to the register name for the provided width (eg: conver_reg(eax, 8) -> rax)""" reg_idx = ida_idp.str2reg(reg_name) if reg_idx > 15: # R15 is index 15. 8-bit registers have indexes above 15 but are 0 indexed, so sub 16 reg_idx -= 16 return ida_idp.get_reg_name(reg_idx, width)
def get_ymm_mreg(xmm_mreg): """ Return the YMM microcode register for a given XMM register. """ xmm_reg = ida_hexrays.mreg2reg(xmm_mreg, XMM_SIZE) xmm_name = ida_idp.get_reg_name(xmm_reg, XMM_SIZE) xmm_number = int(xmm_name.split("mm")[-1]) # compute the ymm mreg id ymm_reg = ida_idp.str2reg("ymm%u" % xmm_number) ymm_mreg = ida_hexrays.reg2mreg(ymm_reg) # sanity check... xmm_name = ida_hexrays.get_mreg_name(xmm_mreg, XMM_SIZE) ymm_name = ida_hexrays.get_mreg_name(ymm_mreg, YMM_SIZE) assert xmm_name[1:] == ymm_name[ 1:], "Reg escalation did not work... (%s, %s)" % (xmm_name, ymm_name) # return the ymm microcode register id return ymm_mreg
def reg2str(register, width=None): """Convert given register index to the register name with the provided width (eg: reg2str(0, 8) -> rax)""" if not width: width = 8 if idc.__EA64__ else 4 return ida_idp.get_reg_name(register, width)
def _expand_locations(arch, pfn, ty, argloc, out_locs): """Expand the locations referred to by `argloc` into a list of `Location`s in `out_locs`.""" reg_names = ida_idp.ph_get_regnames() where = argloc.atype() if where == ida_typeinf.ALOC_STACK: sp_adjust_retaddr = int( ida_frame.frame_off_args(pfn) - ida_frame.frame_off_retaddr(pfn) ) loc = Location() loc.set_memory(arch.stack_pointer_name(), argloc.stkoff() + sp_adjust_retaddr) loc.set_type(ty) out_locs.append(loc) # Distributed across two or more locations. elif where == ida_typeinf.ALOC_DIST: for part in argloc.scattered(): part_ty = ty.extract(arch, part.off, part.size) _expand_locations(arch, pfn, part_ty, part, out_locs) # Located in a single register, possibly in a small part of the register # itself. elif where == ida_typeinf.ALOC_REG1: ty_size = ty.size(arch) reg_name = reg_names[argloc.reg1()].upper() try: reg_offset = argloc.regoff() family = arch.register_family(reg_name) # Try to guess the right name for the register based on the size of the # type that it will contain. For example, IDA will tell us register `ax` # is used, not specify if it's `al`, `ah`, `eax`, or `rax`. # # NOTE: The registers in the family tuple are sorted in descending # order of size. found = False for f_reg_name, f_reg_offset, f_reg_size in family: if f_reg_offset != reg_offset: continue if ty_size == f_reg_size: found = True reg_name = f_reg_name break if not found: raise Exception() except: reg_name = ( ida_idp.get_reg_name(argloc.reg1(), ty_size) or reg_name ).upper() loc = Location() loc.set_register(arch.register_name(reg_name)) loc.set_type(ty) out_locs.append(loc) # Located in a pair of registers. elif where == ida_typeinf.ALOC_REG2: ty_size = ty.size(arch) reg_name1 = reg_names[argloc.reg1()].upper() reg_name2 = reg_names[argloc.reg2()].upper() ty1 = ty.extract(arch, 0, ty_size / 2) ty2 = ty.extract(arch, ty_size / 2, ty_size / 2) try: found = False family1 = arch.register_family(reg_name1) family2 = arch.register_family(reg_name2) # Try to guess which registers IDA actually meant. For example, for # an `EDX:EAX` return value, our `ty_size` will be 8 bytes, but IDA will # report the registers as `ax` and `dx` (due to those being the names in # `ph_get_regnames`). So, we have to scan through the associated family # and try to see if we can guess the right version of those registers. for r1_info, r2_info in itertools.product(family1, family2): f_reg_name1, f_reg_offset1, f_reg_size1 = r1_info f_reg_name2, f_reg_offset2, f_reg_size2 = r2_info if f_reg_offset1 or f_reg_offset2: continue if ty_size == (f_reg_size1 + f_reg_size2): found = True reg_name1 = f_reg_name1 reg_name2 = f_reg_name2 ty1 = ty.extract(arch, 0, f_reg_size1) ty2 = ty.extract(arch, f_reg_size1, f_reg_size2) break if not found: raise Exception() except Exception as e: reg_name1 = ( ida_idp.get_reg_name(argloc.reg1(), ty_size) or reg_name1 ).upper() reg_name2 = ( ida_idp.get_reg_name(argloc.reg2(), ty_size) or reg_name2 ).upper() loc1 = Location() loc1.set_register(arch.register_name(reg_name1)) loc1.set_type(ty1) out_locs.append(loc1) loc2 = Location() loc2.set_register(arch.register_name(reg_name2)) loc2.set_type(ty2) out_locs.append(loc2) # Memory location computed as value in a register, plus an offset. # # TODO(pag): How does this work if the register itself is not # treated as an argument? elif where == ida_typeinf.ALOC_RREL: rrel = argloc.get_rrel() loc = Location() loc.set_memory( _get_address_sized_reg(arch, reg_names[rrel.reg].upper()), rrel.off ) loc.set_type(ty) out_locs.append(loc) # Global variable with a fixed address. We can represent this # as computing a PC-relative memory address. elif where == ida_typeinf.ALOC_STATIC: loc = Location() loc.set_memory(arch.program_counter_name(), argloc.get_ea() - ea) loc.set_type(ty) out_locs.append(loc) # Unsupported. else: raise InvalidLocationException( "Unsupported location {} with type {}".format( str(argloc), ty.serialize(arch, {}) ) )
def value(self, value): """Sets the value of the argument to the cpu context.""" # TODO: Pull value data based on type. argloc = self._funcarg_obj.argloc loc_type = argloc.atype() # This type occurs when we have created an uninitialized argument. if loc_type == ida_typeinf.ALOC_NONE: logger.warning('Argument {} location is of type ALOC_NONE'.format( self.idx)) elif loc_type == ida_typeinf.ALOC_STACK: # read the argument from the stack using the calculated stack offset from the disassembler logger.debug( f'Setting argument {self.idx} at {hex(self.addr)}, stack: {hex(self._cpu_context.sp)}, offset: {hex(argloc.stkoff())}' ) self._cpu_context.memory.write( self.addr, utils.struct_pack(value, width=self._cpu_context.byteness)) elif loc_type == ida_typeinf.ALOC_DIST: # arguments described by multiple locations # TODO: Uses the scattered_aloc_t class, which is a qvector or argpart_t objects # argloc.scattered() raise NotImplementedError( "Argument {} location of type ALOC_DIST".format(self.idx)) elif loc_type == ida_typeinf.ALOC_REG1: # single register reg_name = ida_idp.get_reg_name(argloc.reg1(), self.width) self._cpu_context.registers[reg_name] = value elif loc_type == ida_typeinf.ALOC_REG2: # register pair (eg: edx:eax [reg2:reg1]) reg_width = self.width // 2 reg1_name = ida_idp.get_reg_name(argloc.reg1(), reg_width) reg2_name = ida_idp.get_reg_name(argloc.reg2(), reg_width) self._cpu_context.registers[reg1_name] = value & utils.get_mask( reg_width) self._cpu_context.registers[reg2_name] = value >> (reg_width * 8) elif loc_type == ida_typeinf.ALOC_RREL: # register relative (displacement from address pointed by register # TODO: CURRENTLY UNTESTED logger.info("Argument {} of untested type ALOC_RREL. " "Verify results and report issues.".format(self.idx)) # Obtain the register-relative argument location rrel = argloc.get_rrel() reg_name = ida_idp.get_reg_name(rrel.reg, self.width) self._cpu_context.registers[reg_name] = value - rrel.offset elif loc_type == ida_typeinf.ALOC_STATIC: # global address # TODO: CURRENTLY UNTESTED logger.info("Argument {} of untested type ALOC_STATIC. " "Verify results and report issues.".format(self.idx)) # return argloc.get_ea() raise NotImplementedError( f"Argument {self.idx} location of type ALOC_STATIC") elif loc_type >= ida_typeinf.ALOC_CUSTOM: # custom argloc # TODO: Will need to figure out the functionality and usage for the custloc_desc_t structure # argloc.get_custom() raise NotImplementedError( f"Argument {self.idx} location of type ALOC_CUSTOM")
def value(self): """Retrieves the value of the argument based on the cpu context.""" # TODO: Pull value data based on type. argloc = self._funcarg_obj.argloc loc_type = argloc.atype() # This type occurs when we have created an uninitialized argument. if loc_type == ida_typeinf.ALOC_NONE: logger.warning("Argument {} location is of type ALOC_NONE".format( self.idx)) return None elif loc_type == ida_typeinf.ALOC_STACK: # Get function logger.debug( f'Retrieving argument {self.idx} at {hex(self.addr)}, stack: {hex(self._cpu_context.sp)}, offset: {hex(argloc.stkoff())}' ) # read the argument from the stack using the calculated stack offset from the disassembler value = self._cpu_context.memory.read(self.addr, self._cpu_context.byteness) return utils.struct_unpack(value) elif loc_type == ida_typeinf.ALOC_DIST: # arguments described by multiple locations # TODO: Uses the scattered_aloc_t class, which is a qvector or argpart_t objects # argloc.scattered() raise NotImplementedError( "Argument {} location of type ALOC_DIST".format(self.idx)) elif loc_type == ida_typeinf.ALOC_REG1: # single register # TODO: Determine better way to convert reg1 integer to name. reg_name = ida_idp.get_reg_name(argloc.reg1(), self.width) return self._cpu_context.registers[reg_name] elif loc_type == ida_typeinf.ALOC_REG2: # register pair (eg: edx:eax [reg2:reg1]) # Width is the combination of both registers. reg_width = self.width // 2 reg1_name = ida_idp.get_reg_name(argloc.reg1(), reg_width) reg2_name = ida_idp.get_reg_name(argloc.reg2(), reg_width) logger.debug("Register pair: [%s:%s]", reg1_name, reg2_name) reg1_value = self._cpu_context.registers[reg1_name] reg2_value = self._cpu_context.registers[reg2_name] return (reg2_value << (reg_width * 8)) | reg1_value elif loc_type == ida_typeinf.ALOC_RREL: # register relative (displacement from address pointed by register # TODO: CURRENTLY UNTESTED logger.info("Argument {} of untested type ALOC_RREL. " "Verify results and report issues.".format(self.idx)) # Obtain the register-relative argument location rrel = argloc.get_rrel() reg_name = ida_idp.get_reg_name(rrel.reg, self.width) value = self._cpu_context.registers[reg_name] return value + rrel.offset elif loc_type == ida_typeinf.ALOC_STATIC: # global address # TODO: CURRENTLY UNTESTED logger.info("Argument {} of untested type ALOC_STATIC. " "Verify results and report issues.".format(self.idx)) return argloc.get_ea() elif loc_type >= ida_typeinf.ALOC_CUSTOM: # custom argloc # TODO: Will need to figure out the functionality and usage for the custloc_desc_t structure # argloc.get_custom() raise NotImplementedError( "Argument {} location of type ALOC_CUSTOM".format(self.idx))
def value(self): # TODO: Pull value data based on type. argloc = self._funcarg_obj.argloc loc_type = argloc.atype() # This type occurs when we have created an uninitialized argument. if loc_type == ida_typeinf.ALOC_NONE: logger.warning('Argument {} location is of type ALOC_NONE'.format(self.idx)) return None elif loc_type == ida_typeinf.ALOC_STACK: # read the argument from the stack using the calculated stack offset from the disassembler cur_esp = self._cpu_context.sp + argloc.stkoff() value = self._cpu_context.memory.read(cur_esp, self._cpu_context.byteness) return utils.struct_unpack(value) elif loc_type == ida_typeinf.ALOC_DIST: # arguments described by multiple locations # TODO: Uses the scattered_aloc_t class, which is a qvector or argpart_t objects # argloc.scattered() raise NotImplementedError("Argument {} location of type ALOC_DIST".format(self.idx)) elif loc_type == ida_typeinf.ALOC_REG1: # single register # TODO: Determine better way to convert reg1 integer to name. reg_name = ida_idp.get_reg_name(argloc.reg1(), self.width) return self._cpu_context.registers[reg_name] elif loc_type == ida_typeinf.ALOC_REG2: # register pair (eg: edx:eax [reg2:reg1]) # TODO: CURRENTLY UNTESTED logger.info( "Argument {} of untested type ALOC_REG2. " "Verify results and report issues".format(self.idx)) # TODO: Assuming registers are the same width.. # need to determine if that is an accurate assumption. reg1_name = ida_idp.get_reg_name(argloc.reg1(), self.width) reg2_name = ida_idp.get_reg_name(argloc.reg2(), self.width) reg1_value = self._cpu_context.registers[reg1_name] reg2_value = self._cpu_context.registers[reg2_name] return reg2_value << self.width | reg1_value elif loc_type == ida_typeinf.ALOC_RREL: # register relative (displacement from address pointed by register # TODO: CURRENTLY UNTESTED logger.info("Argument {} of untested type ALOC_RREL. " "Verify results and report issues.".format(self.idx)) # Obtain the register-relative argument location rrel = argloc.get_rrel() reg_name = ida_idp.get_reg_name(rrel.reg, self.width) value = self._cpu_context.registers[reg_name] return value + rrel.offset elif loc_type == ida_typeinf.ALOC_STATIC: # global address # TODO: CURRENTLY UNTESTED logger.info( "Argument {} of untested type ALOC_STATIC. " "Verify results and report issues.".format(self.idx)) return argloc.get_ea() elif loc_type >= ida_typeinf.ALOC_CUSTOM: # custom argloc # TODO: Will need to figure out the functionality and usage for the custloc_desc_t structure # argloc.get_custom() raise NotImplementedError("Argument {} location of type ALOC_CUSTOM".format(self.idx))
def get_reg(op): return ida_idp.get_reg_name(op.reg, 0)
def get_reg_name(reg_value): return ida_idp.get_reg_name(reg_value, BITS // 8)
def get_reg_str(operand): if operand.type == ida_ua.o_reg: # TODO: Find the operand length. Hardcoded to 8 for now. return ida_idp.get_reg_name(operand.reg, 8) return ''
def update_struct_offsets_for_xref(xref, struct_name, pad=0): regs = {} regs['hndl'] = [] regs['ptr'] = [] inst = idautils.DecodeInstruction(xref) # Are we looking at a handle or a pointer? if inst.get_canon_mnem() == "mov": regs['ptr'].append(get_reg_str(inst.Op1)) print "{} - Tracking pointer register {} at 0x{:08x}: {}".format( ' ' * pad, regs['ptr'][0], xref, disasm(xref) ) elif inst.get_canon_mnem() == "lea": regs['hndl'].append(get_reg_str(inst.Op1)) print "{} - Tracking handle register {} at 0x{:08x}: {}".format( ' ' * pad, regs['hndl'][0], xref, disasm(xref) ) items = get_func_items_from_xref(xref) # Iterate through the rest of the instructions in this function looking # for tracked registers with a displacement for item in items: regs['nhndl'] = [] regs['nptr'] = [] inst = idautils.DecodeInstruction(item) # Update any call instruction with a displacement from our register for op_no, op in enumerate(inst.Operands): if op_no == 2: break if op.type == ida_ua.o_displ and \ ida_idp.get_reg_name(op.reg, 8) in regs['ptr']: print("{} - Updating operand {} in instruction at " + "0x{:08x}: {}").format( ' ' * pad, get_op_str(inst.ea, op_no), item, disasm(item) ) apply_struct_offset(inst, op_no, struct_name) # If we find a mov instruction that copies a tracked register, track # the new register if inst.get_canon_mnem() == 'mov': if get_reg_str(inst.Op2) in regs['ptr'] and \ get_reg_str(inst.Op1) not in regs['ptr']: print("{} - Copied tracked register at 0x{:08x}, " + "tracking register {}").format( ' ' * pad, item, get_reg_str(inst.Op1) ) regs['ptr'].append(get_reg_str(inst.Op1)) regs['nptr'].append(get_reg_str(inst.Op1)) if get_reg_str(inst.Op2) in regs['hndl'] and \ get_reg_str(inst.Op1) not in regs['hndl']: print("{} - Copied tracked register at 0x{:08x}, " + "tracking register {}").format( ' ' * pad, item, get_reg_str(inst.Op1) ) regs['hndl'].append(get_reg_str(inst.Op1)) regs['nhndl'].append(get_reg_str(inst.Op1)) # If we find a mov instruction that dereferences a handle, track the # destination register for reg in regs['hndl']: if get_reg_str(inst.Op2) == "[{}]".format(reg) and \ inst.get_canon_mnem() == 'mov' and \ get_reg_str(inst.Op1) not in regs['ptr']: print("{} - Found a dereference at 0x{:08x}, " + "tracking register {}").format( ' ' * pad, item, get_reg_str(inst.Op1) ) regs['ptr'].append(get_reg_str(inst.Op1)) regs['nptr'].append(get_reg_str(inst.Op1)) # If we've found an instruction that overwrites a tracked register, # stop tracking it if inst.get_canon_mnem() in ["mov", "lea"] and \ inst.Op1 == ida_ua.o_reg: if get_reg_str(inst.Op1) in regs['ptr']: if get_reg_str(inst.Op1) not in regs['nptr']: print "{} - Untracking pointer register {}: ".format( ' ' * pad, get_reg_str(inst.Op1) + disasm(item) ) regs['ptr'].remove(get_reg_str(inst.Op1)) elif get_reg_str(inst.Op1) in regs['hndl']: if get_reg_str(inst.Op1) not in regs['nhndl']: print "{} - Untracking handle register {}: ".format( ' ' * pad, get_reg_str(inst.Op1) + disasm(item) ) regs['hndl'].remove(get_reg_str(inst.Op1)) # If we hit a call, just bail if inst.get_canon_mnem() == "call": break # If we're not tracking any registers, bail if len(regs['ptr']) == 0 and len(regs['hndl']) == 0: break
def rename_tables_internal(ea, regs, visited, renamed): names = { 'im': IMAGE_HANDLE_NAME, 'st': SYSTEM_TABLE_NAME, 'bs': BOOT_SERVICES_NAME, 'rs': RUNTIME_SERVICES_NAME } print "Processing function at 0x{:08x}".format(ea) visited.add(ea) for item in idautils.FuncItems(ea): inst = idautils.DecodeInstruction(item) # Bail out if we hit a call if inst.get_canon_mnem() == "call": to_visit = inst.Op1.addr if to_visit in visited: return else: target_types = [ida_ua.o_imm, ida_ua.o_far, ida_ua.o_near] if inst.Op1.type in target_types: rename_tables_internal( to_visit, copy.deepcopy(regs), visited, renamed ) # Keeps saved registers based on UEFI spec regs = keep_saved_regs(regs) else: print " - Can't follow call, bailing!" return if inst.get_canon_mnem() in ["mov", "lea"]: # Rename data for key in names: if get_reg_str(inst.Op2) in regs[key] and \ inst.Op1.type == ida_ua.o_mem: print(" - Found a copy to a memory address for {}, " + "updating: {}").format(names[key], disasm(inst.ea)) to_rename = inst.Op1.addr if to_rename not in renamed: renamed.add(to_rename) label = get_next_unused_label( names[key] ) ida_name.set_name(to_rename, label) break # Eliminate overwritten registers for key in names: if get_reg_str(inst.Op1) in regs[key] and \ get_reg_str(inst.Op2) not in regs[key]: print " - Untracking register {} for {}: {}".format( get_reg_str(inst.Op1), names[key], disasm(inst.ea) ) regs[key].remove(get_reg_str(inst.Op1)) # Keep track of registers containing the EFI tables etc if get_reg_str(inst.Op2) in regs['im'] and \ inst.Op1.type == ida_ua.o_reg and \ get_reg_str(inst.Op1) not in regs['im']: # A tracked register was copied to a new register, # track the new one print " - Tracking register {} for image handle: {}".format( get_reg_str(inst.Op1), disasm(inst.ea) ) regs['im'].append(get_reg_str(inst.Op1)) if get_reg_str(inst.Op2) in regs['st'] and \ inst.Op1.type == ida_ua.o_reg and \ get_reg_str(inst.Op1) not in regs['st']: # A tracked register was copied to a new register, # track the new one print " - Tracking register {} for system table: {}".format( get_reg_str(inst.Op1), disasm(inst.ea) ) regs['st'].append(get_reg_str(inst.Op1)) if inst.Op2.type == ida_ua.o_displ and \ ida_idp.get_reg_name(inst.Op2.reg, 8) in regs['st']: # A tracked register was used in a right operand with a # displacement offset = inst.Op2.addr if offset == 0x60: print " - Tracking register {} for boot services " + \ "table: {}".format(get_reg_str(inst.Op1), disasm(inst.ea)) regs['bs'].append(get_reg_str(inst.Op1)) elif offset == 0x58: print " - Tracking register {} for runtime services " + \ "table: {}".format(get_reg_str(inst.Op1), disasm(inst.ea)) regs['rs'].append(get_reg_str(inst.Op1)) apply_struct_offset(inst, 1, SYSTEM_TABLE_STRUCT)
def rename_tables_internal(ea, regs, stackdepth=0): names = { 'im': IMAGE_HANDLE_NAME, 'st': SYSTEM_TABLE_NAME, 'bs': BOOT_SERVICES_NAME, 'rs': RUNTIME_SERVICES_NAME } print "Processing function at 0x{:08x}".format(ea) for item in idautils.FuncItems(ea): inst = idautils.DecodeInstruction(item) # Bail out if we hit a call if inst.get_canon_mnem() == "call": if stackdepth == MAX_STACK_DEPTH: print " - Hit stack depth limit, bailing!" return else: target_types = [ida_ua.o_imm, ida_ua.o_far, ida_ua.o_near] if inst.Op1.type in target_types: # TODO : Currently assuming that the registry will be # inaffected by calls rename_tables_internal( inst.Op1.addr, copy.deepcopy(regs), stackdepth+1 ) else: print " - Can't follow call, bailing!" return if inst.get_canon_mnem() in ["mov", "lea"]: # Rename data for key in names: if get_reg_str(inst.Op2) in regs[key] and \ inst.Op1.type == ida_ua.o_mem: print(" - Found a copy to a memory address for {}, " + "updating: {}").format(names[key], disasm(inst.ea)) ida_name.set_name(inst.Op1.addr, names[key]) break # Eliminate overwritten registers for key in names: if get_reg_str(inst.Op1) in regs[key] and \ get_reg_str(inst.Op2) not in regs[key]: print " - Untracking register {} for {}: {}".format( get_reg_str(inst.Op1), names[key], disasm(inst.ea) ) regs[key].remove(get_reg_str(inst.Op1)) # Keep track of registers containing the EFI tables etc if get_reg_str(inst.Op2) in regs['im'] and \ inst.Op1.type == ida_ua.o_reg and \ get_reg_str(inst.Op1) not in regs['im']: # A tracked register was copied to a new register, # track the new one print " - Tracking register {} for image handle: {}".format( get_reg_str(inst.Op1), disasm(inst.ea) ) regs['im'].append(get_reg_str(inst.Op1)) if get_reg_str(inst.Op2) in regs['st'] and \ inst.Op1.type == ida_ua.o_reg and \ get_reg_str(inst.Op1) not in regs['st']: # A tracked register was copied to a new register, # track the new one print " - Tracking register {} for system table: {}".format( get_reg_str(inst.Op1), disasm(inst.ea) ) regs['st'].append(get_reg_str(inst.Op1)) if inst.Op2.type == ida_ua.o_displ and \ ida_idp.get_reg_name(inst.Op2.reg, 8) in regs['st']: # A tracked register was used in a right operand with a # displacement offset = inst.Op2.addr if offset == 0x60: print " - Tracking register {} for boot services " + \ "table: {}".format(get_reg_str(inst.Op1), disasm(inst.ea)) regs['bs'].append(get_reg_str(inst.Op1)) elif offset == 0x58: print " - Tracking register {} for runtime services " + \ "table: {}".format(get_reg_str(inst.Op1), disasm(inst.ea)) regs['rs'].append(get_reg_str(inst.Op1)) apply_struct_offset(inst, 1, SYSTEM_TABLE_STRUCT)