#!/bin/bash


# TODO
#-Make a list of available versions of gentoo releases that work with this script
# and make sure the script verifies by the filename it is using an
# appropriate release.
#-Add the option of formatting the pen if necessary


BACK_UP="\033[1K\033[0G"
NORMAL="\033[0m"
WARN="\033[33;1m"
BAD="\033[31;1m"
BOLD="\033[1m"
GOOD="\033[32;1m"


big_msg() {
echo -e '\e[37;44m\033[1m'$1'\e[0m\033[0m'
}

good_msg() {	
	msg_string=$1
	msg_string="${msg_string:-...}"
	echo -e "${GOOD}>>${NORMAL}${BOLD} ${msg_string} ${NORMAL}"
}

normal_msg() {	
	msg_string=$1
	msg_string="${msg_string:-...}"
	echo -e "${GOOD}${NORMAL}${BOLD} ${msg_string} ${NORMAL}"
}

bad_msg() {
	msg_string=$1
	msg_string="${msg_string:-...}"
	echo -e "${BAD}!!${NORMAL}${BOLD} ${msg_string} ${NORMAL}"
} 

leaving(){
	[ -d "${cdmnt}" -a "${isomode}" = "1" ] && (umount cdmnt ; rmdir cdmnt) >/dev/null 2>/dev/null
	[ -d "${utmp}" ] && (umount utmp ; rmdir utmp) >/dev/null 2>/dev/null
}

usage() {
    echo "$0  [--verify] <iso or cd mount point> <usb device>"
    bad_msg "$1"
    exit 1
}

die() {
        bad_msg "$1"
        leaving
        exit 1
}

filesystem() {
    local USBFS=$(/lib/udev/vol_id -t ${USBDEV})
    if [ "$USBFS" != "vfat" -a "$USBFS" != "msdos" ]; then
	die "USB filesystem must be vfat or msdos"
    fi
}

testing() {

if [ $(id -u) != 0 ]; then 
    die "You need to be root to run this script"
fi

if [ \( \( ! -e /usr/lib/syslinux/mbr.bin \) -o \( ! -e /usr/bin/syslinux \) \) \
	-a \( \( ! -e /sbin/grub-install \) -o \( ! -e /sbin/grub \) \) ]
then
	die "You must have syslinux-3 or grub installed"
fi

if [ -z "$ISO" -o -z "$USBDEV" ]
then
	usage
elif [ ! -e "$ISO" ]
then
    usage "Iso or mount point not found"
elif [ ! -b "$USBDEV" ]
then
    usage "USB device not found"
fi

if [ "${verify}" = "1" ]
then
	[ -d "${ISO}" ] && die "You did not provide an ISO file"
	[ ! -e "${ISO}".DIGESTS ] && die "You do not have a valid *.DIGESTS in the same directory as the ISO file"
	good_msg "Verifying md5sum"
	md5sum -c --status "${ISO}".DIGESTS
	[ $? -ne 0 ] && die "md5sum check failed"	
fi

filesystem

}

getdisk() {
    
    local p=$(udevinfo -q path -n ${USBDEV})
    if [ -e /sys/$p/device ]; then
	device=$(basename /sys/$p)
    else
	device=$(basename $(readlink -f /sys/$p/../))
    fi
    if [ ! -e /sys/block/$device -o ! -e /dev/$device ]; then
	die "Error finding block device of $DEV.  Aborting!"
    fi

    device="/dev/${device}"
}

#checkMBR() {
#    local bs=$(mktemp /tmp/bs.XXXXXX)
#    dd if=$device of=$bs bs=512 count=1 2>/dev/null || die "Error checking the MBR"
#    
#    local mbrword=$(hexdump -n 2 $bs |head -n 1|awk {'print $2;'})
#    rm -f $bs
#    if [ "$mbrword" = "0000" ]; then
#	echo "MBR appears to be blank."
#	echo "Do you want to replace the MBR on this device?"
#	echo "Press Enter to continue or ctrl-c to abort"
#	read
#	resetMBR $1
#    fi
#
#    return 0
#}

