#!/bin/sh

# Copyright (c) 2015 Electric Sheep Fencing, LLC. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.
#
# 3. All advertising materials mentioning features or use of this software
#    must display the following acknowledgment:
#    "This product includes software developed by the pfSense Project
#    for use in the pfSense® software distribution. (http://www.pfsense.org/).
#
# 4. The names "pfSense" and "pfSense Project" must not be used to
#    endorse or promote products derived from this software without
#    prior written permission. For written permission, please contact
#    coreteam@pfsense.org.
#
# 5. Products derived from this software may not be called "pfSense"
#    nor may "pfSense" appear in their names without prior written
#    permission of the Electric Sheep Fencing, LLC.
#
# 6. Redistributions of any form whatsoever must retain the following
#    acknowledgment:
#
# "This product includes software developed by the pfSense Project
# for use in the pfSense software distribution (http://www.pfsense.org/).
#
# THIS SOFTWARE IS PROVIDED BY THE pfSense PROJECT ``AS IS'' AND ANY
# EXPRESSED OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE pfSense PROJECT OR
# ITS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
# NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.

usage() {
	echo "Usage: $(basename ${0}) [-46bdyf] [-u|-i PKG_NAME|-r PKG_NAME]" >&2
	echo "	-4          - Force IPv4"
	echo "	-6          - Force IPv6"
	echo "	-b          - Platform is booting" >&2
	echo "	-c          - Check if upgrade is necessary" >&2
	echo "	-d          - Turn on debug" >&2
	echo "	-f          - Force package installation" >&2
	echo "	-h          - Show this usage help" >&2
	echo "	-l          - Logfile path (defaults to /cf/conf/upgrade_log.txt)" >&2
	echo "	-p socket   - Write pkg progress to socket"
	echo "	-R          - Do not reboot (this can be dangerous)"
	echo "	-y          - Assume yes as the answer to any possible interaction" >&2
	echo "" >&2
	echo "The following parameters are mutually exclusive:" >&2
	echo "	-i PKG_NAME - Install package PKG_NAME" >&2
	echo "	-r PKG_NAME - Remove package PKG_NAME" >&2
	echo "	-u          - Update repository information" >&2
}

_echo() {
	local _n=""
	if [ "${1}" = "-n" ]; then
		shift
		_n="-n"
	fi

	if [ -z "${logfile}" ]; then
		logfile=/dev/null
	fi

	echo ${_n} "${1}" | tee -a ${logfile}
}

_exec() {
	local _cmd="${1}"
	local _msg="${2}"
	local _mute="${3}"
	local _ignore_result="${4}"
	local _stdout="${stdout}"

	if [ -z "${_cmd}" -o -z "${_msg}" ]; then
		return 1
	fi

	if [ "${_mute}" != "mute" ]; then
		_stdout=''
	fi

	_echo -n ">>> ${_msg}... "
	if [ -z "${_stdout}" ]; then
		_echo ""
		# Ref. http://stackoverflow.com/questions/1221833/bash-pipe-output-and-capture-exit-status
		exec 4>&1
		local _result=$({ { ${_cmd} 2>&1 3>&-; printf $? 1>&3; } 4>&- | \
			tee -a ${logfile} 1>&4; } 3>&1)
		exec 4>&-
	else
		# Ref. http://stackoverflow.com/questions/1221833/bash-pipe-output-and-capture-exit-status
		exec 4>&1
		local _result=$({ { ${_cmd} >${_stdout} 2>&1 3>&-; printf $? 1>&3; } 4>&- | \
			tee -a ${logfile} 1>&4; } 3>&1)
		exec 4>&-
	fi

	if [ ${_result} -eq 0 -o -n "${_ignore_result}" ]; then
		[ -n "${_stdout}" ] \
			&& _echo "done."
		return 0
	else
		[ -n "${_stdout}" ] \
			&& _echo "failed."
		_exit 1
	fi
}

