#!/bin/bash -e
#
# Script Name: curl-prgrs
#
# Description:
# The curl-prgrs script augments the functionality of curl by adding a custom progress bar
# for file downloads, making it a convenient drop-in replacement for curl with the added
# benefit of visual progress tracking. It accepts the same command-line arguments and options as curl.
#
# Additionally, this script provides a mechanism to mitigate endless data attacks (as defined
# in the TUF threat model) by setting the CURL_PRGRS_MAX_FILE_SIZE_BYTES environment variable.
# Unlike curl's --max-filesize option, which does not have an effect when the file size is unknown prior to download,
# this script ensures the file size restriction is enforced. The limitation of curl's --max-filesize is acknowledged in
# the curl man page as follows:
# "NOTE: The file size is not always known prior to download, and for such files this
# option has no effect even if the file transfer ends up being larger than this given limit."
#
# Usage:
# Substitute curl with curl-prgrs for downloading files:
#     $ curl-prgrs -O http://example.com/file.tar.gz
#     $ curl-prgrs http://example.com/file.tar.gz > file.tar.gz
#
# Features:
# - Displays a custom-drawn progress bar to visualize download progress.
# - Provides error handling and cleanup for various termination scenarios.
# - Conducts preliminary checks for required dependencies before proceeding.
# - Utilizes temporary files for capturing progress information and facilitating process communication.
# - Allows custom configurations via environment variables.
#
# Environment Variables:
# - CURL: Defines the path to the curl binary. Acceptable values are "curl" or "scurl" (secure curl). Default: "curl"
# - CURL_PRGRS_MAX_FILE_SIZE_BYTES: Sets the maximum allowed file size for downloads in bytes. Mandatory.
# - CURL_OUT_FILE: Specifies the path to the output file for download. Mandatory.
# - CURL_PRGRS_EXEC: Designates the command to execute for updating the progress bar.
#   Last argument will be the number in percent. Optional.
#
# Authors:
# - Sam Stephenson <sstephenson@gmail.com>
# - Patrick Schleizer <adrelanos@whonix.org>
#
# License:
# (c) 2013 Sam Stephenson <sstephenson@gmail.com>
# Released into the public domain on 2013-01-21.
# Source: https://gist.github.com/sstephenson/4587282
#
# Modifications:
# - Subsequent modifications by Patrick Schleizer under the same license.
#
# Note: Several parts of this script are broken into functions when it doesn't
# seem necessary. This is for the sake of automated testing. Do not factor out
# functions unless necessary.

#set -x

## provides: was_executed
# shellcheck source=./check_runtime.bsh
source "${HELPER_SCRIPTS_PATH:-}"/usr/libexec/helper-scripts/check_runtime.bsh

## provides: draw_progress_bar
source "${HELPER_SCRIPTS_PATH:-}"/usr/libexec/helper-scripts/progress-bar

## provides: is_whole_number
source "${HELPER_SCRIPTS_PATH:-}"/usr/libexec/helper-scripts/strings.bsh

# shellcheck source=./has.bsh
source "${HELPER_SCRIPTS_PATH:-}"/usr/libexec/helper-scripts/has.bsh

## Allows tests to override the stderr TTY check.
stderr_is_tty() {
  [ -t 2 ]
}

initialize_terminal() {
  ## We want to print the progress bar to stderr, but only if stderr is a
  ## terminal. To avoid a conditional every time we print something, we can
  ## instead print everything to an allocated file descriptor, and then point
  ## that file descriptor to the right place: stderr if it's a TTY, or
  ## /dev/null otherwise.

  if stderr_is_tty; then
    exec {stderr_fd}>&2
  else
    exec {stderr_fd}>/dev/null
  fi
}

