[原创]ida循环搜索

#!/usr/bin/python
#coding=utf-8
#通过深度优先遍历，找到当前函数的循环，需要整个IDA文件搜索请自行修改
from idaapi import *
loop_list = []
#判断是否跳转到其他函数,否则用Rnext会跳转到调用的函数中去
def isSeamFunc(addr, next):
        addr_funcname = GetFunctionName(addr)
        next_funcname = GetFunctionName(next)
        if cmp(addr_funcname, next_funcname)==0:
                return next
        else:
                return BADADDR
               
#使用Rfirst遍历代码路径
#用法：Rfirst(addr), Rnext(addr, Rfirst(addr))
def codePath(addr):
        addr = GetFunctionAttr(addr, FUNCATTR_START)
        addr_stack = [addr] #用来保存代码路径的栈
        ref_stack = [BADADDR] #用来保存下一个引用的栈
        flag = 0
        while 1:
                print '%x'%addr, GetDisasm(addr)
                ref_first = Rfirst(addr)
                ref_next = Rnext(addr, ref_first)
                ref_next = isSeamFunc(addr, ref_next)
                if addr in addr_stack[:-1]: #当前函数已经在路径栈中了，说明有循环，记录循环，并出栈
                        loop_list.append([addr,addr_stack[-2]])
                        flag = 1
                else:
                        flag = 0
                if ref_first==BADADDR or flag==1: #找到函数的返回点或者出现循环了，出栈回退
                        print '--------------------->'
                        while 1:
                                if len(addr_stack)==0:
                                        print 'code path finish'
                                        return
                                tmp_first = addr_stack.pop()
                                tmp_next = ref_stack.pop()
                                if tmp_next!=BADADDR: #上一个分支点的处理
                                        ref_next = Rnext(tmp_first, tmp_next)
                                        addr_stack.append(tmp_first)
                                        ref_stack.append(ref_next)
                                        addr = tmp_next
                                        break
                else:
                        addr_stack.append(ref_first)
                        ref_stack.append(ref_next)
                        addr = ref_first
       
if __name__ == '__main__':
        codePath(ScreenEA())
        print '---------------->loop list:'
        for addr in loop_list:
                print '%x'%addr[0], GetDisasm(addr[0]), '---------->loop size:', abs(addr[1]-addr[0])