_exit() {
	trap "-" 1 2 15 EXIT

	pkg_lock ${kernel_pkg}

	if [ -f "${pid_file}" ]; then
		rm -f ${pid_file}
	fi

	if [ -n "${chroot_dir}" ]; then
		umount -f ${chroot_dir} >/dev/null 2>&1
	fi

	if [ -z "${booting}" -o "${boot_stage}" != "2" ]; then
		/etc/rc.conf_mount_ro
	fi

	local _rc=${1:-"0"}

	# If EVENT_PIPE is defined, GUI is calling
	[ -n "${progress_socket}" ] \
		&& _echo "__RC=${_rc}"

	exit ${_rc}
}

pkg_with_pb() {
	local _event_pipe=""

	if [ -n "${progress_socket}" ]; then
		if [ -e "${progress_socket}" ]; then
			rm -f ${progress_socket}
		fi

		_event_pipe="-o EVENT_PIPE=${progress_socket}"

		nc -lU ${progress_socket} >> ${progress_file} &

		while [ ! -e "${progress_socket}" ]; do
			sleep 0.1
		done
	fi

	pkg ${_event_pipe} $@
	return $?
}

fetch_upgrade_packages() {
	local _pkgs_to_fetch=""
	if [ "${platform}" = "nanobsd" ]; then
		local _pkg=""

		# Check if all non-auto packages installed on 2nd partition are
		# installed on current one, if not, mark them to be deleted by
		# pkg autoremove
		for _pkg in $(pkg ${pkg_chroot} query -e '%a == 0' %n); do
			if ! pkg info -e ${_pkg}; then
				_exec "pkg ${pkg_chroot} set -A 1 ${_pkg}" "Scheduling package ${_pkg} for removal"
			fi
		done

		# Check if all non-auto packages installed on current partition are
		# installed on 2nd one, if not, we need to fetch them
		for _pkg in $(pkg query -e '%a == 0' %n); do
			if ! pkg ${pkg_chroot} info -e ${_pkg}; then
				_pkgs_to_fetch="${_pkgs_to_fetch}${_pkgs_to_fetch:+ }${_pkg}"
			fi
		done

	fi

	_echo ">>> Downloading upgrade packages..."
	if ! pkg_with_pb ${pkg_chroot} upgrade -F 2>&1 | tee -a ${logfile}; then
		_echo "ERROR: It was not possible to download packages"
		_exit 1
	fi

	if [ -n "${_pkgs_to_fetch}" ]; then
		_echo ">>> Fetching packages not present on upgrade partition..."
		if ! pkg_with_pb ${pkg_chroot} fetch -d ${_pkgs_to_fetch} 2>&1 | tee -a ${logfile}; then
			_echo "ERROR: It was not possible to fetch packages"
			_exit 1
		fi
	fi
}

pkg_lock() {
	local _pkg="${1}"

	if [ -z "${_pkg}" ]; then
		return
	fi

	if [ "$(pkg ${pkg_chroot} query %k ${_pkg})" = "0" ]; then
		_exec "pkg ${pkg_chroot} lock ${_pkg}" "Locking package ${_pkg}" mute
	fi
}

pkg_unlock() {
	local _pkg="${1}"

	if [ -z "${_pkg}" ]; then
		return
	fi

	if [ "$(pkg ${pkg_chroot} query %k ${_pkg})" = "1" ]; then
		_exec "pkg ${pkg_chroot} unlock ${_pkg}" "Unlocking package ${_pkg}" mute
	fi
}

pkg_update() {
	local _run_update=1

	local _force=""
	if [ "${1}" = "force" ]; then
		local _force=1
	fi

	if [ -z "${_force}" -a -f ${last_update_file} ]; then
		local _last_update=$(head -n 1 ${last_update_file})
		# Verify if content contain only numbers
		if echo "${_last_update}" | grep -E -q '^[0-9]+$'; then
			local _now=$(date +%s)
			# Only run update hourly, and if last update is in the future
			[ ${_now} -gt ${_last_update} -a $((${_now} - ${_last_update})) -le $((60 * 60)) ] \
				&& unset _run_update
		fi
	fi

	[ -z "${_run_update}" ] \
		&& return 0

	_exec "pkg ${pkg_chroot} update" "Updating repositories" mute
	date +%s > ${last_update_file}
}

