#!/bin/bash
# Gettext internationalization

#"BASED ON CODE OF BUILD_SLACKWARE-LIVE.SH FROM linux-nomad"
AUTHOR='Dimitris Tzemos - djemos@slackel.gr '
LICENCE='GPL v3+'
SCRIPT=$(basename "$0")
SCRIPT=$(readlink -f "$SCRIPT")
VER=1.0

version() {
  echo "Salix USB installer and persistent creator for 32 and EFI/UEFI 64 v$VER"
  echo " by $AUTHOR"
  echo "Licence: $LICENCE"
}
usage() {
  echo 'install-on-USB.sh [-h/--help] [-v/--version]'
  echo ' -h, --help: this usage message'
  echo ' -v, --version: the version author and licence'
  echo ''
  echo "`basename $0` --usb isoname 32|64"
  echo "`basename $0` --persistent 32|64"
  echo ''
  echo '-> Install syslinux on an USB key using an ISO and the USB key itself.'
  echo '-> Create a persistent ext3 file at installation (optional).'
  echo '-> Create a persistent ext3 file after installation.'
  exit 1
}

if [ "$1" = "--version" ] || [ "$1" = "-v" ]; then
version
  exit 0
fi
if [ "$1" = "--help" ] || [ "$1" = "-h" ]; then
usage
fi
#if [ $(id -ru) -ne 0 ]; then
#echo "Error : you must run this script as root" >&2
#  exit 2
#fi

LIVELABEL="LIVE" #edit in init script too
CWD=`pwd`
CMDERROR=1
PARTITIONERROR=2
FORMATERROR=3
BOOTERROR=4
INSUFFICIENTSPACE=5
persistent_file=""
SCRIPT_NAME="$(basename $0)"
NAME="persistent"
SVER=14.1

function check_root(){
# make sure I am root
if [ `id -u` != "0" ]; then
	echo "ERROR: You need to be root to run $SCRIPT_NAME"
	echo "The install_on_usb_uefi.sh script needs direct access to your boot device."
	echo "Use sudo or kdesudo or similar wrapper to execute this."
	exit 1
fi
}

function find_iso(){
isoname=$1
if [ -f "$isoname" ]
then
	isoname=$1
	iso_arch="${isoname##*/}"
	iso_arch="${iso_arch##*live}"
	iso_arch="${iso_arch%%-*}"
	if [ "$iso_arch" != "64" ]; then
		iso_arch=32
	fi
else
	echo "'Sorry, $isoname iso file does not exist"
	exit
fi
}

function check_if_file_iso_exists(){
isoname=$1
iso_arch=$2
if [ -f "$isoname" ]; then
	iso_archf="${isoname##*/}"
	iso_archf="${iso_archf##*live}"
	iso_archf="${iso_archf%%-*}"
	if [ "$iso_archf" != "64" ]; then
		iso_archf=32
	fi
else
	echo "Sorry, $isoname iso file does not exist"
	exit
fi

if [ "$iso_arch" != "$iso_archf" ]; then
	echo "Sorry, $isoname arch is $iso_archf while parameter is $iso_arch"
	echo "You provide the wrong iso image"
    exit
fi    
}

function find_usb(){
disks=`cat /proc/partitions | sed 's/  */:/g' | cut -f5 -d: | sed -e /^$/d -e /[1-9]/d -e /^sr0/d -e /loop/d -e /$bootdevice/d`
installmedia=""
for disk in $disks; do
 if [ "$disk" != "sda" ]; then
	installmedia=/dev/$disk
 fi
done
echo "========================================="
echo "										   "
echo "Removable device is $installmedia        "
echo "										   "
echo "========================================="

if [ "$installmedia" == NULL ]; then
	echo "There is no removable usb device attached to your system"
	exit
fi
}

