#!/bin/sh

# Copyright (c) 2014-2018 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=

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} ]; then
		if [ -n "${PART}" ]; then
			umount ${MNT} && rm -rf ${MNT}
		elif [ -n "${POOL}" ]; then
			zpool export ${POOL} && rm -rf ${MNT}
		fi
	fi

	if [ -z "${ZFS}" ]; then
		kldunload zfs
	fi

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

	exit "${RET}"
}

trap bootstrap_and_exit 2

while getopts b OPT; do
	case ${OPT} in
	b)
		DO_BOOT="-b"
		;;
	*)
		echo "Usage: man opnsense-importer" >&2
		exit 1
		;;
	esac
done

shift $((${OPTIND} - 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 ZFS="yes"
	else
		kldload zfs
	fi
}

probe_for_part()
{
	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 [ "${DEV}" = "zroot" ]; then
		if zpool import | awk '{ print $1 " " $2 }' | grep -q "${DEV} ONLINE"; then
			export POOL="${DEV}"
			return 0
		fi
	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=${1}
DEVS=
PART=

if [ -n "${DEV}" ]; then
	zfs_load

	if ! probe_for_part ${DEV}; then
		echo "No known partition layout was found for '${DEV}'."
		bootstrap_and_exit 1
	fi
else
	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

	DEVS=$(camcontrol devlist; gmirror status -s; graid status -s)
fi

while : ; do
	if [ -z "${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
			DEV=
			continue
		elif ! probe_for_part ${DEV}; then
			echo "No known partition layout was found for '${DEV}'."
			DEV=
			continue
		fi

		DEV=
	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

		break
	elif [ -n "${POOL}" ]; then
		echo "Starting import for ZFS pool '${POOL}'."
		echo

		zpool import -f -R ${MNT} ${POOL}
		if ! zfs mount ${POOL}/ROOT/default; then
			echo "The device could not be mounted."
			POOL=
		fi

		break
	else
		echo "Something went wrong during partition selection."
		bootstrap_and_exit 1
	fi

	if [ -n "${DEV}" ]; then
		echo "Import from partition '${DEV}' failed."
		bootstrap_and_exit 1
	fi
done

if [ -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
else
	echo "No previous configuration was found on this device."
	bootstrap_and_exit 1
fi

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

bootstrap_and_exit 0
