#!/bin/bash
# Project page: http://www.sysresccd.org/
# (C) 2008 Francois Dupoux
# This scipt is available under the GPL-2 license

# it may also be interesting to reset the MBR
# dd if=/usr/lib/syslinux/mbr.bin of=/dev/sdf

## HELP AND BASIC ARGUMENT PROCESSING
#####################################

PROG="${0}"
ISODIR="/livemnt/boot"
USBDIR="/mnt/usbstick"
SYSLINUX="sylinux"
MINSIZEMB=300
cdfiles=('sysrcd.dat' 'sysrcd.md5' 'version' 'isolinux/initram.igz' 
	'isolinux/rescuecd' 'isolinux/rescue64' 'isolinux/f1boot.msg'
	'isolinux/isolinux.bin' 'isolinux/isolinux.cfg')

usage()
{
	cat <<EOF
sysresccd-usbstick: SystemRescueCd installation script for USB-sticks
Syntax:             sysresccd-usbstick <command> ...

Please, read the manual for help about how to use this script.
http://www.sysresccd.org/Online-Manual-EN
This script must be executed outside of the chroot environment.

Commands (execute in that order):
 1) listdev               Show the list of removable media
 2) writembr <devname>    Recreate the MBR + partition table on the stick
 3) format <partname>     Format the USB-stick device (overwrites its data)
 4) copyfiles <partname>  Copy all the files from the cdrom to the USB-stick
 5) syslinux <partname>   Make the device bootable

 -h|--help	          Display this screen
 -i ISODIR                Set the mount directory for the stock SysRescCD iso ($ISODIR)
 -u USBDIR                Set the directory for mounting the USB stick ($USBDIR)
 -m SYSLINUX              Set the syslinux program ($SYSLINUX)
 -s MINSIZEMB             Minimal size required for SysRescCD in mega-bytes
 dialog                   Dialog to make SystemRescueCD-USB-Sticks

Distributed under the GNU Public License version 2 - http://www.sysresccd.org
EOF
}

## MISC FUNCTIONS: Many utilities functions
###########################################

# show the error message ($1 = first line of the message)
help_readman()
{
	echo "$1"
	echo "Please, read the manual for more help about this script"
	echo "Web: http://www.sysresccd.org"
	exit 1
}

## ERROR HANDLING
#####################################

die()
{
	if [ -n "$1" ]
	then
		echo "${PROG}: error: $1"
	else
		echo "${PROG}: aborting."
	fi
	exit 1
}

## MISC FUNCTIONS: Many utilities functions
###########################################

# check that there is one partition and one only on block-device $1
find_first_partition()
{
	devname="$1"
	if [ -z "${devname}" ] || [ ! -d "/sys/block/$(basename ${devname})" ]
	then
		die "${devname} is not a valid device name"
	fi
	
	partcnt=0
	firstpart=0
	for i in $(seq 1 4)
	do
		partname="${devname}${i}"
		if [ -b "${partname}" ]
		then
			[ "${firstpart}" = '0' ] && firstpart="$i"
			partcnt=$((partcnt+1))
		fi
	done
	
	if [ "${partcnt}" = '1' ]
	then
		return ${partcnt}
	else
		return 0
	fi
}

# check $1 is a valid partition name
check_valid_partname()
{
	if [ -z "${partname}" ]
	then
		die "you have to provide a valid partition device-name as argument of this command"
	fi

	if [ -z "${partname}" ] || [ ! -b "${partname}" ]
	then
		die "${partname} is not a valid partition name"
	fi
	
	if ! echo "${partname}" | grep -qE '^/dev/[a-z]*[1-4]+$'
	then
		die "device [${partname}] is not a valid partition. Expect something like [/dev/sdf1]"
	fi

	if is_dev_mounted "${partname}"
	then
		die "${partname} is already mounted, cannot continue"
	fi
	
	return 0
}

# check $1 is a valid block device name
check_valid_blkdevname()
{
	if [ -z "${devname}" ]
	then
		die "you have to provide a valid device name as argument of this command"
	fi
	
	if [ ! -b "${devname}" ] || [ ! -d "/sys/block/$(basename ${devname})" ]
	then
		die "${devname} is not a valid device name"
	fi
	
	if is_dev_mounted "${devname}"
	then
		die "${devname} is already mounted, cannot continue"
	fi
	
	return 0
}

check_sysresccd_files()
{
	rootdir="$1"
	[ -z "${rootdir}" ] && rootdir="${ISODIR}"
	for curfile in ${cdfiles[*]}
	do
		curcheck="${rootdir}/${curfile}"
		if [ ! -f ${curcheck} ]
		then
			die "Cannot find ${curcheck}, cannot continue"
		fi
	done
	return 0
}

