#!/bin/bash
#usbinstall.sh
# Arg 1: usb device for installation
# Arg 2: iso file
# Arg 3: option
#~ # (c) 2009 ercole 'ercolinux' carpanetto - ercole69@gmail.com
#~ # (c) 2009-2022 Silvan Calarco - silvan.calarco@mambasoft.it
# Released under the terms of the GPL v3 license

TEXTDOMAINDIR=/usr/share/locale/
TEXTDOMAIN=usbinstall

# Set progressbar (only if $DBUS_USBINSTALL = 1)
# Arg 1: value%
set_pbar()
{
	if [ "$DBUS_USBINSTALL" == "1" ]; then
		dbus-send --system --print-reply --dest=org.openmamba.usbinstall /org/openmamba/usbinstall org.openmamba.usbinstall.Set_progressBar int32:$1 &>/dev/null
		if [ "$?" != "0" ]; then
			[ "$CPPID" ] && kill -9 $CPPID
			exit 1
		fi
	fi
}

# Set status label (only if $DBUS_USBINSTALL = 1)
# Arg 1: string
set_status()
{
	if [ "$DBUS_USBINSTALL" == "1" ]; then
		dbus-send --system --print-reply --dest=org.openmamba.usbinstall /org/openmamba/usbinstall org.openmamba.usbinstall.Set_description "string:$1" &>/dev/null
		if [ "$?" != "0" ]; then
			[ "$CPPID" ] && kill -9 $CPPID
			exit 1
		fi
	fi
}

# Terminate the frontend (only if $DBUS_USBINSTALL = 1)
exit_frontend()
{
	
	if [ "$DBUS_USBINSTALL" == "1" ]; then
		dbus-send --system --print-reply --dest=org.openmamba.usbinstall /org/openmamba/usbinstall org.openmamba.usbinstall.Exit "string:$1" &>/dev/null
		if [ "$?" != "0" ]; then
			[ "$CPPID" ] && kill -9 $CPPID
			exit 1
		fi
	fi
}

####################################(DBUS)###########################################


function errorAndExit()
{
	echo "$1"
	# Se l'argomento 2 è settato stampa anche l'help
	if [ -n "$2" ]; then
		usage
	fi
        if [ "${SELECT_USB1}" ]; then
           udisksctl unmount -b ${SELECT_USB1} >/dev/null
        fi
        if [ "${SELECT_USB2}" ]; then
           udisksctl unmount -b ${SELECT_USB2} >/dev/null
        fi
	exit_frontend "$1" # Termina il frontend con codice di errore 1
	exit 1
}

function usage()
{
	echo $"usbinstall: install openmamba into your usb key!"
	echo -e "\t"$"Usage: sudo usbinstall Arg1 Arg2 Arg3 [-y]"
	echo -e "\t"$"Arg 1: Usb device path (e.g. /dev/sdc)"
	echo -e "\t"$"Arg 2: ISO file (e.g. ~/openmamba-livecd-it-milestone2-2.0pre8.i586.iso)"
	echo -e "\t"$"Arg 3: "$"-d: pass message through dbus (when invoked by GUI)"
	echo -e "\t       "$"-y: install without asking confirmation [Be careful!]"
}


#=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~
#--------------------------------(INITIAL-CHECK)-------------------------------------
#=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~
# Confirm
if [ "$3" = "-y" ]; then
	CONFIRM=0
elif [ "$3" = "-d" ]; then
	CONFIRM=0
	DBUS_USBINSTALL=1	# First set this
else
	CONFIRM=1
fi

# Help
if [ "$1" = "-h" -o "$1" = "--help" ]; then
	usage
	exit 0
fi

# Check if usbinstall have root permission
if [ "`id -u`" != 0 ]; then
	errorAndExit $"Error: usbinstall must be run as root" 1
fi

# Check if Arg 1 is valid
if [ -b "$1" ]; then	# If is a block device
	ISKEY=`udisksctl info -b $1 | grep HintSystem | sed "s|.*[[:space:]]||"`
	if [ "$ISKEY" = "true" -o "$CONFIRM" = "0" ]; then	# If is a usb
		export SELECT_USB=$1
	else # If not, user can choose
		echo $"Warning: you haven't choosen a removable device ($1). Are you sure to continue? [yes/no]"
		read KEY
		case "$KEY" in
			$yes)
				export SELECT_USB=$1
				;;
			*)
				errorAndExit $"Installation aborted"
				;;
		esac
	fi
