#!/bin/bash
export LC_ALL=C
export TZ=UTC
server="${1:-localhost}"
rounds="${ROUNDS:-5}"
interactive="${INTERACTIVE:-no}"

set -u

for p in df cat head gzip ls awk rm sort ssh scp
do
  type "${p}" >/dev/null || { echo "Please install this missing program." >&2; exit 1; }
done

_cleanup() {
  rm -f /tmp/ssh_cipher_benchmark_input.$$ /tmp/ssh_cipher_benchmark_ouput.$$
  [[ ${rm_in:-0} = 1 ]] && rm /tmp/ssh_cipher_benchmark_input
}

ssh -q -oCompression=no -oBatchMode=yes "${server}" 'exit 42'
if [[ $? -ne 42 ]]
then
  printf 'FATAL: Could not connect to "%s" via SSH with -oBatchMode=yes in effect. Please ensure that this works.\n' "${server}" >&2
  exit 1
fi

rm_in=0
trap _cleanup EXIT

printf 'About to benchmark SSH cipher performance with peer "%s".\n' "${server}"
printf 'Please ensure constant server and client loads for all %d iterations per cipher.\n' "${rounds}"

if ! [[ -f /tmp/ssh_cipher_benchmark_input ]]
then
  printf 'No transfer input file found at /tmp/ssh_cipher_benchmark_input - creating a temporary one...\n'
  echo
  df -m /tmp/
  read -r -p "How many MBytes of transfer input data to generate in /tmp/? " iters
  gzip -c8 < /dev/urandom | head -c 1048576 >/tmp/ssh_cipher_benchmark_input.$$
  for (( i = 0; i < iters; i++ ))
  do
    cat /tmp/ssh_cipher_benchmark_input.$$ >> /tmp/ssh_cipher_benchmark_input
  done
  rm_in=1
fi

if ! ls -lh /tmp/ssh_cipher_benchmark_input
then
  printf 'FATAL: Required input file "/tmp/ssh_cipher_benchmark_input" not present.\n' >&2
  exit 1
fi
echo "Transfer input file ready. Determining peer ciphers..."

echo
IFS=, readarray -t ciphers < \
    <(
	ssh -vv -oBatchMode=yes "${server}" : 2>&1 </dev/null |\
	tr '\r\n' '\n' |\
	awk '/^debug2: peer server KEXINIT proposal/{parse=1} /^debug2: ciphers stoc:/{if(parse){split($NF,C,","); for(c in C){printf("%s\n", C[c])}; exit}} END{print "__END_OF_CIPHERS__"}'
    )

if [[ ${ciphers[0]} = __END_OF_CIPHERS__ ]]
then
  printf 'FATAL: Could not parse supported ciphers for peer "%s". This might be a bug - sorry.\n' "${server}"
  exit 1
else
  printf 'Negotiated %d ciphers for this run, performing %d rounds each.\n' "$(( ${#ciphers[@]} - 1 ))" "${rounds}"
fi

echo
cat /tmp/ssh_cipher_benchmark_input > /dev/null
for c in "${ciphers[@]}"
do
  [[ ${c} = __END_OF_CIPHERS__ ]] && printf '\n' && continue
  printf 'Benchmarking %30s \t' "${c}"
  [[ ${interactive} != no && -t 1 ]] && read -r -p 'Hit ENTER to continue... ' -s _
  for (( i=0; i < rounds; i++ ))
  do
    read -r timeinfo < <(
      set +x;
      TIMEFORMAT='%3R secs  %P %%cpu   %3U user %3S sys';
      { time scp -q -oCompression=no -oBatchmode=yes -oCiphers="${c}" \
        /tmp/ssh_cipher_benchmark_input "${server}":/dev/null; } 2>&1
    )
    printf '.'
    printf '%s   %s\n' "${timeinfo}" "${c}" >> /tmp/ssh_cipher_benchmark_ouput.$$
  done
  printf '\n'
done

echo
echo 'Results (from worst to best, descending):'
echo
sort -r -n /tmp/ssh_cipher_benchmark_ouput.$$
