#!/bin/bash

TIMEOUT_SINGLE_CMD_SEC=60

TIMEOUT_SINGLE_HTTP_REQUEST_SEC=1
AWS_IMDSV2_TOKEN_TTL_SECONDS=21600
AWS_IMDS_TOKEN=

CURL_ARGS=(--silent --max-time $TIMEOUT_SINGLE_HTTP_REQUEST_SEC)
CURL_CODE_PREFIX="CODE:"
CURL_WRITE_OUT=(--write-out "$CURL_CODE_PREFIX%{http_code}")

SCRIPT_GIT_HASH="b1ad6b6487fce6514985079649c24a5ed9e1f4b1"

args=("$@")

# Parameters.
while [[ $# -gt 0 ]]; do
  opt="$1"
  case "$opt" in
  --preflight)
    # Used for checking sudo target user permissions.
    PREFLIGHT=1
    ;;
  --user)
    USER_FLAG=1
    ;;
  --force-imdsv1)
    # Used for testing purposes.
    FORCE_IMDSV1=1
    ;;
  --low-priority)
    LOW_PRIORITY=1
    ;;
  --collect-netstat)
    # Used for enabling netstat collection
    COLLECT_NETSTAT=1
    ;;
  --stats)
    PRINT_STATS=1
    ;;
  --stats-to-file)
    STATS_TO_FILE=1
    ;;
  --output)
    OUTPUT_FILE="$2"
    shift
    ;;
  *)
    echo "Unknown parameter: $opt"
    ;;
  esac
  shift
done

if [[ ${EUID} -ne 0 && -z "${SUDO_COMMAND}" ]]; then
  # We are not root and not running as sudo.
  # Try running this script with sudo.
  sudo -n "$0" --preflight &> /dev/null
  if [[ $? -eq 0 ]]; then
    # The preflight is successful.
    # The sudo target user has permissions to run the script.
    # We re-run the script with sudo and propagate the script exit code.
    sudo -n "$0" "${args[@]}"
    exit $?
  fi
  # We keep running the script with the current user.
fi

# This is a preflight. No need to run the rest of the script.
[[ ${PREFLIGHT} -eq 1 ]] && exit

if [[ ${USER_FLAG} -eq 1 ]]; then
  echo >&2 "--user flag is deprecated, root is no longer strictly required."
fi

if [[ ${EUID} -ne 0 ]]; then
  echo >&2 "This script is run with non-root user, it might fail or collect partial data."
fi

shopt -s nullglob

files_dump_targets=(
  "/sys/class/dmi/id/product_serial"     # Serial from SMBIOS (requires root).
  "/proc/cpuinfo"                        # CPU info.
  "/proc/meminfo"                        # Memory info.
  "/proc/diskstats"                      # Disk usage
  "/proc/net/dev"                        # Network usage
  "/proc/uptime"                         # Machine uptime.
)

# Use the prefix "@outfile filename" to specify the output filename, otherwise the command name will be used.
# To run a command named "@outfile" use "[@outfile filename] -- @outfile [flags]"
commands=(
  "ip addr"                             # Show IP addressed assigned to NICs
  "top -n 1 -b | head -6"               # CPU utilization.
)

if [[ ${COLLECT_NETSTAT} -eq 1 ]]; then
  commands+=(
    "netstat -tuanp"
    "ss -tuanp"
  )
fi

tmp=$(mktemp -d)

if [[ ${LOW_PRIORITY} -eq 1 ]]; then
  current_priority=$(ps -o ni= -p $$)
  renice -n "$(($current_priority + 5))" -p $$ 1> "$tmp/renice" 2> "$tmp/renice.stderr"
  echo $? > "$tmp/renice.exitcode"
fi

function collect_single_filecommand(){
  filecommand=$1
  file_dumps_dir=$2

  fileglob=$(echo "$filecommand" | cut -d " " -f1)
  command=$(echo "$filecommand" | cut -d " " -f1 --complement -s)
  found=0
  # nullglob is set above, so if this doesn't match a file we will skip the loop entirely.
  for file in $fileglob; do
    found=1
    if [[ ! -f "$file" ]]; then
      echo "$file, is not a file" >>"$tmp/collect.log"
      continue
    fi
    targetDir=$(dirname "$file_dumps_dir/$file")
    mkdir -p "$targetDir"
    "$SHELL" -c "cat $file $command" >"$file_dumps_dir/$file"
  done
  if [[ found -eq 0 ]]; then
    echo "file: $fileglob, not found" >>"$tmp/collect.log"
  fi
}

function collect_file_dumps() {
  file_dumps_dir="$tmp/files"
  mkdir "$file_dumps_dir"

  for filecommand in "${files_dump_targets[@]}"; do
    collect_single_filecommand "$filecommand" "$file_dumps_dir"
  done
}

