import csv
import os
def_flags="--no-default-features --features std,snapshot_restore,singlecore,restarting,do_hash_notify_state"
remote="remote/"
RUNTIME=1800
TARGET_REPS_A=2
TARGET_REPS_B=2
NUM_NODES=2
REP_PER_NODE_A=int(TARGET_REPS_A/NUM_NODES)
REP_PER_NODE_B=int(TARGET_REPS_B/NUM_NODES)
NODE_ID= 0 if os.getenv('NODE_ID') == None else int(os.environ['NODE_ID'])
MY_RANGE_A=range(NODE_ID*REP_PER_NODE_A,(NODE_ID+1)*REP_PER_NODE_A)
MY_RANGE_B=range(NODE_ID*REP_PER_NODE_B,(NODE_ID+1)*REP_PER_NODE_B)

rule build_showmap:
    output:
        directory("bins/target_showmap")
    shell:
        "cargo build --target-dir {output} {def_flags},config_stg"

rule build_random:
    output:
        directory("bins/target_random")
    shell:
        "cargo build --target-dir {output} {def_flags},feed_longest"

rule build_frafl:
    output:
        directory("bins/target_frafl")
    shell:
        "cargo build --target-dir {output} {def_flags},config_frafl,feed_longest"

rule build_afl:
    output:
        directory("bins/target_afl")
    shell:
        "cargo build --target-dir {output} {def_flags},config_afl,observer_hitcounts"

rule build_stg:
    output:
        directory("bins/target_stg")
    shell:
        "cargo build --target-dir {output} {def_flags},config_stg"

rule build_stgpath:
    output:
        directory("bins/target_stgpath")
    shell:
        "cargo build --target-dir {output} {def_flags},feed_stg_abbhash,sched_stg_abbhash,mutate_stg"

rule build_showmap_int:
    output:
        directory("bins/target_showmap_int")
    shell:
        "cargo build --target-dir {output} {def_flags},config_stg,fuzz_int"

rule build_random_int:
    output:
        directory("bins/target_random_int")
    shell:
        "cargo build --target-dir {output} {def_flags},feed_longest,fuzz_int"


rule build_frafl_int:
    output:
        directory("bins/target_frafl_int")
    shell:
        "cargo build --target-dir {output} {def_flags},config_frafl,fuzz_int"

rule build_afl_int:
    output:
        directory("bins/target_afl_int")
    shell:
        "cargo build --target-dir {output} {def_flags},config_afl,fuzz_int,"

rule build_stg_int:
    output:
        directory("bins/target_stg_int")
    shell:
        "cargo build --target-dir {output} {def_flags},config_stg,fuzz_int"

rule build_stgpath_int:
    output:
        directory("bins/target_stgpath_int")
    shell:
        "cargo build --target-dir {output} {def_flags},feed_stg_abbhash,sched_stg_abbhash,mutate_stg,fuzz_int"

rule build_feedgeneration1:
    output:
        directory("bins/target_feedgeneration1")
    shell:
        "cargo build --target-dir {output} {def_flags},feed_genetic,gensize_1"

rule build_feedgeneration1_int:
    output:
        directory("bins/target_feedgeneration1_int")
    shell:
        "cargo build --target-dir {output} {def_flags},feed_genetic,fuzz_int,gensize_1"

rule build_feedgeneration10:
    output:
        directory("bins/target_feedgeneration10")
    shell:
        "cargo build --target-dir {output} {def_flags},feed_genetic,gensize_10"

rule build_feedgeneration10_int:
    output:
        directory("bins/target_feedgeneration10_int")
    shell:
        "cargo build --target-dir {output} {def_flags},feed_genetic,fuzz_int,gensize_10"

rule build_feedgeneration100:
    output:
        directory("bins/target_feedgeneration100")
    shell:
        "cargo build --target-dir {output} {def_flags},config_genetic,gensize_100"

rule build_feedgeneration100_int:
    output:
        directory("bins/target_feedgeneration100_int")
    shell:
        "cargo build --target-dir {output} {def_flags},config_genetic,fuzz_int,gensize_100"