function usb_message(){
MSG="ATTENTION!!!!!\n\
This script is going to install liveiso to your $installmedia device\n\
This will take some time. So please be patient ..\n\
If you continue, your drive $installmedia will be erased and repartitioned.\n\
All data will be lost.\n\
\n\
\n\
ARE YOU SURE YOU KNOW WHAT YOU'RE DOING?"

dialog --title "Are you sure you want to do this?" \
	--defaultno \
	--yesno "$MSG" 0 0
retval=$?
if [ $retval -eq 1 ] || [ $retval -eq 255 ]; then
	exit 0
fi
}

function persistent_message(){
MSG="Do you want to create a persistent file on your drive $installmedia ?\n\
\n\
"

dialog --title "Create a Persistent file" \
	--defaultno \
	--yesno "$MSG" 0 0
retval=$?
if [ $retval -eq 1 ] || [ $retval -eq 255 ]; then
	persistent_file="no"
else	
    persistent_file="yes"
    
    answer="$(eval dialog --title \"Select size in MB\" \
	--stdout \
	--menu \"Select the size of persistent file:\" \
	0 0 0 \
	'100' 'o' \
	'300' 'o' \
	'400' 'o' \
	'500' 'o' \
	'700' 'o' \
	'800' 'o' \
	'1000' 'o' \
	'1500' 'o' \
	'2000' 'o' \
	'2500' 'o' \
	'3000' 'o' \
	'3500' 'o' \
	'3998' 'o')"
	retval=$?
	if [ $retval -eq 1 ] || [ $retval -eq 255 ]; then
		persistent_file="no"
	else
		SIZE=$answer
	fi  
fi
}
    
function create_persistent(){
mkdir -p /mnt/install
mount $installmedia /mnt/install
CWD=`pwd`
cd /mnt/install/
FREE=$('df' -k .|tail -n1|awk '{print $4}')
AFTER=$(( $FREE - 1024 * $SIZE ))
if [ $AFTER -gt 0 ]; then
	echo "Creating persistent file 'persistent'. Please wait ..."
	echo ""
	dd if=/dev/zero of="$NAME" bs=1M count=$SIZE
	mkfs.ext3 -F -m 0 -L "persistent" "$NAME" && CHECK='OK'
	if [ -n "$CHECK" ]; then
		echo ""
		echo "The persistent file $NAME is ready."
		cd $CWD
		sleep 5
		umount /mnt/install
		rmdir /mnt/install
	else
		echo "ERROR: $SCRIPT_NAME was not able to create the persistent file $NAME"
		cd $CWD
		sleep 5
		umount /mnt/install
		rmdir /mnt/install	
		exit 1
	fi
else
  echo "ERROR: There is not enough free space left on your device \
for creating the persistent file $NAME"
	cd $CWD
	sleep 5
	umount /mnt/install
	rmdir /mnt/install	
  exit 1
fi
}
    
