#!/bin/bash
#
# Helper script for Ubiquiti UniFi Network
#

# General parameters
NAME="unifi"
BASEDIR="/usr/lib/unifi"
ADDITIONAL_STORAGE_BASE="/srv"

# Constants
APP_STATUS_UP="up"
APP_STATUS_DOWN="down"
APP_STATUS_MIGRATING="migrating"
STATUS_RESPONSE_FILE="/tmp/unifi-network-status-response"

### Helper functions ###

function does_not_have_ubnt_tools() {
	if [[ ! -f /sbin/ubnt-tools ]]; then
		return 0
	else
		return 1
	fi
}

function is_uck () {
	if [[ "${MODEL_SHORTNAME}" == "UCK" ]]; then
		return 0
	else
		return 1
	fi
}

function log() {
	if does_not_have_ubnt_tools || is_uck ; then
		logger -s -t "${NAME}" "${@}"
	else
		logger -s -t "${NAME}" -p 3 "${@}"
	fi
}

function is_link() {
	[ -L "$1" ] && [ -e "$1" ]
}

function is_mountpoint() {
	mountpoint -q "$1"
}

function is_additional_storage_available() {
	is_mountpoint "$ADDITIONAL_STORAGE_BASE" || is_link "$ADDITIONAL_STORAGE_BASE"
}

function create_additional_storage_dir() {
	if is_additional_storage_available && [ ! -d "$ADDITIONAL_STORAGE_DIR" ]; then
		mkdir "$ADDITIONAL_STORAGE_DIR"
		chown "$USER:$GROUP" "$ADDITIONAL_STORAGE_DIR"
	fi
}

function mount_tmpfs_dir() {
	if does_not_have_ubnt_tools || is_uck ; then
		return 0
	fi
	# Support directories
	mkdir -p "$TMPFS_DIR"
	chown -R "$USER:$GROUP" "$TMPFS_DIR"

	# Mount temp filesystem
	mount -t tmpfs -o size=$TMPFS_SIZE tmpfs $TMPFS_DIR || log "Warning: Could not mount tmpfs to $TMPFS_DIR"
}

function unmount_tmpfs_dir() {
	if does_not_have_ubnt_tools || is_uck ; then
		return 0
	fi
	log "tmpfs: $TMPFS_DIR"
	if is_mountpoint "$TMPFS_DIR"; then
		umount "$TMPFS_DIR" || log "Warning: Could not unmount tmpfs from $TMPFS_DIR"
	fi
}

function create_dirs() {
	log "user: $USER, group: $GROUP"
	log "data: $DATADIR, logs: $LOGDIR, run: $RUNDIR, additional storage: $ADDITIONAL_STORAGE_DIR"
	log "tmpfs: $TMPFS_DIR, size: $TMPFS_SIZE"
	link_to_persistent "${DATADIR}" "${BASEDIR}"/data
	link_to_persistent "${LOGDIR}"  "${BASEDIR}"/logs
	link_to_persistent "${RUNDIR}"  "${BASEDIR}"/run
	create_additional_storage_dir
	mount_tmpfs_dir

	check_dir_permissions || log "Warning: Wrong permissions for ${DATADIR}, ${LOGDIR}, ${RUNDIR}"

	return 0
}

function check_dir_permissions() {
	local unifi_uid=$(id -u "${USER}")

	local datadir_uid=$(stat ${DATADIR} -Lc %u)
	local logdir_uid=$( stat ${LOGDIR}  -Lc %u)
	local rundir_uid=$( stat ${RUNDIR}  -Lc %u)

	[ "${unifi_uid}" == "${datadir_uid}" ] && \
	[ "${unifi_uid}" == "${logdir_uid}"  ] && \
	[ "${unifi_uid}" == "${logdir_uid}"  ] && \
	\
	sudo -u "${USER}" test -w "${DATADIR}" && \
	sudo -u "${USER}" test -w "${LOGDIR}"  && \
	sudo -u "${USER}" test -w "${RUNDIR}"  && \
	\
		return 0;

	return 1;
}

function link_to_persistent() {
	local persistent_dir="${1}"
	local orig_dir="${2}"
	install -o "${USER}" -dD "${persistent_dir}"
	chown -R "${USER}": "${persistent_dir}"
	if ! [ -L "${orig_dir}" -a "${persistent_dir}" -ef  "${orig_dir}" ]; then
		rm -rf "${orig_dir}"
		ln -sf "${persistent_dir}" "${orig_dir}"
	fi
}

function is_application_migrating() {
	[[ `cat "$STATUS_RESPONSE_FILE" | grep '"db_migrating":true'` != "" ]]
}

function is_application_up() {
	[[ `cat "$STATUS_RESPONSE_FILE" | grep '"server_running":true'` != "" ]]
}