pkg_upgrade() {
	# figure out which kernel variant is running
	export kernel_pkg=$(pkg query %n $(pkg info ${product}-kernel-\* | grep -v -- -debug-))

	if [ -z "${kernel_pkg}" ]; then
		_echo "ERROR: It was not possible to identify which ${product} kernel is installed"
		_exit 1
	fi

	export next_stage=$(pkg annotate -q -S ${kernel_pkg} next_stage)

	if [ -n "${next_stage}" -a -n "${booting}" -a -n "${boot_stage}" ]; then
		if [ ${boot_stage} != ${next_stage} ]; then
			_exit 0
		fi
	fi

	# If it's booting and first stage didn't run, just exit
	if [ -n "${booting}" -a -z "${next_stage}" ]; then
		_exit 0
	fi

	unset need_reboot
	# First upgrade stage
	if [ -z "${next_stage}" ]; then
		if [ -f "${logfile}" ]; then
			rm -f ${logfile}
		fi

		pkg_update

		if [ "$(compare_pkg_version pkg)" = "<" ]; then
			_exec "pkg upgrade pkg" "Upgrading pkg" mute
			pkg_update force
		fi

		if [ $(pkg upgrade -nq | wc -l) -le 1 ]; then
			_echo "Your packages are up to date"
			_exit 0
		fi

		if [ $(pkg upgrade -r ${product}-core -nq | wc -l) -gt 1 ]; then
			setup_nanobsd_env
			need_reboot=1
		fi

		pkg_unlock ${kernel_pkg}

		if [ "${platform}" = "nanobsd" ] && \
		   [ $(pkg ${pkg_chroot} upgrade -nq | wc -l) -le 1 ]; then
			_echo "**** WARNING ****"
			_echo "Reboot will be required!!"
			_echo "Secondary partition is up to date"
			if [ -z "${yes}" ]; then
				_echo -n "Proceed with upgrade? (y/N) "
				read answer
				if [ "${answer}" != "y" ]; then
					_echo "Aborting..."
					_exit 0
				fi
			fi
			switch_active_nanobsd_partition
			do_reboot
			_exit 0
		fi

		if [ -z "${yes}" ]; then
			# Show user which packages are going to be upgraded
			pkg ${pkg_chroot} upgrade -nq 2>&1 | tee -a ${logfile}

			_echo ""
			if [ -n "${need_reboot}" ]; then
				_echo "**** WARNING ****"
				_echo "Reboot will be required!!"
			fi
			_echo -n "Proceed with upgrade? (y/N) "
			read answer
			if [ "${answer}" != "y" ]; then
				_echo "Aborting..."
				_exit 0
			fi
		fi

		# Download all upgrade packages first
		fetch_upgrade_packages

		if [ $(pkg ${pkg_chroot} upgrade -nq ${kernel_pkg} | wc -l) -gt 1 ]; then
			_exec "pkg ${pkg_chroot} upgrade ${kernel_pkg}" "Upgrading ${product} kernel"
		fi

		pkg ${pkg_chroot} annotate -q -M ${kernel_pkg} next_stage 2
		next_stage=2

		if [ -n "${need_reboot}" -a "${platform}" != "nanobsd" ]; then
			do_reboot
			_exit 0
		fi
	fi

	if [ "${next_stage}" = "2" ]; then
		pkg_lock "${pkg_prefix}*"

		if [ $(pkg ${pkg_chroot} upgrade -nq | wc -l) -gt 1 ]; then
			_echo "Upgrading necessary packages..."
			if ! pkg ${pkg_chroot} upgrade 2>&1 | tee -a ${logfile}; then
				pkg ${pkg_chroot} annotate -q -D ${kernel_pkg} next_stage
				pkg_unlock "${pkg_prefix}*"
				_echo "ERROR: An error occurred when upgrade was running..."
				_exit 1
			fi
		fi

		# XXX: workaround for #5300
		sort -u ${chroot_dir}/usr/local/etc/php/extensions.ini > /tmp/extensions.ini
		mv /tmp/extensions.ini ${chroot_dir}/usr/local/etc/php/extensions.ini

		pkg ${pkg_chroot} annotate -q -M ${kernel_pkg} next_stage 3
		next_stage=3

		pkg_unlock "${pkg_prefix}*"

		if [ -n "${need_reboot}" -a "${platform}" = "nanobsd" ]; then
			switch_active_nanobsd_partition
			do_reboot
			_exit 0
		fi

		if [ -n "${booting}" ]; then
			_exit 0
		fi
	fi

	if [ "${next_stage}" = "3" ]; then
		if [ $(pkg upgrade -nq | wc -l) -gt 1 ]; then
			_echo "Upgrading necessary packages..."
			if ! pkg ${pkg_chroot} upgrade 2>&1 | tee -a ${logfile}; then
				pkg ${pkg_chroot} annotate -q -D ${kernel_pkg} next_stage
				_echo "ERROR: An error occurred when upgrade was running..."
				_exit 1
			fi
		fi

		pkg ${pkg_chroot} annotate -q -D ${kernel_pkg} next_stage

		# cleanup caches
		_exec "pkg ${pkg_chroot} autoremove" "Removing unnecessary packages" mute ignore_result
		_exec "pkg ${pkg_chroot} clean" "Cleanup pkg cache" mute ignore_result
	fi
}