function install_usb() {
	    device=`echo $installmedia | cut -c6-8`
		partition=`echo $installmedia | cut -c6-`
		sectorscount=`cat /sys/block/$device/subsystem/$partition/size`
		sectorsize=`cat /sys/block/$device/queue/hw_sector_size`
		let destinationsize=$sectorscount*$sectorsize/1048576
		installdevice=`echo $installmedia | cut -c1-8`
	 
	 
	 if [ "$installdevice" == "$installmedia" ]; then #install on whole disk: partition and format media
		partnum=2
		installmedia="$installdevice$partnum"
		efipartition="$installdevice""1"
		dd if=/dev/zero of=$installdevice bs=512 count=34 >/dev/null 2>&1
		#hybrid MBR with BIOS boot partition (1007K) EFI partition (32M) and live partition
		#echo -e "2\nn\n\n\n+32M\nef00\nn\n\n\n\n\nn\n128\n\n\nef02\nr\nh\n1 2\nn\n\ny\n\nn\nn\nwq\ny\n" | gdisk $installdevice || return $PARTITIONERROR
		echo -e "2\nn\n\n\n+32M\nef00\nn\n\n\n\n0700\nn\n128\n\n\nef02\nr\nh\n1 2\nn\n\ny\n\nn\nn\nwq\ny\n" | gdisk $installdevice || return $PARTITIONERROR
		partprobe $installdevice; sleep 3
		mkfs.fat -n "efi" $efipartition || return $FORMATERROR
		#mkfs.ext3 -L "$LIVELABEL" $installmedia || return $FORMATERROR
		fat32option="-F 32"
		mkfs.vfat $fat32option -n "$LIVELABEL" $installmedia || return $FORMATERROR
		sleep 3
		
	else #install on partition: filesystem check and format if needed
		partnum=`echo $installmedia | cut -c9-`
		mkdir -p /mnt/tmp
		if mount $installmedia /mnt/tmp >/dev/null 2>&1; then #filesystem check
			sleep 1
			umount /mnt/tmp
			fsck -fy $installmedia >/dev/null 2>&1
		else #format partition
			if echo -e "print\nquit" | parted $installdevice | grep -q ": msdos$"; then
				partitiontype=`gdisk -l $installdevice | grep "^  *$partnum " | sed 's/  */:/g' | cut -f7 -d:`
			else
				partitiontype=`fdisk -l $installdevice 2>/dev/null | grep "^$installmedia " | sed -e 's/\*//' -e 's/  */:/g' | cut -f5 -d:`
			fi
			case $partitiontype in
			83|8300) 
				mkfs.ext3 -L "$LIVELABEL" $installmedia || return $FORMATERROR
				;;
			*)
				partition=`echo $installmedia | cut -c6-`
				size=`cat /proc/partitions | grep " $partition$" | sed 's/  */:/g' | cut -f4 -d:`
				let size=$size/1024
				if (( $size > 1024 )); then
					fat32option="-F 32"
				fi
				mkfs.fat $fat32option -n "$LIVELABEL" $installmedia || return $FORMATERROR
			esac
			sleep 3
		fi
		
		#efi partition detection + filesystem check and format if needed:
		efipartnum=`gdisk -l $installdevice 2>/dev/null | grep " EF00 " | sed 's/  */:/g' | cut -f2 -d:`
		if [ ! -z "$efipartnum" ]; then
			efipartition="$installdevice$efipartnum"
			if [ "$efipartition" != "$installmedia" ]; then
				if mount $efipartition /mnt/tmp >/dev/null 2>&1; then
					sleep 1
					umount /mnt/tmp
					fsck -fy $efipartition >/dev/null 2>&1
				else
					mkfs.fat -n  "efi" $efipartition || return $FORMATERROR
					sleep 3
				fi
			fi
		fi
	fi
	
	#copy
	mkdir -p /mnt/install
	mount $installmedia /mnt/install
	cp -r $livedirectory/boot /mnt/install/
	grubstatus=$GRUBERROR
	
	#grub
	#if [ `uname -m` == "x86_64" ] && [ ! -z "$efipartnum" ]; then #EFI
	if [ "$iso_arch" == "64" ] && [ ! -z "$efipartnum" ]; then #EFI
		mkdir -p /mnt/efi
		mount $efipartition /mnt/efi
		mkdir -p /mnt/efi/EFI/BOOT/
		#~ uuid=`blkid | grep "^$installmedia:" | sed 's/.*UUID=\"\([^\"]*\)\".*/\1/'`
		cat > /tmp/grub.$$.cfg << EOF
