#!/bin/bash

input=$1
if [ -d "$input" ]; then
    pcap_dir=$input
    pcap_files=($(find ${pcap_dir} -type f -name "*.pcap"))
elif [ -f "$input" ]; then
    input_file=$input
    pcap_files=($input_file)
else
    echo "Usage: $0 input_pcap_dir or input_pcap_file"
    exit 1
fi

# remove l7 protocol fields
function preprocess_tshark_ouput() {
    input_file=$1
    output_file=$2
    cp ${input_file} ${output_file}
    kv_array=(
        ":tcp:pptp" ":tcp"
        ":tcp-text-lines" ":tcp"
        ":icmp:ip:udp" ":icmp"
        ":icmpv6:ipv6:udp" ":icmpv6"
        ":ieee8021ad" ":vlan"
        ":x509sat" ""
        ":x509ce" ""
        ":pkix1implicit" ""
        ":pkix1explicit" ""
        ":data-text-lines" ""
        ":http-text-lines" ""
        ":websocket" ""
        ":ssl" ""
        ":ftp-data" ""
        ":x11" ""
        ":ntp" ""
        ":rip" ""
        ":isakmp" ""
        ":esp" ""
        ":udpencap" ""
        ":sip:sdp" ""
        ":sip" ""
        ":sdp" ""
        ":rtcp" ""
        ":rtp" ""
        ":ssh" ""
        ":dns" ""
        ":gquic" ""
        ":http:data-text-lines" ""
        ":http:data" ""
        ":msmms" ""
        ":bfd" ""
        ":ssdp" ""
        ":mdns" ""
        ":radius" ""
        ":pop" ""
        ":smtp" ""
        ":rtmpt" ""
        ":bittorrent" ""
        ":oicq" ""
        ":json" ""
        ":media" ""
        ":telnet" ""
        ":nbdgm:smb:browser" ""
        ":smb2" ""
        ":smb" ""
        ":nbss" ""
        ":memcache" ""
        ":rtspi" ""
        ":rdt" ""
        ":rtsp" ""
        ":nbns" ""
        ":lcp" ""
        ":chap" ""
        ":ipcp" ""
        ":comp_data" ""
        ":ccp" ""
        ":snmp" ""
        ":socks:http:data" ""
        ":socks:http" ""
        ":socks" ""
        ":bgp" ""
        ":eigrp" ""
        ":bootp" ""
        ":xml" ""
        ":echo" ""
        ":vssmonitoring" ""
        ":mndp" ""
        ":websocket-text-lines" ""
        ":image-jfif" ""
        ":png" ""
        ":llmnr" ""
        ":pkcs-1" ""
        ":bitcoin:bitcoin" ""
        ":bitcoin" ""
        ":image-gif" ""
        ":dhcpv6" ""
        ":arp" ""
        ":ccsrl" ""
        ":h245" ""
        ":srp" ""
        ":amr" ""
        ":mp4v-es" ""
        ":ajp13" ""
        ":ocsp" ""
        ":irc" ""
        ":http" ""
        ":ftp" ""
        ":data" ""
        ":sctp:m3ua" ""
        ":sctp" ""
        ":igmp" ""
        ":icmp:ip:tcp" ":icmp"
        ":pwethheuristic:pwethnocw" ""
        ":ospf" ""
        ":teredo" ""
    )
    for ((i = 0; i < ${#kv_array[@]}; i += 2)); do
        key=${kv_array[i]}
        val=${kv_array[i + 1]}
        sed -i "s/$key/$val/g" ${output_file}
    done
}

# When MPLS is nested, packet_tool will output multiple mpls fields, and tshark will only output one mpls field, so we need to preprocess the output
function preprocess_tool_ouput() {
    input_file=$1
    output_file=$2
    cp ${input_file} ${output_file}
    kv_array=(
        ":mpls:mpls" ":mpls"
    )
    for ((i = 0; i < ${#kv_array[@]}; i += 2)); do
        key=${kv_array[i]}
        val=${kv_array[i + 1]}
        sed -i "s/$key/$val/g" ${output_file}
    done
}

err_count=0
pass_count=0
curr_count=0
total_count=${#pcap_files[@]}
tmp_file_dir="cmp_tmp_files/"
err_pcap_dir="cmp_err_pcaps/"

rm -rf ${err_pcap_dir} && mkdir ${err_pcap_dir}
for pcap in "${pcap_files[@]}"; do
    rm -rf ${tmp_file_dir} && mkdir ${tmp_file_dir}

    curr_count=$((curr_count + 1))

    # tshark output
    tshark -r ${pcap} -T fields \
        -e frame.number \
        -e frame.protocols \
        -e eth.src \
        -e eth.dst \
        -e vlan.id \
        -e ip.src \
        -e ip.dst \
        -e ipv6.src \
        -e ipv6.dst \
        -e tcp.srcport \
        -e tcp.dstport \
        -e udp.srcport \
        -e udp.dstport \
        >>${tmp_file_dir}/tshark_output.txt

    # packet_tool output
    ./packet_tool -f ${pcap} -t >>${tmp_file_dir}/tool_output.txt

    # compare tshark and packet_tool output
    preprocess_tshark_ouput ${tmp_file_dir}/tshark_output.txt ${tmp_file_dir}/tshark_format.txt
    preprocess_tool_ouput ${tmp_file_dir}/tool_output.txt ${tmp_file_dir}/tool_format.txt
    diff ${tmp_file_dir}/tshark_format.txt ${tmp_file_dir}/tool_format.txt >>${tmp_file_dir}/diff.txt

    # print result
    line_count=$(cat ${tmp_file_dir}/diff.txt | wc -l)
    if [ "$line_count" -ne 0 ]; then
        printf "\033[31m [${curr_count}/${total_count}] ${pcap} TEST FAILED \033[0m\n"
        cat ${tmp_file_dir}/diff.txt | head -n 100
        cp ${pcap} ${err_pcap_dir}
        err_count=$((err_count + 1))
    else
        printf "\033[32m [${curr_count}/${total_count}] ${pcap} TEST PASSED \033[0m\n"
        pass_count=$((pass_count + 1))
    fi
done

printf "\033[33m Total  : ${total_count} \033[0m\n"
printf "\033[32m Passed : ${pass_count}  \033[0m\n"
printf "\033[31m Failed : ${err_count}   \033[0m\n"

if [ "$err_count" -ne 0 ]; then
    printf "\n\033[31m failed pcap: ${err_pcap_dir} \033[0m\n"
fi
