#!/bin/sh

# Copyright (c) 2014-2020 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.install.media"
MNT="/tmp/hdrescue"
WAIT="........."

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

DO_BOOT=
DO_PASSWORD=
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 bpz OPT; do
	case ${OPT} in
	b)
		DO_BOOT="-b"
		;;
	p)
		DO_PASSWORD="-p"
		;;
	z)
		DO_ZFS="-z"
		;;
	*)
		echo "Usage: man ${0##*/}" >&2
		exit 1
		;;
	esac
done

shift $((${OPTIND} - 1))

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)
	if [ -n "${POOLS}" ]; then
		# insert missing newline due to $()
		export POOLS="${POOLS}
"
	fi
}

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

zfs_probe()
{
	zpool import -aNf

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

import_start()
{
	local DEV=${1}

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

	return 1
}


DEV=
DEVS=
PART=
POOL=
POOLS=

if [ -n "${DO_ZFS}" ]; then
	zfs_load
	zfs_unload
	echo -n "${POOLS}"
	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; gmirror status -s; graid status -s; echo -n "${POOLS}")

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}" ]; then
		echo "Starting import for partition '${PART}'."
		echo

		TYPE="ufs"

		if [ -z "${MSDOS}" ]; then
			echo -n "Running fsck..."
			fsck -t ${TYPE} -y ${PART} > /dev/null
			echo "done."
		else
			TYPE="msdosfs"
		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
	fi

	if [ -z "${PART}" -a -z "${POOL}" ]; then
		if [ -n "${DO_DEV}" ]; then
			bootstrap_and_exit 1
		fi

		continue
	fi

	if [ -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
		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

		# 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}" ]; then
	echo "Please reboot."
fi

bootstrap_and_exit 0