#search --fs-uuid --set=root $uuid
search --label --set=root "$LIVELABEL"
set prefix=(\$root)/boot/grub
source \$prefix/grub.cfg
EOF
		grub-mkimage -p "" -o "/mnt/efi/EFI/BOOT/bootx64.efi" -O x86_64-efi -c /tmp/grub.$$.cfg \
			search part_gpt fat ext2 normal boot configfile linux multiboot all_video font gfxterm png jpeg gfxmenu && grubstatus=0
		rm -f /tmp/grub.$$.cfg
		umount /mnt/efi
		rmdir /mnt/efi
	fi
	
	if echo -e "print\nquit" | parted $installdevice | grep -q ": msdos$" || 
	gdisk -l $installdevice 2>/dev/null | grep -q " EF02 "; then #CSM
		sfdisk $installdevice -A $partnum 2>/dev/null
		grub-install --force --target=i386-pc --boot-directory=/mnt/install/boot $installdevice && grubstatus=0
	fi
	
	umount /mnt/install
	rmdir /mnt/install
	umount /tmp/iso
	rmdir /tmp/iso
	
	return $grubstatus
}

action=$1
case $action in
"--usb")
    isoname=$2
    iso_arch=$3
    check_root
 if  [ "$iso_arch" == "32" ] || [ "$iso_arch" == "64" ]; then   
	check_if_file_iso_exists $isoname $iso_arch
	find_usb
	usb_message
	mkdir -p /tmp/iso
	mount -o loop $isoname /tmp/iso > /dev/null 2>&1
	livedirectory=/tmp/iso/
    if [ -f "$isoname" ] && [ -b "$installmedia" ]; then
		livesystemsize=`du -s -m $livedirectory | sed 's/\t.*//'`
		device=`echo $installmedia | cut -c6-8`
		partition=`echo $installmedia | cut -c6-`
		sectorscount=`cat /sys/block/$device/subsystem/$partition/size`
		sectorsize=`cat /sys/block/$device/queue/hw_sector_size`
		let destinationsize=$sectorscount*$sectorsize/1048576
		if (( $livesystemsize > $destinationsize)); then 
			echo "error: insufficant space on device '$installmedia'"
			umount /tmp/iso
			rm -rf /tmp/iso
			exit $INSUFFICIENTSPACE
		else
			persistent_message
			install_usb $livedirectory $installmedia
			if [ "$persistent_file" == "yes" ]; then
			 create_persistent
			fi 
			exit $!
		fi
	else
		echo "`basename $0` --usb iso_name 32|64"
		exit $CMDERROR
	fi
else
	echo "`basename $0` --usb iso_name 32|64"
		exit $CMDERROR
fi		
;;
	
"--persistent")
	iso_arch=$2
if  [ "$iso_arch" == "32" ] || [ "$iso_arch" == "64" ]; then
	find_usb
	device=`echo $installmedia | cut -c6-8`
	sectorscount=`cat /sys/block/$device/size`
	sectorsize=`cat /sys/block/$device/queue/hw_sector_size`
	let mediasize=$sectorscount*$sectorsize/1048576 #in MB
	#installdevice="/dev/$device"
	installdevice=`echo $installmedia | cut -c1-8`
	#if [ `uname -m` == "x86_64" ] && [ "$iso_arch" == "64" ]; then #EFI/GPT
	#if [ "$iso_arch" == "64" ]; then #EFI/GPT
	if [ "$installdevice" == "$installmedia" ]; then 
			#partitionnumber=2
			#installmedia="$installdevice$partitionnumber"
			partnum=2
			installmedia="$installdevice$partnum"
	else #BIOS/MBR
			#partitionnumber=4
			installmedia="$installdevice$partitionnumber"
			partnum=`echo $installmedia | cut -c9-`
			installmedia="$installdevice$partnum"
	fi		
         echo $installmedia
		persistent_message
		if [ "$persistent_file" == "yes" ]; then
			 create_persistent
		fi 
			exit $!
else
		echo "`basename $0` --persistent 32|64"
		exit $CMDERROR
fi
;;
"--version")
	version
    exit 0
	;;
"--help")
usage
exit 0
;;	
	
*)	echo "`basename $0` --persistent 32|64"
	echo "`basename $0` --usb isoname 32|64"
	exit $CMDERROR
	;;
esac