#!/usr/bin/env bash

## of-flash v1.18 (1st March 2022) by Andrew Davison
##  Script to handle automated reflashing of internal memory on an OpenFrame device.

if [ "$USER" != "root" ]; then
	echo "You need to run this with superuser privileges. Try 'sudo $0'."
	exit 1
fi

if [ "$1" ]; then
	IMGLOC="$1"
else
	IMGLOC="/boot/reflash"
fi

COLOUR=6
HAPPYSTOP=10
SADSTOP=60
DEBUG=0

tput cnorm



amixer -c 0 sset "Master" 20% unmute &>/dev/null

header() {

	setterm -powersave off
	setterm -blank 0

	clear
	echo
	echo
	tput setaf $COLOUR
	echo "                                                                        OpenFrame Reflash v1.17"
	[ $DEBUG -eq 0 ] && echo "==============================================================================================="
	[ $DEBUG -ne 0 ] && echo "====== DEBUG =================================================================================="
	echo "                                                                             birdslikewires.net"
	echo
	tput setab $COLOUR
	tput setaf 0
	echo "  This software has the potential to damage your system and is used entirely at your own risk! "
	echo "  If you find it useful, perhaps you could make a contribution towards hosting & development?  "
	tput sgr0
	echo
	echo

}

countdown() {
	MSG="$1"
	SEC=$2
	tput civis
	while [ $SEC -gt 0 ]; do
		echo -ne "$MSG in $SEC\033[0K\r"
		sleep 1
		: $((SEC--))
	done
	echo -e "$MSG now..."
	tput cnorm
}

droptoshell() {

	counter=15
	counterest=$((counter*2))

	echo -n "Stalling for network (approximately $counterest seconds)..."
	until ping -c 1 -W 1 8.8.8.8 &>/dev/null
	do
		[ $counter -eq 0 ] && break
		sleep 1;
		let counter-=1
	done

	if [ $counter -eq 0 ]; then
		NETWORKUP=0
		echo " not available."
		echo
		countdown "Moving on" $1
		sleep 1
		clear
	else
		NETWORKUP=1
		echo " network ready."
		echo -n "Fetching network time..."
		ntpdate 0.pool.ntp.org 1.pool.ntp.org 2.pool.ntp.org 3.pool.ntp.org 2>&1 >/dev/null
		echo " done."
		echo
		countdown "Dropping to shell" $1
		sleep 1
		clear
		of-ip
		tput sgr0
		echo
	fi

}

shutitdown() {

	tput sgr0
	echo
	echo
	read -s -t $1 -n 1 -p "Press any key for a login prompt, or wait $1 seconds for shutdown..." GOTOLOGIN

	echo > /etc/udev/rules.d/70-persistent-net.rules

	if [[ "$GOTOLOGIN" != "" ]]; then
		echo
		droptoshell 3
		sleep 1
		exit 0
	else
		echo
		echo
		echo "This OpenFrame is about to shut down. When shutdown is complete, turn off the power and"
		echo "disconnect the USB storage device. Restore the power and allow the OpenFrame to boot."
		echo
		countdown "Shutting down" 10
		sleep 1
		echo
		echo
		[ $DEBUG -eq 0 ] && poweroff || exit 0
	fi

}

header

REFLASHCONTENTS=`find $IMGLOC -maxdepth 1 -type f -iname "*.img*" -not -path '*/\.*' -printf '%f\n' | sort`

[ "$REFLASHCONTENTS" == "" ] && FILECOUNT=0 || FILECOUNT=`echo "$REFLASHCONTENTS" | wc -l`

IMGPRESENT=`echo "$REFLASHCONTENTS" | sed -n 1p | grep -c .img.gz`
CHKPRESENT=`echo "$REFLASHCONTENTS" | sed -n 2p | grep -c .img.gz.md5`
EFIPRESENT=`echo "$REFLASHCONTENTS" | grep -c _efi_`

