#!/usr/bin/env bash
#
# 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.

#set -x
set -e
set -o nounset
set -o pipefail
set -o errtrace

## provides: draw_progress_bar
source /usr/libexec/helper-scripts/progress-bar

## provides: is_whole_number
source /usr/libexec/helper-scripts/strings.bsh

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 file descriptor 4, and then point that file
  # descriptor to the right place: stderr if it's a TTY, or /dev/null
  # otherwise.
  if [ -t 2 ]; then
    exec 4>&2
  else
    exec 4>/dev/null
  fi
}

initialize_variables() {
  command -v tput >/dev/null
  command -v curl >/dev/null
  command -v safe-rm >/dev/null
  command -v mktemp >/dev/null

  : "${CURL:="curl"}"
  : "${CURL_PRGRS_MAX_FILE_SIZE_BYTES:=""}"
  : "${CURL_OUT_FILE:=""}"
  : "${CURL_PRGRS_EXEC:=""}"
  : "${curl_prgrs_print_progress:="yes"}"
  percent_last=""

  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"

  expected_header_size=8000
  maximum_http_header_size=32000
}

check_variables() {
  if [ "$CURL_OUT_FILE" = "" ]; then
    stecho "${BASH_SOURCE[0]} ERROR: Variable CURL_OUT_FILE is empty." >&4
    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." >&4
    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.
shutdown() {
  local exit_code="$?"
  local signal="$1"
  local last_err="$BASH_COMMAND"
  if [ "$signal" = "err" ]; then
    stecho "${BASH_SOURCE[0]} ERROR: Signal $signal received. Exiting." >&4
    stecho "${BASH_SOURCE[0]} ERROR: BASH_COMMAND '$BASH_COMMAND' exit code '$exit_code'." >&4
    error_exit_code=110
  elif [ "$signal" = "exit" ]; then
    true "${BASH_SOURCE[0]} INFO: Signal $signal received. Exiting." >&4
  else
    stecho "${BASH_SOURCE[0]} INFO: Signal $signal received. Exiting." >&4
  fi

  trap - SIGHUP SIGINT SIGTERM ERR EXIT

  # If we wrote an exit status to the temporary file, read it. Otherwise,
  # we reached this trap function abnormally; assume a non-zero status.

  sync

  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"
      error_exit_code="111"
    fi
  else
    true "${BASH_SOURCE[0]} ERROR: no status file"
    status="112"
    error_exit_code="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" >&4
  else
    printf '%s\n' '' >&4
  fi

  #stat="$(stcat "$statusfile")"
  #stecho "$stat" >&4

  if [ "$temp_dir_auto_generated" = "true" ]; then
    safe-rm -r -f -- "$temporary_directory"
  fi

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

  processes_list="$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
    fi
  done

  if [ -n "${error_exit_code-}" ]; then
     true "${BASH_SOURCE[0]} INFO: exit with error_exit_code $error_exit_code"
     exit "$error_exit_code"
  fi

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

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
}

# 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=$(( $bytes * 100 / $length ))

  if ! is_whole_number "$percent" ; then
    curl_exit 113
  fi

  if [ "$percent" -ge "100" ]; then
    percent=100
  fi

  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" >&4
    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" >&4
      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"
  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"
}

curl_download() {
  local size_file_downloaded_bytes

  $CURL --no-progress-meter "$@" &
  curl_pid="$!"

  ## 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

  while true ; do
    if [ -f "$CURL_OUT_FILE" ]; then
      size_file_downloaded_bytes="$(stat -c "%s" "$CURL_OUT_FILE")"

      if ! is_whole_number "$size_file_downloaded_bytes" ; then
        curl_exit 113
      fi

      true "size_file_downloaded_bytes: $size_file_downloaded_bytes"
      true "CURL_PRGRS_MAX_FILE_SIZE_BYTES: $CURL_PRGRS_MAX_FILE_SIZE_BYTES"
      if [ "$size_file_downloaded_bytes" -gt "$CURL_PRGRS_MAX_FILE_SIZE_BYTES" ]; then
        curl_exit 81
      fi

      true "size_file_downloaded_bytes: $size_file_downloaded_bytes"
      true "curl_prgrs_content_length: $curl_prgrs_content_length"
      if [ "$size_file_downloaded_bytes" -gt "$curl_prgrs_content_length" ]; then
        curl_exit 114
      fi

      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

    sleep 1
  done

  ## curl already terminated.

  : "${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=()

  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[@]}"
}

main() {
  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}' \
        "${header_arguments[@]}" \
        --output "$header_file" \
    )"

  ## 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

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

  ## Launching into the background required.
  ## If attempting to refactor this, make sure signal sigterm stops downloads.
  header_download="false" curl_download "$@" &

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

traps_enable
initialize_terminal
initialize_variables
check_variables
main "$@"

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