#! /bin/bash
set -ex

#This custom action script only applies to clusters created after July 2016, where managed patching reboot is supported.
#This script updates the config file for reboot.

#Usage:
#This script needs two input arguments. 
#  - Argument 1 must be 
#      0 (enable kernel updates only) or 
#      1 (enable kernel+security updates only) or 
#      2 (enable all updates). 
#    Argument will be defaulted to 0 (enable kernel updates only) if input is not provided.
#
#  - Argument 2 must be 
#      0 (disable reboot) or 
#      1 (enable scheduled reboots) or 
#      2 (enable immediate reboot). 
#    Argument will be defaulted to 0 (disable reboot) if input is not provided.

#Version 1.1
echo "Executing OS patching reboot config script Version 1.1"
logger -p user.info "Executing OS patching reboot config script Version 1.1"

#Usage:
echo "The first input argument to this script must be 0 (enable kernel updates only) or 1 (enable kernel+security updates only) or 2 (enable all updates). Defaulted to 0 if input is not provided."
logger -p user.info "The first input argument to this script must be 0 (enable kernel updates only) or 1 (enable kernel+security updates only) or 2 (enable all updates). Defaulted to 0 if input is not provided."
echo "The second input argument to this script must be 0 (disable reboot) or 1 (enable scheduled reboots) or 2 (enable immediate reboot)."
logger -p user.info "The second input argument to this script must be 0 (disable reboot) or 1 (enable scheduled reboots) or 2 (enable immediate reboot)."

#numerical regex for argument validation
re="^[012]?$"
 
#Validating input variables
if [[ ! $1 =~ $re ]] 
then
  echo "Invalid argument 1 received: $1. Argument must be 0 (enable kernel updates only) or 1 (enable kernel+security updates only) or 2 (enable all updates). Aborting."
  logger -p user.info "Invalid argument 1 received: $1. Argument must be 0 (enable kernel updates only) or 1 (enable kernel+security updates only) or 2 (enable all updates). Aborting."
  exit 1
fi

if [[ ! $2 =~ $re ]]
then
  echo "Invalid argument 2 received: $2. Must be 0 (disable reboot) or 1 (enable scheduled reboots) or 2 (enable immediate reboot). Aborting."
  exit 1
fi

# Function to check if file descriptors exist
check_fd_exists() {
  lsof "$1" >/dev/null 2>&1
}

# Function to check if file descriptors exist for apt and dpkg lock files
check_apt_dpkg_lock_exists() {
  dpkg_lock="/var/lib/dpkg/lock-frontend"
  apt_lists_lock="/var/lib/apt/lists/lock"
  timeout_seconds=60  # 1 minute timeout

  start_time=$(date +%s)

  # Check if file descriptors exist for either lock file
  if (check_fd_exists "$dpkg_lock") || (check_fd_exists "$apt_lists_lock"); then
      echo "check_apt_dpkg_lock_exists status: File descriptors exist. Waiting $timeout_seconds secs for them to finish..."
      logger -p user.info "check_apt_dpkg_lock_exists status: File descriptors exist. Waiting for them to finish..."

      # Get the PID using lsof
      pid_dpkg=$(lsof -t "$dpkg_lock")
      pid_apt=$(lsof -t "$apt_lists_lock")
      echo "pid_dpkg: ${pid_dpkg}"
      echo "pid_apt: ${pid_apt}"
      logger -p user.info "pid_dpkg: ${pid_dpkg}"
      logger -p user.info "pid_apt: ${pid_apt}"

      while [ -n "$pid_dpkg" ] || [ -n "$pid_apt" ]; do
          current_time=$(date +%s)
          elapsed_time=$((current_time - start_time))

          # Get additional information using ps
          if [ -n "$pid_dpkg" ]; then
              lsof_installation_pid=`ps -ef | grep "$pid_dpkg"`
              echo "lsof_installation_pid: ${lsof_installation_pid}"
              logger -p user.info "${lsof_installation_pid}"
          elif [ -n "$pid_apt" ]; then
              lsof_update_pid=`ps -ef | grep "$pid_apt"`
              echo "lsof_update_pid: ${lsof_update_pid}"
              logger -p user.info "${lsof_update_pid}"
          fi

          if [ "$elapsed_time" -ge "$timeout_seconds" ]; then
              echo "check_apt_dpkg_lock_exists status: Timeout reached. Exiting."            
              logger -p user.info "check_apt_dpkg_lock_exists status: Timeout reached. Exiting."
              return 1
          fi

          sleep 5
          pid_dpkg=$(lsof -t "$dpkg_lock")
          pid_apt=$(lsof -t "$apt_lists_lock")
          echo "pid_dpkg after ${elapsed_time}: ${pid_dpkg}"
          echo "pid_apt after ${elapsed_time}: ${pid_apt}"
      done

      # Log the result to the specified file
      lsof_output=`lsof "$dpkg_lock" "$apt_lists_lock"`
      logger -p user.info "${lsof_output}"
      return 0
  else
      echo "check_apt_dpkg_lock_exists status: No file descriptors found. Proceeding with the installation."
      logger -p user.info "check_apt_dpkg_lock_exists status: No file descriptors found. Proceeding with the installation."
      return 0
  fi
}