initialize_variables() {
  has tput
  has curl
  has safe-rm
  has mktemp

  default_if_empty CURL "curl"
  default_if_empty CURL_PRGRS_MAX_FILE_SIZE_BYTES ""
  default_if_empty CURL_OUT_FILE ""
  default_if_empty CURL_PRGRS_EXEC ""
  default_if_empty curl_prgrs_print_progress "yes"
  percent_last=""

  expected_header_size=8000
  maximum_http_header_size=32000
}

## Separate from initialize_variables to avoid creating a tempdir if a config
## error occurs.
initialize_temporary_files() {
  temp_dir_auto_generated=true
  temporary_directory="$(mktemp --directory)"

  ## Compute names for our temporary files by joining the current date and
  ## time with the current process ID. We will need two temporary files: one
  ## for reading progress information from curl, and another for sending the
  ## exit status of curl from the forked child process back to the parent.
  statusfile="${temporary_directory}/status"

  ## curl runs in a subshell, so we have to save its PID to a file so the
  ## parent can read it.
  curl_pid_file="${temporary_directory}/curl.pid"
}

check_variables() {
  if [ "${CURL_OUT_FILE}" = "" ]; then
    stecho "${BASH_SOURCE[0]} ERROR: Variable CURL_OUT_FILE is empty." >&"${stderr_fd}"
    exit 57
  fi
  if [ "${CURL_PRGRS_MAX_FILE_SIZE_BYTES}" = "" ]; then
    stecho "${BASH_SOURCE[0]} ERROR: Variable CURL_PRGRS_MAX_FILE_SIZE_BYTES is empty." >&"${stderr_fd}"
    exit 57
  fi

  is_whole_number "${CURL_PRGRS_MAX_FILE_SIZE_BYTES}"
  is_whole_number "${expected_header_size}"
  is_whole_number "${maximum_http_header_size}"
}

