#!/bin/sh

# /etc/rc.firmware
# originally part of m0n0wall (http://neon1.net/m0n0wall)
# Copyright (C) 2005-2009 Scott Ullrich <sullrich@pfsense.org>.
# Copyright (C) 2003 Manuel Kasper <mk@neon1.net>.
# All rights reserved.

# mount /cf
/etc/rc.conf_mount_rw

# Reset file(s)
echo "" >/conf/upgrade_log.txt
echo "" >/conf/firmware_update_misc_log.txt
echo "" >/conf/fdisk_upgrade_log.txt

exec 3>&2 2>>/conf/firmware_update_misc_log.txt

export ACTION=$1
export IMG=$2
if [ $# -eq 3 ]; then
	export CUSTOMIMG=$3
fi

file_notice() {
	/usr/local/bin/php -q -d auto_prepend_file=config.inc <<ENDOFF
	<?php
		require_once("globals.inc");
		require_once("functions.inc");
		file_notice("$1", "$2", "$1", "");
	?>
ENDOFF
}

output_env_to_log() {
	date >> /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt

	ls -lah /dev/ >> /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt

	ls -lah $IMG >> /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt

	md5 $IMG >> /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt

	mount >> /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt

	top >> /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt
}

backup_chflags() {
	TOPROCESS="bin lib libexec sbin usr"
	for files in $TOPROCESS; do
		/usr/sbin/mtree -Pcp /${files} | bzip2 -9 > /tmp/chflags.dist.${files}.bz2 2>> /conf/upgrade_log.txt
	done
}

restore_chflags() {
	TOPROCESS="bin lib libexec sbin usr"
	for files in $TOPROCESS; do
		cd / && /usr/bin/bzcat /tmp/chflags.dist.${files}.bz2 | /usr/sbin/mtree -P -p /${files} >> /conf/upgrade_log.txt 2>&1
	done
}

remove_chflags() {
	TOPROCESS="bin lib libexec sbin usr"
	for files in $TOPROCESS; do
		/bin/chflags -R noschg /${files}
		/bin/chmod -R u+rw /${files}
	done
}

case $ACTION in
enable)
	touch /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt
	echo "Enable" >> /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt
	/etc/rc.conf_mount_ro
	;;
auto)
	touch /var/run/firmwarelock.dirty
	backup_chflags
	remove_chflags
	/etc/rc.firmware_auto
	restore_chflags
	/etc/rc.conf_mount_ro
	;;
