#!/bin/bash

# bin paths
LSB_RELEASE="/usr/bin/lsb_release -i"
APT_GET="/usr/bin/apt-get -qq"
YUM="/usr/bin/yum"
NTOPNG_CONFIG_TOOL="/usr/bin/ntopng-utils-manage-config"
REDIS_CLI="/usr/bin/redis-cli"
PRODUCT="ntopng"

# redis keys
REDIS_CHECK_FOR_UPDATES_KEY="ntopng.updates.check_for_updates"
REDIS_NEW_VERSION_AVAILABLE_KEY="ntopng.updates.new_version"
REDIS_RUN_UPGRADE_KEY="ntopng.updates.run_upgrade"
REDIS_UPDATE_FAILURE_KEY="ntopng.updates.update_failure"
REDIS_AUTO_UPDATES_KEY="ntopng.prefs.is_autoupdates_enabled"
REDIS_IN_PROGRESS_KEY="ntopng.updates.in_progress"
REDIS_PRODUCT_NAME_KEY="ntopng.product_name"

ACTION=""
QUIET=true

LOGFILE_BASE="/var/log/ntopng-updates"
LOGFILE="${LOGFILE_BASE}.log"
RETAIN_NUM_LINES=1000
APT_SOURCE="ntop@STABLE_SUFFIX@.list"

# Required by apt-get on Ubuntu 18 to run in non interactive mode
export DEBIAN_FRONTEND="noninteractive"

# Workaround for third-party installers not using absolute paths
export PATH="$PATH:/sbin:/usr/sbin"

# Read redis connection info
REDIS_INFO=$(${NTOPNG_CONFIG_TOOL} -a print-redis-info)
HOST=$(echo ${REDIS_INFO} | cut -d'=' -f2 | cut -d'@' -f1 | cut -d':' -f1)
PORT=$(echo ${REDIS_INFO} | cut -d'=' -f2 | cut -d'@' -f1 | cut -d':' -f2 -s)
PSWD=$(echo ${REDIS_INFO} | cut -d'=' -f2 | cut -d'@' -f1 | cut -d':' -f3 -s)
DBID=$(echo ${REDIS_INFO} | cut -d'@' -f2 -s)
REDIS_CLI_OPT=""
if [ ! -z "${HOST}" ]; then
	REDIS_CLI_OPT="${REDIS_CLI_OPT}-h ${HOST} "
fi
if [ ! -z "${PORT}" ]; then
	REDIS_CLI_OPT="${REDIS_CLI_OPT}-p ${PORT} "
fi
if [ ! -z "${PSWD}" ]; then
	REDIS_CLI_OPT="${REDIS_CLI_OPT}-a ${PSWD} "
fi
if [ ! -z "${DBID}" ]; then
	REDIS_CLI_OPT="${REDIS_CLI_OPT}-n ${DBID} "
fi

# Detect OS
OS="DEBIAN"
REDHAT=$(${LSB_RELEASE} | grep -i "CentOS\|RedHat")
if [ ! -z "${REDHAT}" ]; then
	OS="REDHAT"
fi

if [ "${OS}" == "DEBIAN" ]; then
	# Check for nEdge
	if dpkg --get-selections | grep -q "^nedge[[:space:]]*install$" >/dev/null; then
		PRODUCT="nedge"
	fi
fi

# Log init
function logsetup {
	if [ ! -z "$1" ]; then
		LOGFILE="${LOGFILE_BASE}-$1.log"
	fi
	TMP=$(tail -n $RETAIN_NUM_LINES $LOGFILE 2>/dev/null) && echo "${TMP}" > $LOGFILE
	exec > >(tee -a $LOGFILE)
	exec 2>&1
}

# Log
function log {
	echo "$(date +"%b %d %T"): $*"
}

# Help
function print_usage() {
	echo "${PRODUCT} updates utility"
	echo ""
	echo "Usage:"
	echo "`basename $0` -a check-updates"
	echo "`basename $0` -a check-updates-on-demand"
	echo "`basename $0` -a handle-upgrade-requests"
	echo ""
	echo "[-a check-updates]"
	echo "`basename $0` -a check-updates"
	echo "Check for ${PRODUCT} updates."
	echo ""
	echo "[-a check-updates-on-demand]"
	echo "`basename $0` -a check-updates-on-demand"
	echo "Check for requests for checking ${PRODUCT} updates."
	echo ""
	echo "[-a handle-upgrade-requests]"
	echo "`basename $0` -a handle-upgrade-requests"
	echo "Check for update requests and update ${PRODUCT}."
	echo ""
}

# Uncomment to enable logging
#logsetup $ACTION
#log Requested $ACTION

POSITIONAL=()
while [[ $# -gt 0 ]]
do
key="$1"

case $key in
	-a|--action)
	ACTION="$2"
	shift
	shift
	;;
	-v|--verbose)
	QUIET=false
	shift
	;;
	*)
	POSITIONAL+=("$1")
	shift
	;;
esac
done
set -- "${POSITIONAL[@]}" # restore positional parameters

if [[ ( $ACTION != "check-updates" && $ACTION != "check-updates-on-demand" && $ACTION != "handle-upgrade-requests" ) ]]; then
	print_usage
	exit 1
