#!/bin/bash
# live-grub-stick
#
# Copyright (c) 2015 CyberOrg Info
# Copyright (c) 2014 Mindaugas Baranauskas

# 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; either version 2 of the License, or (at your option) any later
# version.
#
# 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 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.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
#
# Original authors:      Jigish Gohil <cyberorg@opensuse.org>
# Contributors:          Mindaugas Baranauskas <opensuse.lietuviu.kalba@gmail.com>
#
# This script creates bootable openSUSE, Fedora or Ubuntu(or clones)
# live usb stick on fat partition
#
# To get POT file for localization, execute:
# bash --dump-po-strings live-grub-stick > live-grub-stick.pot

PATH=$PATH:/sbin:/usr/sbin:/usr/local/sbin:/usr/local/bin:/usr/bin:/bin
ionice -c3 -p$$
umount_everything () {
	sync
	if [[ $distroname != isohybrid ]]; then
                umount $isomount &>/dev/null             
                umount $stickdevpart &>/dev/null
		rmdir $isomount &>/dev/null || echo "unable to remove $isomount, please remove manually"
                umount $stickmount &>/dev/null
		while $(mount | grep $stickmount &>/dev/null); do
			umount $stickdevpart &>/dev/null
			sleep 2
		done
		rmdir $stickmount &>/dev/null || echo "unable to remove $stickmount, please remove manually"
		if [[ -f $liveusbgui ]]; then
			rm $liveusbgui &>/dev/null
		fi
		if [[ -h $iso_symlink ]]; then
			rm "$iso_symlink" &>/dev/null
		fi
	fi
}
clean_up () {
	if [[ $1 == user_killed ]] ; then
		echo "aborting and cleaning up"
		umount_everything
		exit 1
	else
		echo "cleaning up"
		umount_everything
		echo "Your bootable usb stick is now ready"
		echo "have a lot of fun..."
	fi
}
trap "clean_up user_killed" SIGINT SIGTERM
need_help() {
	cat <<EOF
        Create multi boot USB stick/hard disk with whole iso/s on vfat/fat32 partition
        keeping existing data untouched.

        Note: File size greater than 4G is not usable on vfat/fat32 partition so the
        live CD/DVD iso file should not exceed this limit.

        Note2: Install 32bit/x86 iso on the stick first if creating multiboot with both
        x86 and x86_64 arch images.

        Note3: Requires: grub2 and dd_rescue/ddrescue installed on the system running this.

        Run this command as root (su -, not sudo)
                live-grub-stick isopath stickpartition
        e.g.: 
                live-grub-stick /home/geeko/openSUSE-Edu-li-f-e-12.2-1-i686.iso /dev/sdXY

        To add various distribution iso to the stick, run the following:
                For openSUSE    : live-grub-stick --suse /path/to/openSUSE-filename.iso /dev/sdXY
                For openSUSE with persistence    : live-grub-stick --suse-persistent /path/to/openSUSE-filename.iso /dev/sdXY
                For Ubuntu clones     : live-grub-stick --ubuntu /path/to/ubuntu-filename.iso /dev/sdXY
                For Ubuntu clones with persistence      : live-grub-stick --ubuntu-persistent /path/to/ubuntu-filename.iso /dev/sdXY
                For Mint        : live-grub-stick --mint /path/to/mint-filename.iso /dev/sdXY
                For Fedora      : live-grub-stick --fedora /path/to/fedora-filename.iso /dev/sdXY
                For iPXE        : live-grub-stick --ipxe /path/to/ipxe.iso /dev/sdXY

                For isohybrid   : live-grub-stick --isohybrid /path/to/isohybridimage.iso /dev/sdX

        Standard openSUSE installation DVD iso can also be written with --suse option when used
        with live-grub-stick on non-vfat partition.

        isopath should be full absolute path of iso image and the device should be 
        actual partition on the stick like /dev/sdb1, /dev/sdc1,/dev/sdc2...

        The stick partition has to be vfat/fat32 format if the image is not isohybrid.

        Persistent option requires minimum 500M free space on the USB device apart from the spare needed by iso image.
        To allocate more space for cow file use cowsize variable like this example: "export cowsize=1000M" before running the script

        Please note that using isohybrid option will remove all existing data on the USB device
        and create new partitions.

        Run live-grub-stick -l(or --list) to list the possible usb storage devices available.

        openSUSE users can install it via 1-click from here:
        http://software.opensuse.org/package/live-grub-stick

EOF
}
if [[ $(id -u) != 0 ]]; then
	echo "run this command as root"
	need_help
	exit 1