else
	errorAndExit $"Error(Arg 1): \"$1\" isn't a block device" 1
fi

if [ "${SELECT_USB:0:9}" = "/dev/loop" ]; then
   SELECT_USB1=${SELECT_USB}p1
   SELECT_USB2=${SELECT_USB}p2
else
   SELECT_USB1=${SELECT_USB}1
   SELECT_USB2=${SELECT_USB}2
fi

# Check if Arg 2 is valid
if [ -r "$2" ]; then	# Check it is readable
   # Check if Arg 2 is an iso file
   if [ "${2##*.}" != "iso" ]; then
      errorAndExit "Error(Arg 2): you must chose a ISO (.iso) file" 1
   fi
   SELECT_ISO=`readlink -f $2`
else
   errorAndExit "Error(Arg 2): \"$2\" is not readable" 1
fi

# Check if Arg 3 is valid
if [ "$3" ]; then
   if [ "$3" != "-d" -a "$3" != "-y" ]; then	# Check it is readable
      errorAndExit $"Error(Arg 3): \"$3\" is not -d nor -y" 1
   fi
fi

#################################(INITIAL-CHECK)#####################################




#=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~
#--------------------------------(INSTALLATION)--------------------------------------
#=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~
export LOGFILE=/var/log/usbinstall.log
trap "[ \"$CPPID\" ] && kill -9 $CPPID" HUP INT QUIT ABRT KILL TERM
set_status $"Unmount usb device"
echo $"Unmount usb device"

DEVICEMOUNTPOINTS=`udisksctl info -b ${SELECT_USB2} | grep MountPoints: | sed "s|.*[[:space:]]||"`

if [ "$DEVICEMOUNTPOINTS" ]; then
   LANG=C udiskctl unmount -b ${SELECT_USB2} || { 
	errorAndExit $"Error: fail to umount the usb key (${SELECT_USB2})"
   }
fi
set_pbar 2

# Prima di formattare :P
if [ "$CONFIRM" = "1" ]; then
	echo $"Warning: You are going to format this device ($1). Are you sure to continue? [yes/no]"
	read KEY
		case "$KEY" in
			$"yes"|"Y"|"y")
				;;
			*)
				errorAndExit $"Installation aborted"
				;;
		esac
fi

set_status $"Partitioning $SELECT_USB..."
echo $"Partitioning $SELECT_USB..."
# Wipe any old GPT table or grub-install may fail due to multiple partition labels
sudo dd if=/dev/zero of=$SELECT_USB bs=512 count=2047 skip=1
# Start from a new partition table
sudo parted -s $SELECT_USB mktable msdos &>>$LOGFILE
if [ "$?" != "0" ]; then
   errorAndExit $"Error:"" \"parted\" "$"exit with error (see $LOGFILE)"
fi
# EFI partition
sudo parted -s -a optimal -- $SELECT_USB mkpart primary fat32 1M 20M &>>$LOGFILE
if [ "$?" != "0" ]; then
   errorAndExit $"Error:"" \"parted\" "$"exit with error (see $LOGFILE)"
fi
# openmamba partition
sudo parted -s -a optimal -- $SELECT_USB mkpart primary ext4 20M -1 &>>$LOGFILE
if [ "$?" != "0" ]; then
   errorAndExit $"Error:"" \"parted\" "$"exit with error (see $LOGFILE)"
fi
# boot flag on openmamba partition
sudo parted -s $SELECT_USB set 2 boot on &>>$LOGFILE
if [ "$?" != "0" ]; then
   errorAndExit $"Error:"" \"parted\" "$"exit with error (see $LOGFILE)"
fi

set_pbar 5

set_status $"Formatting $SELECT_USB..."
echo $"Formatting $SELECT_USB..."
sudo mkfs.vfat ${SELECT_USB1} -n EFI &>>$LOGFILE
if [ "$?" != "0" ]; then
   errorAndExit $"Error:"" \"mkfs.vfat\" "$"exit with error (see $LOGFILE)"
fi
sudo mkfs.ext4 ${SELECT_USB2} -L openmamba_live &>>$LOGFILE
if [ "$?" != "0" ]; then
   errorAndExit $"Error:"" \"mkfs.ext4\" "$"exit with error (see $LOGFILE)"