# returns 0 if the device is big enough
check_sizeof_dev()
{
	devname="$1"

	if [ -z "${devname}" ]
	then
		die "check_sizeof_dev(): devname is empty"
	fi

	if [ -z "$(which blockdev)" ]
	then
		echo "blockdev not found, assuming the size is ok"
		return 0
	fi
	
	secsizeofdev="$(blockdev --getsz ${devname})"
	mbsizeofdev="$((secsizeofdev/2048))"
	if [ "${mbsizeofdev}" -lt "${MINSIZEMB}" ]
	then
		die "The device [${devname}] is only ${mbsizeofdev} MB. It is too small to copy all the files, an USB-stick of at least ${MINSIZEMB}MB is recommended"
	else
		echo "The device [${devname}] seems to be big enough: ${mbsizeofdev} MB."
		return 0
	fi
}

# say how much freespace there is on a mounted device
check_disk_freespace()
{
	freespace=$(\df -m -P ${1} | grep " ${1}$" | tail -n 1 | awk '{print $4}')
	echo "Free space on ${1} is ${freespace}MB"
	if [ "${freespace}" -lt "${MINSIZEMB}" ]
	then
		die "There is not enough free space on the USB-stick to copy the SystemRescuecd files."
	fi
	return 0
}

# check that device $1 is an USB-stick
is_dev_usb_stick()
{
	curdev="$1"
	
	remfile="/sys/block/${curdev}/removable"
	vendor="$(cat /sys/block/${curdev}/device/vendor 2>/dev/null)"
	model="$(cat /sys/block/${curdev}/device/model 2>/dev/null)"
	if [ -f "${remfile}" ] && cat ${remfile} 2>/dev/null | grep -qF '1' \
		&& cat /sys/block/${curdev}/device/uevent 2>/dev/null | grep -qF 'DRIVER=sd'
	then
		return 0
	else
		return 1
	fi
}

do_writembr()
{
	devname="$1"
	shortname="$(echo ${devname} | sed -e 's!/dev/!!g')"
	
	check_valid_blkdevname "${devname}"
	if ! is_dev_usb_stick "${shortname}"
	then
		die "Device [${devname}] does not seem to be an usb-stick. Cannot continue."
	fi
	
	check_sizeof_dev "${devname}"
	
	if [ -z "$(which install-mbr)" ] || [ -z "$(which parted)" ]
	then
		die "install-mbr and parted must be installed, check these programs first."
	fi
	
	cmd="install-mbr ${devname} --force"
	echo "--> ${cmd}"
	if ! ${cmd}
	then
		die "${cmd} --> failed"
	fi
	
	cmd="parted -s ${devname} mklabel msdos"
	echo "--> ${cmd}"
	if ! ${cmd} 2>/dev/null
	then
		die "${cmd} --> failed"
	fi
	
	cmd="parted -s ${devname} mkpartfs primary fat32 0 100%"
	echo "--> ${cmd}"
	if ! ${cmd} 2>/dev/null
	then
		die "${cmd} --> failed"
	fi
	
	cmd="parted -s ${devname} set 1 boot on"
	echo "--> ${cmd}"
	if ! ${cmd} 2>/dev/null
	then
		die "${cmd} --> failed"
	fi
}

do_format()
{
	partname="$1"
	check_valid_partname "${partname}"

	check_sizeof_dev "${partname}"

	if [ -z "$(which mkfs.vfat)" ]
	then
		die "mkfs.vfat not found on your system, please install dosfstools first."
	fi
	
	if mkfs.vfat -F 32 -n SYSRESC ${partname}
	then
		echo "Partition ${partname} has been successfully formatted"
		return 0
	else
		echo "Partition ${partname} cannot be formatted"
		return 1
	fi
}