if [ $FILECOUNT -eq 2 ] && [ $IMGPRESENT -eq 1 ] && [ $CHKPRESENT -eq 1 ]; then

	IMG=`echo "$REFLASHCONTENTS" | head -1 | awk -F\.img.gz {'print $1'}`
	IMGSUM=`cat $IMGLOC/$IMG.img.gz.md5 | awk -F\   {'print $1'}`

	echo -n "Checking integrity of image..."
	CHECKSUM=`md5sum $IMGLOC/$IMG.img.gz | awk -F\  '{print $1}'`
	if [ "$CHECKSUM" == "$IMGSUM" ]; then

		echo " successful."
		echo
		tput setaf $COLOUR
		echo "   -----------------------------------------------"
		echo "    WARNING: OPENFRAME FLASHING IS ABOUT TO BEGIN "
		echo "   -----------------------------------------------"
		tput sgr0
		echo
		echo "Flashing will permanently erase all files contained on internal memory and replace them with"
		echo -n "the contents of "
		tput setaf $COLOUR
		echo -n "$IMG.img.gz"
		tput sgr0
		echo "."
		echo
		echo "The process should take 5-7 minutes to complete."
		echo
		echo "To stop this, turn off the power to this device BEFORE THE COUNTDOWN ENDS."
		echo
		countdown "Reflash will begin" 20
		echo
		tput setaf 1
		echo "   --------------------------------------------------------------------------------------------"
		echo -n "    Writing the contents of "
		tput setaf 7
		echo -n "'$IMG.img.gz'"
		tput setaf 1
		echo " to internal memory... "
		echo "   --------------------------------------------------------------------------------------------"
		tput setaf 7
		echo
		[ $DEBUG -eq 0 ] && MMCDEST="/dev/mmcblk0" || MMCDEST="/dev/null"
		gzip -dc $IMGLOC/$IMG.img.gz | dd of=$MMCDEST bs=1M
		sleep 1

		aplay -q -f a_law -r 44100 /usr/local/share/sounds/success.au &

		tput setaf $COLOUR
		echo
		echo "   -------------------"
		echo "    FLASHING COMPLETE "
		echo "   -------------------"
		tput sgr0
		echo
		echo "Reflashing of the internal memory has completed."
		sleep 1
		shutitdown $HAPPYSTOP

	else

		echo " checksum mismatch."
		echo
		echo
		tput setaf 3
		echo "    The firmware image you have provided does not match the checksum and may be corrupt."
		echo "    It is possible that there was an error writing the file to this USB device, or that it"
		echo "    did not download properly."
		echo
		echo "    Please try copying the firmware image files to this USB device again and ensure that"
		echo "    the drive is ejected properly before disconnecting it from your PC. If this process fails"
		echo "    again, try downloading new copies of the files or use a different USB storage device."
		echo
		echo "    If the problem persists, contact the provider of the image files."
		shutitdown $SADSTOP

	fi

