#!/bin/sh
# reflash
#  ensure the flash disk is not mounted
#  save configuration files
#  update the kernel
#  update the flashdisk
#  restore the saved configuration files
# the set of configuration files is described by
# /etc/default/conffiles.
#
# /etc/default/functions contains useful utility functions
. /etc/default/functions
load_functions sysconf
#
# NSLU2 flash layout is non-standard.
case "$(machine)" in
nslu2)
	isnslu2=1
	isdsmg600=
	isnas100d=
	imageok=1
	apexpart="Loader"
	usrpart=
	kpart="Kernel"
	ffspart="Flashdisk";;
nas100d)
	isnslu2=
	isdsmg600=
	isnas100d=1
	imageok=1
	apexpart=
	usrpart="usr"
	kpart="kernel"
	ffspart="filesystem";;
dsmg600)
	isnslu2=
	isdsmg600=1
	isnas100d=
	imageok=1
	apexpart=
	usrpart="usr"
	kpart="kernel"
	ffspart="filesystem";;
*)
	isnslu2=
	isdsmg600=
	isnas100d=
	imageok=
	apexpart=
	usrpart=
	kpart="kernel"
	ffspart="filesystem";;
esac
#
# CHECKING FOR INPUT (ARGUMENTS ETC)
# ----------------------------------
#
# find the kernel and the new flash file system, an image file can
# be used to specify both images.
ffsfile=
kfile=
imgfile=
preserve_config=1
while test $# -gt 0
do
	case "$1" in
	-n)	preserve_config=
		shift;;
	-k)	shift
		test $# -gt 0 || {
			echo "reflash: -k: give the file containing the kernel image" >&2
			exit 1
		}
		kfile="$1"
		shift;;
	-[jr])	shift
		test $# -gt 0 || {
			echo "reflash: -j: give the file containing the root jffs2 image" >&2
			exit 1
		}
		ffsfile="$1"
		shift;;
	-i)	shift
		test -n "$imageok" || {
			echo "reflash: -i: only supported on the LinkSys NSLU2,"     >&2
			echo " Iomega NAS 100d and D-Link DSM-G600 systems; use -k and -j" >&2
			echo " to specify the kernel and root file system instead." >&2
			exit 1
		}
		test $# -gt 0 || {
			echo "reflash: -i: give the file containing the complete flash image" >&2
			exit 1
		}
		imgfile="$1"
		shift;;
	*)	if test -n "$imageok"
		then
			echo "reflash: usage: $0 [-n] [-k kernel] [-j rootfs] [-i image]" >&2
		else
			echo "reflash: usage: $0 [-n] [-k kernel] [-j rootfs]" >&2
		fi
		echo "  -n: do not attempt to preserve the configuration" >&2
		echo "  -k file: the new compressed kernel image ('zImage')" >&2
		echo "  -j file: the new root file system (jffs2)" >&2
		test -n "$imageok" &&
			echo "  -i file: a complete flash image (gives both kernel and jffs2)" >&2
		echo " The current jffs2 will be umounted if mounted." >&2
		exit 1;;
	esac
done
#
# Sanity check on the arguments (note that the first case can only fire
# on NSLU2 or DSM-G600 because of the check for -i above.)
if test -n "$imgfile" -a -n "$ffsfile" -a -n "$kfile"
then
	echo "reflash: specify at most two files" >&2
	echo "  -i has both a kernel and rootfs, the kernel from -k and" >&2
	echo "  the rootfs from -j override the one in the image (if given)" >&2
	exit 1
elif test -z "$imgfile" -a -z "$ffsfile" -a -z "$kfile"
then
	echo "reflash: specify at least one file to flash" >&2
	exit 1
fi
#
# Perform basic checks on the input (must exist, size must be ok).
if test -n "$imgfile"
then
	if test -r "$imgfile" -a -n "$isnslu2"
	then
		# read the partition table and from this find the offset
		# and size of $kpart and $ffspart partitions.  The following
		# devio command just dumps the partition table in a format
		# similar to /proc/mtd (but it outputs decimal values!)
		#NOTE: this uses a here document because this allows the while
		# loop to set the variables, a pipe would put the while in
		# a sub-shell and the variable settings would be lost.  This
		# works in ash, no guarantees about other shells!
		while read size base name
		do
			if test "$name" = "$apexpart"
			then
				imgapexsize="$size"
				imgapexoffset="$base"
			elif test "$name" = "$kpart"
			then
				imgksize="$size"
				imgkoffset="$base"
			elif test "$name" = "$ffspart" -o "$name" = "rootfs"
			then
				imgffssize="$size"
				imgffsoffset="$base"
			fi
		done <<EOI