pfSenseNanoBSDupgrade)

	# Sanity check - bail early if there's no firmware file!
	if [ ! -r $IMG ]; then
		echo "2nd parameter has not been passed or file does not exist. Exiting." >> /conf/upgrade_log.txt 2>&1
		/etc/rc.conf_mount_ro
		exit 1
	fi

	# Prevent full upgrade file from being used to upgrade
	if [ `echo $IMG | grep "full"` ]; then
		echo "You cannot use a full file for upgrade.  Please use a file labelled nanobsd upgrade."
		file_notice "NanoBSDUpgradeFailure" "You have attempted to use a full NanoBSD installation file as an upgrade.  Please use a NanoBSD file labelled 'upgrade' instead."
		rm -f $IMG
		/etc/rc.conf_mount_ro
		exit 1
	fi

	touch /var/run/firmwarelock.dirty

	echo "NanoBSD Firmware upgrade in progress..." >> /conf/upgrade_log.txt 2>&1
	echo "NanoBSD Firmware upgrade in progress..." | wall
	/etc/rc.notify_message -e -g -m "NanoBSD Firmware upgrade in progress..."

	# backup config
	/bin/mkdir -p /tmp/configbak
	cp -Rp /conf/* /tmp/configbak 2>/dev/null

	# Remove logs from backup dir to avoid clobbering upon restore.
	rm /tmp/configbak/*_log.txt 2>/dev/null

	echo "" >> /conf/upgrade_log.txt

	echo "Installing ${IMG}." >> /conf/upgrade_log.txt 2>&1
	echo "Installing ${IMG}." | wall

	# resolve glabel label that we booted from
	BOOT_DEVICE=`/sbin/mount | /usr/bin/grep pfsense | /usr/bin/cut -d'/' -f4 | /usr/bin/cut -d' ' -f1`
	# resolve glabel to the real boot dev entry
	REAL_BOOT_DEVICE=`/sbin/glabel list | /usr/bin/grep -B2 ufs/${BOOT_DEVICE} | /usr/bin/head -n 1 | /usr/bin/cut -f3 -d' '`
	# grab the boot device, example ad1, ad0
	BOOT_DRIVE=`/sbin/glabel list | /usr/bin/grep -B2 ufs/pfsense | /usr/bin/head -n 1 | /usr/bin/cut -f3 -d' ' | /usr/bin/cut -d's' -f1`
	# test the slice.  if we are on slice 1 we need to flash 2 and vice versa
	if [ `echo $REAL_BOOT_DEVICE | /usr/bin/grep "s1"` ]; then
		SLICE="2"
		OLDSLICE="1"
		TOFLASH="${BOOT_DRIVE}s${SLICE}"
		COMPLETE_PATH="${BOOT_DRIVE}s${SLICE}a"
		GLABEL_SLICE="pfsense1"
		UFS_ID="1"
		OLD_UFS_ID="0"
	else
		SLICE="1"
		OLDSLICE="2"
		TOFLASH="${BOOT_DRIVE}s${SLICE}"
		COMPLETE_PATH="${BOOT_DRIVE}s${SLICE}a"
		GLABEL_SLICE="pfsense0"
		UFS_ID="0"
		OLD_UFS_ID="1"
	fi

	# Output specific information that this script is using
	echo "SLICE         ${SLICE}" >> /conf/upgrade_log.txt
	echo "OLDSLICE      ${OLDSLICE}" >> /conf/upgrade_log.txt
	echo "TOFLASH       ${TOFLASH}" >> /conf/upgrade_log.txt
	echo "COMPLETE_PATH ${COMPLETE_PATH}" >> /conf/upgrade_log.txt
	echo "GLABEL_SLICE  ${GLABEL_SLICE}" >> /conf/upgrade_log.txt

	# First ensure the new file can fit inside the
	# slice that we are going to be operating on.
	NEW_IMG_SIZE=`echo $((\`gzip -l ${IMG} | grep -v compressed | awk '{ print $2}'\` / 1024 / 1024))`
	SIZE=`/sbin/fdisk ${COMPLETE_PATH} | /usr/bin/grep Meg | /usr/bin/awk '{ print $5 }' | /usr/bin/cut -d"(" -f2`
	# USB slices are under-reported even more than CF slices when viewed
	# directly, instead of when looking at the entire disk. Compensate
	# by adding exactly 6MB. 4MB was consistently 2MB too few, and
	# was resulting in failing upgrades on USB Flash based installs.
	SIZE=`expr $SIZE + 6`
	if [ "$SIZE" -lt "$NEW_IMG_SIZE" ]; then
		file_notice "UpgradeFailure" "Upgrade failed due to the upgrade image being larger than the partition that is configured on disk.  Halting. Size on disk: $SIZE < Size of new image: $NEW_IMG_SIZE"
		echo "Upgrade failed.  Please check the system log file for more information" | wall
		rm -f $IMG
		rm -f /var/run/firmwarelock.dirty
		rm -f /var/run/firmware.lock
		rm -f ${IMG}
		/etc/rc.conf_mount_ro
		exit 1
	fi

	# Output environment information to log file
	output_env_to_log

	# Grab a before upgrade look at fdisk
	echo "" >> /conf/fdisk_upgrade_log.txt
	echo "Before upgrade fdisk/bsdlabel" >> /conf/fdisk_upgrade_log.txt
	fdisk $BOOT_DRIVE >> /conf/fdisk_upgrade_log.txt
	bsdlabel -A ${BOOT_DRIVE}s1 >> /conf/fdisk_upgrade_log.txt
	bsdlabel -A ${BOOT_DRIVE}s2 >> /conf/fdisk_upgrade_log.txt
	bsdlabel -A ${BOOT_DRIVE}s3 >> /conf/fdisk_upgrade_log.txt
	echo "---------------------------------------------------------------" >> /conf/fdisk_upgrade_log.txt
	echo "" >> /conf/fdisk_upgrade_log.txt

	# Log that we are really doing a NanoBSD upgrade
	echo "" >> /conf/upgrade_log.txt
	echo "NanoBSD upgrade starting" >> /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt

	# Remove TOFLASH and get ready for new flash image
	echo "" >> /conf/upgrade_log.txt
	echo "dd if=/dev/zero of=/dev/${TOFLASH} bs=1m count=1" >> /conf/upgrade_log.txt
	dd if=/dev/zero of=/dev/${TOFLASH} bs=1m count=1 >> /conf/upgrade_log.txt 2>&1

	# Stream gzipped image to dd and explode image to new area
	echo "" >> /conf/upgrade_log.txt
	echo "/usr/bin/gzip -dc $IMG | /bin/dd of=/dev/${TOFLASH} obs=64k" >> /conf/upgrade_log.txt
	/usr/bin/gzip -dc $IMG | /bin/dd of=/dev/${TOFLASH} obs=64k >> /conf/upgrade_log.txt 2>&1

	# Grab an after upgrade look at fdisk
	echo "" >> /conf/fdisk_upgrade_log.txt
	echo "After upgrade fdisk/bsdlabel" >> /conf/upgrade_log.txt
	fdisk $BOOT_DRIVE >> /conf/fdisk_upgrade_log.txt
	bsdlabel -A ${BOOT_DRIVE}s1 >> /conf/fdisk_upgrade_log.txt
	bsdlabel -A ${BOOT_DRIVE}s2 >> /conf/fdisk_upgrade_log.txt
	bsdlabel -A ${BOOT_DRIVE}s3 >> /conf/fdisk_upgrade_log.txt
	echo "---------------------------------------------------------------" >> /conf/fdisk_upgrade_log.txt
	echo "" >> /conf/fdisk_upgrade_log.txt

	# Ensure that our new system is sound and bail if it is not and file a notice
	echo "" >> /conf/upgrade_log.txt
	echo "/sbin/fsck_ufs -y /dev/${COMPLETE_PATH}" >> /conf/upgrade_log.txt
	/sbin/fsck_ufs -y /dev/${COMPLETE_PATH} >> /conf/upgrade_log.txt 2>&1
	if [ $? != 0 ]; then
		file_notice "UpgradeFailure" "{\$g['product_name']} upgrade has failed.   Your system has been left in a usable state."
		rm -f $IMG
		rm -f /var/run/firmwarelock.dirty
		rm -f /var/run/firmware.lock
		/etc/rc.conf_mount_ro
		exit 1
	fi

	# Enable foot shooting
	sysctl kern.geom.debugflags=16

	# Add back the corresponding glabel
	echo "" >> /conf/upgrade_log.txt
	echo "/sbin/tunefs -L ${GLABEL_SLICE} /dev/${COMPLETE_PATH}" >> /conf/upgrade_log.txt
	/sbin/tunefs -L ${GLABEL_SLICE} /dev/${COMPLETE_PATH} >> /conf/upgrade_log.txt 2>&1

	# restore config
	cp -Rp /tmp/configbak/* /conf 2>/dev/null

	# Remove upgrade file
	rm -f $IMG

	# Mount newly prepared slice
	mkdir /tmp/$GLABEL_SLICE
	mount /dev/ufs/$GLABEL_SLICE /tmp/$GLABEL_SLICE

	# If /boot/loader.conf.local exists
	# copy to the other slice.
	if [ -f /boot/loader.conf.local ]; then
		cp /boot/loader.conf.local /tmp/$GLABEL_SLICE/boot/loader.conf.local
	fi

	# If /tmp/$GLABEL_SLICE/usr/local/share/pfSense/post_upgrade_command exists
	# after update then execute the command.
	echo "Checking for post_upgrade_command..." >> /conf/upgrade_log.txt
	if [ -f /tmp/$GLABEL_SLICE/usr/local/share/pfSense/post_upgrade_command ]; then
		echo "Found post_upgrade_command, executing ($GLABEL_SLICE)..." >> /conf/upgrade_log.txt
		sh /tmp/$GLABEL_SLICE/usr/local/share/pfSense/post_upgrade_command $GLABEL_SLICE >> /conf/upgrade_log.txt 2>&1
	fi

	# Update fstab
	cp /etc/fstab /tmp/$GLABEL_SLICE/etc/fstab
	sed -i "" "s/pfsense${OLD_UFS_ID}/pfsense${UFS_ID}/g" /tmp/$GLABEL_SLICE/etc/fstab
	if [ $? != 0 ]; then
		echo "Something went wrong when trying to update the fstab entry.  Aborting upgrade."
		file_notice "UpgradeFailure" "Something went wrong when trying to update the fstab entry.  Aborting upgrade."
		rm -f $IMG
		rm -f /var/run/firmwarelock.dirty
		rm -f /var/run/firmware.lock
		umount /tmp/$GLABEL_SLICE
		/etc/rc.conf_mount_ro
		exit 1
	fi
	echo "" >> /conf/upgrade_log.txt
	cat /tmp/$GLABEL_SLICE/etc/fstab >> /conf/upgrade_log.txt

	echo "" >> /conf/upgrade_log.txt
	find /tmp/$GLABEL_SLICE >/conf/file_upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt

	# Unmount newly prepared slice
	umount /tmp/$GLABEL_SLICE

	sync

	# Set active mount slice in fdisk
	echo "" >> /conf/upgrade_log.txt
	echo "gpart set -a active -i ${SLICE} ${BOOT_DRIVE}" >> /conf/upgrade_log.txt
	gpart set -a active -i ${SLICE} ${BOOT_DRIVE} >> /conf/upgrade_log.txt 2>&1

	sync

	# Set active boot source - NanoBSD does not do this but otherwise we
	# end up with the wrong partition being active.
	echo "" >> /conf/upgrade_log.txt
	echo "/usr/sbin/boot0cfg -s ${SLICE} -v /dev/${BOOT_DRIVE}" >> /conf/upgrade_log.txt
	/usr/sbin/boot0cfg -s ${SLICE} -v /dev/${BOOT_DRIVE} >> /conf/upgrade_log.txt 2>&1

	# Disable foot shooting
	sysctl kern.geom.debugflags=0

	# Grab a final look at fdisk
	echo "" >> /conf/fdisk_upgrade_log.txt
	echo "Final upgrade fdisk/bsdlabel" >> /conf/fdisk_upgrade_log.txt
	fdisk $BOOT_DRIVE >> /conf/fdisk_upgrade_log.txt
	bsdlabel -A ${BOOT_DRIVE}s1 >> /conf/fdisk_upgrade_log.txt
	bsdlabel -A ${BOOT_DRIVE}s2 >> /conf/fdisk_upgrade_log.txt
	bsdlabel -A ${BOOT_DRIVE}s3 >> /conf/fdisk_upgrade_log.txt
	echo "---------------------------------------------------------------" >> /conf/fdisk_upgrade_log.txt
	echo "" >> /conf/fdisk_upgrade_log.txt

	# Remove extra stuff
	rm -rf /usr/savecore/*

	date >> /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt

	# Trigger a package reinstallation on reboot
	touch /conf/needs_package_sync

	# remount /cf ro
	/etc/rc.conf_mount_ro
	/bin/sync

	echo "NanoBSD Firmware upgrade is complete.  Rebooting in 10 seconds." >> /conf/upgrade_log.txt 2>&1
	echo "NanoBSD Firmware upgrade is complete.  Rebooting in 10 seconds." | wall
	/etc/rc.notify_message -e -g -m "NanoBSD Firmware upgrade is complete.  Rebooting in 10 seconds."

	sleep 10

	rm -f /var/run/firmwarelock.dirty
	rm -f /var/run/firmware.lock
	. /etc/rc.reboot

	;;
pfSenseupgrade)

	# Sanity check - bail early if there's no firmware file!
	if [ ! -r $IMG ]; then
		echo "2nd parameter has not been passed or file does not exist. Exiting." >> /conf/upgrade_log.txt 2>&1
		/etc/rc.conf_mount_ro
		exit
	fi

	# wait 1 second before beginning
	sleep 1

	# Log that we are really doing a pfSense upgrade
	echo "" >> /conf/upgrade_log.txt
	echo "pfSenseupgrade upgrade starting" >> /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt

	touch /var/run/firmwarelock.dirty

	if [ -f /tmp/perform_full_backup.txt ]; then
		echo "Performing full backup" >> /conf/upgrade_log.txt
		/etc/rc.create_full_backup
		rm /tmp/perform_full_backup.txt
	fi

	touch /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt

	# Output environment information to log file
	output_env_to_log

	backup_chflags
	remove_chflags

	# Do we have a pre-upgrade hook in the update file?
	if [ `tar tvzf $IMG | grep /usr/local/share/pfSense/pre_upgrade_command | wc -l` -gt 0 ]; then
		tar xzvf $IMG -C / ./usr/local/share/pfSense/pre_upgrade_command >> /conf/upgrade_log.txt 2>&1
		chmod a+rx /usr/local/share/pfSense/pre_upgrade_command >> /conf/upgrade_log.txt 2>&1
		sh /usr/local/share/pfSense/pre_upgrade_command >> /conf/upgrade_log.txt 2>&1
	fi

	echo "Firmware upgrade in progress..." >> /conf/upgrade_log.txt 2>&1
	echo "Firmware upgrade in progress..." | wall
	/etc/rc.notify_message -e -g -m "Firmware upgrade in progress..."

	# backup config
	[ -d /tmp/configbak ] && rm -rf /tmp/configbak
	/bin/mkdir -p /tmp/configbak
	cp -Rp /conf/* /tmp/configbak 2>/dev/null

	# Remove logs from backup dir to avoid clobbering upon restore.
	rm /tmp/configbak/*_log.txt 2>/dev/null

	# tar explode image onto hd
	killall -9 check_reload_status
	killall -9 check_reload_status
	echo "Installing $IMG." >> /conf/upgrade_log.txt 2>&1
	cd / && /usr/bin/tar --exclude=./dev -xzUPf $IMG >> /conf/upgrade_log.txt 2>&1
	/usr/local/sbin/check_reload_status
	echo "Image installed $IMG." >> /conf/upgrade_log.txt 2>&1

	# process custom image if its passed
	if [ $# -eq 3 ]; then
		if [ -f $CUSTOMIMG ]; then
			echo "Custom image $CUSTOMIMG found." >> /conf/upgrade_log.txt 2>&1
			echo "Custom image ($CUSTOMIMG) found." >> /conf/upgrade_log.txt 2>&1
			PWD_DIR=`pwd`
			cd / && /usr/bin/tar xzPUf $CUSTOMIMG >> /conf/upgrade_log.txt 2>&1
			cd $PWD_DIR
			echo "Custom image $CUSTOMIMG installed." >> /conf/upgrade_log.txt 2>&1
		fi
	fi

	# restore config
	cp -Rp /tmp/configbak/* /conf 2>/dev/null

	# restore /etc symlinks
	rm /etc/hosts
	ln -s /var/etc/hosts /etc/hosts

	restore_chflags

	# Remove upgrade file
	rm -f $IMG

	if [ -e /etc/init_bootloader.sh ]; then
		if [ ! -x /etc/init_bootloader.sh ]; then
			chmod ug+x /etc/init_bootloader.sh
		fi
		/etc/init_bootloader.sh >> /conf/upgrade_log.txt 2>&1
	fi

	# Remove saved commit ID for gitsync
	rm -f /etc/version.gitsync

	# If /usr/local/share/pfSense/post_upgrade_command exists after update
	# then execute the command.
	if [ -f /usr/local/share/pfSense/post_upgrade_command ]; then
		if [ ! -x /usr/local/share/pfSense/post_upgrade_command ]; then
			chmod ug+x /usr/local/share/pfSense/post_upgrade_command
		fi
		/usr/local/share/pfSense/post_upgrade_command >> /conf/upgrade_log.txt 2>&1
	fi

	# remove unused files
	rm -rf /etc/rc.conf
	rm -rf /usr/savecore/*

	date >> /conf/upgrade_log.txt
	echo "" >> /conf/upgrade_log.txt

	# remount /cf ro
	/etc/rc.conf_mount_ro

	# release the firmware lock
	rm -f /var/run/firmwarelock.dirty
	rm -f /var/run/firmware.lock
	/bin/sync

	echo "Firmware upgrade is complete. Rebooting in 10 seconds." >> /conf/upgrade_log.txt 2>&1
	echo "Firmware upgrade is complete. Rebooting in 10 seconds." | wall
	/etc/rc.notify_message -e -g -m "Firmware upgrade is complete. Rebooting in 10 seconds."

	# Sleep and allow disks to catch up
	sleep 10

	# If the archive has unpacked a file called
	# /tmp/no_upgrade_reboot_required then do
	# not reboot after upgrade.
	if [ -f /tmp/no_upgrade_reboot_required ]; then
		rm /tmp/no_upgrade_reboot_required
	else
		. /etc/rc.reboot
	fi

	;;
esac