ida scripts---path search

ida-path-search.py

from idautils import *
import idc
import idaapi
import sark
import networkx as nx
import Queue


print 'Start'

ivar = {'cmp': ['cmp', 'test'], 'call': ['call']}

f = open("output.txt", 'w')

def check(inst):
    if GetMnem(inst) in ['cmp', 'test', 'call']:
        return True
    if GetOpType(inst, 1) == 2: 
        return True
    return False


def get_end_bb(func):
    start = func.startEA
    ends = []
    while start < func.endEA:
        if GetMnem(start) == 'retn':
            ends.append(sark.get_block_start(start))
        start = idc.NextHead(start)
    return ends

def getPathInsts(path):
    insts = []
    for bb_ea in path:
        for inst in sark.CodeBlock(bb_ea).lines:
            if check(inst.ea):
                # print(inst)
                insts.append(inst)
    return insts

def dumpBB(bb):
    print 'BB Start'
    for inst in bb.lines:
        print inst
    print 'BB End'


def checkInst(insts, seq, depth):
    count = 0
    for inst in insts:
        # print depth, count
        if not check(inst.ea):
            continue
        # print inst
        if depth + count == len(seq):
            return count
        if inst.insn.mnem in ivar[seq[depth + count]]:
            # print inst
            count += 1
        else:
            return 0
        
    return count

def BFS(bb, seq, depth, path):
    # dumpBB(bb)
    num = checkInst(bb.lines, seq, depth)
    succs = list(bb.next)
    result = False
    if (len(seq) - 1) <= (depth + num):
        return True
    if len(succs) == 0:
        return False
    for succ in succs:
        result = result or BFS(succ, seq, depth + num, path)
        if result:
            path.append(succ)
            # dumpBB(succ)
        # print result
    return result

def findFirstBB(bb, seq):
    insts = getPathInsts([bb.start_ea])
    for index, inst in enumerate(insts):
        count =  checkInst(insts[index:], seq, 0)
        if count != 0:
            return count

    return 0

def isMatched(start, seq):
    queue = Queue.Queue()
    queue.put(start)
    result = False
    path = []
    while not queue.empty():
        bb = queue.get()
        for succ in bb.next:
            queue.put(succ)
        count = findFirstBB(bb, seq)
        if count != 0:
            # print 'Find first bb'
            # dumpBB(bb)
            path *= 0
            path.append(bb)
            result = BFS(bb, seq, count, path)
            if result:
                for bb in path:
                    dumpBB(bb)
                
                print "Path Matched"
                insts = getPathInsts([bb.start_ea for bb in path])
                print "".join("%s\n" % inst for inst in insts)
                return result
    
    getPathInsts([bb.start_ea for bb in path])
    return result

def isFindSwitch(func):
    for (startea, endea) in Chunks(func):
        for head in Heads(startea, endea):
            switch_info = idaapi.get_switch_info_ex(head)
            if switch_info != None:
                num_cases = switch_info.get_jtable_size()
                return True
    return False



def run(start_addr, end_addr):
    print(hex(start_addr), hex(end_addr))
    func = sark.Function(start_addr)
    print(func.name)
    # func = idaapi.get_func(start_addr)
    graph = sark.get_nx_graph(func.ea)
    # viewer = sark.ui.NXGraph(graph, title="My Graph", handler=sark.ui.AddressNodeHandler())
    # viewer.Show()
    start = sark.get_block_start(start_addr)
    end = sark.get_block_start(end_addr)
    if isFindSwitch(func.ea):
        print "Find Switch case and skip func"
        return
    seq = ['call', 'call', 'cmp', 'call', 'cmp', 'call']
    print isMatched(sark.CodeBlock(start), seq)
    f.write('%s: 0x%X-0x%X\n' % (func.name, start_addr, end_addr))
    try:
        for path in nx.shortest_simple_paths(graph, start, end):
            print 'New Path (%d)' % len(path)
            insts = getPathInsts(path)
            f.write('New Path (%d)\n' % len(path))
            f.write("".join("%s\n" % inst for inst in insts))
            # f.write()
    except Exception as e:
        print('Error:', e)
    print(hex(start_addr), hex(end_addr))
    

def readLine():
    data = []
    with open("data.txt", 'r') as f:
        lines = f.readlines()
        for line in lines:
            line = line.strip()
            line = line.split(' ')
            if len(line) == 2 and line[0].isdigit() and line[1].isdigit():
                data.append((int(line[0]), int(line[1])))
    return data

data = readLine()
# for d in data:
#     run(d[0], d[1])

# run(990715, 990893)
run(447970, 448010)
f.close()
print 'End'