#!/bin/bash
set -e
_help() {
    echo '
SSH performance test

args:  [-s size] [-r repeats] [-t timeout] [-d/-u] user@host

Options:

  -s size: default: 100MB, acceptable suffixes: M,G,T,P
  -r repeats: default: 1
  -t timeout: default: 60
  -d  Download only
  -u  Upload only
'
    exit
}

_size_in_bytes() {
  size=${size^^}
  if [[ ! ${size: -1} = "B" ]]; then
    size="${size}B"
  fi
  numeric=${size%??}
  suffix=${size: -2}
  onemb=$(( 1024**2 ))
  case $suffix in
    MB)  bytes=$(( 1024**2 * $numeric )) ;;
    GB)  bytes=$(( 1024**3 * $numeric )) ;;
    TB)  bytes=$(( 1024**4 * $numeric )) ;;
    PB)  bytes=$(( 1024**5 * $numeric )) ;;
    *) echo "Cant parse size"; exit 1; ;;
  esac
  megabytes=$(( $bytes / (1024**2) ))

}
_msg() {
    echo "$1 $size repeated $repeats times, timeout ${timeout}s:"
}
size=100MB
repeats=1
upload=1
download=1
timeout=60

shift_arg=0
for (( i=1; i<=$#; i++ )); do
  value=${!i}
  j=$(( i + 1 ))
  if [[ $shift_arg -eq 1 ]]; then
      shift_arg=0
      continue
  fi
  if [[ "$value" = "--help" ]]; then _help; fi
  if [[ "${value}" = "-"* ]]; then
    [[ "$value" = -h ]] && { _help; }
    [[ "$value" = -s ]] && { size=${!j}; }
    [[ "$value" = -r ]] && { repeats=${!j}; }
    [[ "$value" = -u ]] && { download=0; continue; }
    [[ "$value" = -d ]] && { upload=0; continue; }
    [[ "$value" = -t ]] && { timeout=${!j}; }
    shift_arg=1
    continue
  else
    host="$value"
  fi
done


if [[ -z "$host" ]]; then
  _help
fi
_size_in_bytes


if [[ $download -eq 1 ]]; then
    _msg "Downloading"
    for i in $( seq $repeats ); do
        timeout $timeout ssh -o 'Compression no' "$host" "dd if=/dev/zero bs=$onemb count=$megabytes 2>/dev/null" | \
          pv -W -f -s ${bytes} > /dev/null
    done
fi
if [[ $upload -eq 1 ]]; then
    _msg "Uploading"
    for i in $( seq $repeats ); do
        timeout $timeout dd if=/dev/zero bs=$onemb count=$megabytes 2>/dev/null | \
          ssh -o 'Compression no' "$host" "cat | pv -W -f -s ${bytes} > /dev/null"
    done
fi