fi

sleep 1
LANG=C udisksctl mount -b ${SELECT_USB1} &>>$LOGFILE
if [ "$?" != "0" ]; then
   errorAndExit $"Error:"" \"udisks --mount\" "$"exit with error (see $LOGFILE)"
fi
LANG=C udisksctl mount -b ${SELECT_USB2} &>>$LOGFILE
if [ "$?" != "0" ]; then
   errorAndExit $"Error:"" \"udisks --mount\" "$"exit with error (see $LOGFILE)"
fi
set_pbar 30

MOUNTDIR_EFI=`udisksctl info -b ${SELECT_USB1} | grep -m1 MountPoints: | sed "s|.*[[:space:]]||"`
if [ ! -d "$MOUNTDIR_EFI" ]; then
   errorAndExit $"Error: unable to found usb mount path"
fi

# create EFI grub 32 and 64 bit images
mkdir -p $MOUNTDIR_EFI/EFI/BOOT/
grub-mkimage -o $MOUNTDIR_EFI/EFI/BOOT/bootx64.efi -O x86_64-efi -p "(,msdos2)/boot/grub" \
   part_gpt part_msdos ntfs ntfscomp hfsplus fat ext2 normal chain boot linux echo \
   help gfxterm gettext png efi_gop efi_uga search search_label search_fs_uuid || {
   errorAndExit $"Error: unable to create GRUB x86_64-efi image" &>>$LOGFILE
}

# 32bit EFI optional because untested and curently missing in x86_64 arch
if [ -d /usr/lib/grub/i386-efi ]; then
      grub-mkimage -o $MOUNTDIR_EFI/EFI/BOOT/bootia32.efi -O i386-efi -p "(,msdos2)/boot/grub" \
      part_gpt part_msdos ntfs ntfscomp hfsplus fat ext2 normal chain boot linux echo \
      help gfxterm gettext png efi_gop efi_uga search search_label search_fs_uuid || {
      errorAndExit $"Error: unable to create GRUB i386-efi image" &>>$LOGFILE
   }
fi

MOUNTDIR=`udisksctl info -b ${SELECT_USB2} | grep -m1 MountPoints: | sed "s|.*[[:space:]]||"`
if [ ! -d "$MOUNTDIR" ]; then
   errorAndExit $"Error: unable to found usb mount path"
fi

pushd $MOUNTDIR >/dev/null

sudo chmod 777 ./
if [ "$SELECT_BOOT" ]; then
   set_status $"Extracting bootusb archive..."
   echo $"Extracting bootusb archive..."
   gunzip -c $SELECT_BOOT | cpio -i
   if [ ! -d "./boot" ]; then
      errorAndExit $"Error: Invalid bootusb archive"
   fi
else
   set_status $"Extracting bootloader data from ISO image..."
   echo $"Extracting bootloader data from ISO image..."
   ISOMOUNTDIR=`mktemp -d --suffix=.usbinstall`
   mount -o loop,ro $SELECT_ISO $ISOMOUNTDIR || {
      errorAndExit $"Error: unable to mount ISO image"
   }

   mkdir LiveOS boot rwroot ovlwork
   cp $ISOMOUNTDIR/boot/{initrd*,vmlinuz*} boot/

   set_status $"Copying SquashFS image from ISO image..."
   echo $"Copying SquashFS image..."
   # Copy code
   ORIG_SIZE=`stat -c %s $ISOMOUNTDIR/LiveOS/squashfs.img`
   cp $ISOMOUNTDIR/LiveOS/squashfs.img LiveOS/ &
   export CPPID=$!
   DEST_SIZE=0
   DEST_FILE=LiveOS/squashfs.img
   ps $CPPID >/dev/null
   while [ $? -eq 0 ]; do
       if [ -e $DEST_FILE ]; then
           DEST_SIZE=$(stat -c %s $DEST_FILE)
           percentuale=$((30 + ( 50 * $DEST_SIZE ) / $ORIG_SIZE ))
           set_pbar $percentuale
       fi
       sleep 2
       ps $CPPID >/dev/null
   done

   if [ -e $ISOMOUNTDIR/boot/isolinux/isolinux.cfg ]; then
       # basic extlinux.conf -> grub.cfg conversion
       # Legacy for old openmamba iso images with isolinux bootloader
       echo "\
#
# GRUB configuration generated by openmamba usbinstall tool
#
font=unicode
set gfxmode=auto
insmod all_video
insmod gfxterm
insmod gettext
terminal_output=gfxterm
timeout=10
# theme
insmod=png
set menu_color_normal=white/black
set menu_color_highlight=white/green
search --no-floppy --label --set=root openmamba_live