## Define our `shutdown` function, which will be responsible for cleaning
## up when the program terminates, either normally or abnormally.
# shellcheck disable=SC2317
shutdown() {
  local exit_code="$?"
  local signal="$1"
  local last_err="${BASH_COMMAND}"
  if [ "${signal}" = "err" ]; then
    stecho "${BASH_SOURCE[0]} ERROR: Signal ${signal} received while running '${last_err}'. Exiting." >&"${stderr_fd}"
    stecho "${BASH_SOURCE[0]} ERROR: BASH_COMMAND '${BASH_COMMAND}' exit code '${exit_code}'." >&"${stderr_fd}"
  elif [ "${signal}" = "exit" ]; then
    true "${BASH_SOURCE[0]} INFO: Signal ${signal} received. Exiting." >&"${stderr_fd}"
  else
    stecho "${BASH_SOURCE[0]} INFO: Signal ${signal} received. Exiting." >&"${stderr_fd}"
  fi

  trap - SIGHUP SIGINT SIGTERM ERR EXIT
  sync

  ## 'status' is the single source of truth for the exit code: curl_exit is the
  ## only writer of the status file, recording the specific outcome (0, or a
  ## code such as 81 'file too large' / 114 / 115 / 116 from the endless-data
  ## mitigation). That recorded code MUST survive to the caller -- masking every
  ## abort as a generic code would defeat the documented exit-code contract.
  ##
  ## Fallbacks, only when the recorded code cannot be trusted:
  ##   111 - the status file holds a non-number (corrupt),
  ##   112 - no status file (shutdown reached before any curl_exit),
  ##   110 - the recorded status is 0 but the shutdown was NOT a clean success:
  ##         any signal but 'exit' (the ERR trap on a command that was NOT a
  ##         curl_exit -- which never records 0 on failure; or a
  ##         SIGTERM/SIGINT/SIGHUP mid-run, after a header curl_exit recorded 0
  ##         and before the body completed), OR the normal 'exit' path carrying a
  ##         NON-zero code (a failed 'wait' -> exit "${wait_exit_code}", which
  ##         does NOT trip ERR). Only a clean 'exit' with exit code 0 may report
  ##         success; a generic error beats a false success.
  local status

  if [ -f "${statusfile}" ]; then
    true "${BASH_SOURCE[0]} INFO: got status file"
    status="$(stcat "${statusfile}")"
    if ! is_whole_number "${status}"; then
      true "${BASH_SOURCE[0]} ERROR: status is not a number! status: '${status}'"
      status="111"
    elif [ "${status}" -eq 0 ] && { [ "${signal}" != "exit" ] || [ "${exit_code}" -ne 0 ]; }; then
      true "${BASH_SOURCE[0]} ERROR: abnormal termination (signal '${signal}', exit code '${exit_code}') with a 0 status file!"
      status="110"
    fi
  else
    true "${BASH_SOURCE[0]} ERROR: no status file"
    status="112"
  fi

  ## If we are exiting normally, jump back to the beginning of the line
  ## and clear it. Otherwise, print a newline.
  if [ "${status}" -eq 0 ]; then
    printf '%b' "\x1B[0G\x1B[0K" >&"${stderr_fd}"
  else
    printf '%s\n' '' >&"${stderr_fd}"
  fi

  #stat="$(stcat "$statusfile")"
  #stecho "$stat" >&"${stderr_fd}"

  ## Read the PID curl_download published before removing the temp dir.
  local published_curl_pid=""
  default_if_empty curl_pid_file ""
  if [ -n "${curl_pid_file}" ] && [ -f "${curl_pid_file}" ]; then
    published_curl_pid="$(cat -- "${curl_pid_file}" 2>/dev/null || true)"
  fi

  ## Only clean up the temp dir if it is ephemeral and we are the parent shell.
  if [ "${temp_dir_auto_generated}" = "true" ] && [ "${BASHPID}" = "$$" ]; then
    safe-rm -r -f -- "${temporary_directory}"
  fi

  true curl_pid
  : "${curl_pid:=""}"

  ## Kill curl by both the local curl_pid (set when shutdown runs inside the
  ## worker) and the published PID (the only one the main shell knows).
  processes_list="${curl_pid} ${published_curl_pid}"
  for processes_item in ${processes_list} ; do
    if kill -0 -- "${processes_item}" 2>/dev/null ; then
      #ps -p "$processes_item" || true
      kill -s sigkill -- "${processes_item}" &>/dev/null || true
      ## There is only ever one curl process we want to kill, so we're done
      ## now.
      break
    fi
  done

  true "${BASH_SOURCE[0]} INFO: exit with status ${status}"
  exit "${status}"
}

# shellcheck disable=SC2317
shutdown_sigint() {
  shutdown sigint
}
# shellcheck disable=SC2317
shutdown_sigterm() {
  shutdown sigterm
}
# shellcheck disable=SC2317
shutdown_err() {
  shutdown err
}
# shellcheck disable=SC2317
shutdown_exit() {
  shutdown exit
}
# shellcheck disable=SC2317
shutdown_sighup() {
  shutdown sighup
}

traps_enable() {
  ## Register our `shutdown` function to be invoked when the process dies.
  trap shutdown_sigint SIGINT
  trap shutdown_sigterm SIGTERM
  trap shutdown_err ERR
  trap shutdown_exit EXIT
  trap shutdown_sighup SIGHUP
}

## Map downloaded bytes and expected total to a clamped 0..100 percentage.
## Both arguments are expected to be already-validated whole numbers.
##
## A total of 0 (an empty file / 'Content-Length: 0') means there is nothing
## left to fetch, so it maps to 100.
compute_percent() {
  local bytes="$1"
  local length="$2"
  local percent

  if [ "${length}" -le 0 ]; then
    printf '%s' 100
    return 0
  fi

  percent=$(( bytes * 100 / length ))
  if [ "${percent}" -ge 100 ]; then
    percent=100
  fi
  printf '%s' "${percent}"
}