fi
while [ $# -gt 2 ]; do
        case $1 in
                --fedora)
                distroname=fedora
                ;;
                --suse|--opensuse)
                distroname=suse
                ;;
                --ubuntu|--mint)
                distroname=ubuntu
                ;;
                --suse-persistent|--opensuse-persistent)
                distroname=suse
		PERSISTENT_IMAGE=true
                ;;
                --ubuntu-persistent|--mint-persistent)
                distroname=ubuntu
                PERSISTENT_IMAGE=true
                ;;
		--isohybrid)
		distroname=isohybrid
		;;
                --ipxe)
                distroname=ipxe
                ;;
                *)
                echo "invalid arg -- $1, check spelling or if the distribution is supported"
		need_help
                exit 1
                ;;
        esac
        shift
done
if [ "$1" == "-h" -o "$1" == "--help" ]; then
        need_help
        exit
fi
if [ "$1" == "-l" -o "$1" == "--list" ]; then
        echo "Here is the list of possible usb storage devices on your computer,"
        echo "use the correct one, usually the one with a number at the end:"
        echo "$(for i in $(find /dev/disk/by-path/ |grep usb); do readlink -f $i;done)"
        exit
fi
if [[ x"$1" == x ]]; then
	echo "Requires first argument as iso image path"
	need_help
	exit 1
fi
if [[ x"$2" == x ]]; then
	echo "Requires second arguement as device partition path, /dev/sdb1 for example"
	need_help
	echo "Here is the list of possible usb storage devices on your computer, "
	echo "use the correct one, usually the one with a number at the end: "
	echo "$(for i in $(find /dev/disk/by-path/ |grep usb); do readlink -f $i;done)"
	exit 1
fi
if lsb_release -i | grep -qi suse; then
	grubinstall='grub2-install'
	if [[ ! -e $(which qemu-img) ]]; then
        	echo "qemu-img not found, please install qemu-tools package"
        	exit 1
	fi
else
	grubinstall='grub-install'
fi
if [[ ! -e $(which $grubinstall) ]]; then
	echo "$grubinstall command not found, please install grub2 package"
	exit 1
fi
if [[ ! -e $1 ]]; then
	echo "File $1 does not exist"
	exit 1
fi
if [[ ! -e $2 ]]; then
        echo "Device $2 does not exist"
        exit 1
fi
#variables
distroname=${distroname:-suse}
isomount=$(mktemp -d)
isopath=$(readlink -f "$1")
iso_symlink=$(mktemp) ; rm $iso_symlink ; ln -s "$isopath" $iso_symlink
isosize=$(stat -c%s "$isopath")
stickmount=$(mktemp -d)
stickdevice=$(echo $2 | sed 's/[0-9]*//g')
stickbase=$(basename $2)
isoname=$(basename "$1")
isonametr=$(echo $isoname | tr " " _)
isonameshort=$(echo $isoname | cut -d "-" -f 1,2,3 | tr " " _)
stickuuid=$(blkid -s UUID -o value $2)
isolabel=$(blkid -s LABEL -o value $1)
stickpart=$(basename $2 | sed 's/[a-z]*//g')
stickdevpart=$2
liveusbgui=/tmp/liveusbgui
benice="nice ionice -c3"
cowsize=${cowsize:-500M}
if [[ $distroname != isohybrid ]]; then
	if [[ $isosize -gt 4187593113 ]]; then
		if [[ $(blkid -s TYPE -o value "$2") == vfat ]];then
			echo "ISO size is larger than that allowed on vfat partition, exiting"
			echo "use grub supported partition format(ext3/4) that allow larger file size"
			exit 1
		fi
	fi
fi
if [[ $isoname != $isonametr ]]; then
	echo "removing space/s from iso name copied to the destination disk"
	export isoname=$isonametr
fi

if [[ ! -h /dev/disk/by-uuid/$stickuuid ]]; then
        echo "wrong uuid detected, please file bug with output of: blkid $2"
fi

if echo $isoname | grep -qi "Li-f-e" | grep -qi "suse"; then
	isonameshort="openSUSE-Li-f-e"
