Python idautils 模块，Heads() 实例源码

def get_list_of_function_instr(addr):
    f_start = addr
    f_end = idc.FindFuncEnd(addr)
    chunks = enumerate_function_chunks(f_start)
    list_of_addr = list()
    list_of_calls = list()
    image_base = idaapi.get_imagebase(addr)
    for chunk in chunks:
        for head in idautils.Heads(chunk[0], chunk[1]):
            # If the element is an instruction
            if head == hex(0xffffffffL):
                raise Exception("Invalid head for parsing")
            if idc.isCode(idc.GetFlags(head)):
                call_name = get_call_name(head)
                if call_name != None:
                    list_of_calls.append(call_name)
                head = head - image_base
                head = str(hex(head))
                head = head.replace("L", "")
                head = head.replace("0x", "")
                list_of_addr.append(head)
    return list_of_addr, list_of_calls

def get_list_of_function_instr(addr):
    '''
    The function returns a list of instructions from a function
    @addr - is function entry point
    @return - list of instruction's addresses
    '''
    f_start = addr
    f_end = idc.FindFuncEnd(addr)
    chunks = enumerate_function_chunks(f_start)
    list_of_addr = list()
    image_base = idaapi.get_imagebase(addr)
    for chunk in chunks:
        for head in idautils.Heads(chunk[0], chunk[1]):
            # If the element is an instruction
            if head == hex(0xffffffffL):
                raise Exception("Invalid head for parsing")
            if idc.isCode(idc.GetFlags(head)):
                head = head - image_base
                head = str(hex(head))
                head = head.replace("L", "")
                head = head.replace("0x", "")
                list_of_addr.append(head)
    return list_of_addr

def send_comments(self):
        """
            Initial sync of comments
        """
        cmt_types = [idc.Comment, idc.RptCmt]  # Maybe GetFunctionCmt also?
        for head in idautils.Heads():
            send_cmt = ""
            for cmt_type in cmt_types:
                cmt = None
                cmt = cmt_type(head)
                if cmt and not SkelUtils.filter_coms_blacklist(cmt):
                    if cmt not in send_cmt:
                        send_cmt += cmt
            if len(send_cmt) > 0:
                try:
                    self.skel_conn.push_comment(head, send_cmt)
                except Exception as e:
                    g_logger.exception(e)

def _get_blocks_codes_per_func_iter():
    """
    Iterative function to generate all blocks and opcodes
    :return: N/A
    """
    all_blocks = {}
    all_codes = {}
    all_opcodes = []
    for func in idautils.Functions():
        # blocks_in_func contains
        #   <idaapi.BasicBlock object at 0x0545F3F0>, ...
        blocks_in_func = idaapi.FlowChart(idaapi.get_func(func))
        blocks = []
        for block in blocks_in_func:

            # IDA BUG! block.startEA == block.endEA??
            # Should be in the range "block.startEA <= block < block.endEA"
            if block.startEA != block.endEA:
                blocks.append((block.startEA, block.endEA))
            for head in idautils.Heads(block.startEA, block.endEA):
                ibytes = idc.GetManyBytes(head, idc.ItemEnd(head) - head)
                spd = idc.GetSpd(head)
                all_codes[head] = insn.Instruction(head, ibytes, spd)

                # IDA BUG! all_codes[head].bytes == 0??
                # The size of the code should be positive
                if all_codes[head].bytes != 0:
                    all_opcodes.append((all_codes[head].addr, all_codes[head].addr + len(all_codes[head].bytes)))
        all_blocks[func] = blocks
    yield (all_blocks, all_opcodes)

def get_code_heads():
    """Returns a set with all the recognized code heads from all the
    code sections."""
    global code_heads

    if len(code_heads) == 0:
        for begin, end, name in code_segments_iter():
            code_heads |= set(filter(lambda x: idaapi.isCode(idc.GetFlags(x)), idautils.Heads(begin, end)))

    return code_heads