elif [ $FILECOUNT -eq 2 ] && [ $EFIPRESENT -eq 2 ]; then

		if [ -e /dev/fh ] && [ -f /tmp/openframe.ver ]; then

			OFVER=`cat /tmp/openframe.ver`

			EFI=`echo "$REFLASHCONTENTS" | head -1 | awk -F\.img {'print $1'}`
			EFISUM=`cat $IMGLOC/$EFI.img.md5 | awk -F\   {'print $1'}`
			EFIVER=`echo $EFI | head -1 | awk -F\of {'print $2'} | awk -F\_ {'print $1'}`

			echo -n "Checking integrity of EFI image..."

			if [ "$EFIVER" == "$OFVER" ]; then

				CHECKSUM=`md5sum $IMGLOC/$EFI.img | awk -F\  '{print $1}'`
				if [ "$CHECKSUM" == "$EFISUM" ]; then

					echo " successful."
					echo
					tput setaf $COLOUR
					echo "   -----------------------------------------------------"
					echo "    WARNING: OPENFRAME $OFVER EFI FLASHING IS ABOUT TO BEGIN "
					echo "   -----------------------------------------------------"
					tput setaf 7
					echo
					echo "Flashing will PERMANENTLY OVERWRITE the contents of the boot EFI chip, replacing the data with"
					echo -n "the "
					tput setaf $COLOUR
					echo -n "$EFI.img"
					tput setaf 7
					echo " file. This is a potentially dangerous change to your device."
					echo
					tput setaf 1
					echo "   ---------------------------------------------------------------------------"
					echo "    REFLASHING THE EFI HAS THE POTENTIAL TO RENDER YOUR OPENFRAME UNBOOTABLE! "
					[ "$OFVER" -eq 2 ] && echo "           YOUR DEVICE HAS A SOLDERED EFI CHIP WHICH CANNOT BE CHANGED!"
					echo "   ---------------------------------------------------------------------------"
					tput setaf 7
					echo
					
					echo "The process should take approximately 1 minute to complete."
					echo
					echo "To stop this, turn off the power to this device BEFORE THE COUNTDOWN ENDS."
					echo
					countdown "Reflash will begin" 30
					echo
					sleep 3
					tput setaf 1
					echo "   --------------------------------------------------------------------------------------------"
					echo -n "    Reprogramming EFI boot chip using "
					tput setaf 7
					echo -n "'$EFI.img'"
					tput setaf 1
					echo " - do not interrupt! "
					echo "   --------------------------------------------------------------------------------------------"
					tput setaf 7
					echo
					[ $DEBUG -eq 0 ] && EFIDEST="/dev/fh" || EFIDEST="/dev/null"
					dd if=$IMGLOC/$EFI.img of=$EFIDEST bs=1024k count=1
					echo
					sleep 2
					tput setaf 3
					echo "   --------------------------------------------------------------------------------------------"
					echo "    Checking integrity of new EFI boot chip contents... "
					echo "   --------------------------------------------------------------------------------------------"

					EFIPOSTSUM=`md5sum /dev/fh | awk -F\  {'print $1'}`

					if [ "$EFIPOSTSUM" == "$EFISUM" ]; then

						[ ! -d $IMGLOC/complete ] && mkdir $IMGLOC/complete
						mv $IMGLOC/$EFI.img $IMGLOC/$EFI.img.md5 $IMGLOC/complete/

						aplay -q -f a_law -r 44100 /usr/local/share/sounds/success.au &

						tput setaf $COLOUR
						echo
						echo "   -----------------------"
						echo "    EFI FLASHING COMPLETE "
						echo "   -----------------------"
						tput sgr0
						echo
						echo "Reprogramming of the EFI chip is complete. The EFI image files have been moved into a"
						echo "directory named 'complete' to ensure they are not written again accidentally."
						sleep 1
						shutitdown $HAPPYSTOP

					else

						aplay -q -f a_law -r 44100 /usr/local/share/sounds/failure.au &

						tput setaf 3
						echo
						echo "   ----------------------------------"
						echo "    EFI FLASHING FAILED VERIFICATION "
						echo "   ----------------------------------"
						tput sgr0
						echo
						echo "The new contents of the EFI chip compared to the file you were reprogramming it with do not"
						echo "appear to match. It is possible that there was a write error or a problem with the chip."
						echo
						echo "This system will shortly drop you to a shell prompt and advise you on what to do next."
						echo
						sleep 1
						droptoshell 3

						tput setab 3
						tput setaf 0
						echo "                                                                                                    "
						echo "        EFI FLASHING FAILED VERIFICATION - DO NOT POWER OFF OR REBOOT YOUR OPENFRAME DEVICE!        "
						echo "                                                                                                    "

						echo
						tput sgr0
						echo "If the EFI is improperly written, switching off may lead to a situation where the OpenFrame can"
						echo "no longer boot. It is recommended that you seek advice from the O2 Joggler Forums."
						echo
						tput setaf $COLOUR
						echo "    http://www.jogglerwiki.com/forum"
						tput setaf 7
						echo
						echo "This Reflash Tool contains all of the components necessary to attempt a manual EFI flash. If you"
						echo "wish to repeat the process with the same image files, enter the following command:"
						echo
						tput setaf $COLOUR
						echo "    sudo offlash"
						tput setaf 7
						echo
						echo "In case you do not have a keyboard attached, a Secure Shell (SSH) server is running which will"
						echo "allow you to connect to your OpenFrame from a PC using a terminal application such as PuTTY."

						if [ $NETWORKUP -eq 1 ]; then

							echo
							echo "Your IP address should be displayed at the top of the screen."

						else

							tput setaf 3
							echo
							echo "   -----------------------------------------------------------------"
							echo "    Your system does not appear to have a wired network connection. "
							echo "   -----------------------------------------------------------------"
							echo
							tput setaf 7
							echo "In case you do not have a keyboard attached, the Reflash Tool will wait until a wired network"
							echo "connection can be made. Plug an ethernet cable into the socket on the back of your OpenFrame"
							echo "and connect it to your network. Alternatively, press CTRL-C to exit to the shell prompt."
							echo

							sleep 10
							until ping -c 1 -W 1 8.8.8.8 &>/dev/null
							do
								sleep 1;
							done

							echo "Network ready."
							echo
							ifconfig eth0 | grep inet | grep -v inet6 | sed 's/^ *//'

						fi

						echo
						sleep 1
						exit 0

					fi

				else

					echo " checksum mismatch."
					echo
					echo
					tput setaf 3
					echo "    The firmware image you have provided does not match the checksum and may be corrupt."
					echo "    It is possible that there was an error writing the file to this USB device, or that it"
					echo "    did not download properly."
					echo
					echo "    Please try copying the firmware image files to this USB device again and ensure that"
					echo "    the drive is ejected properly before disconnecting it from your PC. If this process fails"
					echo "    again, try downloading new copies of the files or use a different USB storage device."
					echo
					echo "    If the problem persists, contact the provider of the image files."

				fi

			else

				echo " potentially incompatible firmware."
				echo
				echo
				tput setaf 3
				echo "    The image you have provided does not appear to have been designed for this OpenFrame $OFVER."
				echo 
				echo "    The filename specifies that it was written for the OpenFrame $EFIVER. Compatibility with this"
				echo "    device cannot be guaranteed. If you are intent on writing this image file then the"
				echo "    workaround should be obvious, but doing so is not recommended."

			fi

		else

			if [ ! -e /dev/fh ]; then

				echo "Firmware Hub driver unavailable."
				echo
				echo
				tput setaf 3
				echo "    The driver for the Intel Firmware Hub has not been loaded."
				echo 
				echo "    This is a problem with the Reflash System, the cause of which cannot be determined."
				echo "    Please try rebooting. If the problem continues, please rebuild your Reflash System"
				echo "    USB stick or seek advice from the O2 Joggler Forums."
				echo
				tput setaf $COLOUR
				echo "        http://www.jogglerwiki.com/forum"
				tput setaf 7

			fi

			if [ ! -f /tmp/openframe.ver ]; then

				echo "Unable to determine OpenFrame type."
				echo
				echo
				tput setaf 3
				echo "    The type of OpenFrame device you are using could not be determined."
				echo 
				echo "    It is not safe to continue with an EFI reprogramming attempt on unknown hardware."
				echo "    If you wish to attempt this manually you may do so, but this automated system will"
				echo "    not help you to proceed."

			fi

		fi

		shutitdown $SADSTOP