rule run_bench:
    input:
        "build/{target}.elf",
        "bins/target_{fuzzer}"
    output:
        multiext("timedump/{fuzzer}/{target}#{num}", ".time", ".log") # , ".case"
    run:
        with open('target_symbols.csv') as csvfile:
            reader = csv.DictReader(csvfile)
            line = next((x for x in reader if x['kernel']==wildcards.target), None)
            if line == None:
                return False
            kernel=line['kernel']
            fuzz_main=line['main_function']
            fuzz_input=line['input_symbol']
            fuzz_len=line['input_size']
            bkp=line['return_function']
            if wildcards.fuzzer.find('random') >= 0:
                script="""
                    export RUST_BACKTRACE=1
                    mkdir -p $(dirname {output[0]})
                    set +e
                    echo $(pwd)/{input[1]}/debug/fret -n $(pwd)/timedump/{wildcards.fuzzer}/{wildcards.target}#{wildcards.num} -t -a -g -k {input[0]} -c ./target_symbols.csv fuzz --random -t {RUNTIME} -s {wildcards.num}
                    $(pwd)/{input[1]}/debug/fret -n $(pwd)/timedump/{wildcards.fuzzer}/{wildcards.target}#{wildcards.num} -t -a -g -k {input[0]} -c ./target_symbols.csv fuzz --random -t {RUNTIME} -s {wildcards.num} > {output[1]} 2>&1
                    exit 0
                    """
            else:
                script="""
                    export RUST_BACKTRACE=1
                    mkdir -p $(dirname {output[0]})
                    set +e
                    echo $(pwd)/{input[1]}/debug/fret -n $(pwd)/timedump/{wildcards.fuzzer}/{wildcards.target}#{wildcards.num} -t -a -g -k {input[0]} -c ./target_symbols.csv fuzz -t {RUNTIME} -s {wildcards.num}
                    $(pwd)/{input[1]}/debug/fret -n $(pwd)/timedump/{wildcards.fuzzer}/{wildcards.target}#{wildcards.num} -t -a -g -k {input[0]} -c ./target_symbols.csv fuzz -t {RUNTIME} -s {wildcards.num} > {output[1]} 2>&1
                    exit 0
                    """
            shell(script)

rule run_showmap:
    input:
        "{remote}build/{target}.elf",
        "bins/target_showmap",
        "bins/target_showmap_int",
        "{remote}timedump/{fuzzer}/{target}#{num}.case"
    output:
        "{remote}timedump/{fuzzer}/{target}#{num}_case.trace.ron",
        "{remote}timedump/{fuzzer}/{target}#{num}_case.time",
    run:
        with open('target_symbols.csv') as csvfile:
            reader = csv.DictReader(csvfile)
            line = next((x for x in reader if x['kernel']==wildcards.target), None)
            if line == None:
                return False
            kernel=line['kernel']
            fuzz_main=line['main_function']
            fuzz_input=line['input_symbol']
            fuzz_len=line['input_size']
            bkp=line['return_function']
            script=""
            if wildcards.fuzzer.find('_int') > -1:
                script="export FUZZER=$(pwd)/{input[2]}/debug/fret\n"
            else:
                script="export FUZZER=$(pwd)/{input[1]}/debug/fret\n"
            script+="""
                mkdir -p $(dirname {output})
                set +e
                echo $(pwd)/{input[1]}/debug/fret -n $(pwd)/{remote}/timedump/{wildcards.fuzzer}/{wildcards.target}#{wildcards.num}_case -t -a -r -g -k {input[0]} -c ./target_symbols.csv showmap -i {input[3]}
                $(pwd)/{input[1]}/debug/fret -n $(pwd)/{remote}/timedump/{wildcards.fuzzer}/{wildcards.target}#{wildcards.num}_case -t -a -r -g -k {input[0]} -c ./target_symbols.csv showmap -i {input[3]}
                exit 0
                """
            if wildcards.fuzzer.find('random') >= 0:
                script="export FUZZ_RANDOM=1\n"+script
            shell(script)

rule tarnsform_trace:
    input:
        "{remote}timedump/{fuzzer}/{target}.{num}.trace.ron"
    output:
        "{remote}timedump/{fuzzer}/{target}.{num}.trace.csv"
    shell:
        "$(pwd)/../../../../state2gantt/target/debug/state2gantt {input} > {output[0]}"

rule trace2gantt:
    input:
        "{remote}timedump/{fuzzer}/{target}.{num}.trace.csv"
    output:
        "{remote}timedump/{fuzzer}/{target}.{num}.trace.csv.png"
    shell:
        "Rscript --vanilla $(pwd)/../../../../state2gantt/gantt.R {input}"

