vd

#!/usr/bin/python2

"""
This file is part of VDISCOVER.

VDISCOVER is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

VDISCOVER is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with VDISCOVER. If not, see <http://www.gnu.org/licenses/>.

Copyright 2014 by G.Grieco
"""

import os
import os.path
import argparse
import sys
import csv
import random

csv.field_size_limit(sys.maxsize)
sys.setrecursionlimit(1024 * 1024 * 1024)

#from vdiscover.Detection import WriteTestcase
from vdiscover.Process import Process
from vdiscover.Printer import TypePrinter
#from vdiscover.Cluster  import ClusterScikit, ClusterConv
from vdiscover.Utils import update_progress
from vdiscover.Sampling import cluster_sampler

if __name__ == "__main__":

    if open("/proc/sys/kernel/randomize_va_space").read().strip() != "0":
        print("Address space layout randomization (ASLR) is enabled, disable it before continue to use the cache")
        print("Hint: # echo 0 > /proc/sys/kernel/randomize_va_space")
        sys.exit(-1)

    # Arguments
    parser = argparse.ArgumentParser(description='')
    parser.add_argument("-i", help="", type=str,
                        default=None, required=True, dest="seeds")
    parser.add_argument("-o", help="", type=str,
                        default=None, required=True, dest="out")
    parser.add_argument("-m", help="", type=str, nargs='+', dest="mods")

    #parser.add_argument("-v", help="", type=str, default=None, required=False, dest="vectorizer")
    #parser.add_argument("-m", help="", type=str, default="afl", dest="fuzzer")
    parser.add_argument("cmd", help="", type=str, default=None)

    options = parser.parse_args()
    seeds = options.seeds
    outfile = options.out
    mods = options.mods
    #fuzzer = options.fuzzer
    cmd = options.cmd
    #vectorizer = options.vectorizer
    program = cmd.split(" ")[0]
    programf = program.replace("/", "__")
    main_module = program.split("/")[-1]

    timeout = 15
    envs = dict()
    traces_path = outfile  # outdir+"/traces.raw"

    if os.path.exists(traces_path):
        print traces_path, "exists. I will not overwritte it. Aborting"
    else:

        modules_to_trace = [main_module]
        if mods is not None:
            modules_to_trace = modules_to_trace + mods

        if "LD_LIBRARY_PATH" in os.environ:
            libs = os.environ["LD_LIBRARY_PATH"]
            for _, _, files in os.walk(libs):
                for f in files:
                    modules_to_trace.append(f)

        print "Tracing", modules_to_trace
        app = Process(program, envs, timeout, modules_to_trace, [], True)
        prt = TypePrinter(traces_path, program, 0)
        traces = []
        all_files = []

        print "Extracting traces.."
        for x, y, files in os.walk(seeds):
            nfiles = len(files)
            # print "Processing directory ","./"++("/".join(y)), "with",
            # nfiles, "seeds"
            for f in files:
                all_files.append(x + "/".join(y) + "/" + f)

        random.shuffle(all_files)
        nfiles = len(all_files)

        for progress, testcase in enumerate(all_files):
            # print testcase
            progress = round(float(progress) / nfiles, 4)
            update_progress(progress)
            prepared_cmd = cmd.replace(program, "")
            if "@@" in prepared_cmd:
                prepared_cmd = prepared_cmd.split("@@")
                prepared_cmd = prepared_cmd[0].split(
                    " ") + [testcase] + prepared_cmd[1].split(" ")
            else:
               prepared_cmd = [prepared_cmd ,"< "+testcase]
               #print(prepared_cmd)
 
            prepared_cmd = filter(lambda x: x != '', prepared_cmd)
            # print "Getting data.."
            events = app.getData(prepared_cmd)
            # print "Printing data.. ", len(events)
            traces.append(prt.print_events(testcase, events))
            # print prepared_cmd
            # print traces[-1]

    # if vectorizer is None:
    #  clustered_traces = ClusterScikit(vectorizer, traces, None, "dynamic", None)
    # else:
    #  clustered_traces = ClusterConv(vectorizer, traces, None, "dynamic", None, None)
    #  cluster_sampler(clustered_traces,1)
    #  #print clusters