#!/bin/sh

# Copyright (c) 2014-2022 Franco Fichtner <franco@opnsense.org>
# Copyright (c) 2004-2009 Scott Ullrich <sullrich@gmail.com>
#
# 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.
#
# THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS 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
# AUTHOR 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.

INSTALL="/.probe.for.readonly"
MNT="/tmp/hdrescue"
SHELL="/bin/sh"
WAIT="........."

if [ "$(id -u)" != "0" ]; then
	echo "Must be root."
	exit 1
fi

DO_BOOT=
DO_MANUAL=
DO_PASSWORD=
DO_VERBOSE=
DO_ZFS=

bootstrap_and_exit()
{
	RET=${1}

	# ensure config directory structure
	mkdir -p /conf/backup
	mkdir -p /conf/sshd

	# create initial config.xml if necessary
	if [ ! -f /conf/config.xml ]; then
		echo -n "Bootstrapping config.xml..."
		cp /usr/local/etc/config.xml /conf/config.xml
		echo "done."
	fi

	# clean up after a finished import
	if [ -d ${MNT} -a "$(mount | grep -cw ${MNT})" != "0" ]; then
		if [ -n "${PART}" ]; then
			umount ${MNT}
		elif [ -n "${POOL}" ]; then
			zpool export ${POOL}
		fi
	fi

	zfs_unload

	# error code given or assumed ok (trap)
	if [ -z "${RET}" ]; then
		RET=0
	fi

	exit "${RET}"
}

while getopts bmpzV OPT; do
	case ${OPT} in
	b)
		DO_BOOT="-b"
		;;
	m)
		DO_MANUAL="-m"
		;;
	p)
		DO_PASSWORD="-p"
		;;
	V)
		DO_VERBOSE="-V"
		;;
	z)
		DO_ZFS="-z"
		;;
	*)
		echo "Usage: man ${0##*/}" >&2
		exit 1
		;;
	esac
done

shift $((OPTIND - 1))

if [ -n "${DO_VERBOSE}" ]; then
	set -x
fi

DO_DEV=${1}

timeout_prompt()
{
	OUTPUT=$(echo ${2} | sed 's/./& /g')
	MESSAGE=${1}
	RETURN=1

	echo -n "${MESSAGE} "

	stty cbreak -echo
	for NEXT in ${OUTPUT}; do
		echo -n ${NEXT}
		if timeout 1 dd of=/dev/null count=1 status=none; then
			RETURN=0
			break
		fi
	done
	stty -cbreak echo

	echo

	return ${RETURN}
}

zfs_load()
{
	# we need to load ZFS to list pools
	if ! kldstat -qm zfs; then
		export UNLOAD_ZFS="yes"
		kldload zfs
	fi

	export POOLS=$(zfs_probe)
}

zfs_unload()
{
	if [ -n "${UNLOAD_ZFS}" ]; then
		kldunload zfs
	fi
}

zfs_probe()
{
	zpool import -aNf > /dev/null 2> /dev/null

	zpool get -H cachefile | while read ZPOOL ZMORE; do
		ZGUID=$(zpool get -H guid ${ZPOOL} | awk '{ print $3 }')
		ZSIZE=$(zpool get -H size ${ZPOOL} | awk '{ print $3 }')
		if [ "$(mount | grep -w / | grep -c ${ZPOOL})" = "0" ]; then
			zpool export ${ZPOOL} > /dev/null 2> /dev/null
		fi
		echo "${ZPOOL} ${ZGUID} ${ZSIZE}"
	done
}

import_start()
{
	local DEV=${1}

	# We are going to make an educated guess about the enclosed
	# file system.  The types "msdos" and "cd9660" are meant as
	# fallbacks to try, but by no means meant to be correct.
	# Correctness stems from the selection of the device given,
	# so that e.g. a CD will not have any attached partitions,
	# but we do not actually check for the device node name.

	export PART=
	export TYPE=
	export POOL=

	if [ -e "/dev/${DEV}s1a" ]; then
		# MBR UFS
		export PART="/dev/${DEV}s1a"
		export TYPE="ufs"
		return 0
	elif [ -e "/dev/${DEV}p3" ]; then
		# GPT UFS
		export PART="/dev/${DEV}p3"
		export TYPE="ufs"
		return 0
	elif [ "$(echo ${POOLS} | grep -c "^${DEV} ")" != "0" ]; then
		# ZFS POOL
		export POOL="${DEV}"
		return 0
	elif [ -e "/dev/${DEV}s1" ]; then
		# MBR MSDOS
		export PART="/dev/${DEV}s1"
		export TYPE="msdos"
		return 0
	elif [ -e "/dev/${DEV}p1" ]; then
		# GPT MSDOS
		export PART="/dev/${DEV}p1"
		export TYPE="msdos"
		return 0
	elif [ -e "/dev/${DEV}" ]; then
		# ISO 9660
		export PART="/dev/${DEV}"
		export TYPE="cd9660"
		return 0
	fi

	return 1
}


DEVS=
POOLS=