def enum_heads():
    for segment in enum_segments():
        for head in idautils.Heads(segment.start, segment.end):
            yield head

def _feature_functiontype(self, f_ea):
        '''
        functiontype here is to identify the type of the function, now we just identify whether the function is doing memory
        operation like memcpy. later maybe we will extend the types.
        for memory operation type,the way we identify is:
            a. There're loops
            b. There're index change
                how to identify index change?
            c. Memory operation include but not limited to:
                a. Mov [eax....], ecx, lea....
                b. Stos, movs, lods
                    for 8-bit, 16-bot 
                c. Call library memory function, strcpy, ...      

        prior feature: loopcount
        ''' 
#         lflag = 0
        imflag = 0
        for loop in self.loops.values():
#             lflag = 1
            for block in loop:
                for l_ea in idautils.Heads(block[0],block[1]):
                    inst = idautils.DecodeInstruction(l_ea)
                    if inst == None:
                        continue
                    if inst.itype in [122]: # mov
                        # mov 122
                        if 3 == inst[0].type or 4 == inst[0].type: 
                            imflag = 1         
                    elif inst.itype in [124,207,107]: #movs/movsd, stos lods
                        # 124 movs 207 stos 107 lods
                        imflag = 1
                    elif inst.itype in [16]: # call library function
                        # 16 call
                        istr = GetInstruction(l_ea)
                        if 'strcpy' in istr or 'memcpy' in istr or 'alloc' in istr or 'free' in istr:
                            imflag = 1
        if  imflag:#lflag and    
            return 1
        else:
            return 0

def get_code_and_blocks(ea):
    """Extracts the control flow graph for the function at the given address.
    Returns a dictionary with the instructions (ea->insn.Instruction) and a list
    of the basic blocs (bbl.BasicBlock)."""
    code = {}
    blocks = {}
    ida_blocks = set(idaapi.FlowChart(idaapi.get_func(ea)))

    for bb in ida_blocks:

        # XXX: it seems that it's not a bug but inter-function jumps!
        if bb.startEA == bb.endEA: # skip that .. it's IDA's bug
            #print "skipping block %x : %x in func %x"%(bb.startEA, bb.endEA, ea)
            continue

        blocks[bb.startEA] = bbl.BasicBlock(bb.startEA, bb.endEA, {})

        for head in idautils.Heads(bb.startEA, bb.endEA):
            ibytes = idc.GetManyBytes(head, idc.ItemEnd(head) - head)
            spd = idc.GetSpd(head)
            code[head] = insn.Instruction(head, ibytes, spd)
            blocks[bb.startEA].instrs.append(code[head])
            next_head = idc.NextHead(head, bb.endEA)

            if idaapi.isFlow(idc.GetFlags(next_head)):
                code[head].succ.add(next_head)

        for suc_bb in (s for s in bb.succs() if s.startEA != s.endEA):
            #assume head is the last instruction of the block
            code[head].succ.add(suc_bb.startEA)

    for bb in (b for b in ida_blocks if b.startEA != b.endEA):
        for suc_bb in (s for s in bb.succs() if s.startEA != s.endEA):
            # a jump with zero offset (like, jz 0) gives two succs to the same bb
            if blocks[suc_bb.startEA] not in blocks[bb.startEA].successors:
                blocks[bb.startEA].successors.append(blocks[suc_bb.startEA])
        blocks[bb.startEA].successors.sort(key=lambda x: x.begin, reverse=True)

    #FIXME: find a better way ..
    for block in blocks.itervalues():
        if block.instrs[0].addr == ea:
            #print "found the entry!:", block.instrs
            block.instrs[0].f_entry = True
            block.type |= bbl.BasicBlock.ENTRY
            break
    else:
        print "BUG: could not find function entry in instrs!!"
    #print "blocks:", blocks

    return code, blocks.values()


#XXX: good test function in 0x070016E7 (BIB.dll)