#!/bin/sh
# by Dominik Stanisław Suchora <hexderm@gmail.com>
# License: GNU GPLv3

#this script should be run in screen session

set -e

if [ "$#" -lt 3 ]
then
    echo "$0 <DIR> <THREADS> <COMMAND...>" >/dev/stderr
    exit 1
fi

src="$1"
shift
if [ \! -e "$src" ]
then
    echo "$src doesn't exist" >/dev/stderr
    exit 1
fi
if [ \! -d "$src" ]
then
    echo "$src is not a directory" >/dev/stderr
    exit 1
fi
if [ \! -r "$src" ]
then
    echo "$src is not readable" >/dev/stderr
    exit 1
fi

threads="$1"
[ -z "$threads" ] && threads=1
if echo "$threads" | grep -vq '^0*[1-9][0-9]*$'
then
    echo "number of threads '$threads' is too small"
    exit 1
fi
shift

ramfs="/tmp/ramfs"

if [ \! -d "$ramfs" ]
then
    if [ -e "$ramfs" ]
    then
        echo "$ramfs exists and is not a directory" >/dev/stderr
        exit 1
    fi
    echo "$ramfs does not exist" >/dev/stderr
    sudo sh -c "mkdir \"$ramfs\"; mount -t ramfs ramfs \"$ramfs\"; chmod -R 777 \"$ramfs\""
fi

if [ \! -w "$ramfs" ]
then
    echo "$ramfs is not writeable" >/dev/stderr
    exit 1
fi

dir=$(mktemp -d -p "$ramfs")

cp -r "$src" "$dir/input"
mkdir "$dir/output"

comm=""
while [ "$#" -gt 0 ]
do
    add="$1"
    if echo "$1" | grep -q "/" && [ -r "$1" ]
    then
        cp -r "$1" "$dir"
        add="$dir/$(basename "$1")"
    fi

    comm="$comm $add"

    shift
done

if [ "$threads" == '1' ]
then
    afl-fuzz -i "$dir/input" -o "$dir/output" $comm
else
    screen afl-fuzz -i "$dir/input" -M fuzzer1 -o "$dir/output" $comm &
    for i in $(seq 2 "$threads")
    do
        screen afl-fuzz -i "$dir/input" -S fuzzer"$i" -o "$dir/output" $comm &
    done

    wait
fi