## The `print_progress` function draws our progress bar to the screen. It
## takes two arguments: the number of bytes read so far, and the total
## number of bytes expected.
print_progress() {
  local bytes="$1"
  local length="$2"

  if ! is_whole_number "${bytes}" ; then
    curl_exit 113
  fi
  if ! is_whole_number "${length}" ; then
    curl_exit 113
  fi

  ## If we are expecting less than 8 KB of data, don't bother drawing a
  ## progress bar. (This helps avoid a flicker when following redirects.)
  #[ "$length" -gt 8192 ] || return 0

  ## Calculate the progress percentage and the size of the filled and
  ## unfilled portions of the progress bar.
  local percent
  true "${BASH_SOURCE[0]} INFO: bytes: '${bytes}'"
  true "${BASH_SOURCE[0]} INFO: length: '${length}'"
  percent="$(compute_percent "${bytes}" "${length}")"

  if [ "${percent_last}" = "${percent}" ]; then
    true "${BASH_SOURCE[0]} INFO: percentage number unchanged. Not re-drawing progress bar to avoid flicker."
  else
    draw_progress_bar "${percent}" >&"${stderr_fd}"
    if [ "${CURL_PRGRS_EXEC}" = "" ]; then
      true "${BASH_SOURCE[0]} INFO: CURL_PRGRS_EXEC is empty. Not executing CURL_PRGRS_EXEC."
    else
      true "${BASH_SOURCE[0]} INFO: CURL_PRGRS_EXEC is set. Executing CURL_PRGRS_EXEC..."
      true "${BASH_SOURCE[0]} INFO: ${CURL_PRGRS_EXEC} '${percent}'"
      ${CURL_PRGRS_EXEC} "${percent}" >&"${stderr_fd}"
      true "${BASH_SOURCE[0]} INFO: CURL_PRGRS_EXEC success."
    fi
  fi

  percent_last="${percent}"
}

curl_exit() {
  curl_exit_code="$1"
  true "${BASH_SOURCE[0]} INFO: write ${curl_exit_code} to ${statusfile}"
  stecho "${curl_exit_code}" > "${statusfile}"
  ## curl is either already shut down or is about to be killed, so clear its
  ## PID file so the parent doesn't try to kill it again.
  if [ -n "${curl_pid_file:-}" ]; then
    printf '%s' '' > "${curl_pid_file}" 2>/dev/null || true
  fi
  if [ "${curl_exit_code}" = "0" ]; then
    return 0
  fi
  : "${curl_pid:=""}"
  if [ "${curl_pid}" != "" ]; then
    if kill -0 -- "${curl_pid}" 2>/dev/null; then
      kill -s SIGKILL -- "${curl_pid}" &>/dev/null || true
    fi
  fi
  return "${curl_exit_code}"
}

## Given the bytes seen on disk so far and the two ceilings (the hard
## CURL_PRGRS_MAX_FILE_SIZE_BYTES cap and the advertised content length), echo
## the curl exit code the caller must raise, or 0 when the size is still within
## bounds:
##   113 - size is not a whole number
##    81 - size exceeded the hard maximum file-size cap
##   114 - size exceeded the advertised content length
classify_download_size() {
  local downloaded="$1"
  local max_bytes="$2"
  local content_length="$3"

  if ! is_whole_number "${downloaded}" ; then
    printf '%s' 113
    return 0
  fi
  if [ "${downloaded}" -gt "${max_bytes}" ]; then
    printf '%s' 81
    return 0
  fi
  if [ "${downloaded}" -gt "${content_length}" ]; then
    printf '%s' 114
    return 0
  fi
  printf '%s' 0
}

## The content-length ceiling classify_download_size must enforce for the
## current phase. The body phase uses the content length, the header phase uses
## the maximum allowable download size since we don't know how big the header
## will be in advance.
content_length_ceiling_for_phase() {
  local header_download="$1"
  local advertised="$2"
  local max_bytes="$3"

  if [ "${header_download}" = "true" ]; then
    printf '%s' "${max_bytes}"
    return 0
  fi
  printf '%s' "${advertised}"
}