fi

function set_failure() {
	STATUS=$1
	# Report failure status
	${REDIS_CLI} ${REDIS_CLI_OPT} SET ${REDIS_UPDATE_FAILURE_KEY} "${STATUS}" >/dev/null 2>&1
	# Cleanup requests
	${REDIS_CLI} ${REDIS_CLI_OPT} DEL ${REDIS_CHECK_FOR_UPDATES_KEY} >/dev/null 2>&1 
	${REDIS_CLI} ${REDIS_CLI_OPT} DEL ${REDIS_RUN_UPGRADE_KEY} >/dev/null 2>&1
}

function stall_check() {
	INSTANCES=$(pgrep -fc "$0 -a ${ACTION}")
	if [ $INSTANCES -gt 2 ]; then
		echo "Another instance is running with the same action"
		exit 0
	fi
}

function maintenance_check() {

	APP_BIN="/usr/bin/${PRODUCT}"

	NOW=`date +'%s'`
	EXPIRY_EPOCH=`${APP_BIN} --check-maintenance | cut -d' ' -f 1`

	if [ "${EXPIRY_EPOCH}" = "Invalid" ]; then
		# Missing or expired or invalid license

		# Allow updates on missing license file as this does not 
		# invalidate the license and handle community mode
		if [ `${APP_BIN} --check-license | grep Empty | wc -l ` -eq 0 ]; then

			# Expired or invalid license
			set_failure "no-license"
			exit 0
		fi
	elif echo "${EXPIRY_EPOCH}" | grep -q "^[0-9]\+$"; then
		if [ "${EXPIRY_EPOCH}" -lt "${NOW}" ]; then
			# Expired maintenance
			set_failure "expired-maintenance"
			exit 0
		fi
	else
		# Unable to read maintenance
		set_failure "no-license"
		exit 0
	fi
}

function service_enabled_check() {
	SERVICE_ENABLED=$(/bin/systemctl is-enabled ${PRODUCT} 2>/dev/null)
	if [ ! "${SERVICE_ENABLED}" == "enabled" ]; then
		set_failure "service-not-enabled"
		exit 0
	fi
}

# Check if the user requested a ntopng update, and upgrade it in case
function run_upgrade() {
	RESULT=1

	# Check if upgrade is in progress (it can take more than 1 min)
	IN_PROGRESS=$(${REDIS_CLI} ${REDIS_CLI_OPT} GET ${REDIS_IN_PROGRESS_KEY} 2>/dev/null | grep 1)
	if [ "${IN_PROGRESS}" == "1" ]; then
		return
	fi

	# Set "in progress" flag
	${REDIS_CLI} ${REDIS_CLI_OPT} SET ${REDIS_IN_PROGRESS_KEY} "1" EX 3600 >/dev/null 2>&1 

	if [ "${OS}" == "DEBIAN" ]; then
		# Debian or Ubuntu

		# Update repo index
		${APT_GET} update -o Dir::Etc::sourcelist="sources.list.d/${APT_SOURCE}" -o Dir::Etc::sourceparts="-" -o APT::Get::List-Cleanup="0"

		# Install
		${APT_GET} upgrade --assume-yes --fix-broken --allow-unauthenticated --with-new-pkgs ${PRODUCT}
		RESULT=$?

		# Check if installation is successful (we do not trust $?)
		NTOPNG_VERSION=$(${APT_GET} --just-print upgrade 2>&1 | grep "Inst ${PRODUCT} " | cut -d'(' -f2 | cut -d' ' -f1)
		if [ -z "${NTOPNG_VERSION}" ]; then
			RESULT=0
		fi
	else
		# CentOS

		# Install
		${YUM} clean all
		${YUM} check-update ${PRODUCT}
		${YUM} -y install ${PRODUCT}
		RESULT=$?
	fi

	if [ $RESULT -eq 0 ]; then
		[ $QUIET = false ] && echo "${PRODUCT} updated successfully"
		# Cleanup new version available key
		${REDIS_CLI} ${REDIS_CLI_OPT} DEL ${REDIS_NEW_VERSION_AVAILABLE_KEY} >/dev/null 2>&1 
		${REDIS_CLI} ${REDIS_CLI_OPT} DEL ${REDIS_UPDATE_FAILURE_KEY} >/dev/null 2>&1

		# Restart the service if not running (e.g. old prerm was disabling it)
		/bin/systemctl -q is-active ${PRODUCT} || /bin/systemctl restart ${PRODUCT}
	else
		[ $QUIET = false ] && echo "Unable to update"
		${REDIS_CLI} ${REDIS_CLI_OPT} SET ${REDIS_UPDATE_FAILURE_KEY} "upgrade-failure" >/dev/null 2>&1
	fi

	# Cleanup "in progress" flag
	${REDIS_CLI} ${REDIS_CLI_OPT} DEL ${REDIS_IN_PROGRESS_KEY} >/dev/null 2>&1

	# Cleanup upgrade request key if any
	${REDIS_CLI} ${REDIS_CLI_OPT} DEL ${REDIS_RUN_UPGRADE_KEY} >/dev/null 2>&1
}