#
# Boot entries taken from ISO bootloader
#
" > $MOUNTDIR_EFI/EFI/BOOT/grub.cfg
      while read line; do
         case $line in
            "MENU LABEL "*) [ "$CURR_BRACKET_OPEN" = "1" ] && echo "}" >> $MOUNTDIR_EFI/EFI/BOOT/grub.cfg
                           echo "menuentry \"${line/MENU LABEL }\" {" >> $MOUNTDIR_EFI/EFI/BOOT/grub.cfg
                           CURR_BRACKET_OPEN=1
                           ;;
            "kernel "*) CURR_KERNEL="${line/kernel }"
                        echo "echo \"Loading ${CURR_KERNEL}\"" >> $MOUNTDIR_EFI/EFI/BOOT/grub.cfg
                        ;;
            "append "*) set -- ${line/append initrd=}
                        CURR_INITRD="$1"
                        CURR_CMDLINE="${line/append initrd=$CURR_INITRD}"
                        CURR_CMDLINE="${CURR_CMDLINE/vga=792}"
                        echo "linux ${CURR_KERNEL} ${CURR_CMDLINE}" >> $MOUNTDIR_EFI/EFI/BOOT/grub.cfg
                        echo "echo \"Loading ${CURR_INITRD}\"" >> $MOUNTDIR_EFI/EFI/BOOT/grub.cfg
                        echo "initrd ${CURR_INITRD}" >> $MOUNTDIR_EFI/EFI/BOOT/grub.cfg
                        echo "}" >> $MOUNTDIR_EFI/EFI/BOOT/grub.cfg
                        CURR_BRACKET_OPEN=0
                        ;;
         esac
      done < $ISOMOUNTDIR/boot/isolinux/isolinux.cfg
      [ "$CURR_BRACKET_OPEN" = "1" ] && echo "}" >> $MOUNTDIR_EFI/EFI/BOOT/grub.cfg
      sed -i "s|root=live:LABEL=.*splash|root=live:LABEL=openmamba_live rd.live.overlay=/dev/disk/by-label/openmamba_live:/rwroot selinux=0 quiet splash rw|" $MOUNTDIR_EFI/EFI/BOOT/grub.cfg

      mkdir boot/grub
      cp $MOUNTDIR_EFI/EFI/BOOT/grub.cfg boot/grub/grub.cfg
   else
      cp -a $ISOMOUNTDIR/boot/grub boot/
      sed -i "s|root=live:LABEL=[^[:space:]]*|root=live:LABEL=openmamba_live rd.live.overlay=/dev/disk/by-label/openmamba_live:/rwroot|" boot/grub/grub.cfg
   fi
   umount $ISOMOUNTDIR
   rmdir $ISOMOUNTDIR
fi
set_pbar 85

set_status $"Writing bootloader..."
echo $"Writing bootloader..."
# Install mbr bootloder (i386-pc)
grub-install --modules="part_msdos part_gpt ext2" --removable --target=i386-pc --boot-directory=boot/ $SELECT_USB &>>$LOGFILE || {
      errorAndExit $"Error: failed to install bootloader into device (${SELECT_USB})"
}
set_pbar 90

popd >/dev/null
sync
set_pbar 95

set_status $"Unmount usb device"
echo $"Unmount usb device"
LANG=C udisksctl unmount -b ${SELECT_USB1} &>>$LOGFILE || {
	errorAndExit $"Error: fail to umount the usb key (${SELECT_USB1})"
}
LANG=C udisksctl unmount -b ${SELECT_USB2} &>>$LOGFILE || {
	errorAndExit $"Error: fail to umount the usb key (${SELECT_USB2})"
}

## WARNING: calling install-mbr after unmounting to prevent problems of FAT since EFI partition introduction
#install-mbr $SELECT_USB &>>$LOGFILE

set_pbar 100
exit_frontend ""
echo $"Done!"
#################################(INSTALLATION)######################################