function is_application_connected_to_mongo() {
	grep -q '"db_connected":true' "$STATUS_RESPONSE_FILE"
}

function get_application_status() {
	local status="$APP_STATUS_DOWN"
	local http_code=$(curl -s --connect-timeout 1 --max-time 5 -o "$STATUS_RESPONSE_FILE" -w "%{http_code}" "$STATUS_URL")
	if [ "${http_code}" == "200" ]; then
		if is_application_migrating; then
		  status="$APP_STATUS_MIGRATING"
		elif is_application_up && is_application_connected_to_mongo; then
			status="$APP_STATUS_UP"
		fi
	fi
	echo "$status"
}

function health_check_loop() {
	log "health-check max retry count: $HEALTHCHECK_MAX_RETRIES"
	local retry_count=0
	local migration_retry_count=0
	while :; do
		sleep 3
		local application_status="$(get_application_status)"
		if [[ $application_status == $APP_STATUS_UP ]]; then
			log "${NAME} is up and running. Health check finished successfully."
			exit 0
		else
			if [[ $application_status == $APP_STATUS_MIGRATING ]]; then
				migration_retry_count=$((migration_retry_count+1))
			else
				retry_count=$((retry_count+1))
			fi
		fi
		if [ ${migration_retry_count} -ge ${HEALTHCHECK_MAX_RETRIES} ]; then
			log "${NAME} is still running DB migration. Health check timed out."
			exit 1
		elif [ ${retry_count} -ge ${HEALTHCHECK_MAX_RETRIES} ]; then
			log "${NAME} is not running. Health check timed out."
			exit 1
		fi
	done
}

function get_java_version () {
	java -version 2>&1 | head -1 | cut -d'"' -f2 | sed '/^1\./s///' | cut -d'.' -f1
}

function set_java_version () {
	sudo update-alternatives --set java $1
}

function log_unsupported_java_version() {
	log "Warning: Could not find alternative/supported Java version. Make sure '/usr/bin/java' points to supported version - Java $REQUIRED_JAVA_VERSION"
}

function set_java_home () {
	supported_java_version=$REQUIRED_JAVA_VERSION
	log "Setting Java home to supported version - $supported_java_version"
	arch=`dpkg --print-architecture 2>/dev/null`
	if [ ! "$arch" ]; then
		arch=`uname -m`
		if [ "$arch" = aarch64 ]; then
			arch=arm64
		fi
	fi
	java_list=''
	java_list=`echo ${java_list} java-${supported_java_version}-openjdk-${arch}`
	java_list=`echo ${java_list} java-${supported_java_version}-openjdk`

	for a in i386 amd64 armhf arm64; do
		java_list=$(echo ${java_list} oracle-java${supported_java_version}-jdk-${a}/jre)
	done
	for a in i386 amd64; do
		java_list=$(echo ${java_list} oracle-java${supported_java_version}-jre-${a})
	done
	for a in x64 i586 arm32-vfp-hflt arm64-vfp-hflt; do
		java_list=$(echo ${java_list} jdk-${supported_java_version}-oracle-${a}/jre)
	done
	for a in x64 i586; do
		java_list=$(echo ${java_list} jre-${supported_java_version}-oracle-${a})
	done
	for a in i386 amd64 armhf arm64; do
		java_list=$(echo ${java_list} adoptopenjdk-${supported_java_version}-hotspot-${a})
	done

	java_list=$(echo ${java_list} java-${supported_java_version}-oracle/jre)

	cur_java=`update-alternatives --query java 2>/dev/null | awk '/^Value: /{print $2}'`
	cur_real_java=`readlink -f ${cur_java} 2>/dev/null`
	for jvm in ${java_list}; do
		jvm_real_java=`readlink -f /usr/lib/jvm/${jvm}/bin/java 2>/dev/null`
		[ "${jvm_real_java}" != "" ] || continue
		if [ "${jvm_real_java}" == "${cur_real_java}" ]; then
			set_java_version $cur_java
			return 0
		fi
	done

	alts_java=`update-alternatives --query java 2>/dev/null | awk '/^Alternative: /{print $2}'`
	for cur_java in ${alts_java}; do
		cur_real_java=`readlink -f ${cur_java} 2>/dev/null`
		for jvm in ${java_list}; do
			jvm_real_java=`readlink -f /usr/lib/jvm/${jvm}/bin/java 2>/dev/null`
			[ "${jvm_real_java}" != "" ] || continue
			if [ "${jvm_real_java}" == "${cur_real_java}" ]; then
				set_java_version $cur_java
				return 0
			fi
		done
	done

	log_unsupported_java_version
}

