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

. "/usr/sbin/unifi-network-service-common"

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

# Constants

### Helper functions ###

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 ; 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 ; 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 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 [${SUPPORTED_JAVA_VERSIONS[*]}]"
}

function set_java_home () {
	for supported_version in "${SUPPORTED_JAVA_VERSIONS[@]}"; do
		if set_java_home_version "$supported_version"; then
			return 0
		fi
	done
	log_unsupported_java_version
}
function set_java_home_version() {
	supported_java_version=$1
	log "Attempting to set 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`
	java_list=`echo ${java_list} temurin-${supported_java_version}-jre-${arch}`
	java_list=`echo ${java_list} temurin-${supported_java_version}-jre`

	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
	return 1
}

function get_unifi_property () {
	if [ -f ${DATADIR}/system.properties ]; then
		property_name=$1
		cut -d "=" -f2 <<< $(grep "^[^#;]" ${DATADIR}/system.properties | grep $property_name) | cut -d ' ' -f1
	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)
	for supported_version in "${SUPPORTED_JAVA_VERSIONS[@]}"; do
		if [[ "java-${current_java_version}" == "java-${supported_version}" ]]; then
			return 0
		fi
	done
	return 1
}

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 remove_system_env_file_if_exists() {
	[ -f "${DATADIR}/system_env" ] && rm "${DATADIR}/system_env"
}

function load_environment () {
	remove_system_env_file_if_exists
	if 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
}

function shutdown_unifi_mongo_service_if_needed() {
	if [ "${UNIFI_MONGODB_SERVICE_ENABLED}" == "true" ] && [ $(systemctl is-active --quiet unifi-mongodb && echo $?) ]; then
		runningMongoDbVersion=$(mongo localhost:27117/ace --quiet --eval "db.version()")
		echo "Running MongoDB version - $runningMongoDbVersion"
		installedMongoDbVersion=$(dpkg-query --showformat="\${VERSION}" --show mongodb-server)
		echo "Installed MongoDB version - $installedMongoDbVersion"
		if [[ $installedMongoDbVersion == *"$runningMongoDbVersion"* ]]; then
			echo "MongoDB version has not changed"
		else
			log "MongoDB version has changed, shutting down unifi-mongodb.service"
			systemctl stop unifi-mongodb
			log "Finished shutting down unifi-mongodb.service"
		fi
	fi
}

### 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}
MODEL_SHORTNAME=$(get_model_shortname)
UNIFI_CORE_ENABLED=${UNIFI_CORE_ENABLED:-"false"}
UNIFI_MONGODB_SERVICE_ENABLED=${UNIFI_MONGODB_SERVICE_ENABLED:-false}
SUPPORTED_JAVA_VERSIONS=("17" "21" "25")

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
	;;
	cleanup)
		unmount_tmpfs_dir
	;;
	set-java-home)
		set_java_home_if_needed
	;;
	load-environment)
		load_environment
	;;
	shutdown-unifi-mongo-if-needed)
		shutdown_unifi_mongo_service_if_needed
	;;
	service-started)
		log_unifi_service_cpu_time
	;;
	*)
		log "Usage: $0 {init|init-uos|create-dirs|healthcheck|cleanup|set-java-home|load-environment}"
		exit 1
	;;
esac

exit 0
