#!/bin/bash
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; version 2 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Library General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
#
# $Id$
 
export PATH=/sbin:/usr/sbin:/bin:/usr/bin
 
err() {
  echo "$*" 1>&2
  exit 1
}
 
usage() {
  err "Usage: $0 [--reset-mbr ] <isopath> <usb device>"
}
 
if [ $(id -u) != 0 ]; then
  err "You need to be root to run this script"
fi
 
if [ ! -x /usr/bin/syslinux ]; then
  err "syslinux needs to be installed for this script to run"
fi
 
while [ $# -gt 2 ]; do
  case $1 in
    --reset-mbr)
      resetmbr="yes"
      ;;
    *)
      usage
      ;;
  esac
  shift
done
 
cleanup() {
  [ -d "${ISOMNT}" ] && umount "${ISOMNT}" && rmdir "${ISOMNT}"
  [ -d "${USBMNT}" ] && umount "${USBMNT}" && rmdir "${USBMNT}"
}
 
checkfs() {
  device=$1
  usbfs=$(/lib/udev/vol_id -t $device)
  case "${usbfs}" in
    vfat|msdos|ext[23])
      ;;
    *)
      err "USB filesystem must be vfat/ext2/ext3"
      ;;
  esac
}
 
resetmbr() {
  device=$1
  sysdev=$(udevadm info -q path -n ${device})
  if [ -e "/sys/${sysdev}/device" ]; then
    dev=$(basename /sys${sysdev})
  else
    dev=$(basename $(readlink -f /sys${sysdev}/../))
  fi
  if [ ! -e /sys/block/${dev} -o ! -e /dev/${dev} ]; then
    err "Error finding ${dev} block device. Bailing out"
  fi
  dev=/dev/$dev
  if [ -f /usr/lib/syslinux/mbr.bin ]; then
    cat /usr/lib/syslinux/mbr.bin > $dev
  elif [ -f /usr/share/syslinux/mbr.bin ]; then
    cat /usr/share/syslinux/mbr.bin > $dev
  else
    err "Unable to find syslinux mbr"
  fi
}
 
copyfiles() {
  echo "Copying files..."
  for item in ${COPYLIST}; do
    cp -af ${ISOMNT}/${item} ${USBMNT}
  done
  if [ ! -d ${USBMNT}/syslinux ]; then mkdir ${USBMNT}/syslinux; fi
  cp ${ISOMNT}/isolinux/* ${USBMNT}/syslinux
}
 
sanitychecks() {
  if [ -z "${ISOIMAGE}" -o ! -f "${ISOIMAGE}" ]; then
    usage
  fi
  if [ -z "${USBDEV}" -o ! -b "${USBDEV}" ]; then
    usage
  fi
}
 
install_syslinux() {
  device=$1
  echo "Updating boot config"
  # Figure out the filesystem and setup (sys|ext)linux accordingly
  usbfs=$(/lib/udev/vol_id -t $device)
  echo "Installing bootloader on ${device}"
  case $usbfs in
    msdos|vfat)
      mv ${USBMNT}/syslinux/isolinux.cfg ${USBMNT}/syslinux/syslinux.cfg
      syslinux -d syslinux ${device}
      ;;
    ext[23])
      mv ${USBMNT}/syslinux/isolinux.cfg ${USBMNT}/syslinux/extlinux.conf
      extlinux -i ${USBMNT}/syslinux
      ;;
  esac
  # TODO: Mark the partition bootable.
  # parted $device toggle N boot
}
 
ISOIMAGE=$1
USBDEV=$2
ISOMNT=$(mktemp -d /tmp/iso.XXXXXXXXXX)
USBMNT=$(mktemp -d /tmp/usb.XXXXXXXXXX)
 
# Files needed from the iso
COPYLIST="casper dists pool preseed .disk README.diskdefines md5sum.txt install"
 
sanitychecks
checkfs ${USBDEV}
if [ "$resetmbr" == "yes" ]; then resetmbr ${USBDEV}; fi
mount -o loop $ISOIMAGE $ISOMNT
mount $USBDEV $USBMNT
 
trap cleanup EXIT
copyfiles
install_syslinux ${USBDEV}
 
echo "Done!"
exit 0