# Function to run command with retry
function run_command_with_retry() {
  # Usage run_command_with_retry 'command' maxRetries sleepInSeconds (gets doubled every retry, except when timeout ocurs)
  # Parameters with provided defaults.
  local command=$1
  local maxRetries=${2:-5}
  local sleepInSeconds=${3:-10}

  # Check if command is provided
  if [ -z "$command" ]; then 
    echo "No command supplied"
    exit 1
  fi

  # Initialize retry count
  local retryCount=0
  local START=$(date "+%s.%N")
  local END
  set +ex

  # Run the command with retries
  while [ $retryCount -lt $maxRetries ]; do
    echo "ATTEMPT(${retryCount}): '$command' Started..."
    # Check if apt or apt-get is used and if so, check if apt/dpkg lock exists
    if [[ $command =~ "apt" ]] || [[ $command =~ "apt-get" ]] || [[ $command =~ "dpkg" ]]; then
      echo "Check if command has apt and apt/dpkg lock exists..."
      logger -p user.info "Check if command has apt and apt/dpkg lock exists..."
      check_apt_dpkg_lock_exists
      local ap_lock_status=$?
      echo "apt/dpkg lock status: $ap_lock_status"
      if [ $ap_lock_status -ne 0 ]; then
        echo "Apt/dpkg lock exists. Exiting."
        logger -p user.info "Apt/dpkg lock exists. Exiting."
        echo "ATTEMPT(${retryCount}): $command Skipped/Exited with $ap_lock_status, sleeping $sleepInSeconds seconds before retry..."
        sleep $sleepInSeconds
        sleepInSeconds=$((sleepInSeconds * 2))
        retryCount=$((retryCount+1))
        continue
      fi
    fi
    echo "ATTEMPT(${retryCount}): '$command' Running..."
    bash -c "$command"
    local status=$?
    if [ $status -eq 0 ]; then
      END=$(date "+%s.%N")
      echo "ATTEMPT(${retryCount}): $command Succeeded, time spent: " | awk "{ printf  \$0\"%.6g sec.\n\", $END - $START }"
      set -ex
      return 0
    else
      echo "ATTEMPT(${retryCount}): '$command' Exited with $status, sleeping $sleepInSeconds seconds before retry..."
      sleep $sleepInSeconds
      sleepInSeconds=$((sleepInSeconds * 2))
    fi
    retryCount=$((retryCount+1))
  done
  END=$(date "+%s.%N")
  echo ": Failed $maxRetries times in total, giving up! Time spent: " | awk "{ printf  \$0\"%.6g sec.\n\", $END - $START }"
  set -ex
  exit 3
}

OS_VER=`lsb_release -sr`
OS_PATCH_SCRIPT=os.xenial.patching.sources.list
SECURITY_PATCH_SCRIPT=security.xenial.sources.list

if [ $OS_VER = '18.04' ]; then
  OS_PATCH_SCRIPT=os.bionic.patching.sources.list
  SECURITY_PATCH_SCRIPT=security.bionic.sources.list
fi

# Install updates
# Get source list from remote, based on user input. Using update to re-synchronize the package index files from their sources.
# Using upgrade to install the newest versions of packages currently installed on the system.
# if first argument is 2 (enable all updates), download os.patching.sources.list 
if [[ $1 -eq 2 ]]
then
  echo "Downloading full os updates list..."
  echo "(Incase the execution halts here, retrying might yield success)" 
  logger -p user.info "Downloading full os updates list..."
  wget -O /tmp/userdefined.sources.list -nv https://hdiconfigactions.blob.core.windows.net/linuxospatchingrebootconfigv02/${OS_PATCH_SCRIPT}
# if first argument is 1 (enable kernel+security updates only), download security.sources.list
elif [[ $1 -eq 1 ]]
then
  echo "Downloading security updates list only..."
  echo "(Incase the execution halts here, retrying might yield success)"
  logger -p user.info "Downloading security updates list only..."
  wget -O /tmp/userdefined.sources.list -nv https://hdiconfigactions.blob.core.windows.net/linuxospatchingrebootconfigv02/${SECURITY_PATCH_SCRIPT}
fi

if [[ $1 -eq 1 ]] || [[ $1 -eq 2 ]]
then
  if [ ! -s "/tmp/userdefined.sources.list" ]
  then
    echo "Failed to download updates list from remote. Aborting."
    logger -p user.info "Failed to download updates list from remote. Aborting."
    exit 1
  fi
  chmod 750 /tmp/userdefined.sources.list
  mv /tmp/userdefined.sources.list /etc/apt/userdefined.sources.list
  run_command_with_retry 'DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" update'
  echo "Installing updates..."
  logger -p user.info "Installing updates..."
  run_command_with_retry 'DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" upgrade -o Dir::Etc::SourceList=/etc/apt/userdefined.sources.list'
  run_command_with_retry 'DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" install linux-azure'
# if first argument is 0 (enable kernel updates only), install kernel updates only
else
  run_command_with_retry 'DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" update'
  echo "Installing kernel updates only..."
  logger -p user.info "Installing kernel updates only..."  
  run_command_with_retry 'DEBIAN_FRONTEND=noninteractive apt-get -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" install linux-azure'
fi

if [[ $2 -eq 1 ]] || [[ $2 -eq 2 ]]
then
  echo "Downloading the schedule-reboots script..."
  echo "(Incase the execution halts here, retrying might yield success)" 
  logger -p user.info "Downloading the schedule-reboots script..."
  wget -O /tmp/schedule-reboots.sh -nv https://hdiconfigactions.blob.core.windows.net/linuxospatchingrebootconfigv02/schedule-reboots.sh
  if [ ! -s "/tmp/schedule-reboots.sh" ]
  then
    echo "Failed to download schedule-reboots script from remote. Aborting."
    logger -p user.info "Failed to download schedule-reboots script from remote. Aborting."
    exit 1
  fi
  chmod 750 /tmp/schedule-reboots.sh
  #run the script to update the config file with reboot option enabled
  /tmp/schedule-reboots.sh $2
  echo "Updated config file with reboot option enabled"
  logger -p user.info "Updated config file with reboot option enabled"
fi

echo "Installed updates successfully"
logger -p user.info "Installed updates successfully"