#!/bin/bash
#
# USAGE: ./ditl_anonymization.sh [OPTIONS] file1 [file2 [... ] ]
#

INP_COMP=xzcat
FMT_PCAP=false
STREAMING=false
OUT_COMP=bzip2
OUT_SUFF=.bz2
OUT_FMT=pcap
OUT_DIR=.
OUT_PARALLEL=false

function usage() {
	cat <<EOF
USAGE:
  $0 [options] <input_glob>

WHERE
  <input_glob> is in FS (or NFS) 

OPTIONS are
  [-i input_decompressor] defaults to $INP_COMP
  [-k keyfile]            defaults to $OUT_DIR/keyfile
  [-f output_format]      defaults to $OUT_FMT (could be erf,pcap; assumes libtrace-tools installed)
  [-F FILTER]             specify the tcpdump filter passed to dag_scrubber (default: port 53)
	                          Note that payloads are kept ONLY for port 53.
  [-o output_compressor]  defaults to $OUT_COMP
  [-s output suffix]      defaults to $OUT_SUFF
  [-O output dir]         defaults to $OUT_DIR
  [-P]                    input/output is in PCAP format (ERF by default)
  [-p]                    output a list of jobs to be run by parallel
  [-S]                    streaming mode: read from stdin, write to stdout

EOF
	exit $1
}


PROG=$0
DONTSCRAMBLE='192.228.79.201/32 199.9.14.201/32 2001:500:84::b/128 2001:500:200::b/128'
while getopts "fF:hi:k:o:O:pPs:S" opt; do
    case $opt in
        ( i ) inp_comp=$OPTARG ;;
        ( k ) key_file=$OPTARG ;;
        ( o ) out_comp=$OPTARG ;;
	( s ) out_suffix=$OPTARG ;;
	( S ) STREAMING=true ;;
        ( f ) out_fmt=$OPTARG ;;
        ( F ) FILTER_PASS="$OPTARG" ;;
        ( O ) out_dir=$OPTARG ;;
        ( P ) fmt_pcap=true ;;
        ( p ) out_parallel=true ;;

        ( h|* ) usage 0 ;;
    esac
done
shift $((OPTIND-1))

if $STREAMING; then
    [ $# == "0" ] || usage 1
else
    [ $# -ge 1 ] || usage 1
fi

INP_GLOB="$@"

inp_comp=${inp_comp:-$INP_COMP}
out_comp=${out_comp:-$OUT_COMP}
fmt_pcap=${fmt_pcap:-$FMT_PCAP}
out_parallel=${out_parallel:-$OUT_PARALLEL}

out_suffix=${out_suffix:-$OUT_SUFF}
out_fmt=${out_fmt:-$OUT_FMT}
out_dir=${out_dir:-$OUT_DIR}

key_file=${key_file:-$out_dir/keyfile}

set -o pipefail

FILTER_PAYL='( port 53 || port 853 )'
FILTER_PASS=${FILTER_PASS:-"(tcp || udp) && $FILTER_PAYL"}
if $STREAMING; then
    if $fmt_pcap; then
        $inp_comp | 
            dag_scrubber -P -s $key_file -m --pass4=24 --pass6=64 \
                         --dont-scramble="$DONTSCRAMBLE" \
	                 -F "$FILTER_PASS" -n "$FILTER_PAYL" | 
            $out_comp || {
                echo "Streaming mode (pcap) failed" >&2
                exit 1
            }
    else
        #erf
        $inp_comp |
            dag_scrubber -s $key_file -m --pass4=24 --pass6=64 \
                         --dont-scramble="$DONTSCRAMBLE" \
                         -F "$FILTER_PASS" -n "$FILTER_PAYL" | 
            tracesplit erf:- $out_fmt:- |
            $out_comp || {
                echo "Streaming mode (erf) failed" >&2
                exit 1
            }
    fi
    exit 0
fi

for f in $INP_GLOB; do
    echo "# Working with file $f"
    fn=$(basename $f)
    ofn=${fn%.*}
    ofn=${ofn%.*}.${out_fmt}${out_suffix}

    [ -e "$out_dir/$ofn" ] && continue

    if $fmt_pcap; then
        #pcap
	if $out_parallel; then
            echo "$inp_comp $f | dag_scrubber -P -s $key_file -m --pass4=24 --pass6=64 \
                                              --dont-scramble='$DONTSCRAMBLE' \
                                              -F \"$FILTER_PASS\" -n \"$FILTER_PAYL\" | $out_comp >$out_dir/${ofn}"
	else
	    $inp_comp $f |
		dag_scrubber -P -s $key_file -m --pass4=24 --pass6=64 \
                             --dont-scramble="$DONTSCRAMBLE" \
			     -F "$FILTER_PASS" -n "$FILTER_PAYL" | 
                $out_comp >$out_dir/${ofn} || {
                echo "FAILED with $f" >&2
                exit 1
            }
	fi
    else
        #erf
	if $out_parallel; then
            echo "$inp_comp $f | dag_scrubber -s $key_file -m --pass4=24 --pass6=64 \
                                     --dont-scramble='$DONTSCRAMBLE' \
                                     -F \"$FILTER_PASS\" -n \"$FILTER_PAYL\" | \
                                     tracesplit erf:- $out_fmt:- | $out_comp >$OUT_DIR/${ofn}"
	else
            $inp_comp $f |
                dag_scrubber -s $key_file -m --pass4=24 --pass6=64 \
                         --dont-scramble="$DONTSCRAMBLE" \
                         -F "$FILTER_PASS" -n "$FILTER_PAYL" | 
                    tracesplit erf:- $out_fmt:- |
                    $out_comp >$OUT_DIR/${ofn} || {
                        echo "FAILED with $f" >&2
                        exit 1
            }
	fi
    fi

done