function collect_command_outputs() {
  command_outputs_dir="$tmp/commands"
  mkdir "$command_outputs_dir"

  for raw_command_line in "${commands[@]}"; do
    # convert command line to array
    command_line_array=($raw_command_line)

    relativeTargetFile=
    # parse the command line, look for exsitence of "@outfile" or "--" prefixes
    while [[ -n "${command_line_array[0]}" ]]; do
      case "${command_line_array[0]}" in
        "@outfile")
          relativeTargetFile=${command_line_array[1]}
          command_line_array=(${command_line_array[@]:2})
        ;;
        "--")
          command_line_array=(${command_line_array[@]:1})
          break
        ;;
        *)
          break
        ;;
      esac
    done

    # no "@outfile" prefix found, use the command name as outfile
    if [[ -z "$relativeTargetFile" ]]; then
      relativeTargetFile=${command_line_array[0]}
    fi

    targetFile="$command_outputs_dir/$relativeTargetFile"

    cmd=${command_line_array[*]}

    echo "- cmd: $cmd -" >>"$tmp/collect.log"
    (time timeout ${TIMEOUT_SINGLE_CMD_SEC} "$SHELL" -c "$cmd" 2>"$targetFile.stderr") >"$targetFile" 2>>"$tmp/collect.log"
    echo "$?" >"$targetFile.exitcode"
    cat "$targetFile.stderr" >>"$tmp/collect.log"
  done
}

function collect_nics_information() {
  # Collect the NICs symlinks which indicates which of them are physical.
  "$SHELL" -c 'find /sys/class/net/* -type l -ls' >"$tmp/nics_symlinks" 2>>"$tmp/collect.log"

  # Collect the operstate which indicates the status of the NIC.
  "$SHELL" -c 'grep "" /sys/class/net/*/operstate' >"$tmp/nics_operstate" 2>>"$tmp/collect.log"
}

# Collect disk IO stats.
function collect_diskstats_information() {
  output_dir="$tmp/diskstats"
  mkdir "$output_dir"

  cat /proc/diskstats > "$output_dir/1"
  sleep 1
  cat /proc/diskstats > "$output_dir/2"
}

function collect_common_metadata() {
  full_collect_date="$(date --iso-8601=sec -u | sed 's/+00:00/.000Z/g')"
  hostname >"$tmp"/machinename
  echo "$full_collect_date" >"$tmp"/timestamp
  echo Linux >"$tmp"/ostype
  echo "$SCRIPT_GIT_HASH" >"$tmp"/version
  echo "performance" >"$tmp"/collection_mode
  whoami | tr -d '\n' >"$tmp"/whoami
}

function pack_collected_data() {
  collect_date="$(date +%Y-%m-%d-%H-%M)"
  if [[ -z "$OUTPUT_FILE" ]]; then
    result_filename="mcdc-collect-performance-$(hostname)-$collect_date.tar"
  else
    out_dir=$(dirname "$OUTPUT_FILE")
    mkdir -p "$out_dir"
    result_filename="$OUTPUT_FILE"
  fi

  tar czf "$result_filename" -C "$tmp" .
  if [[ $? -ne 0 ]]; then
    echo >&2 "Couldn't create $result_filename"
    exit 1
  fi

  echo "Collected info saved to: $result_filename"

  # $tmp/collect.log file might not exist by the time this fails
  # so make sure it's completely silent and error free.
  rm -rf "$tmp" >/dev/null 2>&1 || true
}

function print_stats() {
  local stage
  stage="$1"

  # Check if statistics collection is enabled
  if [[ -z "$PRINT_STATS" && -z "$STATS_TO_FILE" ]]; then
    return
  fi

  # Get CPU and memory usage of the script. See `man 5 proc` for fields definitions.
  proc_stat=($(sed -E 's/\([^)]+\)/X/' "/proc/$$/stat"))
  utime=${proc_stat[15]} # user space time
  stime=${proc_stat[16]} # kernel space time
  clock_tick=$(getconf CLK_TCK)
  let cpu_time="($utime + $stime) / $clock_tick"
  memory_usage=$(grep "VmPeak" /proc/$$/status | awk '{print $2}')
  stat_line="ScriptStats: $(date '+%Y-%m-%d %H:%M:%S'),$stage,$cpu_time,$memory_usage"

  # Print stats to console or log to file
  if [[ -n "$PRINT_STATS" ]]; then
    echo "$stat_line"
  elif [[ -n "$STATS_TO_FILE" ]]; then
    local log_file
    log_file=$(dirname "$OUTPUT_FILE")/script-stats.log
    echo "$stat_line" >> "$log_file"
  fi
}

main() {
  print_stats "start"
  echo "--- file dumps -----" >>"$tmp/collect.log"
  (time collect_file_dumps) 2>>"$tmp/collect.log"
  print_stats "file_dumps"
  echo "--- commands -----" >>"$tmp/collect.log"
  (time collect_command_outputs) 2>>"$tmp/collect.log"
  print_stats "commands"
  echo "--- nics information -----" >>"$tmp/collect.log"
  (time collect_nics_information) 2>>"$tmp/collect.log"
  print_stats "nics_information"
  echo "--- disk stats -----" >>"$tmp/collect.log"
  (time   collect_diskstats_information) 2>>"$tmp/collect.log"
  print_stats "disk_stats"
  echo "--- common metadata -----" >>"$tmp/collect.log"
  (time collect_common_metadata) 2>>"$tmp/collect.log"
  print_stats "common_metadata"
  echo "--- pack data -----" >>"$tmp/collect.log"
  time pack_collected_data 2>>"$tmp/collect.log"
  print_stats "pack_data"
}

main