else

	header

	if [ -d $IMGLOC ]; then

		if [ $FILECOUNT -ne 2 ]; then

			if [ $FILECOUNT -eq 0 ]; then

				echo "No firmware image found."
				echo
				echo
				tput setaf 3
				echo "    When your USB flash memory device is connected to a computer, there will be a volume"
				echo "    named '"`cat /etc/fstab | grep vfat | awk -F\  {'print $1'} | awk -F\= {'print $2'}`"' which should contain a directory named 'reflash'."
				echo
				echo "    Please copy the two firmware image files into this 'reflash' directory."
				echo
				echo "    Firmware must be in '.img.gz' format and be accompanied by a '.img.gz.md5' checksum file."

			elif [ $FILECOUNT -eq 1 ] ; then

				echo "Only one file has been found in the 'reflash' directory."
				echo
				echo
				tput setaf 3
				echo "    Please ensure that you have copied both files into the 'reflash' directory."
				echo
				echo "    Firmware must be in '.img.gz' format and be accompanied by a '.img.gz.md5' checksum file."

			elif [ $FILECOUNT -gt 2 ]; then

				echo "There are $FILECOUNT files in 'reflash' directory."
				echo
				echo
				tput setaf 3
				echo "    Please ensure that you have copied ONLY TWO files into the 'reflash' directory."
				echo
				echo "    Firmware must be in '.img.gz' format and be accompanied by a '.img.gz.md5' checksum file."

			fi

		else

			if [ $IMGPRESENT -eq 0 ]; then

				IMGDECOMP=`ls -1 $IMGLOC | sed -n 1p | grep -c .img`

				if [ $IMGDECOMP -eq 1 ]; then

					echo "Decompressed image file located. Cannot perform integrity check."
					echo
					echo
					tput setaf 3
					echo "    An image file has been located, but it cannot be verified against its checksum because"
					echo "    it has been decompressed. Certain web browsers will expand compressed files automatically"
					echo "    after downloading them in an attempt to be helpful."
					echo
					echo "    Please locate the original compressed image file (check your Trash or Recycle Bin) and"
					echo "    copy this to your USB device in place of the expanded version."
					echo
					echo "    Firmware must be in '.img.gz' format and be accompanied by a '.img.gz.md5' checksum file."

				fi

			else

				echo "File confusion."
				echo
				echo
				tput setaf 3
				echo "    Two files have been located in the 'reflash' directory, but we cannot determine what they"
				echo "    are for. Please allow your OpenFrame to shut down, then reconnect this USB device to"
				echo "    your PC in order to check that the files are correct and properly named."
				echo
				echo "    Firmware must be in '.img.gz' format and be accompanied by a '.img.gz.md5' checksum file."

			fi

		fi

	else

		echo "No $IMGLOC directory found."
		echo
		echo
		tput setaf 3
		echo "    When your USB flash memory device is connected to your PC, there will be a volume"
		echo "    named 'rfl-boot' which should contain a directory named 'reflash'. Please recreate this"
		echo "    directory and copy the firmware image files into it."
		echo
		echo "    Firmware must be in '.img.gz' format and be accompanied by a '.img.gz.md5' checksum file."

	fi

	shutitdown $SADSTOP

fi