# Check for new ntopng updates available
function check_updates() {
	CRON_UPDATES=$1
	NTOPNG_VERSION=""
	RESULT=1
	if [ "${OS}" == "DEBIAN" ]; then
		# Debian or Ubuntu

		# Update repo index
		${APT_GET} update -o Dir::Etc::sourcelist="sources.list.d/${APT_SOURCE}" -o Dir::Etc::sourceparts="-" -o APT::Get::List-Cleanup="0"

		# Check for broken packages
		${APT_GET} --just-print upgrade >/dev/null 2>&1
		if [ ! $? -eq 0 ]; then
			# Something went wrong, trying to fix it
			${APT_GET} --assume-yes --fix-broken install

			${APT_GET} --just-print upgrade >/dev/null 2>&1
		fi

		if [ ! $? -eq 0 ]; then
			# Something went wrong, unable to fix it
			${REDIS_CLI} ${REDIS_CLI_OPT} SET ${REDIS_UPDATE_FAILURE_KEY} "update-failure" >/dev/null 2>&1 
		else
			# Check update and get version
			NTOPNG_VERSION=$(${APT_GET} --just-print --assume-yes --with-new-pkgs upgrade 2>&1 | grep "Inst ${PRODUCT} " | cut -d'(' -f2 | cut -d' ' -f1)
			RESULT=0
		fi
	else
		# CentOS

		# Check update and get version
		${YUM} clean all
		NTOPNG_VERSION=$(${YUM} check-update ${PRODUCT} | grep ${PRODUCT} | tr -s ' ' | cut -d' ' -f2)
		RESULT=0
	fi

	# If there is an update, set new version on redis, otherwise delete the current version (if any)
	if [ ! -z "${NTOPNG_VERSION}" ]; then
		${REDIS_CLI} ${REDIS_CLI_OPT} SET ${REDIS_NEW_VERSION_AVAILABLE_KEY} ${NTOPNG_VERSION} EX 86400 >/dev/null 2>&1 

		# Reset failure reason if any
		${REDIS_CLI} ${REDIS_CLI_OPT} DEL ${REDIS_UPDATE_FAILURE_KEY} >/dev/null 2>&1

		# Check if this is run by cron (not a manual check)
		if [ "${CRON_UPDATES}" = "1" ]; then
			# Check if automatic updates are enabled by the user
			AUTO_UPDATES=$(${REDIS_CLI} ${REDIS_CLI_OPT} GET ${REDIS_AUTO_UPDATES_KEY} 2>/dev/null | grep 1)
			if [ "${AUTO_UPDATES}" == "1" ]; then
				run_upgrade
			fi
		fi
	else
		${REDIS_CLI} ${REDIS_CLI_OPT} DEL ${REDIS_NEW_VERSION_AVAILABLE_KEY} >/dev/null 2>&1 
		if [ $RESULT -eq 0 ]; then
			${REDIS_CLI} ${REDIS_CLI_OPT} DEL ${REDIS_UPDATE_FAILURE_KEY} >/dev/null 2>&1 
		fi
	fi
}

# Check if the user requested a 'Check for updates' manually
function check_updates_on_demand() {
	CHECK_FOR_UPDATES=$(${REDIS_CLI} ${REDIS_CLI_OPT} GET ${REDIS_CHECK_FOR_UPDATES_KEY} 2>/dev/null | grep 1)
	if [ "${CHECK_FOR_UPDATES}" == "1" ]; then
		stall_check
		check_updates 0
		${REDIS_CLI} ${REDIS_CLI_OPT} DEL ${REDIS_CHECK_FOR_UPDATES_KEY} >/dev/null 2>&1 
	fi
}

# Check if the user requested a ntopng update, and upgrade it in case
function handle_upgrade_requests() {
	# Check redis for upgrade requests
	UPGRADE=$(${REDIS_CLI} ${REDIS_CLI_OPT} GET ${REDIS_RUN_UPGRADE_KEY} 2>/dev/null | grep 1)
	if [ "${UPGRADE}" == "1" ]; then
		stall_check
		run_upgrade
	fi
}

# Run this cron script when the ntopng service is enabled only
service_enabled_check

# Check if the product is under maintenance
maintenance_check

# Check OEM mode
OEM_MODE=$(${REDIS_CLI} ${REDIS_CLI_OPT} EXISTS ${REDIS_PRODUCT_NAME_KEY} 2>/dev/null | grep 1)
if [ "${OEM_MODE}" == "1" ]; then
	# Exit unless there is an OEM source file
	APT_SOURCE="ntop-oem.list"
	eval $(apt-config shell APT_ETC Dir::Etc)
	[ -e "/${APT_ETC}sources.list.d/${APT_SOURCE}" ] || exit 0
fi

if [ $ACTION == "check-updates" ]; then

	check_updates 1

elif [ $ACTION == "check-updates-on-demand" ]; then

	check_updates_on_demand

elif [ $ACTION == "handle-upgrade-requests" ]; then

	handle_upgrade_requests

else
	# never reached
	echo "Unknown action $ACTION"
	exit 1
fi

exit 0