if [ -n "${DO_ZFS}" ]; then
	zfs_load
	zfs_unload
	if [ -n "${POOLS}" ]; then
		echo "${POOLS}"
	fi
	exit 0
fi

trap bootstrap_and_exit 2

if [ -n "${DO_BOOT}" ]; then
	touch ${INSTALL} 2> /dev/null
	if [ -f ${INSTALL} -a -f /conf/config.xml ]; then
		bootstrap_and_exit 0
	fi

	if ! timeout_prompt \
	    'Press any key to start the configuration importer:' ${WAIT}; then
		bootstrap_and_exit 0
	fi
fi

zfs_load

if [ -n "${DO_DEV}" ]; then
	if ! import_start ${DO_DEV}; then
		echo "No known partition layout was found for '${DO_DEV}'."
		bootstrap_and_exit 1
	fi
fi

DEVS=$(
	camcontrol devlist
	echo "${POOLS}" | while read ZPOOL ZGUID ZSIZE ZMORE; do
		if [ -z "${ZPOOL}" ]; then
			continue
		fi
		printf "%-35s%s\n" "<${ZGUID} ${ZSIZE}>" \
		    "ZFS pool (${ZPOOL})"
	done
	gmirror status -s
	graid status -s
)

while : ; do
	if [ -z "${DO_DEV}" ]; then
		echo
		echo "${DEVS}"
		echo
		read -p "Select device to import from (e.g. ada0) or leave blank to exit: " DEV
		echo

		if [ -z "${DEV}" ]; then
			bootstrap_and_exit 0
		elif [ "${DEV}" = "!" ]; then
			# secret escape! (not so secret now, is it?)
			csh
			continue
		elif ! import_start ${DEV}; then
			echo "No known partition layout was found for '${DEV}'."
			continue
		fi
	fi

	mkdir -p ${MNT}

	if [ -n "${PART}" -a -n "${TYPE}" ]; then
		echo "Starting import for partition '${PART}'."
		echo

		if [ "${TYPE}" = "ufs" ]; then
			echo -n "Running fsck..."
			fsck -t ${TYPE} -y ${PART} > /dev/null
			echo "done."
		fi

		if ! mount -t ${TYPE} ${PART} ${MNT}; then
			echo "The device could not be mounted."
			PART=
		fi
	elif [ -n "${POOL}" ]; then
		echo "Starting import for ZFS pool '${POOL}'."
		echo

		zpool import -fNR ${MNT} ${POOL}
		if ! mount -t zfs ${POOL}/ROOT/default ${MNT}; then
			echo "The pool could not be mounted."
			POOL=
		fi
	else
		if [ -n "${DO_DEV}" ]; then
			bootstrap_and_exit 1
		fi

		continue
	fi

	if [ -n "${DO_MANUAL}" ]; then
		echo "# Manual shell mode, selected file system available at ${MNT}"
		${SHELL}
		break
	elif [ -n "${DO_PASSWORD}" ]; then
		if [ -f "${MNT}/usr/local/sbin/opnsense-shell" ]; then
			mount -t devfs devfs ${MNT}/dev
			chroot ${MNT} /bin/sh /etc/rc.d/ldconfig start
			chroot ${MNT} /usr/local/sbin/opnsense-shell password root -x 0
			umount ${MNT}/dev

			echo "The password was reset successfully."

			break
		else
			echo "The installed version does not yet support recovery."

			if [ -n "${DO_DEV}" ]; then
				bootstrap_and_exit 1
			fi
		fi
	elif [ -f "${MNT}/conf/config.xml" ]; then
		if [ "$(grep -cx -- '---- BEGIN config.xml ----' ${MNT}/conf/config.xml)" != "0" ]; then
			echo "The file /conf/config.xml is encrypted."

			if [ -n "${DO_DEV}" ]; then
				bootstrap_and_exit 1
			fi

			continue
		fi

		rm -rf /conf/*

		for FILE in captiveportal.sqlite config.xml dhcpleases.tgz dhcp6c_duid netflow.tgz rrd.tgz; do
			if [ -f "${MNT}/conf/${FILE}" ]; then
				echo -n "Restoring ${FILE}..."
				cp "${MNT}/conf/${FILE}" /conf
				echo "done."
			fi
		done

		for DIR in backup sshd; do
			if [ -d "${MNT}/conf/${DIR}" ]; then
				echo -n "Restoring ${DIR}..."
				cp -r "${MNT}/conf/${DIR}" /conf
				echo "done."
			else
				mkdir -p "/conf/${DIR}"
			fi
		done

		for FILE in $(find /conf/sshd -type f -name '*key'); do
			# fixup possibly unsupported file system permissions
			chmod 600 ${FILE}
		done

		# hooray, we're done!
		break
	else
		echo "The file /conf/config.xml could not be found."

		if [ -n "${DO_DEV}" ]; then
			bootstrap_and_exit 1
		fi
	fi
done

if [ -z "${DO_BOOT}${DO_MANUAL}${DO_PASSWORD}" ]; then
	echo "Please reboot."
fi

bootstrap_and_exit 0
