#!/bin/bash

IPTABLES_SAVE_FILE_NAME=/tmp/udp-client-server-iptables.save
PLATFORM=`uname -s`
PACKET_LOSS_RATE=$1

# BEGIN: Platform-specific.

function check_preconditions {
  case $PLATFORM in
  Linux)
    if [ -f $IPTABLES_SAVE_FILE_NAME ]
    then
      echo "iptables backup file already exists. Please clean up and try again."
      return 1
    fi
    ;;
  Darwin)
    ;;
  *)
    echo "Unsupported platform \"$PLATFORM\"."
    return 1
    ;;
  esac
}

function set_up {
  case $PLATFORM in
  Linux)
    iptables-save -c > $IPTABLES_SAVE_FILE_NAME
    iptables -A INPUT -m statistic --mode random --probability $PACKET_LOSS_RATE -j DROP
    ;;
  Darwin)
    ipfw add pipe 1 in proto udp
    ipfw pipe 1 config plr $PACKET_LOSS_RATE
    ;;
  *)
    echo "Unsupported platform \"$PLATFORM\"."
    return 1
    ;;
  esac
}

function tear_down {
  case $PLATFORM in
  Linux)
    if [ -f $IPTABLES_SAVE_FILE_NAME ]
    then
      iptables-restore -c < $IPTABLES_SAVE_FILE_NAME
      rm $IPTABLES_SAVE_FILE_NAME
    fi
    ;;
  Darwin)
    ipfw -q flush
    ;;
  *)
    echo "Unsupported platform \"$PLATFORM\"."
    return 1
    ;;
  esac
}

function display_diagnostics {
  echo "--------------------"
  case $PLATFORM in
  Linux)
    iptables --list
    ;;
  Darwin)
    ipfw list
    ;;
  *)
    echo "Unsupported platform \"$PLATFORM\"."
    return 1
    ;;
  esac
  echo "--------------------"
}

# END: Platform-specific.

if [[ `whoami` != "root" ]]
then
  echo "Please run this script as root or using sudo."
  exit 1
fi

if [[ "$PACKET_LOSS_RATE" == "" ]]
then
  echo "Please specify packet loss rate (in range 0.0 to 1.0)."
  exit 1
fi

function on_exit {
  tear_down
  exit $?
}
trap on_exit EXIT

check_preconditions
if [ $? != 0 ]
then
  exit 1
fi

set_up
if [ $? != 0 ]
then
  exit 1
fi

display_diagnostics
if [ $? != 0 ]
then
  exit 1
fi

while [ true ]
do
  echo -n "Press Q to quit and restore firewall settings >"
  read text
  if [[ "$text" == "Q" || "$text" == "q" ]]
  then
    break
  fi
done

tear_down
if [ $? != 0 ]
then
  exit 1
fi