#!/bin/bash
# Author: Steven Shiau <steven _at_ nchc org tw>
# License: GPL
# Description: This program will put Debian Live minimal + DRBL/Clonezilla program into a bootable device, such as pendrive/usb stick.
# Some notes and programs about make bootable USB device in M$ windows are modified from http://www.pendrivelinux.com/, Thanks to PDLA.

# Load DRBL setting and functions
if [ ! -f "/opt/drbl/sbin/drbl-conf-functions" ]; then
  echo "Unable to find /opt/drbl/sbin/drbl-conf-functions! Program terminated!" 
  exit 1
fi
. /opt/drbl/sbin/drbl-conf-functions

# Load some necessary setting from drbl-ocs.conf
. /opt/drbl/conf/drbl-ocs.conf

# Load basic functions of ocs
. /opt/drbl/sbin/ocs-functions

# other default settings
# If insert_mode=prog_only, only copy DRBL/Clonezilla programs only, no ocs images.
insert_mode="prog_and_img"
img_size_sum=0
# default is to create a bootable device instead of a released zip file.
# Two options: boot_dev or release_file
target_mode="boot_dev"

#
prog="$(basename $0)"

# functions
USAGE() {
    echo "Usage:"
    echo "To put clonezilla image into a bootable device (such as pendrive/usb stick):"
    echo "$prog [OPTION] DEV CLONEZILLA_IMAGE_NAME"
    echo "CLONEZILLA_IMAGE_NAME  The image exists in the system"
    echo "OPTION:"
    language_help_prompt_by_idx_no
    echo "-b, --bg-mode  [text|graphic]  Assign the back ground of boot menu. Default is graphic"
    echo "-p, --image-path   Assign the clonezilla image source path where CLONEZILLA_IMAGE_NAME exists in this server.  Default = $ocsroot"
    echo "-d, --dev DEV      Write the output to DEV (such as /dev/sda1)"
    echo "-a, --boot-loader [grub|syslinux]  Force to use grub or syslinux as boot loader in target device"
    echo "-s, --skip-image   Do not include any clonezilla image. The is used to created a live CD or USB flash drive with DRBL/Clonezilla programs only."
    echo "-c, --create-release  Create a USB package release file, this will not create a bootable device, instead it will create a zip file contain all necessary files to make it boot and run clonezilla live."
    echo "$prog will download a template Debian live CD for clonezilla iso file from $DEBIAN_ISO_URL if ncecessary. You can also download it by yourself, and put it in the working directory when you run $prog. If you want to create that template iso file in Debian Etch, run $DRBL_SCRIPT_PATH/sbin/create-debian-live."
    echo "NOTE! You have to prepare the target partition first, and the filesystem should be ready (FAT or ext2/3)."
    echo "Ex:"
    echo "To put clonezilla image sarge-ocs (located in /home/partimag in clonezilla server) to USB device /dev/sda1, you can run:"
    echo "  $prog -d /dev/sda1 sarge-ocs"
    echo "To put more images, just append them, such as:"
    echo "  $prog -d /dev/sda1 sarge-ocs etch-ocs"
    echo "To create a Live USB device with DRBL/Clonezilla programs, which can be used to save image:"
    echo "  $prog -s -d /dev/sda1"
    echo "To create a zip file of general purpose Clonezilla live. Later it can be put in an USB flash drive or similar device:"
    echo "  $prog -s -c"
}
#

check_if_root

# Parse command-line options
while [ $# -gt 0 ]; do
  case "$1" in
    -l|--language)
            shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
              specified_lang="$1"
              shift
            fi
	    [ -z "$specified_lang" ] && USAGE && exit 1
            ;;
    -a|--boot-loader)
            shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
              boot_loader="$1"
              shift
            fi
	    [ -z "$boot_loader" ] && USAGE && exit 1
            ;;
    -b|--bg-mode)
            shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
              bg_mode="$1"
              shift
            fi
	    [ -z "$bg_mode" ] && USAGE && exit 1
            ;;
    -p|--image-path)
            shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
              image_path="$1"
              shift
            fi
	    [ -z "$image_path" ] && USAGE && exit 1
            ;;
    -s|--skip-image)
            insert_mode="prog_only"
            shift ;;
    -d|--dev)
            shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
              output_dev="$1"
              shift
            fi
	    [ -z "$output_dev" ] && USAGE && exit 1
            ;;
    -c|--create-release)
            target_mode="release_file"
            shift ;;
    -*)     echo "${0}: ${1}: invalid option" >&2
            USAGE >& 2
            exit 2 ;;
    *)      break ;;
  esac