check_upgrade() {
	# figure out main meta package name
	if is_pkg_installed ${product}-vmware; then
		local _meta_pkg="${product}-vmware"
	elif is_pkg_installed ${product}; then
		local _meta_pkg="${product}"
	else
		_echo "ERROR: It was not possible to identify which ${product} meta package is installed"
		_exit 1
	fi

	pkg_update

	if [ "$(compare_pkg_version ${_meta_pkg})" = "<" ]; then
		local _new_version=$(pkg rquery %v ${_meta_pkg})
		_echo "${_new_version} version of ${product} is available"
		_exit 2
	else
		for _pkg in $(pkg query -e "%n ~ ${product}-*" %n); do
			# Ignore additional packages
			if echo "${_pkg}" | grep -q "^${pkg_prefix}"; then
				continue
			fi
			if [ "$(compare_pkg_version ${_pkg})" = "<" ]; then
				local _new_version=$(pkg rquery %v ${_pkg})
				_echo "${_new_version} version of ${_pkg} is available"
				_exit 2
			fi
		done
	fi

	_echo "Your system is up to date"
	_exit 0
}

setup_nanobsd_env() {
	if [ "${platform}" != "nanobsd" ]; then
		return;
	fi

	chroot_dir=/tmp/nanobsd_upgrade
	mkdir -p ${chroot_dir} 2>/dev/null
	local _cur_partition=$(mount -p / | cut -f1)
	local _update_partition=$(echo ${_cur_partition} | sed -e 's,0$,2,; s,1$,0,; s,2$,1,')

	if [ ! -e "${_update_partition}" ]; then
		_echo "Secondary partition (${_update_partition}), used for upgrade not found"
		_exit 1
	fi

	_exec "mount ${_update_partition} ${chroot_dir}" "Mounting second partition to run upgrade" mute

	pkg_chroot="-c ${chroot_dir}"

	pkg_update force

	if [ "$(compare_pkg_version pkg)" = "<" ]; then
		_exec "pkg ${pkg_chroot} upgrade pkg" "Upgrading pkg" mute
		pkg_update force
	fi

}