function get_unifi_property () {
	if [ -f ${DATADIR}/system.properties ]; then
		property_name=$1
		cut -d "=" -f2 <<< $(grep "^[^#;]" ${DATADIR}/system.properties | grep $property_name)
	fi
}

function get_environment_property () {
	if [ -f /etc/default/${NAME} ]; then
		property_name=$1
		cut -d "=" -f2 <<< $(grep "^[^#;]" "/etc/default/${NAME}" | grep $property_name)
	fi
}

function get_model_shortname () {
	if [ -f /sbin/ubnt-tools ]; then
		cut -d "=" -f2 <<< $(sudo /sbin/ubnt-tools id | grep "board.shortname")
	fi
}

function load_jvm_opts () {
	heap_min=$(get_unifi_property 'unifi.xms')
	heap_max=$(get_unifi_property 'unifi.xmx')
	stack_size=$(get_unifi_property 'unifi.xss')
	g1gc_enabled=$(get_unifi_property 'unifi.G1GC.enabled')

	set_jvm_opts=""
	if [ "${heap_min}" != "" ]; then
		set_jvm_opts=$(echo ${set_jvm_opts} -Xms${heap_min}M)
	elif [ "${JVM_INIT_HEAP_SIZE}" != "" ]; then
		set_jvm_opts=$(echo ${set_jvm_opts} -Xms${JVM_INIT_HEAP_SIZE}M)
	fi
	if [ "${heap_max}" != "" ]; then
		set_jvm_opts=$(echo ${set_jvm_opts} -Xmx${heap_max}M)
	elif [ "${JVM_MAX_HEAP_SIZE}" != "" ]; then
		set_jvm_opts=$(echo ${set_jvm_opts} -Xmx${JVM_MAX_HEAP_SIZE}M)
	else
		set_jvm_opts=$(echo ${set_jvm_opts} -Xmx1024M)
	fi
	if [ "${stack_size}" != "" ]; then
		set_jvm_opts=$(echo ${set_jvm_opts} -Xss${stack_size}K)
	fi
	if [[ "${g1gc_enabled}" != "" && $g1gc_enabled ]]; then
		set_jvm_opts=$(echo ${set_jvm_opts} -XX:+UseG1GC)
	else
		set_jvm_opts=$(echo ${set_jvm_opts} -XX:+UseParallelGC)
	fi
	if [ "${JAVA_ENTROPY_GATHER_DEVICE}" != "" ]; then
		set_jvm_opts=$(echo ${set_jvm_opts} -Djava.security.egd=${JAVA_ENTROPY_GATHER_DEVICE})
	fi
	if [ "${UNIFI_JVM_EXTRA_OPTS}" != "" ]; then
		set_jvm_opts=$(echo ${set_jvm_opts} ${UNIFI_JVM_EXTRA_OPTS})
	fi

	if [ "${set_jvm_opts}" != "-Xmx1024M -XX:+UseParallelGC" ]; then
		set_jvm_opts="\""$set_jvm_opts"\""
		log "Adding UNIFI_JVM_OPTS=$set_jvm_opts to ${DATADIR}/system_env"
		echo UNIFI_JVM_OPTS="$set_jvm_opts" >> "${DATADIR}/system_env"
	fi
}

function supported_java_version_set() {
	current_java_version=$(get_java_version)
	[ "java-${current_java_version}" == "java-${REQUIRED_JAVA_VERSION}" ]
}

function set_java_home_if_needed() {
	supported_java_version_set || set_java_home
}

function create_system_env_file() {
	[ -f "${DATADIR}/system_env" ] && echo -n > "${DATADIR}/system_env" || touch "${DATADIR}/system_env"
}

function load_environment () {
	if is_uck ; then
		create_system_env_file
		uck_jvm_opts='"-Xmx768M -XX:+UseParallelGC"'
		log "Adding UNIFI_JVM_OPTS=$uck_jvm_opts to ${DATADIR}/system_env"
		echo UNIFI_JVM_OPTS="$uck_jvm_opts" >> "${DATADIR}/system_env"
	elif does_not_have_ubnt_tools ; then
		create_system_env_file
		load_jvm_opts
	else
		log "Skipping load-environment..."
	fi
}

function dir_symlink_fix () {
	local DSTDIR=$1
	local SYMLINK=$2
	local MYUSER=$3
	local MYGROUP=$4
	local MYMODE=$5

	[ -d ${DSTDIR} ] || install -o ${MYUSER} -g ${MYGROUP} -m ${MYMODE} -d ${DSTDIR}
	[ -d ${SYMLINK} -a ! -L ${SYMLINK} ] && mv ${SYMLINK} `mktemp -u ${SYMLINK}.XXXXXXXX`
	[ "$(readlink ${SYMLINK})" = "${DSTDIR}" ] || (rm -f ${SYMLINK} && ln -sf ${DSTDIR} ${SYMLINK})
}

