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

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

## success: whole numbers such as 1 2 3 456 etc.
## fail:
## 5.5
## 5,500
## trailing spaces
## leading spaces
## empty input
## space only
## scientific notation such as 1.23e10
## negative numbers such as -1
is_whole_number() {
    [[ $1 =~ ^[0-9]+$ ]]
}

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 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
    echo "$BASH_SOURCE ERROR: Variable CURL_OUT_FILE is empty." >&4
    exit "57"
  fi
  if [ "$CURL_PRGRS_MAX_FILE_SIZE_BYTES" = "" ]; then
    echo "$BASH_SOURCE 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
    echo "$BASH_SOURCE ERROR: Signal $signal received. Exiting." >&4
    echo "$BASH_SOURCE ERROR: BASH_COMMAND '$BASH_COMMAND' exit code '$exit_code'." >&4
    error_exit_code=110
  elif [ "$signal" = "exit" ]; then
    true "$BASH_SOURCE INFO: Signal $signal received. Exiting." >&4
  else
    echo "$BASH_SOURCE 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 INFO: got status file"
    status="$(cat "$statusfile")"
    if ! is_whole_number "$status"; then
      true "$BASH_SOURCE ERROR: status is not a number! status: '$status'"
      status="111"
      error_exit_code="111"
    fi
  else
    true "$BASH_SOURCE 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 "\x1B[0G\x1B[0K" >&4
  else
    echo >&4
  fi

  #stat="$(cat "$statusfile")"
  #echo "$stat" >&4

  if [ "$temp_dir_auto_generated" = "true" ]; then
    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 INFO: exit with error_exit_code $error_exit_code"
     exit "$error_exit_code"
  fi

  true "$BASH_SOURCE INFO: exit with status $status"
  exit "$status"
}

# 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 INFO: bytes: '$bytes'"
  true "$BASH_SOURCE 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 INFO: percentage number unchanged. Not re-drawing progress bar to avoid flicker."
  else
    draw_progress_bar "$percent"
    if [ "$CURL_PRGRS_EXEC" = "" ]; then
      true "$BASH_SOURCE INFO: CURL_PRGRS_EXEC is empty. Not executing CURL_PRGRS_EXEC."
    else
      true "$BASH_SOURCE INFO: CURL_PRGRS_EXEC is set. Executing CURL_PRGRS_EXEC..."
      true "$BASH_SOURCE INFO: $CURL_PRGRS_EXEC '$percent'"
      $CURL_PRGRS_EXEC "$percent" >&4
      true "$BASH_SOURCE INFO: CURL_PRGRS_EXEC success."
    fi
  fi

  percent_last="$percent"
}

curl_exit() {
  curl_exit_code="$1"
  true "$BASH_SOURCE INFO: write $curl_exit_code to $statusfile"
  echo "$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
}

curl_download() {
  local size_file_downloaded_bytes

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

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

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

      if [ "$size_file_downloaded_bytes" -ge "$curl_prgrs_content_length" ]; then
        curl_exit 114
      fi

      if [ "$curl_prgrs_print_progress" = "yes" ]; then
        ## Need to echo to stderr to avoid confusing the stdout output of this command.
        #echo "$BASH_SOURCE INFO: print_progress '$size_file_downloaded_bytes' '$curl_prgrs_content_length'" >&2
        print_progress "$size_file_downloaded_bytes" "$curl_prgrs_content_length" >&4
      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
    if [ "$size_file_downloaded_bytes" -le "$curl_prgrs_content_length" ]; then
      curl_exit 115
    fi
  fi

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

remove_argument_for_header_request() {
    local args=()
    local skip_next=false

    for arg in "$@"; do
        if [ "$skip_next" = true ]; then
            skip_next=false
            continue
        fi

        if [ "$arg" = "--continue-at" ]; then
            skip_next=true
            continue
        fi

        if [ "$arg" = "-C" ]; then
            skip_next=true
            continue
        fi

        if [ "$arg" = "--output" ]; then
            skip_next=true
            continue
        fi

        args+=("$arg")
    done

    echo "${args[@]}"
}

main() {
  local header_arguments header_file

  header_arguments=($(remove_argument_for_header_request "$@"))
  header_file="$temporary_directory/header"

  true "$BASH_SOURCE 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 zenity.
  ##
  ## CURL_OUT_FILE and
  ## --output "$header_file" to avoid overwriting files when using "--continue-at -".
  curl_prgrs_content_length=$( \
    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 INFO: Header downloaded done."

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

  true "$BASH_SOURCE INFO: Download file..."

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

  true "$BASH_SOURCE INFO: File download done."

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

initialize_terminal
initialize_variables
check_variables
main "$@"

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