Python idc.NextHead() Examples

def find_dispatch_by_struct_index():
    """Attempts to locate the dispatch function based off it being loaded in a structure
    at offset 70h, based off of https://github.com/kbandla/ImmunityDebugger/blob/master/1.73/Libs/driverlib.py """
    
    out = set()
    for function_ea in idautils.Functions():
        flags = idc.get_func_flags(function_ea)
        # skip library functions
        if flags & idc.FUNC_LIB:
            continue
        func = idaapi.get_func(function_ea)
        addr = func.startEA
        while addr < func.endEA:
            if idc.GetMnem(addr) == 'mov':
                if '+70h' in idc.GetOpnd(addr, 0) and idc.GetOpType(addr, 1) == 5:
                    out.add(idc.GetOpnd(addr, 1))
            addr = idc.NextHead(addr)
    return out 

def chunk_from_path(path):
    assert_ida_available()
    import idc
    chunk = chunk_t()
    for i in xrange(len(path)):
        body = chunk.body.add()
        body.typeid = body.INSTRUCTION
        inst = body.instruction
        inst.thread_id = 0
        addr = path[i]
        inst.address = addr
        inst.opcode = idc.GetManyBytes(addr, idc.NextHead(addr)-addr)
        try:
            next_a = path[i+1]
            inf1 = inst.concrete_infos.add()
            inf1.next_address = next_a
            inf1.typeid = inf1.NEXT_ADDRESS
        except IndexError:
            pass
        inf2 = inst.concrete_infos.add()
        inf2.typeid = inf2.NOT_RETRIEVED
    return chunk 

def create_call_map(self, ftype):
        assert_ida_available()
        import idc
        import idautils
        seg_mapping = {idc.SegName(x): (idc.SegStart(x), idc.SegEnd(x)) for x in idautils.Segments()}
        imports = seg_mapping[".idata"] if ftype == PE else seg_mapping['.plt']
        start, stop = seg_mapping[".text"]
        current = start
        while current <= stop:
            inst = current
            if idc.GetMnem(inst) in ["call", "jmp"]:
                value = idc.GetOperandValue(inst, 0)
                name = idc.GetOpnd(inst, 0)
                if imports[0] <= value <= imports[1]:
                    entry = self.config.call_map.add()
                    entry.address = inst
                    entry.name = name
            current = idc.NextHead(current, stop) 

def enum_function_addrs(fva):
    '''
    yield the effective addresses of each instruction in the given function.
    these addresses are not guaranteed to be in any order.
    
    Args:
      fva (int): the starting address of a function
    
    Returns:
      sequence[int]: the addresses of each instruction
    '''
    f = idaapi.get_func(fva)
    if not f:
        raise ValueError('not a function')

    for block in idaapi.FlowChart(f):
        ea = block.startEA
        while ea <= block.endEA:
            yield ea
            ea = idc.NextHead(ea) 

def addRetBP(self, ea):
        """
        Add a breakpoint for a return instruction.
        This will place a breakpoint on the next instruction,
        e.g - the instruction that will be hit after the function returns.
        @param ea: effective address of a CALL instruction
        @return: True if breakpoint was successfully added, otherwise return False
        """
        if not is_call(ea):
            self.logger.error("The instruction at address %s is not recognized as a CALL instruction", hex(ea))
            raise UnrecognizedCallInstruction()

        next_inst = idc.NextHead(ea)

        # Add breakpoint on next instruction
        if next_inst not in self.ret_bps:
            idc.AddBpt(next_inst)
            self.ret_bps[next_inst] = 0
            return True

        return False 

def _fill_instrs(self):
        cur_addr = self.startEA
        while cur_addr != idc.BADADDR:
            self.instrs.append(cur_addr)
            cur_addr = idc.NextHead(cur_addr, self.endEA) 

def decode_here_clicked(self):
        inst = idc.here()
        if not idc.isCode(idc.GetFlags(inst)):
            print "Not code instruction"
        else:
            raw = idc.GetManyBytes(inst, idc.NextHead(inst)-inst)
            s = to_hex(raw)
            self.decode_ir(s) 

def generate_dead_alive_dump(self):
        f = Path("dead_or_alive_dump.txt")
        handle = f.open("w")
        for cfg in self.functions_cfg.values():
            for bb in cfg.values():
                for i in bb.instrs:
                    status = bb.instrs_status[i] if bb.is_alive() else Status.DEAD
                    size = idc.NextHead(i)-i
                    handle.write(u"%x,%d,%s\n" % (i, size, status))
        handle.close() 

def iter_lines():
    """
    Iterate through all line addresses in the IDB
    Yields addresses of all lines.
    """
    for ea in idautils.Segments():
        seg_start = idc.SegStart(ea)
        seg_end = idc.SegEnd(ea)

        cur_addr = seg_start
        while (cur_addr < seg_end) and (cur_addr != idaapi.BADADDR):
            yield cur_addr
            cur_addr = idc.NextHead(cur_addr) 

def get_function_args_count(self, function_ea, local_vars):
        """
        The function returns count of function arguments
        @function_ea - function entry point
        @local_vars - local variables dictionary
        @return - function arguments count
        """
        # i#9 Now, we can't identify fastcall functions.

        function_args_count = 0
        args_dict = dict()
        for local_var in local_vars:
            usage_list = local_vars.get(local_var, None)
            if usage_list == None:
                print "WARNING: empty usage list for ", local_var
                continue
            for head in usage_list:
                ops = self.get_instr_operands(int(head, 16))
                for idx, (op,type) in enumerate(ops):
                    if op.count("+") == 1:
                        value = idc.GetOperandValue(int (head, 16), idx)
                        if value < (15 * ARGUMENT_SIZE) and "ebp" in op:
                            args_dict.setdefault(local_var, []).append(head)
                    elif op.count("+") == 2:
                        if "arg" in local_var:
                            args_dict.setdefault(local_var, []).append(head)
                    else:
                        continue

        function_args_count = len(args_dict)
        if function_args_count:
            return function_args_count, args_dict

        #TODO Check previous algorithm here
        f_end = idc.FindFuncEnd(function_ea)
        f_end = idc.PrevHead(f_end, 0)
        instr_mnem = idc.GetMnem(f_end)
        #stdcall ?
        if "ret" in instr_mnem:
            ops = self.get_instr_operands(f_end)
            if len(ops) == 1:
                for op,type in ops:
                    op = op.replace("h", "")
                    function_args_count = int(op,16)/ARGUMENT_SIZE
                    return function_args_count, args_dict
        #cdecl ?
        refs = idautils.CodeRefsTo(function_ea, 0)
        for ref in refs:
            #trying to find add esp,x signature after call
            head = idc.NextHead(ref, 0xFFFFFFFF)
            if head:
                disasm = idc.GetDisasm(head)
                if "add" in disasm and "esp," in disasm:
                    ops = self.get_instr_operands(head)
                    op,type = ops[1]
                    if op:
                        op = op.replace("h", "")
                        function_args_count = int(op,16)/ARGUMENT_SIZE
                        return function_args_count, args_dict
        return function_args_count, args_dict