$(devio "<<$imgfile" '
	<= $ 0x20000 -
	L= 0x1000
	$( 1
		# 0xff byte in name[0] ends the partition table
		$? @ 255 =
		# output size base name
		<= f15+
		.= b 0xfffffff &
		<= f4+
		.= b
		pf "%lu %lu "
		<= f28-
		cp 16
		pn
		<= f240+
		L= L256-
	$) L255>')
EOI
		# check the result
		test "$imgksize" -gt 0 -a "$imgkoffset" -ge 0 || {
			echo "reflash: $imgfile: failed to find $kpart partition in image" >&2
			exit 1
		}
		# the kernel is after a 16 byte header which holds the
		# values length,0,0,0  Get the true size.
		ktmp="$(devio "<<$imgfile" "L=$imgksize" "O=$imgkoffset" '
			$( OL+$>!
				<= O
				A= b
				$( AL>!
					pr A
				$) 0
			$) 0')"
		test "$ktmp" -gt 0 || {
			echo "reflash: $imgfile($imgkoffset,$imgksize): invalid kernel offset/size" >&2
			exit 1
		}
		# update the size and offset to these values (the offset is 16+ because
		# of the header).
		imgksize="$ktmp"
		imgkoffset="$(devio "O=$imgkoffset" 'pr O16+')"
		# just test the size for the rootfs
		test "$imgffssize" -gt 0 -a "$imgffsoffset" -ge 0 || {
			echo "reflash: $imgfile: failed to find $ffspart" >&2
			exit 1
		}
	elif test -r "$imgfile" -a \( -n "$isdsmg600" -o -n "$isnas100d" \)
	then
		#
		# For the DSM-G600, this is really easy - the image is just
		# a tar file.  So, extract the contents of the tar file, and
		# set the kernel and filesystem variables (if not already set)
		# to point to the extracted content.  Content will look like:
		# 
		# drwxr-xr-x 500/500         0 2006-11-25 23:47:59 firmupgrade
		# -rw-r--r-- 500/500   4718592 2006-12-02 16:32:51 firmupgrade/rootfs.gz
		# -rw-r--r-- 500/500        40 2006-11-25 22:15:41 firmupgrade/version.msg
		# -rw-r--r-- 500/500         0 2006-11-25 23:47:59 firmupgrade/usr.cramfs
		# -rw-rw-r-- 500/500   1306872 2006-12-02 16:33:37 firmupgrade/ip-ramdisk
		# 
		# Heuristic: if the size of usr.cramfs is zero, the firmware
		# is not a D-Link firmware for the device.  (The version.msg
		# file is not useful for this purpose; it describes the hardware,
		# not the firmware version in the image!)
		# 
		# TODO: If usr.cramfs is non-zero, we should flash that, too, just
		#       to make sure that it matches the native firmware's kernel
		#       and rootfs that we're now flashing back onto the device.

		echo "reflash: unpacking DSM-G600/NAS-100d image file" >&2
		tar -x -f "$imgfile" -C /var/tmp || {
			echo "reflash: unable to unpack image file to be flashed" >&2
			exit 1
		}

		if test -z "$kfile"
		then
			kfile="/var/tmp/firmupgrade/ip-ramdisk"
		fi

		if test -z "$ffsfile"
		then
			ffsfile="/var/tmp/firmupgrade/rootfs.gz"
		fi

		if test -s "/var/tmp/firmupgrade/usr.cramfs"
		then
			echo "reflash: Native flash being restored" >&2
			usrfile="/var/tmp/firmupgrade/usr.cramfs"
			preserve_config=
		fi

	else
		echo "reflash: $imgfile: image file not found" >&2
		exit 1
	fi
fi
if test -n "$kfile"
then
	if test ! -r "$kfile"
	then
		echo "reflash: $kfile: kernel file not found" >&2
		exit 1
	fi
	# the file values override anything from the image.
	imgksize="$(devio "<<$kfile" 'pr$')"
	imgkoffset=0
else
	kfile="$imgfile"
fi
if test -n "$ffsfile"
then
	if test ! -r "$ffsfile"
	then
		echo "reflash: $ffsfile: root file system image file not found" >&2
		exit 1
	fi
	# values override those from the image
	imgffssize="$(devio "<<$ffsfile" 'pr$')"
	imgffsoffset=0
else
	ffsfile="$imgfile"
fi
if test -n "$usrfile"
then
	if test ! -r "$usrfile"
	then
		echo "reflash: $usrfile: usr file system image file not found" >&2
		exit 1
	fi
	# values override those from the image
	imgusrsize="$(devio "<<$usrfile" 'pr$')"
	imgusroffset=0
else
	usrfile=
fi
#
# INPUTS OK, CHECKING THE ENVIRONMENT
# -----------------------------------
# basic setup.  This could be parameterised to use different partitions!
#
kdev=
ksize=0
if test -n "$kfile"
then
	# we have a new kernel
	kdev="$(mtblockdev $kpart)"
	test -n "$kdev" -a -b "$kdev" || {
		echo "reflash: $kpart($kdev): cannot find $kpart mtd partition." >&2
		echo "  check /proc/mtd, either the partition does not exist or there is no" >&2
		echo "  corresponding block device." >&2
		exit 1
	}
	ksize="$(devio "<<$kdev" 'pr$')"
	#
	# check the input file size
	test -n "$imgksize" -a "$imgksize" -gt 0 -a "$imgksize" -le "$ksize" || {
		echo "reflash: $kfile: bad $kpart size ($imgksize, max $ksize)" >&2
		exit 1
	}
fi
#
ffsdev=
ffssize=0
if test -n "$ffsfile"
then
	ffsdev="$(mtblockdev $ffspart)"
	[ -n "$ffsdev" ] || \
	ffsdev="$(mtblockdev rootfs)"
	test -n "$ffsdev" -a -b "$ffsdev" || {
		echo "reflash: $ffspart($ffsdev): cannot find $ffspart mtd partition." >&2
		echo "  check /proc/mtd, either the partition does not exist or there is no" >&2
		echo "  corresponding block device." >&2
		exit 1
	}
	ffssize="$(devio "<<$ffsdev" 'pr$')"
	#
	# check the input file size
	test -n "$imgffssize" -a "$imgffssize" -gt 0 -a "$imgffssize" -le "$ffssize" || {
		echo "reflash: $ffsfile: bad $ffspart size ($imgffsize, max $ffssize)" >&2
		exit 1
	}
fi
#
usrdev=
usrsize=0
if test -n "$usrfile"
then
	usrdev="$(mtblockdev $usrpart)"
	test -n "$usrdev" -a -b "$usrdev" || {
		echo "reflash: $usrpart($usrdev): cannot find $usrpart mtd partition." >&2
		echo "  check /proc/mtd, either the partition does not exist or there is no" >&2
		echo "  corresponding block device." >&2
		exit 1
	}
	usrsize="$(devio "<<$usrdev" 'pr$')"
	#
	# check the input file size
	test -n "$imgusrsize" -a "$imgusrsize" -gt 0 -a "$imgusrsize" -le "$usrsize" || {
		echo "reflash: $usrfile: bad $usrpart size ($imgusrsize, max $usrsize)" >&2
		exit 1
	}
fi

#
# INPUTS OK, ENVIRONMENT OK, UMOUNT ANY EXISTING MOUNT OF THE FLASHDISK
# ---------------------------------------------------------------------
# This is only required if the device is going to be used
if test -n "$ffsdev"
then
	# -r causes this to fail if the flash device is mounted on /
	umountflash -r "$ffsdev" || exit 1
	#
	# Everything is umounted, now remount on a temporary directory.
	ffsdir="/tmp/flashdisk.$$"
	mkdir "$ffsdir" || {
		echo "reflash: $ffsdir: failed to create temporary directory" >&2
		exit 1
	}
	#
	mountflash "$ffsdev" "$ffsdir" -o ro || {
		rmdir "$ffsdir"
		exit 1
	}
	#
	# this is a utility function to make the cleanup easier
	errorexit() {
		umount "$ffsdir" && rmdir "$ffsdir" ||
			echo "reflash: $ffsdir: temporary directory cleanup failed" >&2
		exit 1
	}
	#
	test -r "$ffsdir/etc/default/conffiles" || {
		echo "reflash: [/initrd]/etc/default/conffiles: file not found" >&2
		errorexit
	}
else
	errorexit() {
		exit 1
	}
fi
#
# PRESERVE EXISTING CONFIGURATION
# -------------------------------
# Only required if the flash partition will be written
if test -n "$ffsdev" -a -n "$preserve_config"
then
	echo "reflash: preserving existing configuration file" >&2
	#
	# This step produces /tmp/preserve.$$ and /tmp/cpio.$$, the former is
	# a list of the preserved configuration files together with the processing
	# option, the latter is a directory tree of the preserved files (a directory
	# tree makes the restore step easier.)
	saved=/tmp/cpio.$$
	list=/tmp/preserve.$$
	mkdir "$saved" || {
		echo "reflash: $saved: could not create save directory" >&2
		errorexit
	}
	#
	# sysconf_save_conffiles <flash-directory> <dest> <list>
	sysconf_save_conffiles "$ffsdir" "$saved" "$list" || {
		echo "reflash: $saved: copy of saved configuration files failed" >&2
		rm -rf "$saved"
		rm "$list"
		errorexit
	}
	#
	# If this umount fails do not try to continue...
	umount "$ffsdir" || {
		echo "reflash: $ffsdir: temporary mount point umount failed" >&2
		echo "  No changes have been made." >&2
		rm -rf "$saved"
		rm "$list"
		exit 1
	}
fi
#
# FLASH THE NEW IMAGES
# --------------------
echo "reflash: about to flash new image" >&2
#
# There are four possibilities here - kernel only, flashdisk only, then
# kernel&flashdisk in either one or two different files.  The code used
# to attempt to do everything in one step, but this complicates it,
# so two steps are used (as required).  A failure between the two
# steps is a problem, but then so is a failure in a partial write.
# Write the flashdisk first because this is larger (most likely to
# fail).
#
# -p causes the progress indicator to be displayed
progress=-p
do_kernel() {
	local cmd
	if test -n "$isnslu2"
	then
		# NSLU2: write length,0,0,0 header, then fill
		cmd="wb L,4; fb 12,0; cpL"
	else
		# Other: just write the kernel bytes
		cmd="cpL"
	fi
	devio $progress "$@" "<<$kfile" ">>$kdev" '
		# kernel is at imgkoffset[imgksize]
		' "<= $imgkoffset" "L=$imgksize" "$cmd" '
		# fill with 255
		fb #t-,255'
}
#
do_ffs() {
	devio $progress "$@" "<<$ffsfile" ">>$ffsdev" '
		# rootfs is at imgffsoffset[imgffssize]
		' "<= $imgffsoffset" "cp $imgffssize" '
		# fill with 255
		fb #t-,255'
}
#
do_usr() {
	devio $progress "$@" "<<$usrfile" ">>$usrdev" '
		# usrfs is at imgusroffset[imgusrsize]
		' "<= $imgusroffset" "cp $imgusrsize" '
		# fill with 255
		fb #t-,255'
}
#
# check_status $? type file(offset,size) device
#  check the devio status code (given in $1)
check_status() {
	case "$1" in
	0)	echo " done" >&2;;
	1)	echo " failed" >&2
		echo "reflash: $3: flash $2 failed, no changes have been made to $4" >&2
		if test "$2" = rootfs
		then
			rm -rf "$saved"
			rm "$list"
			exit 1
		else
			echo "reflash: $2: continuing with rootfs changes" >&2
			echo "  NOTE: the old kernel is still in $4!" >&2
		fi;;
	3)	echo " failed" >&2
		echo "reflash: $3: WARNING: partial flash of $2 to $4 the system is unbootable" >&2
		echo "  Reflash from RedBoot or correct the problem here." >&2
		if test "$2" = rootfs
		then
			exit 3
		else
			echo "reflash: $2: continuing with rootfs changes" >&2
			echo "  NOTE: the kernel in $4 must be reflashed!" >&2
		fi;;
	*)	echo " failed" >&2
		echo "reflash($1): $3: internal error flashing $2 to $4" >&2
		exit $1;;
	esac
}
#
if test -n "$usrdev"
then
	echo -n "reflash: writing usrfs to $usrdev " >&2
	do_usr
	check_status $? usrfs "$usrfile($imgusroffset,$imgusrsize)" "$usrdev"