switch_active_nanobsd_partition() {
	if [ "${platform}" != "nanobsd" ]; then
		return;
	fi

	local _cur_partition=$(mount -p / | cut -f1 | sed 's,^/dev/,,')
	local _disk=$(glabel status -s | \
		awk "\$1 == \"${_cur_partition}\" { print substr(\$3, 0, length(\$3)-3)}")
	local _i=$(echo ${_cur_partition} | cut -c ${#_cur_partition})

	if ! echo "${_i}" | egrep -q '^[0-9]$'; then
		_echo "Invalid partition label ${_cur_partition}"
		_exit 1
	fi

	# pfsense0 == part 1 / pfsense1 == part 2
	if [ ${_i} -eq 0 ]; then
		_i=2
	else
		_i=1
	fi

	_exec "gpart set -a active -i ${_i} ${_disk}" "Setting secondary partition as active" mute
}

is_pkg_installed() {
	local _pkg_name="${1}"
	shift
	local _pkg_chroot="$@"

	pkg ${_pkg_chroot} info -e ${_pkg_name}
	return $?
}

compare_pkg_version() {
	local _pkg_name="${1}"

	if ! is_pkg_installed ${_pkg_name} ${pkg_chroot}; then
		echo '!'
		return -1
	fi

	local _lver=$(pkg ${pkg_chroot} query %v ${_pkg_name})

	if [ -z "${_lver}" ]; then
		_echo "ERROR: It was not possible to determine ${_pkg_name} local version"
		_exit 1
	fi

	local _rver=$(pkg ${pkg_chroot} rquery %v ${_pkg_name})

	if [ -z "${_rver}" ]; then
		_echo "ERROR: It was not possible to determine ${_pkg_name} remote version"
		_exit 1
	fi

	local _version=$(pkg version -t ${_lver} ${_rver})

	if [ $? -ne 0 ]; then
		_echo "ERROR: Error comparing ${_pkg_name} local and remote versions"
		_exit 1
	fi

	echo ${_version}
	return 0
}

pkg_install() {
	local _pkg_name="${1}"

	local _force=""
	if [ -n "${2}" ]; then
		_force="-f"
	fi

	if [ -z "${_pkg_name}" ]; then
		_echo "ERROR: Blank package name"
		_exit 1
	fi

	pkg_update

	if is_pkg_installed ${_pkg_name}; then
		local _cversion=$(compare_pkg_version ${_pkg_name})

		if [ -z "${_force}" ]; then
			if [ "${_cversion}" = "=" ]; then
				_echo "Package ${_pkg_name} is up to date"
				_exit 0
			elif [ "${_cversion}" = ">" ]; then
				_echo "Installed ${_pkg_name} version is newer than remote"
				_exit 0
			fi
		fi
		local _cmd="upgrade ${_force}"
		local _msg="Upgrading"
	else
		local _cmd="install"
		local _msg="Installing"
	fi

	_exec "pkg_with_pb ${_cmd} ${_pkg_name}" "${_msg} ${_pkg_name}"
	_exec "pkg clean" "Cleaning up cache" mute ignore_result
}

pkg_delete() {
	local _pkg_name="${1}"

	if [ -z "${_pkg_name}" ]; then
		_echo "ERROR: Blank package name"
		_exit 1
	fi

	if ! is_pkg_installed ${_pkg_name}; then
		_echo "ERROR: Package ${_pkg_name} is not installed"
		_exit 1
	fi

	_exec "pkg_with_pb delete ${_pkg_name}" "Removing ${_pkg_name}"
	_exec "pkg autoremove" "Removing stale packages" mute ignore_result
}

# Reinstall every pfSense-pkg-* package
pkg_reinstall_all() {
	for _pkg in $(pkg query -e '%a == 0' %n); do
		case ${_pkg} in "${pkg_prefix}"* )
			_echo "Reinstalling ${_pkg}"
			pkg_install ${_pkg} 1
			;;
		esac
	done
}

do_reboot() {
	_echo "Upgrade is complete.  Rebooting in 10 seconds."
	echo "Upgrade is complete.  Rebooting in 10 seconds." | wall
	/etc/rc.notify_message -e -g -m "Upgrade is complete.  Rebooting in 10 seconds."
	if [ -z "${dont_reboot}" ]; then
		(sleep 10 && /etc/rc.reboot) &
	fi
}

pid_file="/var/run/$(basename $0).pid"
last_update_file="/var/run/$(basename $0)-last-update"
logfile="/cf/conf/upgrade_log.txt"
stdout='/dev/null'

# pkg should not ask for confirmations
export ASSUME_ALWAYS_YES=true

# Disable automatic update
export REPO_AUTOUPDATE=false

export product=$(/usr/local/bin/php -n /usr/local/sbin/read_global_var product_name pfSense)
export pkg_prefix=$(/usr/local/bin/php -n /usr/local/sbin/read_global_var pkg_prefix pfSense-pkg-)
export platform=$(cat /etc/platform)

USE_MFS_TMPVAR=$(/usr/local/sbin/read_xml_tag.sh boolean system/use_mfs_tmpvar)
if [ "${platform}" = "nanobsd" ] || [ "${USE_MFS_TMPVAR}" = "true" ]; then
	export PKG_DBDIR=/root/var/db/pkg
	export PKG_CACHEDIR=/root/var/cache/pkg
fi

# Upgrade process on nanobsd will happen in chroot
export pkg_chroot=""
export chroot_dir=""

unset dont_reboot
unset booting
unset boot_stage
unset force
unset yes
unset progress_file
unset progress_socket
unset action
unset action_pkg
unset force_ipv4
unset force_ipv6
while getopts 46b:cdfi:hp:l:r:Ruy opt; do
	case ${opt} in
		4)
			if [ -n "${force_ipv6}" ]; then
				usage
				_exit 1
			fi
			force_ipv4=1
			;;
		6)
			if [ -n "${force_ipv4}" ]; then
				usage
				_exit 1
			fi
			force_ipv6=1
			;;
		b)
			booting=1
			boot_stage="${OPTARG}"
			;;
		c)
			action="check"
			;;
		d)
			stdout=''
			;;
		f)
			force=1
			;;
		i)
			if [ -n "${action}" ]; then
				usage
				_exit 1
			fi
			action="install"
			action_pkg="${OPTARG}"
			;;
		h)
			usage
			_exit 0
			;;
		l)
			logfile="${OPTARG}"
			if [ -z "${logfile}" ]; then
				usage
				_exit 1
			fi
			;;
		p)
			progress_socket="${OPTARG}"
			if [ -z "${progress_socket}" ]; then
				usage
				_exit 1
			fi
			;;
		r)
			if [ -n "${action}" ]; then
				usage
				_exit 1
			fi
			action="delete"
			action_pkg="${OPTARG}"
			;;
		R)
			dont_reboot=1
			;;
		u)
			if [ -n "${action}" ]; then
				usage
				_exit 1
			fi
			action="update"
			;;
		y)
			yes=1
			;;
		*)
			usage
			_exit 1
			;;
	esac
