#!/bin/bash

# dd-resume version #VERSION#

usage() {
  >&2 printf 'Usage:  dd-resume [options]\n'
  >&2 printf '  copy resumably with dd via network\n\n'
  >&2 printf 'Options (all options are mandatory unless explicitely stated):\n'
  >&2 printf '  --input $input-device      name of input device\n'
  >&2 printf '  --host $input-host         name of input device'"'"'s host (target only)\n'
  >&2 printf '  --log-file $log-file       name of log file (target only)\n'
  >&2 printf '  --output $output-device    name of output device (target only)\n'
  >&2 printf '  --port $remote-port        port of target host (source only)\n'
  >&2 printf '  --start $start-position    start position for transfer (source only)\n'
  >&2 printf '  --required                 print list of required programs and exit\n'
  >&2 printf '#HELPTEXT#                  #\n'
  [ -n "$1" ] && exit $1
  exit 1
}

block_size=$((1024*1024))
dd_options='bs='"${block_size}"
needed_programs='dd nc ss screen pgrep kill ssh'

eval set -- "$(
  getopt -o '' \
    --long input: \
    --long host: \
    --long log-file: \
    --long output: \
    --long port: \
    --long start: \
    --long required \
    --long help \
    --long version \
    -n "$(basename "$0")" \
    -- "$@" \
    || echo "usage"
)"

while true; do
  case $1 in
    --input)
      if [ -n "${input_device}" ]; then
        >&2 printf 'cannot handle multiple input-devices\n\n'
        usage
      fi
      shift
      input_device="$1"
    ;;
    --host)
      if [ -n "${host}" ]; then
        >&2 printf 'cannot handle multiple hosts\n\n'
        usage
      fi
      shift
      host="$1"
    ;;
    --log-file)
      if [ -n "${log_file}" ]; then
        >&2 printf 'cannot handle multiple log files\n\n'
        usage
      fi
      shift
      log_file=$(readlink -f "$1")
    ;;
    --output)
      if [ -n "${output_device}" ]; then
        >&2 printf 'cannot handle multiple output devices\n\n'
        usage
      fi
      shift
      output_device="$1"
    ;;
    --port)
      if [ -n "${remote_port}" ]; then
        >&2 printf 'cannot handle multiple ports\n\n'
        usage
      fi
      shift
      remote_port="$1"
    ;;
    --start)
      if [ -n "${start}" ]; then
        >&2 printf 'cannot handle multiple starts\n\n'
        usage
      fi
      shift
      start="$1"
    ;;
    --required)
      printf '%s\n' ${needed_programs}
      exit 0
    ;;
    --help)
      usage 0
    ;;
    --version)
      >&2 printf '%s\n' '#VERSION#'
      exit 0
    ;;
    --)
      shift
      break
    ;;
    *)
      >&2 printf 'oops, option "%s" is unkknown\n' "$1"
      exit 1
  esac
  shift
done

if [ $# -ne 0 ]; then
  >&2 printf 'too many arguments\n\n'
  usage
fi

for needed in ${needed_programs}; do
  if ! which "${needed}" >/dev/null; then
    >&2 printf 'please install required program %s\n' "${needed}"
    exit 1
  fi
done

if [ -z "${input_device}" ]; then
  >&2 printf 'no input device specified\n\n'
  usage
fi

if [ -n "${host}" ]; then
  # was run on the receiver side

  if [ -z "${log_file}" ] || \
    [ -z "${output_device}" ] || \
    [ -n "${remote_port}" ] || \
    [ -n "${start}" ]; then
    >&2 printf 'conflicting options\n\n'
    usage
  fi

  if [ ! -w "${log_file%/*}" ] && [ ! -w "${log_file}" ]; then
    >&2 printf 'cannot write/create log-file %s\n\n' "${log_file}"
    usage
  fi

  if [ ! -w "${output_device}" ]; then
    >&2 printf 'cannot write to output device %s\n\n' "${output_device}"
    usage
  fi

  start=0
  if [ -f "${log_file}" ]; then
    start=$(
      sed -n '
        1 {
          s/^start: //
          p
          q
        }
      ' "${log_file}"
    )
    copied=$(
      sed '
        s/^\([0-9]\+\) bytes .\+ copied, .*$/\1/
        t
        d
      ' "${log_file}" | \
      sort -n | \
      tail -n1
    )
    if [ -n "${copied}" ]; then
      start=$((start + copied/block_size))
    fi
  fi

  previously_running=$(
    pgrep -x 'dd|nc'
  )

  printf 'start: %s\n' "${start}" > \
    "${log_file}"

  screen -d -m bash -c '
    set +o pipefail
    nc -l -c | \
      LC_ALL=C dd of="'"${output_device}"'" '"${dd_options}"' seek='"${start}"' 2>>"'"${log_file}"'" >/dev/null
    printf '"'"'\nwrite-pipe-exit: %s\n'"'"' $? >>"'"${log_file}"'"
  '

  for i in {1..20}; do
    if [ -n "${port}" ]; then
      break
    elif [ ${i} -gt 1 ]; then
      sleep 0.1
    fi
    nc_pid=$(
      pgrep -x nc | \
        grep -vxF "$(printf '%s\n' "${previously_running}")"
    )
    dd_pid=$(
      pgrep -x dd | \
        grep -vxF "$(printf '%s\n' "${previously_running}")"
    )
    port=$(
      ss -nlp | \
        grep -F ",pid=${nc_pid}," | \
        awk '{print $5}' | \
        sed 's/^.*://'
    )
  done
  if [ -z "${port}" ]; then
    >&2 printf 'could not determine listening port\n'
    exit 1
  fi
  screen -d -m bash -c '
    while kill -SIGUSR1 '"${dd_pid}"'; do
      sleep 1
    done
  '

  ssh "${host}" "$0 --input ${input_device} --start ${start} --port ${port}"
  printf '\nssh-exit: %s\n' $? >>"${log_file}"

else
  # was run on the sender side

  if [ -n "${log_file}" ] || \
    [ -n "${output_device}" ] || \
    [ -z "${remote_port}" ] || \
    [ -z "${start}" ]; then
    >&2 printf 'conflicting options\n\n'
    usage
  fi

  remote_ip="${SSH_CLIENT%% *}"
  if [ -z "${remote_ip}" ]; then
    >&2 printf 'cannot obtain remote ip\n'
    exit 1
  fi
  if [ ! -r "${input_device}" ]; then
    >&2 printf 'cannot read from %s\n\n' "${input_device}"
    usage
  fi
  set +o pipefail
  dd if="${input_device}" ${dd_options} skip="${start}" 2>/dev/null | \
    nc -cn "${remote_ip}" "${remote_port}"
  exit $?

fi