fi
#
if test -n "$ffsdev"
then
	echo -n "reflash: writing rootfs to $ffsdev " >&2
	do_ffs
	check_status $? rootfs "$ffsfile($imgffsoffset,$imgffssize)" "$ffsdev"
fi
#
if test -n "$kdev"
then
	echo -n "reflash: writing kernel to $kdev " >&2
	do_kernel
	check_status $? kernel "$kfile($imgkoffset,$imgksize)" "$kdev"
fi
#
# verify - this just produces a warning
if test -n "$usrdev"
then
	echo -n "reflash: verifying new usr image " >&2
	if do_usr -v
	then
		echo " done" >&2
	else
		echo " failed" >&2
		echo "reflash: WARNING: usrfs flash image verification failed" >&2
		echo "  The system is may be bootable." >&2
	fi
fi
#
if test -n "$ffsdev"
then
	echo -n "reflash: verifying new flash image " >&2
	if do_ffs -v
	then
		echo " done" >&2
	else
		echo " failed" >&2
		echo "reflash: WARNING: rootfs flash image verification failed" >&2
		echo "  The system is probably unbootable." >&2
		echo "  System configuration files will be restored but this may fail" >&2
		echo "  Starting a shell for user recovery (exit to continue)" >&2
		PS1='badflash$ ' sh -i <>/dev/tty >&0 2>&0
	fi
