#!/bin/bash

DEFPORT=1001
DEFTTY=/dev/ttyUSB0
DEFBAUD=115200

TRANSPORT=${1}
FILESIZE=${2}
IP=${3}
PORT=${4:-$DEFPORT}
TTY=${5:-$DEFTTY}
BAUD=${6:-$DEFBAUD}

SRCFILE=/dev/shm/data.dump
DSTFILE=/tmp/data.dump

#Max size of supported UDP datagram size.
#Higher size packets will be partially lost by application.
MAXDGRAMMSIZE=512

[[ -z $TRANSPORT ]] && { echo "No transport: [tcp|udp]"; exit 1; }
[[ -z $IP ]] && { echo "No IP!"; exit 1; }
[[ -z $FILESIZE ]] && { echo "No file size!"; exit 1; }


trap abort INT


function cleanup {
    rm -rf $SRCFILE
    rm -rf $DSTFILE
    [[ -z $bg_proc ]] && exit 1
    kill $bg_proc
    wait $bg_proc 2>/dev/null
    #kill $(jobs -pr)
}

function abort {
    echo "Aborted!"
    cleanup
    exit 2
}

function setup {
    echo "Transfer size: $FILESIZE"

    #Check if port is available
    [[ $TRANSPORT == "udp" ]] && NC_ARGS="-u"
    nc -vz $NC_ARGS $IP $PORT    
    [[ $? != 0 ]] && { echo "Port $PORT is not available"; cleanup; exit 1; }
    #Wait UDP port ping
    [[ $TRANSPORT == "udp" ]] && sleep 1

    stty -F $TTY $BAUD # -icanon
    dd if=/dev/urandom of=$SRCFILE bs=$FILESIZE count=1 2>&1 &>/dev/null
    md5_sent=$(md5sum $SRCFILE | cut -d$' ' -f1)
    echo "Out MD5: $md5_sent"
}

function check_transfer {
    md5_rcvd=$(md5sum $DSTFILE | cut -d$' ' -f1)
    echo -e " In MD5: $md5_rcvd"
    echo "$(stat -c%s $DSTFILE)/$(stat -c%s $SRCFILE) bytes received"
    echo "--------------------------------------------"
    [[ $md5_sent != $md5_rcvd ]] && { echo -e "Test failed, transferred file MD5 missmatch!\r\n"; cleanup; exit 1; }
    echo -e "Test OK\r\n"
}

function test_serial2tcp {
    setup

    socat -u tcp:$IP:$PORT $DSTFILE &
    bg_proc=$!
    [[ -z $bg_proc ]] && { echo "Socat not started!"; cleanup; exit 1; }

    socat -u "SYSTEM:dd if=$SRCFILE bs=$FILESIZE count=1" $TTY,raw

    #Wait data to be transferred
    sleep 1

    check_transfer
    cleanup
}

function test_serial2udp {
    setup

    #Send initial request
    echo "." | nc -u $IP $PORT > $DSTFILE &
    bg_proc=$!
    [[ -z $bg_proc ]] && { echo "Socat not started!"; cleanup; exit 1; }

    #Wait initial request transferred
    sleep 1

    socat -u "SYSTEM:dd if=$SRCFILE bs=$FILESIZE count=1" $TTY,raw

    #Wait data to be transferred
    sleep 1

    check_transfer
    cleanup
}

function test_tcp2serial {
    setup

    socat -u $TTY,raw $DSTFILE &
    bg_proc=$!
    [[ -z $bg_proc ]] && { echo "Socat not started!"; cleanup; exit 1; }

    #socat -u "SYSTEM:dd if=$SRCFILE bs=$FILESIZE count=1" tcp:$IP:$PORT
    dd if=$SRCFILE | nc $IP $PORT

    #Wait data to be transferred
    sleep 1

    check_transfer
    cleanup
}

function test_udp2serial {
    setup

    socat -u $TTY,raw $DSTFILE &
    bg_proc=$!
    [[ -z $bg_proc ]] && { echo "Socat not started!"; cleanup; exit 1; }

    #Split data to the $MAXDGRAMMSIZE size blocks
    socat -u -b$MAXDGRAMMSIZE "SYSTEM:dd if=$SRCFILE" udp:$IP:$PORT

    #Wait data to be transferred
    sleep 1

    check_transfer
    cleanup
}


if [[ $TRANSPORT == "tcp" ]]
then
    # Serial -> TCP
    echo -e "\r\nSerial->TCP"
    echo "============================================"
    test_serial2tcp

    # TCP -> Serial
    echo -e "\r\nTCP->Serial"
    echo "============================================"
    test_tcp2serial
fi

if [[ $TRANSPORT == "udp" ]]
then
    # Serial -> UDP
    echo -e "\r\nSerial->UDP"
    echo "============================================"
    test_serial2udp

    # UDP -> Serial
    echo -e "\r\nUDP->Serial"
    echo "============================================"
    test_udp2serial
fi

exit 0