from pydbg import *
from pydbg.defines import *

import psutil
import utils
import time
import threading
import os
import copy
import random

PATH_TO_PROCESS = "C:\\class\\cdf\\cdf_reader.exe"
PROCESS_NAME = "cdf_reader.exe"
ORIGINAL_DOCUMENT_PATH = "C:\\class\\cdf\\haiku2.cdf"
FUZZY_DOCUMENT_PATH = "C:\\class\\fuzzing\\fuzzy.cdf"
CRASH_LOG_FILE = "C:\\class\\fuzzing\\crashlog.txt"
BAD_DOCUMENT_STORAGE = "C:\\class\\fuzzing\\crash_docs\\"

crash_log = open(CRASH_LOG_FILE,'w')
crash_counter = 0
iteration_count = 0

def checkAccessViolation(dbg):
	global crash_counter
	crash_counter = crash_counter + 1

        crash_bin = utils.crash_binning.crash_binning()
	crash_bin.record_crash(dbg)

        crash_analysis =  "\nCRASH DETECTED\n"
	crash_analysis += crash_bin.crash_synopsis()
	crash_analysis += "\nCRASH COUNT: %d\n"%crash_counter

        crashlog = open(CRASH_LOG_FILE, 'a')
	crashlog.write(crash_analysis)
	crashlog.close()

        print "crash: %d, iteration: %d"%(crash_counter, iteration_count)
        
	dbg.terminate_process()
        cmd = "copy %s %scrash_%d.cdf"%(FUZZY_DOCUMENT_PATH,BAD_DOCUMENT_STORAGE,crash_counter)
        os.system(cmd)
	return DBG_EXCEPTION_NOT_HANDLED

#got start and stop thread code from:
#http://www.pepelux.org/papers/Using%20Sulley%20on%20local%20applications.pdf
class StartProcessThread(threading.Thread):
    def __init__(self,t_proc_name, t_name):
        threading.Thread.__init__(self)
        self.t_proc_name = t_proc_name
        self.t_name = t_name

    def run(self):
        thread_dbg = pydbg()
        thread_dbg.load(self.t_proc_name, self.t_name)
        thread_dbg.set_callback(EXCEPTION_ACCESS_VIOLATION,checkAccessViolation)
        thread_dbg.run()

class KillProcessThread(threading.Thread):
    def __init__(self,t_proc_name):
        threading.Thread.__init__(self)
        self.t_proc_name = t_proc_name

    def run(self):
        time.sleep(1)
        os.system("taskkill /F /T /IM " + self.t_proc_name + " > blah.out") 
        
def createFuzzyDocument(original_document):
    """
    This is what you should be modifying depending on how you want to 
    generate a fuzzy document based on the original
    """
    fuzzy_document = copy.copy(original_document)
    fuzzy_document[random.randint(0,len(original_document)-1)] = random.randint(0,0xff)

    return fuzzy_document

def main():
    file_handle = open(ORIGINAL_DOCUMENT_PATH,'rb')
    original_document = bytearray(file_handle.read())
    file_handle.close()

    global iteration_count
    
    while True:

        iteration_count += 1
        fuzzy_document = createFuzzyDocument(original_document)

        file_handle = open(FUZZY_DOCUMENT_PATH,'wb')
        file_handle.write(fuzzy_document)
        file_handle.close()

        reader_thread = StartProcessThread(PATH_TO_PROCESS, FUZZY_DOCUMENT_PATH)
        reader_thread.start()

        killer_thread = KillProcessThread(PROCESS_NAME)
        killer_thread.start()

        #limit ourselves to one reader thread at a time to keep things simple
        while reader_thread.isAlive() or killer_thread.isAlive():
            time.sleep(.25)
        

if __name__ == "__main__":
    main()