fi
#
if test -n "$kdev"
then
	echo -n "reflash: verifying new kernel image " >&2
	if do_kernel -v
	then
		echo " done" >&2
	else
		echo " failed" >&2
		echo "reflash: WARNING: kernel flash image verification failed" >&2
		echo "  The system is probably unbootable." >&2
		echo "  System configuration files will be restored in the rootfs." >&2
	fi
fi
#
# RESTORE THE OLD CONFIGURATION
# -----------------------------
# If not write the rootfs none of the following is required - exit now.
test -n "$ffsdev" -a -n "$preserve_config" || exit 0
#
echo "reflash: restoring saved configuration files" >&2
#
# the file /etc/.configured is the datestamp file used to ensure that
# changed configuration files can be recognised.  It is created by
# /etc/rcS.d/S99finish on first boot (if it does not exist).  We need
# a timestamp earlier than any files we create so touch it here, this
# also acts as a test on the mounted file system
mountflash "$ffsdev" "$ffsdir" && :>"$ffsdir/etc/.configured" || {
	rmdir "$ffsdir"
	echo "reflash: mount of new flash root file system failed" >&2
	if test -d "$ffsdir/etc"
	then
		echo "    The file system does not seem to be writeable." >&2
		echo "    The mounted file system is in $ffsdir" >&2
	fi
	echo "  WARNING: the kernel and root file system have been reflashed," >&2
	echo "  HOWEVER the new root file system seems to be unuseable." >&2
	echo "  Saved configuration files are in $saved" >&2
	echo "  The list of saved configuration files is in $list" >&2
	echo " You should determine the reason for the failed mount, mount the new" >&2
	echo " file system and restore the configuration from $saved - it's just a" >&2
	echo " matter of copying the saved files where required." >&2
	exit 1
}
#
# sysconf_restore_conffiles <flash-directory> <source-dir> <restore>
restore="/tmp/restore.$$"
sysconf_restore_conffiles "$ffsdir" "$saved" "$restore" <"$list" || {
	echo "reflash: $saved: restore of saved configuration files failed" >&2
	echo "  The new flash file system is mounted on $ffsdir" >&2
	echo "  The saved files are in $saved and the list in $list, the list of" >&2
	echo "  files selected for restore is in $restore" >&2
	echo "  You should restore any required configuration from $saved," >&2
	echo "  then umount $ffsdir and reboot." >&2
	exit 1
}
#
# remove the copied files (i.e. the ones which were preserved)
(	cd "$saved"
	exec rm $(cat "$restore")
)
rm "$restore"
#
# clean up, files left in $saved need to be handled by the user
files="$(find "$saved" ! -type d -print)"
if test -n "$files"
then
	echo "reflash: the following saved configuration files remain:" >&2
	echo "$files" >&2
	echo "The full list of preserved files is in $list.  To alter the" >&2
	echo "new root file system use the command:" >&2
	echo "" >&2
	echo "  mount -t jffs2 $ffsdev /mnt" >&2
	echo "" >&2
	echo "The saved files are in the temporary directory, they will not" >&2
	echo "be retained across a system boot.  Copy them elsewhere if you" >&2
	echo "are unsure whether they are needed" >&2
else
	rm -rf "$saved"
	rm "$list"
fi
#
# now the final umount
if umount "$ffsdir"
then
	rmdir "$ffsdir"
	echo "reflash: system upgrade complete.  Reboot to continue." >&2
	exit 0
else
	echo "reflash: $ffsdir: temporary mount point umount failed" >&2
	echo "  ALL changes have been made successfully, however the umount of" >&2
	echo "  the new root file system has failed.  You should determine the" >&2
	echo "  cause of the failure, umount $ffsdir, then reboot the system (this" >&2
	echo "  will use the upgraded kernel and root file system)" >&2
	exit 1
fi