done

if [ -n "${force_ipv4}" ]; then
	export IP_VERSION="4"
elif [ -n "${force_ipv6}" ]; then
	export IP_VERSION="6"
fi

# Set default action when no parameter is set
: ${action:="upgrade"}

if pgrep -qF ${pid_file} >/dev/null 2>&1; then
	echo "Another instance is already running... Aborting!"
	exit 1
fi

if [ -z "${booting}" -o "${boot_stage}" != "2" ]; then
	/etc/rc.conf_mount_rw
fi

echo $$ > ${pid_file}

trap _exit 1 2 15 EXIT

if [ "${action}" != "upgrade" -a -f "${logfile}" ]; then
	rm -f ${logfile}
fi

progress_file=${logfile%.*}.json

if [ -e "${progress_file}" ]; then
	rm -f ${progress_file}
fi

case "${action}" in
	check)
		check_upgrade
		;;
	upgrade)
		pkg_upgrade
		;;
	update)
		pkg_update force
		;;
	install)
		if [ ${action_pkg} == "ALL_PACKAGES" ] && [ -n ${force} ]; then
			pkg_reinstall_all
		else
			pkg_install ${action_pkg} ${force}
		fi
		;;
	delete)
		pkg_delete ${action_pkg}
		;;
	*)
		_echo "ERROR: Invalid action!"
		_exit 1
esac

_exit 0