done
ocs_image="$*"
# strip the / to avoid problem
ocs_image="$(echo $ocs_image | sed -e "s|/||g")"
      
if [ -z "$ocs_image" ]; then
  [ "$insert_mode" = "prog_and_img" -a "$target_mode" = "boot_dev" ] && USAGE && exit 1
fi
[ -z "$image_path" ] && image_path=$ocsroot
[ -z "$bg_mode" ] && bg_mode="$BG_MODE_DEF"
[ -z "$output_dev" -a "$target_mode" = "boot_dev" ] && USAGE && exit 1
#
ask_and_load_lang_set $specified_lang

# we need zip to create the release file when target_mode is release_file
if ! type zip &>/dev/null && [ "$target_mode" = "release_file" ]; then
  [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
  echo "mkisofs not found!"
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
  echo "$msg_program_stop"
  exit 1
fi

if [ "$target_mode" = "boot_dev" ]; then
  # Prepare output dev, name, grub dev name...
  # /dev/sda1 -> /dev/sda, sda1 (output_dev_hd, output_dev_name)
  output_dev_hd="$(echo $output_dev | sed -e "s/[[:digit:]]*$//g")"
  output_dev_name="$(basename $output_dev)"
  
  # check if the target exists
  # Todo: if it's devfs ?
  if [ -z "$(grep -Ew "$output_dev_name" /proc/partitions)" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "$output_dev $msg_NOT_found!"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo "$msg_program_stop"
    exit 1
  fi

  # check if the target is busy or not
  if [ -n "$(grep -Ew "^$output_dev" /proc/mounts)" ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "$msg_is_mounted_u_must_unmount_it: $output_dev"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo "$msg_this_is_disk_usage_status:"
    df -h
    echo "$msg_program_stop"
    exit 1
  fi
  [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
  echo "$msg_create_live_device_warning: $output_dev"
  echo "$msg_this_is_disk_usage_status:"
  fdisk -l $output_dev_hd
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
  echo "$msg_are_u_sure_u_want_to_continue"
  echo -n "[y/N] "
  read cont_ans
  case "$cont_ans" in
    y|Y|[yY][eE][sS])
       echo $msg_ok_let_do_it
       ;;
    *)
       echo "Abort!"
       exit 2
  esac
fi


iso_file="$(basename $DEBIAN_ISO_URL)"
if [ ! -f "$iso_file" ]; then
  echo "We need Debian live CD for Clonnezilla to create clonezilla live."
  echo "Downloadling the iso file from $DEBIAN_ISO_URL..."
  wget $DEBIAN_ISO_URL
  get_iso_rc=$?
  # validate it
  if [ "$get_iso_rc" -eq 0 ]; then
    echo -n "Validating $DEBIAN_ISO... "
    md5_tmp="$(mktemp -d /tmp/isomd5.XXXXX)"
    #  wget http://opensource.nchc.org.tw/drbl-core/iso/MD5SUMS
    echo -n "Downloading $md5_file_url... "
    wget -r -l1 --no-parent --passive-ftp -e robots=off -q -nd --retr-symlinks -P "$md5_tmp" $md5_file_url
    if ! grep -w "$DEBIAN_ISO" $md5_tmp/MD5SUMS | md5sum -c; then
      [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
      echo "$DEBIAN_ISO is broken! Something went wrong! Try to remove $DEBIAN_ISO and run this program again."
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
      echo "$msg_program_stop"
      exit 1
    fi
  else
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "Unable to download $DEBIAN_ISO_URL!"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo "$msg_program_stop"
    exit 1
  fi
fi

if [ "$insert_mode" = "prog_and_img" ]; then
  echo "Creating clonezilla live with image(s) $ocs_image from $image_path..."
  # use the 1st image name as iso filename
  ocs_imgs_with_abs_path=""
  for im in $ocs_image; do
    if [ ! -d "$image_path/$im" ]; then
      [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
      echo "$image_path/$im $msg_NOT_found!"
      [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
      echo "$msg_program_stop"
      exit 1
    fi
    ocs_imgs_with_abs_path="$ocs_imgs_with_abs_path $image_path/$im=$image_path/$im"
    convert_ocs_format_from_1.5_to_2.0_or_newer $image_path/$im/
    img_size="$(du -ms $image_path/$im/ | awk -F" " '{print $1}')" # unit: MB
    img_size_sum="$((img_size_sum+img_size))"
  done
else
  echo "Creating clonezilla Live without any clonezilla image embedded..."
fi

DEBIAN_ISO_TMP="$(mktemp -d /tmp/ocs-iso.XXXXXX)"
trap "[ -d "$DEBIAN_ISO_TMP" ] && umount $DEBIAN_ISO_TMP &>/dev/null && rm -rf $DEBIAN_ISO_TMP" HUP INT QUIT TERM EXIT
trap "[ -d "$USB_TMP" ] && umount $USB_TMP &>/dev/null && rm -rf $USB_TMP" HUP INT QUIT TERM EXIT
WD="$(pwd)"
USB_TMP="$(mktemp -d /tmp/usb-dev.XXXXXX)"
WD_TMP="$(mktemp -d /tmp/iso_wd.XXXXXX)"
mount -o loop $DEBIAN_ISO $DEBIAN_ISO_TMP
if [ "$target_mode" = "release_file" ]; then
  output_filename="clonezilla-live-usb-`date +%Y%m%d`.zip"
  # assume boot loader as syslinux for release_file mode.
  boot_loader="syslinux"
  # clean the stale zip file
  rm -f $output_filename
else
  # mount the target boot device
  mount -o loop $output_dev $USB_TMP
  rc=$?
  if [ "$rc" -gt 0 ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "$msg_unable_to_mount_this_dev $output_dev!"
    echo "$msg_format_as_FAT_16_32:"
    echo "mkfs.vfat -F 16 $output_dev"
    echo "$msg_or"
    echo "mkfs.vfat -F 32 $output_dev"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    umount $DEBIAN_ISO_TMP
    exit 1
  fi
fi
echo -n "Creating DRBL/Clonezilla program tarball from $DRBL_SCRIPT_PATH... "
create_live_media_opt_drbl_tarball $WD_TMP
echo "done!"
#
if [ ! -e "$uni_font_file" ]; then
  echo "Download unifont.bgf from $uni_font_url..."
  wget $wget_opt $uni_font_url
  rc=$?
  if [ "$rc" -gt 0 ]; then
    [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
    echo "Unable to download unifont.bgf from $uni_font_url!"
    [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
    echo "$msg_program_stop"
    exit 1
  fi
fi

template_iso_size="$(du -Lms $DEBIAN_ISO | awk -F" " '{print $1}')"
drbl_prog_size="$(du -Lms $WD_TMP/opt_drbl.tgz | awk -F" " '{print $1}')" 
uni_font_size="$(du -Lms $uni_font_file | awk -F" " '{print $1}')"
target_files_size="$((template_iso_size+drbl_prog_size+img_size_sum+uni_font_size))"

echo "Estimated necessary space size in target dev: $target_files_size MB"
#
echo "Copying files from Debian Live CD... "
rsync -a --exclude isolinux --exclude md5sum.txt $DEBIAN_ISO_TMP/*  $USB_TMP/
mkdir -p $USB_TMP/pkg
cp -f $WD_TMP/opt_drbl.tgz $USB_TMP/pkg/
cp -a $uni_font_file $USB_TMP/pkg/

mkdir -p $USB_TMP/ocs.d
cp -a $DRBL_SCRIPT_PATH/setup/files/ocs/ocs.d $USB_TMP/
cp -f $DRBL_SCRIPT_PATH/doc/COPYING $USB_TMP/
# replace the original syslinux and related files.
etherboot_zlilo="$($query_pkglist_cmd drbl-etherboot | grep -E "eb-.*-etherboot-pci.zlilo$")"
cp -f $pxelinux_simple_vesamenu $pxelinux_simple_menu $pxelinux_bg_img $pxelinux_chain_file $DEBIAN_ISO_TMP/isolinux/memtest $USB_TMP/
if [ -n "$etherboot_zlilo" ]; then
  # we have to force it name as etherboot.zdsk, since isolinux only uses the "plain" ISO 9660 filenames, i.e. it does not support Rock Ridge or Joliet filenames.
  # ref: http://syslinux.zytor.com/archives/2006-October/007440.html
  # "-" will be regards as "_" if you want to use "-" for isolinux.
  # In syslinux on vfat, etherboot.zlilo is too long, make it ever shorter as eb.zli
  cp -f $etherboot_zlilo $USB_TMP/eb.zli
fi

# output setting for clonezilla live
mkdir -p $USB_TMP/etc/ocs
if [ "$insert_mode" = "prog_and_img" ]; then
  cat <<-OCS_CONF > $USB_TMP/etc/ocs/ocs-live.conf
ocs_live_run="$DRBL_SCRIPT_PATH/sbin/ocs-live-restore" 
DIA="whiptail"
OCS_CONF
else
  cat <<-OCS_CONF > $USB_TMP/etc/ocs/ocs-live.conf
ocs_live_run="$DRBL_SCRIPT_PATH/sbin/ocs-live-save-for-rw"
DIA="whiptail"
OCS_CONF
fi
# create dir $ocsroot in target dev
mkdir -p $USB_TMP/$ocsroot
[ -n "$ocs_image" ] && echo -n "Copying clonezilla image to $output_dev... "
for im in $ocs_image; do
  echo -n "$im "
  # we just copy images to $ocsroot in usb debice. since $ocsroot in drbl.conf will be read when ocs-live-save-*/ocs-live-restore is run
  cp -rfL $image_path/$im $USB_TMP/$ocsroot/
done

# put boot loader in MBR if not assign
[ -z "$boot_loader" ] && set_boot_loader $output_dev
# check & format boot_loader
case "$boot_loader" in
  grub|GRUB)         boot_loader="grub" ;;
  syslinux|SYSLINUX) boot_loader="syslinux" ;;
  *) 
     [ "$BOOTUP" = "color" ] && $SETCOLOR_FAILURE
     echo "Unknown boot loader $boot_loader!"
     [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
     echo "$msg_program_stop"
     exit 1
esac

echo "Copying kernel, initrd, etc..."
case "$boot_loader" in
  grub)
      mkdir -p $USB_TMP/boot/grub
      cp -a $DEBIAN_ISO_TMP/isolinux/{$kernel_file,$initrd_file} $USB_TMP/boot/
      cp -f $ocs_logo_img_grub $USB_TMP/boot/grub
      ;;
  syslinux)
      cp -a $DEBIAN_ISO_TMP/isolinux/{$kernel_file,$initrd_file} $USB_TMP/
      cp -f $ocs_logo_img_syslinux $USB_TMP/
      ;;
esac

if [ "$target_mode" = "boot_dev" ]; then
  echo -n "Waiting for data to be saved in $output_dev before unmounting it..." && umount $USB_TMP &>/dev/null
  echo
  echo "Making $output_dev_hd bootable..."
  [ -n "$boot_loader" ] && boot_loader_opt="-b $boot_loader"
  $DRBL_SCRIPT_PATH/sbin/makeboot.sh -l $lang -f $boot_loader_opt $output_dev
else
  ocs-live-boot-menu -l $lang_answer -f $VGA_MODE_DEF -b $bg_mode -k $kernel_file -i $initrd_file syslinux $USB_TMP/
  echo "Downloading syslinux.exe and Makeboot.exe..."
  wget $wget_opt -P "$USB_TMP" $syslinux_exe_url
  wget $wget_opt -P "$USB_TMP" $makeboot_exe_url
  # just store it. since big files, like squash flie and opt_drbl.tgz are compressed, it's not necessary to compress it again.
  (cd $USB_TMP; zip -0 -r $WD/$output_filename *)
  echo "The created release file is $output_filename. You can extract all the files into your pendrive, and run makeboot.exe from pendrive in M$ windows."
  [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
  echo "Warning: DO NOT RUN makeboot.exe from your local hard drive!! It is intended to be run from your USB device."
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
fi

# unmount all iso file
umount $DEBIAN_ISO_TMP &>/dev/null

# Clean the tmp working directory
echo "Cleaning tmp dirs..."
[ ! -z "$DEBIAN_ISO_TMP" ] && rm -rf $DEBIAN_ISO_TMP
[ ! -z "$USB_TMP" ] && rm -rf $USB_TMP

echo "Done!"
if [ "$target_mode" = "boot_dev" ]; then
  [ "$BOOTUP" = "color" ] && $SETCOLOR_WARNING
  echo "$msg_boot_clonezilla_live_dev: $output_dev_hd"
  [ "$BOOTUP" = "color" ] && $SETCOLOR_NORMAL
fi