rule all_main:
    input:
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['random','afl','feedgeneration10','state'], target=['waters','watersv2'],num=range(0,3))

rule all_main_int:
    input:
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['random_int','afl_int','feedgeneration10_int','state_int'], target=['waters_int','watersv2_int'],num=range(0,4))

rule all_compare_feedgeneration:
    input:
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['feedgeneration1','feedgeneration10','feedgeneration100'], target=['waters_int','watersv2'],num=range(0,10))

rule all_compare_feedgeneration_int:
    input:
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['feedgeneration1_int','feedgeneration10_int','feedgeneration100_int'], target=['waters_int','watersv2_int'],num=range(0,10))

rule all_compare_afl:
    input:
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['afl','frafl','feedlongest'], target=['waters','watersv2'],num=range(0,10))

rule all_compare_afl_int:
    input:
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['afl_int','frafl_int','feedlongest_int'], target=['waters_int','watersv2_int'],num=range(0,10))

rule all_images:
    input:
        expand("{remote}timedump/{fuzzer}/{target}.{num}.trace.csv.png",remote=remote, fuzzer=['afl','feedgeneration10','state'], target=['waters','watersv2'],num=range(0,3))

rule all_images_int:
    input:
        expand("{remote}timedump/{fuzzer}/{target}.{num}.trace.csv.png",remote=remote, fuzzer=['afl_int','feedgeneration10_int','state_int'], target=['waters_int','watersv2_int'],num=range(0,3))

rule clusterfuzz:
    input:
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['random','afl','feedgeneration10','state'], target=['waters','watersv2'],num=MY_RANGE_A),
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['random_int','afl_int','feedgeneration10_int','state_int'], target=['waters_int','watersv2_int'],num=MY_RANGE_A),
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['feedgeneration1','feedgeneration10','feedgeneration100'], target=['waters_int','watersv2'],num=MY_RANGE_B),
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['feedgeneration1_int','feedgeneration10_int','feedgeneration100_int'], target=['waters_int','watersv2_int'],num=MY_RANGE_B),
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['afl','frafl','feedlongest'], target=['waters','watersv2'],num=MY_RANGE_B),
        expand("timedump/{fuzzer}/{target}.{num}", fuzzer=['afl_int','frafl_int','feedlongest_int'], target=['waters_int','watersv2_int'],num=MY_RANGE_B),

rule all_new:
    input:
        expand("timedump/{fuzzer}/{target}#{num}.time", fuzzer=['feedgeneration100', 'frafl', 'stg'], target=['waters', 'watersv2'],num=range(0,3)),
        expand("timedump/{fuzzer}/{target}#{num}.time", fuzzer=['feedgeneration100_int', 'frafl_int', 'stg_int'], target=['waters_int', 'watersv2_int'],num=range(0,3)),
        expand("timedump/{fuzzer}/{target}#{num}.time", fuzzer=['random', 'stgpath'], target=['waters', 'watersv2'],num=range(0,3)),
        expand("timedump/{fuzzer}/{target}#{num}.time", fuzzer=['random_int', 'stgpath_int'], target=['waters_int', 'watersv2_int'],num=range(0,3))

rule all_showmap:
    input:
        expand("{remote}timedump/{fuzzer}/{target}#{num}_case.trace.ron",remote=remote, fuzzer=['frafl', 'stg'], target=['watersv2'],num=range(2,3)),
        expand("{remote}timedump/{fuzzer}/{target}#{num}_case.trace.ron",remote=remote, fuzzer=['frafl_int', 'stg_int'], target=['watersv2_int'],num=range(0,3)),
        expand("{remote}timedump/{fuzzer}/{target}#{num}_case.trace.ron",remote=remote, fuzzer=['random', 'stgpath'], target=['watersv2'],num=range(0,1)),
        expand("{remote}timedump/{fuzzer}/{target}#{num}_case.trace.ron",remote=remote, fuzzer=['random_int', 'stgpath_int'], target=['watersv2_int'],num=range(0,1))

    

rule all_bins:
    input:
        expand("bins/target_{target}{flag}",target=['random','frafl','stg','stgpath','feedgeneration100'],flag=['','_int'])