fi
are_you_sure ()  {
        echo  -n "$1 [$2/$3]? "
        while true; do
                read answer
                case $answer in
                        y | Y | yes | YES ) answer="y"; break;;
                        n | N | no | NO ) exit;;
                        *) echo "Please answer (y)es or (n)o.";;
                esac
        done
}
try_cp_with_progress () {
	sync
        if [ -x /usr/bin/dd_rescue ]; then
        	$benice dd_rescue -A -b 4M -y 4M "$iso_symlink" "$2"
	elif [ -x /usr/bin/ddrescue ]; then
		$benice ddrescue --force -b 4M "$iso_symlink" "$2"
	else
		echo "dd_rescue or ddrescue not found, please install one of them, using dd for now"
		$benice dd if="$iso_symlink" of="$2" bs=4M
	fi
}
cfg_setup() {
	mount -o loop $isopath $isomount &>/dev/null
        if [[ -d $stickmount/boot/grub2 ]]; then
                grub2path=$stickmount/boot/grub2
                grub2pathrealusb=/boot/grub2
        else
                grub2path=$stickmount/boot/grub
                grub2pathrealusb=/boot/grub
        fi
        grubcfgpath=$grub2path/grub.cfg
        if [[ ! -f $stickmount/grubstick ]]; then
                echo "copying grub files to $stickmount/boot/"
		if [[ -d "$isomount/boot/grub2" ]];then
			cp -r $isomount/boot/grub2/* $grub2path/
		elif [[ -d "$isomount/boot/grub" ]];then
			cp -r $isomount/boot/grub/* $grub2path/
		else
	                cp -r $grub2pathreal/* $grub2path/
		fi
        fi
	if [[ $distroname == suse ]]; then
		if [[ $PERSISTENT_IMAGE == true ]]; then
                        echo "Creating image file for persistent data"
			if ! qemu-img create $stickmount/boot/$isonameshort-cowfile $cowsize &>/dev/null;then
				echo "Can't create cow file"
			else
	                        persistopts="kiwi_cowdevice=/dev/disk/by-uuid/$stickuuid kiwi_cowsystem=/boot/$isonameshort-cowfile"
			fi
		fi
		if [[ -e $isomount/boot/x86_64/loader/linux ]]; then
			isoarch=x86_64
		else
			isoarch=ix86
		fi
		if [[ -e $isomount/LiveOS/squashfs.img ]]; then
                        appendsection="\$linux (loop)/boot/$isoarch/loader/linux iso-scan/filename=/$isoname root=live:CDLABEL=$isolabel $persistopts install=hd:/$isoname loader=syslinux splash=silent quiet showopts"
		else
			appendsection="\$linux (loop)/boot/$isoarch/loader/linux isofrom=/dev/disk/by-uuid/$stickuuid:/$isoname isofrom_device=/dev/disk/by-uuid/$stickuuid isofrom_system=/$isoname $persistopts install=hd:/$isoname loader=syslinux splash=silent quiet showopts"
		fi
		appendsection2="\$initrd (loop)/boot/$isoarch/loader/initrd"
	fi
	if [[ $distroname == ubuntu ]]; then
                if [[ $PERSISTENT_IMAGE == true ]]; then
			cowfolder=$isonameshort-cow
			mkdir -p $stickmount/$cowfolder
                        echo "Creating image file for persistent data"
                        if ! qemu-img create $stickmount/$cowfolder/casper-rw $cowsize &>/dev/null;then
                                echo "Can't create cow file"
			else
				echo "Creating filesystem for persistent data"
				mkfs.ext4 -F -O^has_journal -L casper-rw $stickmount/$cowfolder/casper-rw
				persistopts="persistent persistent-path=/$cowfolder/"
                        fi
		fi
		if [[ -e $isomount/casper/vmlinuz ]]; then
			vmlinuzfile=/casper/vmlinuz
		else
			vmlinuzfile=/casper/vmlinuz.efi
		fi
                if [[ -e $isomount/casper/initrd.lz ]]; then
                        initrdfile=/casper/initrd.lz
                else
                        initrdfile=/casper/initrd
                fi
                appendsection="linux (loop)/$vmlinuzfile boot=casper iso-scan/filename=/$isoname $persistopts quiet splash"
                appendsection2="initrd (loop)/$initrdfile"

	fi
	if [[ $distroname == fedora ]]; then
		appendsection="linux (loop)/isolinux/vmlinuz0 iso-scan/filename=/$isoname root=LABEL=$isolabel rootfstype=auto ro rd.live.image rd.luks=0 rd.md=0 rd.dm=0 quiet rhgb"
		appendsection2="initrd (loop)/isolinux/initrd0.img"
	fi
	if [[ $distroname == ipxe ]]; then
		appendsection="linux16 (loop)/ipxe.krn"
	fi
        if ! grep -q $isonameshort $grubcfgpath | grep -qi menuentry; then
                 echo "adding new image to boot menu"
                 add_menu_section
        fi
	umount $isomount &>/dev/null
}
grub2_stuff () {
	echo "installing grub2 on $stickdevice"
	mkdir -p $stickmount/boot
	$grubinstall --force --target=i386-pc --boot-directory=$stickmount/boot $stickdevice
	echo "setting $stickdevice partition $stickpart active"
	parted $stickdevice set $stickpart boot on &>/dev/null
	cfg_setup
}
add_harddisk_section () {
	if ! grep -q harddisk $grubcfgpath; then
	        cat <<EOF >>$grubcfgpath

menuentry 'Boot from harddisk' {
	set root=(hd1)
	chainloader +1
}

menuentry 'Reboot' {
	reboot
}

EOF
	fi
}
create_grub_cfg() {
cp /usr/share/grub*/unicode.pf2 $grub2path/ 2>/dev/null
export themename=$(echo /$grub2path/themes/*|rev | cut -d / -f1|rev)
        cat <<EOF >$grubcfgpath
insmod fat
insmod ext2
insmod part_gpt
insmod efi_gop
insmod iso9660
insmod chain
insmod linux
insmod echo
insmod configfile
insmod boot
insmod search_label
insmod search_fs_file
insmod search
insmod search_fs_uuid
insmod ls
insmod video
insmod video_fb
insmod normal
insmod test
insmod sleep
insmod png
insmod gettext
insmod gzio
insmod efi_uga

if [ \$grub_platform = "efi" ]; then
    set linux=linuxefi
    set initrd=initrdefi
else
    set linux=linux
    set initrd=initrd
fi
set default=0
set font=$grub2pathrealusb/unicode.pf2
set gfxmode=auto
insmod gfxterm
insmod gfxmenu
terminal_input gfxterm
if terminal_output gfxterm; then true; else
	terminal gfxterm
fi
set theme=$grub2pathrealusb/themes/$themename/theme.txt
export theme
set timeout=10

EOF

}
add_menu_section () {
        cat <<EOF >>$grubcfgpath

menuentry "$isonameshort" --class $distroname --class os {
	loopback loop /$isoname
	echo Loading linux...
	set gfxpayload=keep
	$appendsection
	echo Loading initrd...
	$appendsection2
}

EOF
add_harddisk_section
}

echo "Please make sure the following information is correct:"
echo "iso name: $isoname distro: $distroname stick device: $stickdevice"
echo "stick uuid: /dev/disk/by-uuid/$stickuuid stick partition: $stickpart"
if [[ ! -f $liveusbgui ]]; then
	are_you_sure "continue ?" "y" "n"
fi
if [[ $distroname == isohybrid ]];then
	try_cp_with_progress "$1" "$2"
else
	mkdir $isomount $stickmount &>/dev/null
	if mount | grep $stickdevpart &>/dev/null; then
		if ! umount $stickdevpart &>/dev/null; then
			echo "$stickdevpart in use, please umount it before continuing"
			exit 1
		fi
	fi
	if ! mount $2 $stickmount &>/dev/null; then
		echo "unable to mount the usb stick partition"
		exit 1
	fi
	if [[ -e $stickmount/$isoname ]];then
		isopath=$stickmount/$isoname
		echo "using existing $isoname on the stick"
	else
		isopath=$iso_symlink
		echo "copying $isoname to usb stick"
		try_cp_with_progress "$1" "$stickmount/$isoname"
		if [[ $isosize != $(stat -c%s "$stickmount/$isoname") ]]; then
			echo "Source and destination image does not match"
			exit 1
		fi
	fi
        if [[ -d /boot/grub2 ]]; then
                grub2pathreal=/boot/grub2
	else
		grub2pathreal=/boot/grub
	fi
	if [[ -f $stickmount/grubstick ]]; then
		echo "the stick is already bootable stick"
		cfg_setup
	else
		grub2_stuff
		echo "creating menu entries"
		create_grub_cfg
		add_menu_section
	fi
	touch $stickmount/grubstick
	thisscript=$(readlink -f $0)
	cp $thisscript $stickmount/ &>/dev/null || true
fi
clean_up