do_copyfiles()
{
	partname="$1"
	check_valid_partname "${partname}"
	
	# check the important files are available in ${ISODIR}
	check_sysresccd_files "${ISODIR}"
	
	check_sizeof_dev "${partname}"
	
	mkdir -p ${USBDIR} 2>/dev/null
	if ! mount -t vfat ${partname} ${USBDIR}
	then
		die "cannot mount ${partname} on ${USBDIR}"
	fi
	echo "${partname} successfully mounted on ${USBDIR}"
	
	check_disk_freespace "${USBDIR}"
	
	if cp -r --remove-destination ${ISODIR}/* ${USBDIR}/ && sync
	then
		echo "Files have been successfully copied to ${partname}"
	else
		echo "Cannot copy files to ${partname}"
	fi
	
	for curfile in '${USBDIR}/isolinux/isolinux.cfg'
	do
		if [ ! -f "${curfile}" ]
		then
			umount ${USBDIR}
			die "${curfile} not found, cannot continue"
		fi
	done
	
	# check the important files have been copied
	check_sysresccd_files "${USBDIR}"
	
	# move isolinux files to syslinux files
	rm -rf ${USBDIR}/syslinux
	if ! mv ${USBDIR}/isolinux/isolinux.cfg ${USBDIR}/isolinux/syslinux.cfg \
		|| ! mv ${USBDIR}/isolinux ${USBDIR}/syslinux
	then
		umount ${USBDIR}
		die "cannot move isolinux to syslinux, failed"
	fi
	
	# remove the last lines which produces error messages 'bad keyword' with syslinux
	sed -i -e '/label disk[1-2]$/d' -e '/label floppy$/d' -e '/label nextboot$/d' -e '/localboot/d' ${USBDIR}/syslinux/syslinux.cfg
	
	# add scandelay option which allows the usb devices to be detected
	sed -i -e 's!initrd=initram.igz!initrd=initram.igz scandelay=5!g' ${USBDIR}/syslinux/syslinux.cfg
	
	umount ${USBDIR}
}

do_syslinux()
{
	partname="$1"
	check_valid_partname "${partname}"
	
	if [ -z "$(which ${SYSLINUX})" ]
	then
		die "${SYSLINUX} not found on your system, please install syslinux first."
	fi
	
	if ${SYSLINUX} ${partname} && sync
	then
		echo "${SYSLINUX} has successfully prepared ${partname}"
	else
		echo "${SYSLINUX} failed to prepare ${partname}"
	fi
}

is_dev_mounted()
{
	curdev="$1"
	
	if cat /proc/mounts | grep -q "^${curdev}"
	then
		return 0
	else
		return 1
	fi
}

do_dialog()
{
	lwselection="`mktemp /tmp/lwselection.XXXX`"
	selection='dialog --backtitle "Select USB-Stick" --checklist "Select USB-Stick" 20 61 5'
	devcnt=0
	for curpath in /sys/block/*
	do
		curdev="$(basename ${curpath})"
		devname="/dev/${curdev}"
		if is_dev_usb_stick ${curdev}
		then
			if [ -n "$(which blockdev)" ]
			then
				secsizeofdev="$(blockdev --getsz /dev/${curdev})"
				mbsizeofdev="$((secsizeofdev/2048))"
				sizemsg=" and size=${mbsizeofdev}MB"
			fi	
			echo "Device [${devname}] detected as [${vendor} ${model}] is removable${sizemsg}"
			if is_dev_mounted "${devname}"
			then
				echo "Device [${devname}] is mounted"
			else
				echo "Device [${devname}] is not mounted"
			fi
			find_first_partition ${devname}
			firstpart="$?"
			if [ "${firstpart}" != '0' ]
			then
				echo "Device [${devname}] has one partition: ${devname}${firstpart}"
				selection="$selection \"${devname}\" \"\" off"
			else
				echo "Cannot identify which partition to use on ${devname}"
			fi
			devcnt=$((devcnt+1))
		fi
	done
	if [ "${devcnt}" = '0' ]
	then
		echo "No USB-stick has been detected."
	else
		eval $selection 2>$lwselection
			if test -s $lwselection; then
				for devname2 in `cat $lwselection  | tr -d \" | sort`; do
				do_writembr ${devname2}
				sleep 5
				find_first_partition ${devname2}
				devname2="${devname2}$?"
				do_format ${devname2}
				do_copyfiles ${devname2}
				do_syslinux ${devname2}

				done
			fi
	fi
	rm -f $lwselection
}

do_listdev()
{
	devcnt=0
	for curpath in /sys/block/*
	do
		curdev="$(basename ${curpath})"
		devname="/dev/${curdev}"
		if is_dev_usb_stick ${curdev}
		then
			if [ -n "$(which blockdev)" ]
			then
				secsizeofdev="$(blockdev --getsz /dev/${curdev})"
				mbsizeofdev="$((secsizeofdev/2048))"
				sizemsg=" and size=${mbsizeofdev}MB"
			fi	
			echo "Device [${devname}] detected as [${vendor} ${model}] is removable${sizemsg}"
			if is_dev_mounted "${devname}"
			then
				echo "Device [${devname}] is mounted"
			else
				echo "Device [${devname}] is not mounted"
			fi
			find_first_partition ${devname}
			firstpart="$?"
			if [ "${firstpart}" != '0' ]
			then
				echo "Device [${devname}] has one partition: ${devname}${firstpart}"
			else
				echo "Cannot identify which partition to use on ${devname}"
			fi
			devcnt=$((devcnt+1))
		fi
	done
	if [ "${devcnt}" = '0' ]
	then
		echo "No USB-stick has been detected."
	fi
}

## MAIN SHELL FUNCTION
########################################################

if [ $# -eq 0 ]    # Script invoked with no command-line args?
then
	usage
	exit 1
fi

while getopts ":ui:r:w:m:" Option
do
	case $Option in
		i ) ISODIR="$OPTARG";;
		r ) ROOTDIR="$OPTARG";;
		w ) WORKDIR="$OPTARG";;
		* ) usage; exit 1;;  # Default, handles -h
	esac
done
shift $(($OPTIND - 1))

COMMAND="${1}"
shift

if [ "$(whoami)" != "root" ]
then
	help_readman "$0: This script requires root privileges to operate."
fi

if ! cat /proc/mounts | awk '{print $2}' | grep -q -F '/memory'
then
	help_readman "$0: This script must be executed from SystemRescueCd"
	exit 1
fi

case "${COMMAND}" in
	listdev)
		do_listdev
		;;
	writembr)
		do_writembr "$@"
		;;
	format)
		do_format "$@"
		;;
	copyfiles)
		do_copyfiles "$@"
		;;
	syslinux)
		do_syslinux "$@"
		;;
	dialog)
		do_dialog "$@"
		;;
	*)
		usage 
		exit 1
		;;
esac
exit 0