checkPartActive() {
    local dev="${USBDEV}"
    
    # if we're installing to whole-disk and not a partition, then we 
    # don't need to worry about being active
    if [ "$dev" = "$device" ]; then
	return
    fi

    if [ "$(/sbin/fdisk -l $device 2>/dev/null |grep $dev |awk {'print $2;'})" != "*" ]; then
	echo "Partition isn't marked bootable!"
	echo "You can mark the partition as bootable with "
        echo "    # /sbin/parted $device"
	echo "    (parted) toggle N boot"
	echo "    (parted) quit"
	exit 1
    fi
}

[ \( $# -ne 2 \) -a \( $# -ne 3 \) ] && usage

while [ $# -gt 2 ]; do
    case $1 in
	--verify)
	    verify=1
	    ;;
	*)
	    usage
	    ;;
    esac
    shift
done

ISO=$1
USBDEV=$2

testing
getdisk
#checkMBR
checkPartActive

echo ''	
big_msg "Gentoo LiveUSB configuration script"
big_msg "Beta 1"
echo ''
normal_msg "Creation process started"

#see if ISO is mounted cd or iso file
if [ -d ${ISO} ]
then
	cdmnt=$1
	isomode=0
elif [ -f ${ISO} ]
then
	isomode=1
	cdmnt=$(mktemp -d) || die "Couldn't create temporary ISO dir"
	mount -o loop $1 ${cdmnt} || die "ISO not mounted"
else
	die "ISO not found"
fi

utmp=$(mktemp -d) || die "Couldn't create directory for the key ${utmp}"

good_msg "Unmounting USB device"
umount ${USBDEV} >/dev/null 2>/dev/null

mount -t auto ${USBDEV} ${utmp} || die "Failed mounting ${USBDEV} in ${utmp}"

good_msg "Copying files (may take a long time)..."
cp -R ${cdmnt}/* ${utmp} || die "Failed copying cd contents"
if [ ${isomode} -eq 1 ]; then
	umount ${cdmnt} || die "Failed unmounting ${cdmnt}"
fi

bootloader=""

if [ -d "${utmp}/isolinux" ]; then
  bootloader="isolinux"
  [ \( ! -e /usr/lib/syslinux/mbr.bin \) -o \( ! -e /usr/bin/syslinux \) ] &&
    die "Syslinux is required (in /usr/bin)"
  mv ${utmp}/isolinux/* ${utmp} || die "Failed moving contents"
  mv ${utmp}/isolinux.cfg ${utmp}/syslinux.cfg
  rm -rf ${utmp}/isolinux*
  mv ${utmp}/memtest86 ${utmp}/memtest \
	|| die "Failed moving ${utmp}/memtest86 to ${utmp}/memtest"
  sed -i \
      -e "s:cdroot:cdroot slowusb:" \
      -e "s:kernel memtest86:kernel memtest:" \
      ${utmp}/syslinux.cfg || die "Failed using sed"
  good_msg "Installing new MBR..."
  dd if=/usr/lib/syslinux/mbr.bin of=${device} || die "Failed installing MBR in ${device}"
  syslinux ${USBDEV} || die "Couldn't install the syslinux bootloader"
elif [ -d "${utmp}/boot/grub" ]; then
  bootloader="grub"
  [ \( ! -e /sbin/grub-install \) -o \( ! -e /sbin/grub \) ] && \
    die "Grub is required (in /sbin)"
  device=$(echo ${USBDEV} | sed -e 's:p\?[0-9]\+$::')
  partition=$(echo ${USBDEV} | sed -e 's:^.*p\?\([0-9]\+\)$:\1:')
  grub_partition=$((${partition} - 1))
  echo "(hd0) ${device}" > ${utmp}/boot/grub/device.map
  cp /lib/grub/*/* ${utmp}/boot/grub/
  good_msg "Installing Grub"
  echo -e "root (hd0,${grub_partition})\nsetup (hd0)\nquit" | grub --batch --device-map=${utmp}/boot/grub/device.map || die "Couldn't install the grub bootloader"
  good_msg "Installing new MBR..."
  grub-install --root-directory=${utmp} ${device} || die "MBR install error"
else
 die "No valid bootloader found"
fi

umount ${utmp} || die "Failed unmounting ${utmp}"

normal_msg "Gentoo LiveUSB created"