## Ensure the file being downloaded has not grown larger than either ceiling.
enforce_file_size() {
  local content_length_ceiling="$1"
  if [ ! -f "${CURL_OUT_FILE}" ]; then
    return 0
  fi
  size_file_downloaded_bytes="$(stat -c "%s" -- "${CURL_OUT_FILE}")"
  true "size_file_downloaded_bytes: ${size_file_downloaded_bytes}"
  true "CURL_PRGRS_MAX_FILE_SIZE_BYTES: ${CURL_PRGRS_MAX_FILE_SIZE_BYTES}"
  true "content_length_ceiling: ${content_length_ceiling}"
  local size_check_code
  size_check_code="$(classify_download_size "${size_file_downloaded_bytes}" "${CURL_PRGRS_MAX_FILE_SIZE_BYTES}" "${content_length_ceiling}")"
  if [ "${size_check_code}" != "0" ]; then
    curl_exit "${size_check_code}"
  fi
}

curl_download() {
  local size_file_downloaded_bytes

  ${CURL} --no-progress-meter "$@" &
  curl_pid="$!"
  printf '%s\n' "${curl_pid}" > "${curl_pid_file}"

  ## Additional validation.
  ## Already validated earlier, but:
  ## /usr/libexec/helper-scripts/curl-prgrs: line 266: [: : integer expression expected
  if ! is_whole_number "${curl_prgrs_content_length}" ; then
    curl_exit 116
  fi

  ## Default to false so truncation isn't accidentally permitted.
  default_if_empty header_download 'false'

  local content_length_ceiling
  content_length_ceiling="$(content_length_ceiling_for_phase \
    "${header_download}" "${curl_prgrs_content_length}" "${CURL_PRGRS_MAX_FILE_SIZE_BYTES}")"

  while true ; do
    if [ -f "${CURL_OUT_FILE}" ]; then

      enforce_file_size "${content_length_ceiling}"

      if [ "${curl_prgrs_print_progress}" = "yes" ]; then
        ## Need to print to stderr to avoid confusing the stdout output of this command.
        #stecho "${BASH_SOURCE[0]} INFO: print_progress '$size_file_downloaded_bytes' '$curl_prgrs_content_length'" >&2
        print_progress "${size_file_downloaded_bytes}" "${curl_prgrs_content_length}"
      fi
    fi

    if ! kill -0 -- "${curl_pid}" 2>/dev/null; then
      break
    fi

    ## Poll interval is overridable so tests can speed through this.
    sleep "${curl_prgrs_poll_interval:-1}"
  done

  ## curl already terminated.
  enforce_file_size "${content_length_ceiling}"

  default_if_empty size_file_downloaded_bytes ""
  if is_whole_number "${size_file_downloaded_bytes}" ; then
    true "size_file_downloaded_bytes: ${size_file_downloaded_bytes}"
    true "curl_prgrs_content_length: ${curl_prgrs_content_length}"
    if [ "${header_download}" = "false" ]; then
      if [ "${size_file_downloaded_bytes}" -lt "${curl_prgrs_content_length}" ]; then
        curl_exit 115
      fi
    fi
  fi

  curl_exit_code=0
  wait "${curl_pid}" || { curl_exit_code=$? ; true; };
  curl_exit "${curl_exit_code}"
}

remove_argument_for_header_request() {
  local arg_item
  local arg_list=()
  local skip_next=false
  header_arguments=()

  ## Due to how curl's argument format works, this might strip things it
  ## shouldn't (for instance if `--output` is passed as the argument to a
  ## different option). This shouldn't be an issue in practice.
  for arg_item in "$@"; do
    if [ "${skip_next}" = true ]; then
      skip_next=false
      continue
    fi

    if [ "${arg_item}" = "--continue-at" ]; then
      skip_next=true
      continue
    fi
    if [ "${arg_item}" = "-C" ]; then
      skip_next=true
      continue
    fi

    if [ "${arg_item}" = "--output" ]; then
      skip_next=true
      continue
    fi
    if [ "${arg_item}" = "-o" ]; then
      skip_next=true
      continue
    fi

    arg_list+=("${arg_item}")
  done

  header_arguments=("${arg_list[@]}")

  ## Cannot use. Collapses newlines.
  #stecho "${args[@]}"
}

