#!/bin/bash

EXPECTED_OMI_VERSION="OMI-1.6.9-1"
DOWNLOAD_TIMEOUT=600
LOG_PREFIX="OmsSecurityPatch:"

# Log final failure message.
log_failure() {
  logger -p user.err "$LOG_PREFIX Failed to upgrade OMI to $1".
  exit 1
}

log_info() {
	logger -p user.info "$LOG_PREFIX $1"
}

log_warn() {
  logger -p user.warn "$LOG_PREFIX $1"
}

log_error() {
  logger -p user.err "$LOG_PREFIX $1"
}

log_info "Starting activity to patch the OMS security vulnerability."
START_TIME=$(date +%s)

# Exit if OMI is not installed on the node.
sudo /opt/omi/bin/omiserver --version &>/dev/null
ret=$?
if [ $ret -ne 0 ];then
  log_info "OMI not installed on the node. Skipping activity to apply OMS security patch."
  exit 0
fi

# Check if the node is already patched.
OLD_OMI_VERSION=`sudo /opt/omi/bin/omiserver --version | awk '{ print $2 }'`
log_info "OMI version pre-patch: $OLD_OMI_VERSION"
if [ "${EXPECTED_OMI_VERSION,,}" = "${OLD_OMI_VERSION,,}" ];
then
   log_info "OMI version on the node is already up-to-date: $EXPECTED_OMI_VERSION. Skipping activity to apply OMS security patch."
  exit 0
fi

# Download patch and store it in the /mnt directory.
PATCH_FILE="omsagent-1.14.12-0.universal.x64.sh"
CDN_ENDPOINT="https://hdinsight-oms.azureedge.net"
STORAGE_ACCOUNT_NAME=${1:-hdiconfigactions}
if [ "$STORAGE_ACCOUNT_NAME" == "null" ];
then
        STORAGE_ACCOUNT_NAME=hdiconfigactions
fi

echo "Downloading patch file $PATCH_FILE from $CDN_ENDPOINT"
sudo wget -T $DOWNLOAD_TIMEOUT $CDN_ENDPOINT/$PATCH_FILE -O /mnt/$PATCH_FILE
ret=$?
if [ $ret -ne 0 ];then
    log_warn "Failed to download the patch from $CDN_ENDPOINT. Attempting to download the patch from storage account $STORAGE_ACCOUNT_NAME"
    sudo wget -T $DOWNLOAD_TIMEOUT https://$STORAGE_ACCOUNT_NAME.blob.core.windows.net/clustermonitoringconfigactionv01/omsagent/$PATCH_FILE -P /mnt/
    ret=$?
    if [ $ret -ne 0 ];then
      log_error "Failed to download the patch from storage account $STORAGE_ACCOUNT_NAME"
        log_failure $EXPECTED_OMI_VERSION
    fi
fi

# Exeute log_failure() when encountering an error.
trap 'log_failure $EXPECTED_OMI_VERSION' ERR

sudo chmod +x /mnt/$PATCH_FILE
log_info "Patch downloaded successfully."
# Apply the patch.
log_info "Applying patch: sudo sh /mnt/$PATCH_FILE --upgrade"
sudo sh /mnt/$PATCH_FILE --upgrade

NEW_OMI_VERSION=`sudo /opt/omi/bin/omiserver --version | awk '{ print $2 }'`
log_info "OMI version post-patch: $NEW_OMI_VERSION"

END_TIME=$(date +%s)

# Validate patch.
if [ "${EXPECTED_OMI_VERSION,,}" = "${NEW_OMI_VERSION,,}" ];
then
  log_info "Successfully upgraded OMI to version $EXPECTED_OMI_VERSION"
  log_info "Finished acitvity to patch the OMS security vulnerability. Took $(($END_TIME - $START_TIME)) seconds."
else
   log_error "Something went wrong while applying the patch! Took $(($END_TIME - $START_TIME)) seconds."
   log_failure $EXPECTED_OMI_VERSION
fi