function file_symlink_fix () {
	local DSTFILE=$1
	local SYMLINK=$2

	if [ -f ${DSTFILE} ]; then
		[ -f ${SYMLINK} -a ! -L ${SYMLINK} ] && mv ${SYMLINK} `mktemp -u ${SYMLINK}.XXXXXXXX`
		[ "$(readlink ${SYMLINK})" = "${DSTFILE}" ] || (rm -f ${SYMLINK} && ln -sf ${DSTFILE} ${SYMLINK})
	fi
}

function nested_symlinks_fix () {
	# Fix issue with nested symbolic links
	[ -L "${BASEDIR}"/data/data ] && rm "${BASEDIR}"/data/data
	[ -L "${BASEDIR}"/logs/logs ] && rm "${BASEDIR}"/logs/logs
	[ -L "${BASEDIR}"/run/run ] && rm "${BASEDIR}"/run/run
	[ -L "${BASEDIR}"/data/unifi ] && rm "${BASEDIR}"/data/unifi
	[ -L "${BASEDIR}"/logs/unifi ] && rm "${BASEDIR}"/logs/unifi
	[ -L "${BASEDIR}"/run/unifi ] && rm "${BASEDIR}"/run/unifi
}

function init () {
	local UMASK=027
	umask ${UMASK}
	local DIR_MODE=$(printf '%x' $((0x7777 - 0x${UMASK} & 0x0777)))

	nested_symlinks_fix
	dir_symlink_fix ${DATADIR} "${BASEDIR}"/data ${USER} ${GROUP} ${DIR_MODE}
	dir_symlink_fix ${LOGDIR} "${BASEDIR}"/logs ${USER} ${GROUP} ${DIR_MODE}
	dir_symlink_fix ${RUNDIR} "${BASEDIR}"/run ${USER} ${GROUP} ${DIR_MODE}
	[ -z "${UNIFI_SSL_KEYSTORE}" ] || file_symlink_fix ${UNIFI_SSL_KEYSTORE} "${BASEDIR}"/data/keystore
	[ supported_java_version_set ] || log_unsupported_java_version
	load_environment
}

function init_uos () {
	nested_symlinks_fix
	create_dirs
	set_java_home_if_needed
	load_environment
}

### Helper functions end ###

# Load variables
[ -f "/etc/default/${NAME}" ] && . "/etc/default/${NAME}"

USER=${UNIFI_USER:-unifi}
UNIFI_GROUP=$(id -gn ${USER})
GROUP=${UNIFI_GROUP:-unifi}
DATADIR=${UNIFI_DATA_DIR:-/var/lib/${NAME}}
LOGDIR=${UNIFI_LOG_DIR:-/var/log/${NAME}}
RUNDIR=${UNIFI_RUN_DIR:-/var/run/${NAME}}
ADDITIONAL_STORAGE_DIR=${UNIFI_ADDITIONAL_STORAGE_DIR:-$ADDITIONAL_STORAGE_BASE/$NAME}
TMPFS_DIR=${UNIFI_TMPFS_DIR:-/var/opt/$NAME/tmp}
TMPFS_SIZE=${UNIFI_TMPFS_SIZE:-64m}
HEALTHCHECK_MAX_RETRIES=${UNIFI_HEALTHCHECK_MAX_RETRIES:-100}
UNIFI_HTTP_PORT=$(get_unifi_property 'unifi.http.port')
STATUS_URL="http://localhost:${UNIFI_HTTP_PORT:-8080}/status"
MODEL_SHORTNAME=$(get_model_shortname)
UNIFI_CORE_ENABLED=${UNIFI_CORE_ENABLED:-"false"}
REQUIRED_JAVA_VERSION='8'

case "$1" in
	init)
		if [[ "${UNIFI_CORE_ENABLED}" != "true" ]]; then
			init
			log "init complete..."
		else
			log "Skipping init..."
		fi
	;;
	init-uos)
		if [[ "${UNIFI_CORE_ENABLED}" == "true" ]]; then
			init_uos
			log "init-uos complete..."
		else
			log "Skipping init-uos..."
		fi
	;;
	create-dirs)
		create_dirs
	;;
	healthcheck)
		health_check_loop
	;;
	cleanup)
		unmount_tmpfs_dir
	;;
	set-java-home)
		set_java_home_if_needed
	;;
	load-environment)
		load_environment
	;;
	*)
		log "Usage: $0 {init|init-uos|create-dirs|healthcheck|cleanup|set-java-home|load-environment}"
		exit 1
	;;
esac

exit 0