run_body_download() {
  header_download="false" curl_download "$@"
}

run_download() {
  local header_file

  ## {{{ Debugging.
#   local i arg
#   printf '%s\n' "Before number of args: $#"
#   i=0
#   for arg in "$@"; do
#     i=$(( i + 1 ))
#     printf '  [%d]=%q\n' "$i" "$arg"
#   done
  ## }}}

  ## sets: header_arguments
  remove_argument_for_header_request "${@}"

  ## {{{ Debugging.
#   printf '%s\n' ""
#   printf '%s\n' "After number of args: ${#header_arguments[@]}"
#   i=0
#   for arg in "${header_arguments[@]}"; do
#     i=$(( i + 1 ))
#     printf '  [%d]=%q\n' "$i" "$arg"
#   done
  ## }}}

  header_file="${temporary_directory}/header"

  true "${BASH_SOURCE[0]} INFO: Download header..."

  ## Determine curl_prgrs_content_length.
  ## While we don't know the expected size of the header,
  ## curl_prgrs_content_length and
  ## CURL_PRGRS_MAX_FILE_SIZE_BYTES are set to reasonable values.
  ##
  ## CURL_PRGRS_EXEC="" to avoid a progress bar for the header download.
  ## That would confuse yad.
  ##
  ## CURL_OUT_FILE and
  ## --output "$header_file" to avoid overwriting files when using "--continue-at -".
  ## '--write-out' will echo.
  curl_prgrs_content_length="$(
    header_download="true" \
    curl_prgrs_content_length="${expected_header_size}" \
    CURL_PRGRS_MAX_FILE_SIZE_BYTES="${maximum_http_header_size}" \
    CURL_PRGRS_EXEC="" \
    CURL_OUT_FILE="${header_file}" \
      curl_download \
        --head \
        --write-out '%header{Content-Length}' \
        --output "${header_file}" \
        "${header_arguments[@]}" \
    )"

  ## Reset from previews invocation of curl_download, which calls print_progress.
  percent_last=""

  true "${BASH_SOURCE[0]} INFO: Header download done."

  if ! is_whole_number "${curl_prgrs_content_length}" ; then
    curl_exit 116
  fi

  ## Reject an implausibly large advertised Content-Length.
  if [ "${#curl_prgrs_content_length}" -gt 16 ]; then
    curl_exit 116
  fi

  ## Reset the status file between phases.
  safe-rm -f -- "${statusfile}"

  true "${BASH_SOURCE[0]} INFO: Download file..."

  ## Launching into the background is required so a SIGTERM to this script can
  ## interrupt the in-flight download. If attempting to refactor this, make
  ## sure signal sigterm stops downloads.
  run_body_download "$@" &

  wait_exit_code=0
  wait "$!" &>/dev/null || wait_exit_code=$?
  true "${BASH_SOURCE[0]} INFO: File download done."
  true "${BASH_SOURCE[0]} INFO: END."
  exit "${wait_exit_code}"
}

main() {
  set -o errexit
  set -o nounset
  set -o pipefail
  set -o errtrace
  shopt -s inherit_errexit
  shopt -s shift_verbose
  export LC_ALL=C

  initialize_terminal
  initialize_variables
  check_variables
  initialize_temporary_files
  traps_enable
  run_download "$@"
}

if was_executed "${BASH_SOURCE[0]}"; then
  main "$@"
fi

## Debugging.
#print_progress_bar "$1" "$2"
