#!/bin/bash

clean=
if [ "$1" == "--clean" ]; then
	clean=y
	shift
fi
img=$1
device=$2

usage() {
	echo "Usage: $0 [--clean] <image filename> [USB device]"
	echo -e "\t--clean\t\tremove LEAF persistent user data"
}

if [ -z "$img" ]; then
	echo "I need an image file!"
	usage
	exit 1
fi

if [ ! -f "$img" ]; then
	echo "Image file $img does not exist!"
	usage
	exit 1
fi

if [ -z "$device" ]; then
	for dev in /sys/block/*; do
		removable=`cat $dev/removable`
		usb=`readlink -e $dev/device | grep usb`
		if [ -n "$usb" ]; then
			dev_bn=`basename $dev`
			size=`parted --script -m /dev/$dev_bn "unit GB print" | grep "^/dev/$dev_bn" | awk -F: '{print $2}'`
			desc=`parted --script -m /dev/$dev_bn "unit GB print" | grep "^/dev/$dev_bn" | awk -F: '{print $7}'`
			echo "Found USB device..."
			echo "/dev/$dev_bn: $desc $size"
			until [ "$writeto" = "y" ] || [ "$writeto" = "Y" ] || [ "$writeto" = "n" ] || [ "$writeto" = "N" ]; do
				if [ "$removable" = "0" ]; then
					echo "WARNING: /dev/$dev_bn reported itself as a non-removable device!"
				fi
				echo -n "Write to this device (y/n)? "
				read writeto
			done
			if [ "$writeto" = "y" ] || [ "$writeto" = "Y" ]; then
				device=/dev/$dev_bn
				break
			fi
			writeto=
		fi
	done
fi

if [ -z "$device" ]; then
	echo "I need a device name!"
	usage
	exit 1
fi

if [ ! -b "$device" ]; then
	echo "$device is not a valid block device!"
	exit 1
fi

device_bn=`basename $device`
removable=`cat /sys/block/$device_bn/removable`
if [ "$removable" != "1" ]; then
	echo "WARNING: $device reported itself as a non-removable device!"
fi

ro=`cat /sys/block/$device_bn/ro`
if [ "$ro" = "1" ]; then
	echo "$device is a read-only device!"
	exit 1
fi

if [ -n "$clean" ]; then
	echo "!! You have chosen to clean $device!"
	echo "!! This will remove all persistent LEAF user data!"
	echo "!! Hit Control-C to quit.  Press Enter to continue."
	read
fi

echo "* Checking for mounted $device filesystems..."
mounts=`mount | grep "^$device" | awk '{ print $1 }'`
for mount in $mounts; do
        echo "- Unmounting $mount..."
        umount $mount
done
sleep 1

if [ -n "$clean" ]; then
	device_size=`cat /sys/block/$device_bn/size`
	echo "* Cleaning USB device $device..."
	dd if=/dev/zero of=$device bs=1b count=$device_size
fi

echo "* Writing $img to $device..."
dd if=$img of=$device

if [ "$?" = "0" ]; then
	echo
	echo "* LEAF USB stick is now ready for use!"
	echo
else
	echo
	echo "!! Error writing to the USB device!"
	echo
fi