#! /bin/bash
set -ex

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

#Usage:
#This script needs one input arguments. 
#  - Argument 1 must be 
#      1 (enable scheduled reboots - updates the config file for reboot) or 
#      2 (enable immediate reboot). 
#    Argument will be defaulted to 1 (enable scheduled reboots) if input is not provided.


#Version 1.1
echo "Executing schedule-reboots script Version 1.1"
logger -p user.info "Executing schedule-reboots script Version 1.1"

#numerical regex for argument validation
re="^[12]?$"

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

configuration_file="/usr/hdinsight/.managed_patching_reboot_config"

#if immediate reboot is selected
if [[ $1 -eq 2 ]]
then
  if [ -e "$configuration_file" ]
  then
    echo "Managed patching reboot configuration exists. Updating contents of file..."
    logger -p user.info "Managed patching reboot configuration exists. Updating contents of file..."
    #Updating the config file to disable schedule reboot if set.
    sed -i 's/reboot_enabled=1/reboot_enabled=0/g' $configuration_file
  fi
  #Give a final 5 minute warning before shutdown
  logger -p user.info "Rebooting in 5 minutes."
  sudo /sbin/shutdown -r +5
else
  #If the configuration_file does not exist, the cluster is using V1 of the managed_patching_reboot script, 
  #in which case the script itself needs to be patched first to support custom configurations.
  if [ -e "$configuration_file" ]
  then
    echo "Managed patching reboot configuration exists. Updating contents of file..."
    logger -p user.info "Managed patching reboot configuration exists. Updating contents of file..."
    #Updating the config file to make this a one time reboot.
    sed -i 's/reboot_enabled=0/reboot_enabled=1/g' $configuration_file
  else
    echo "Managed patching reboot configuration does not exist. Initializing managed patching reboot configuration..."
    logger -p user.info "Managed patching reboot configuration does not exist. Initializing managed patching reboot configuration..."
    echo "reboot_enabled=1" > $configuration_file
    echo "count_reboot_slots=48" >> $configuration_file
    echo "total_reboot_cycle_duration_hours=24" >> $configuration_file
     
    chmod 666 $configuration_file
  fi

  logger -p user.info "Updated config file contents."
  logger -p user.info -f $configuration_file

  #Updating managed_patching_reboot script
  echo "Updating managed_patching_reboot script..."
  logger -p user.info "Updating managed_patching_reboot script..."
  wget -O /tmp/managed_patching_reboot.sh -nv https://hdiconfigactions.blob.core.windows.net/linuxospatchingrebootconfigv02/managed_patching_reboot.sh
  if [ ! -s "/tmp/managed_patching_reboot.sh" ]
  then
    echo "Failed to download managed_patching_reboot script from remote. Aborting."
    logger -p user.info "Failed to download managed_patching_reboot script from remote. Aborting."
    exit 1
  fi
  chmod 750 /tmp/managed_patching_reboot.sh
  mv /tmp/managed_patching_reboot.sh /usr/hdinsight/managed_patching_reboot.sh

  #Checking for managed patching reboot cron job. The default frequency in the cron job file is 30 minutes
  cronjob_file="/etc/cron.d/managed_patching_cron_job"
  if [ ! -e "$cronjob_file" ]
  then
    echo "Managed patching reboot cron job does not exist. Updating managed_patching_cron_job..."
    logger -p user.info "Managed patching reboot cron job does not exist. Updating managed_patching_cron_job..."
    #updating managed patching reboot cron job file with slot duration
    echo "*/30 * * * * root /usr/hdinsight/managed_patching_reboot.sh" > $cronjob_file
  fi

  echo "Updated Managed patching reboot configuration successfully"
  logger -p user.info "Updated Managed patching reboot configuration successfully"
fi