IDA Function coverage script

잠깐 출장다녀오면서 회사 동료와 얘기하던 중, corpus 최소화 방법에 대해서 얘기가 나왔었다. 그래서 간단하게 아래와 같이 ida에서 동작하는 function coverage기반으로 중복되는 corpus를 걸러낼 수 있도록 만들었다. 많은 샘플 파일을 구해서 해당 script를 돌리고 덤퍼징만 해도 어느정도 효율적으로 퍼징이 가능할 거 같다.

from idautils import *
import ida_dbg
import idaapi
import idc
import ida_nalt
import struct, md5, glob, os, shutil
from collections import OrderedDict

info = idaapi.get_inf_structure()
fmt = "<L"
if info.is_64bit():
    fmt = "<Q"

class MyDbgHook(ida_dbg.DBG_Hooks):
    coverage = []
    filepath = ""
    covpath  = ""
    corpus   = ""

    def dbg_bpt(self, tid, ea):
        self.coverage += [ ea ]
        idaapi.continue_process()
        return 0

    def dbg_process_exit(self, pid, tid, ea, code):
        allcov = OrderedDict()
        for cv in self.coverage:
            allcov[cv] = 1
        
        cvdata = "\n".join(map(hex, allcov.keys()))
        cov = md5.md5(cvdata).hexdigest()
        covs = []
        for path in glob.glob(self.covpath + "\\*"):
            covs += [ os.path.basename(path) ]

        if cov not in covs:
            with open(self.covpath + "\\" + cov, "w") as f:
                f.write("Coverage:\n" + cvdata)
            
            ext = os.path.splitext(self.filepath)
            shutil.copy(self.filepath, self.corpus + "\\" + cov + ext[1])
        
        self.coverage = []

def getAllFuncAddr():
    result = []
    for segea in Segments():
        for funcea in Functions(segea, SegEnd(segea)):
            for (startea, _) in Chunks(funcea):
                result += [ startea ]
    
    return result

addrs = getAllFuncAddr()

for addr in addrs:
    idc.AddBpt(addr)

basepath = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"

dbg = MyDbgHook()
dbg.covpath = basepath + "\\cov"
dbg.corpus = basepath + "\\corpus"
dbg.hook()

for inp in glob.glob(basepath + "\\input\\*"):
    dbg.filepath = inp
    ida_dbg.start_process(ida_nalt.get_input_file_path(), inp)
    ida_dbg.wait_for_next_event(ida_dbg.WFNE_CONT, 5)
    ida_dbg.exit_process()