#!/bin/dash
# UEFI Installer
# Copyright (C) James Budiono 2017, 2018
# License: GNU GPL Version 3 or later
#
# Jan 2018 - step: sparse gettexting

### configuration
TEXTDOMAIN=fatdog OUTPUT_CHARSET=UTF-8
export TEXTDOMAIN OUTPUT_CHARSET
. gettext.sh
APPTITLE="$(gettext 'UEFI Installer')"
APPVERSION=0.2
MAX_MBR_PART=4    # highest MBR partition number
MAX_GPT_PART=100  # actually unlimited, but we only check this far
REQUIRED_SIZE=512 # in MB, minimum free size to be able to install Fatdog

### run-time variables
#DRY_RUN=    # if non blank, nothing will be changed
DEVICES=     # filled by get_device_info
PARTITIONS=  # filled by get_device_info
USBDEV=      # filled by get_device_info
GPTDEV=      # filled by get_device_info
MBRDEV=      # filled by get_device_info
BLKID_CACHE= # blkid is slow, cache its output

TARGET_DEV=
TARGET_PART= # "new", "flp", or partition-to-install-to
TARGET_PART_NEW= # new partition when "new"
TARGET_PART_NEW_NUMBER= # partition number when "new"
INSTALL_SRC= # "usb", "iso", or "disc"
INSTALL_DEV= # devices or filepath to source
INSTALL_MNT=/tmp/uefi-installer.$$ # source mountpoint

### system caps
SORT="sort" # default

trap 'rmdir $INSTALL_MNT/* $INSTALL_MNT 2>/dev/null; exit' 0 INT HUP TERM

################# helper routines #################

### step movements
step=
last_cond=
prev_step=
go_step() { step=$1; }
go_exit() { step=exit; }
go_cancel() { step=cancel; }

# $1-next, $2-prev, $3-condition
next_wizard_step() {
	local next=$1 prev=$2 cond=$?
	[ "$3" ] && cond=$3
	prev_step=$step
	case $cond in
		0)  # Xdialog next
			go_step $next ;;
		3)  # Xdialog previous
			go_step $prev ;;
		1|255) # Xdialog Esc/window kill
			go_cancel ;;	
	esac
	last_cond=$cond
}
# $1-next, $2-prev, $3-condition
next_yesno_step() {
	local next=$1 prev=$2 cond=$?
	[ "$3" ] && cond=$3
	prev_step=$step
	case $cond in
		0)  # Xdialog yes
			go_step $next ;;
		1)  # Xdialog no
			go_step $prev ;;
		255) # Xdialog Esc/window kill
			go_cancel ;;	
	esac
	last_cond=$cond
}

system_capabilities() {
	echo | sort -V > /dev/null 2>&1 && SORT="sort -V"
}

### get available devices
# out: DEVICES USBDEV GPTDEV MBRDEV PARTITIONS
find_devices() {
	local p
	# find devices
	DEVICES=$(cd /sys/block; ls | grep -vE "loop|nbd|ram|sr" | $SORT)
	USBDEV=
	for p in $DEVICES; do
		readlink /sys/block/$p | grep -q usb && USBDEV="$USBDEV $p"
	done
	GPDDEV=
	for p in $DEVICES; do
		fdisk -l /dev/$p | grep -q gpt && GPTDEV="$GPTDEV $p"
	done
	MBRDEV=
	for p in $DEVICES; do
		fdisk -l /dev/$p | grep -q dos && MBRDEV="$MBRDEV $p"
	done
	PARTITIONS=$(cd /sys/class/block; ls | grep -vE "loop|nbd|ram|sr" | $SORT)
	BLKID_CACHE="$(blkid)"
	#echo devices $DEVICES
	#echo usb $USBDEV
	#echo gpt $GPTDEV
	#echo dos $MBRDEV
	#echo partitions $PARTITIONS	
}

# $1-device
is_device_usb() {
	case "$USBDEV" in
		*${1}*) return 0;
	esac
	return 1
}

# $1-device
is_device_gpt() {
	case "$GPTDEV" in
		*${1}*) return 0;
	esac
	return 1
}

# $1-device
get_volabel() {
	local mine
	mine=$(echo "$BLKID_CACHE" | sed "\_/dev/$1_!d;")
	echo "$mine" | grep -q " LABEL=" &&
	echo "$mine" | sed 's/.* LABEL="//;s/".*//'
}

# $1-device
get_fs() {
	local mine
	mine=$(echo "$BLKID_CACHE" | sed "\_/dev/$1_!d;")	
	echo "$mine" | sed 's/.* TYPE="//;s/".*//'
}

# $1-device
get_target_device_details() {
	local dev="$1"
	is_device_usb $1 && dev="$dev (USB)" || dev="$dev (Internal)"
	is_device_gpt $1 && dev="$dev (GPT)" || dev="$dev (MBR)"
	echo $dev
}
#$1-no blank => usb only
# out: radiolist for Xdialog
# note: list USB devices first
get_target_devices_list() {
	local p
	for p in $USBDEV; do
		printf '%s "%s" off\n' $p "$(get_target_device_details $p)"
	done
	[ "$1" ] && return
	for p in $DEVICES; do
		case "$USBDEV" in
			*$p*) ;; # already done
			*) printf '%s "%s" off\n' $p "$(get_target_device_details $p)"
		esac
	done
}

# $1-device
get_target_part_details() {
	local part=$1
	part="$part ($(get_volabel $1)) ($(get_fs $1))"
	echo $part
}
# $1=dev, out: radiolist for Xdialog
get_target_part_list() {
	local p
	for p in $PARTITIONS; do
		case $p in 
			${1}) continue ;; # ignore own device
			*${1}*) printf '%s "%s" off\n' $p "$(get_target_part_details $p)"
		esac
	done
}

# out: radiolist for Xdialog
get_optical_list() {
	cd /sys/block
	ls -d sr* | $SORT | while read -r p; do
		echo "$p $p off"
	done
}


### wizard actions
introduction() {
	Xdialog --title "$APPTITLE" --no-cancel --textbox - 0 0 << EOF
Fatdog64 Installer for UEFI-based systems version $APPVERSION

This installer is meant to install Fatdog64 to USB flash drive.
Please make sure that your target flash driver is **UN-MOUNTED**.

The installation will only work on UEFI systems. If you need to
install to BIOS systems, please use standard Fatdog64 Installer.

==========================================
Warning: As with any Operating System installer, using this software
can cause data loss. *BACKUP* your data first. 
When you proceed, you agree that you are using it at your own risk.
==========================================

EOF
	go_step choose_devices
}

choose_devices() {
	## note: list USB devices first
	find_devices
	TARGET_DEV=$(eval Xdialog --title "'$APPTITLE'" --wizard --stdout --no-tags \
	--radiolist "'Choose device to install to:'" 00 00 5 \
	$(get_target_devices_list))
	next_wizard_step check_usb intro
}

check_is_usb_device() {
	# warn if not
	if ! is_device_usb $TARGET_DEV; then
		Xdialog --title "$APPTITLE" --yesno \
"\"$TARGET_DEV\" is not a USB device. Installing to an internal device
is possible **BUT** it is unsupported, and risky. 

If you already have an operating system on there, it **MAY FAIL** 
to boot after installation. If not sure, please don't proceed.

Are you sure you want to proceed?" 0 0
		next_yesno_step choose_part choose_devices
	else 
		go_step choose_part
	fi
}

choose_partitions() {
	TARGET_PART=$(eval Xdialog --title "'$APPTITLE'" --wizard --stdout --no-tags \
	--radiolist "'From device \"$TARGET_DEV\", choose partition to install to:'" 00 00 5 \
	$(get_target_part_list $TARGET_DEV) \
	"new" "'Make a new partition on $TARGET_DEV'" off \
	"flp" "'Superfloppy mode on $TARGET_DEV'" off)
	next_wizard_step check_part choose_devices
	[ "$TARGET_PART" = "flp" ] && TARGET_PART=$TARGET_DEV
}

check_partition_type() {
	# must be FAT, skip for "new"
	if [ "$(guess_fstype /dev/$TARGET_PART)" = vfat ] || [ "$TARGET_PART" = new ]; then
		go_step check_new_part
		return 
	fi
	
	Xdialog --title "$APPTITLE" --infobox \
"\"$TARGET_PART\" is not a FAT partition. UEFI needs to be installed in 
a FAT partition. Please choose to install in a new partition, 
or quit this installer, re-format your \"$TARGET_PART\", and re-run 
this installer again." 0 0 10000
	go_step choose_part
}

check_make_partition() {
	# only for "new"
	# if partition is DOS, see if we can make another primary partition
	# if partition is GPT, it's all good
	# also fill in TARGET_PART_NEW	
	local max_part_no=$MAX_MBR_PART n=1 partcheck
	TARGET_PART_NEW=
	TARGET_PART_NEW_NUMBER=
	if [ $TARGET_PART = new ]; then
		is_device_gpt $TARGET_DEV && max_part_no=$MAX_GPT_PART # should be enough
		while [ $n -le $max_part_no ]; do
			partcheck=$TARGET_DEV$n
			case $PARTITIONS in
				*${partcheck}*) : $(( n = n + 1 )) 
				   ;; # next
				*) TARGET_PART_NEW=$partcheck 
				   TARGET_PART_NEW_NUMBER=$n
				   go_step check_space
				   break 
				   ;; # got it
			esac
		done
		if [ -z "$TARGET_PART_NEW" ]; then
			Xdialog --title "$APPTITLE" --infobox \
"Cannot find free partition on \"$TARGET_DEV\".
Please choose another device." 0 0 10000
			go_step choose_part
		fi
	else
		go_step check_space
	fi
}

# $1-free size, in MB
check_avail_space_helper() {
	if [ $REQUIRED_SIZE -gt $1 ]; then
		Xdialog --title "$APPTITLE" --infobox \
"\"$TARGET_PART\" only has $1 MB. You need at least $REQUIRED_SIZE.
Please choose another device." 0 0 10000
		go_step choose_devices
	else
		go_step choose_source
	fi
}
check_avail_space() {
	# "new" : if creating new partition, check freespace in disk
	# else : if using existing partition, check freespace in partition
	local USED TOTAL FREE BLOCKS
	if [ "$TARGET_PART" = new ]; then
		USED=$(cat /sys/class/block/${TARGET_DEV}?/size | awk '{s+=$1} END {printf "%.0f", s}')
		read TOTAL < /sys/class/block/$TARGET_DEV/size
		: $(( USED = USED / (2 * 1024) )) # convert to MB
		: $(( TOTAL = TOTAL / (2 * 1024) )) # convert to MB
		: $(( FREE = TOTAL - USED ))
		check_avail_space_helper $FREE
	else
		# have to mount it first
		mkdir -p $INSTALL_MNT
		if mount /dev/$TARGET_PART $INSTALL_MNT; then
			FREE=$(df -m | awk -vd=/dev/$TARGET_PART '$1==d { print $4-$5 }')
			umount $INSTALL_MNT
			check_avail_space_helper $FREE
		else
			Xdialog --title "$APPTITLE" --infobox \
"Bad disk: I cannot read \"$TARGET_DEV\".
Please choose another device." 0 0 10000		
			go_step choose_devices		
		fi
	fi
}

choose_install_source() {
	INSTALL_SRC=$(eval Xdialog --title "'$APPTITLE'" --wizard --stdout --no-tags \
	--radiolist "'Choose installation source:'" 00 00 5 \
	"usb"  "'Fatdog flash drive made by \"dd\"'" off \
	"disc" "'Fatdog DVD/CD'" off \
	"iso"  "'Fatdog ISO file'" off)
	next_wizard_step choose_source_dev choose_part
}

choose_install_device() {
	# "usb" - choose from $USBDEV
	# "disc" - list from sr*
	# "iso" - get a filename
	case $INSTALL_SRC in
		usb)
			INSTALL_DEV=$(eval Xdialog --title "'$APPTITLE'" --wizard --stdout --no-tags \
			--radiolist "'Choose source USB device:'" 00 00 5 \
			$(get_target_devices_list only-usb) )
			;;
			
		disc)
			INSTALL_DEV=$(eval Xdialog --title "'$APPTITLE'" --wizard --stdout --no-tags \
			--radiolist "'Choose source DVD/CD drive:'" 00 00 5 \
			$(get_optical_list) )
			;;
			
		iso)
			INSTALL_DEV=$(eval Xdialog --title "'$APPTITLE'" --wizard --stdout \
			--fselect "''" 00 00)
			;;
	esac
	next_wizard_step check_source choose_source
}

check_source() {
	case $INSTALL_SRC in
		usb|disc) INSTALL_DEV="/dev/$INSTALL_DEV"
	esac
	# "usb" and "disc" - both mount as iso9660
	# "iso" - mount as iso9660, and as loop
	# check if we can find efiboot.img, vmlinuz, initrd, grub.cfg
	local good=no
	mkdir -p $INSTALL_MNT
	if mount -t iso9660 -o loop,ro $INSTALL_DEV $INSTALL_MNT; then
		[ -e $INSTALL_MNT/efiboot.img ] &&
		[ -e $INSTALL_MNT/vmlinuz ] &&
		[ -e $INSTALL_MNT/initrd ] &&
		[ -e $INSTALL_MNT/grub.cfg ] &&
		good=yes
		umount -d $INSTALL_MNT
	fi
	
	if [ $good = yes ]; then
		go_step confirm
	else 
		Xdialog --title "$APPTITLE" \
			--infobox "Failed to access ${INSTALL_DEV#/dev/}, or some files are missing." 0 0 10000
		go_step choose_source_dev
	fi
}

confirmation() {
	local PART_TYPE
	# informative notes
	Xdialog --fixed --title "$APPTITLE" --wizard --stdout --textbox - 0 0 << EOF
Installation options
====================
Target device/partition: $(
if [ $TARGET_PART = new ]; then 
	echo "New partition - $TARGET_PART_NEW"
	echo
	is_device_gpt $TARGET_DEV && PART_TYPE=1 || PART_TYPE=c
	printf "n\np\n${TARGET_PART_NEW_NUMBER}\n\n+${REQUIRED_SIZE}M\nt\n${TARGET_PART_NEW_NUMBER}\n${PART_TYPE}\np\nq\n" |
	fdisk /dev/$TARGET_DEV | sed -n '\_^Device_p; \_^/dev/_ p'
else 
	echo $TARGET_PART
fi
)

Installation Source: $INSTALL_SRC
Installation Source device/path: ${INSTALL_DEV#/dev/}

Please read the above carefully and make sure they are correct.
Nothing has been done so far. You can cancel if you wish, or go back
to change your previous answers. 

Once you 'click' next, irreversible actions will be carried out 
**WITHOUT FURTHER QUESTION** to perform the installation.
EOF
	next_wizard_step prepare_part choose_source_dev
}

prepare_partition() {
	if [ "$DRY_RUN" ]; then
		go_step copy_files
		return
	fi
	
	local PART_TYPE
	# for "new" only
	# make partition for GPT, for MSDOS
	if [ $TARGET_PART = new ]; then
		is_device_gpt $TARGET_DEV && PART_TYPE=1 || PART_TYPE=c
		printf "n\np\n${TARGET_PART_NEW_NUMBER}\n\n+${REQUIRED_SIZE}M\nt\n${TARGET_PART_NEW_NUMBER}\n${PART_TYPE}\np\nw\n" |
		fdisk /dev/$TARGET_DEV
		sleep 1
		blockdev --rereadpt /dev/$TARGET_DEV
		mkdosfs -n "FATDOG_LIVE" /dev/$TARGET_PART_NEW
	fi
	go_step copy_files
}

copy_files() {
	if [ "$DRY_RUN" ]; then
		go_step success
		return
	fi
	
	local TARGET_PART=$TARGET_PART XPID good=no	
	# mount source, as listed in check_source
	# mount target
	# mount efiboot.img as loop, copy contents
	# copy grub.cfg, vmlinuz, initrd
	[ "$TARGET_PART" = "new" ] && TARGET_PART=$TARGET_PART_NEW
	Xdialog --title "$APPTITLE" --no-buttons --infobox "\
Copying files, this will take a while.
Please wait ..." 0 0 10000000 &
	XPID=$!
	mkdir -p $INSTALL_MNT/src $INSTALL_MNT/efi $INSTALL_MNT/tgt
	if mount /dev/$TARGET_PART $INSTALL_MNT/tgt; then
		if mount -t iso9660 -o loop,ro $INSTALL_DEV $INSTALL_MNT/src; then
			if mount -o loop,ro $INSTALL_MNT/src/efiboot.img $INSTALL_MNT/efi; then

				# this is heart of the installation, really
				cp -R $INSTALL_MNT/efi/* $INSTALL_MNT/tgt &&
				cp $INSTALL_MNT/src/vmlinuz  \
				$INSTALL_MNT/src/initrd  \
				$INSTALL_MNT/src/grub.cfg \
				$INSTALL_MNT/tgt &&
				good=yes

				umount -d $INSTALL_MNT/efi
			fi
			umount -d $INSTALL_MNT/src
		fi
		umount $INSTALL_MNT/tgt
	fi
	kill $XPID
	if [ "$good" = yes ]; then
		go_step success
	else	
		go_step failed
	fi
}

success_msg() {
	Xdialog --title "$APPTITLE" --infobox "\
Installation completed successfully.
" 0 0 10000	
	go_exit
}

failed_msg() {
	Xdialog --title "$APPTITLE" --infobox "\
Installation failed. Please run from terminal to capture error messages.
" 0 0 10000	
	go_exit
}

warn_cancel() {
	if Xdialog --title "$APPTITLE" --yesno "\
Do you really want to quit?
Nothing has been modified.
" 0 0; then

		Xdialog --title "$APPTITLE" --infobox "\
Installation Cancelled. No modification has been done.
" 0 0 10000
		go_exit

	else
		go_step $prev_step
	fi
}


################## main ###################
! [ -e /sys/firmware/efi ] && Xdialog --icon "/usr/share/mini-icons/mini-stop.xpm" --title "Warning" \
        --infobox "\
Installation created by this installer only boots on UEFI systems.
Use Fatdog64 standard installer if you want to boot on BIOS systems.
" 0 0 10000

### preparation
[ "$1" = "-n" ] && DRY_RUN=yes
[ "$DRY_RUN" ] && APPTITLE="$APPTITLE (Dryrun)"

### main loop
system_capabilities
while true; do 
	case $step in 
		""|intro) introduction ;;
		
		# target device selection and checks
		choose_devices)    choose_devices ;;
		check_usb)         check_is_usb_device ;;		
		choose_part)       choose_partitions ;;		
		check_part)        check_partition_type ;;
		check_new_part)    check_make_partition ;;
		check_space) check_avail_space ;;
		
		# source selection and checks
		choose_source)     choose_install_source ;;
		choose_source_dev) choose_install_device ;;
		check_source)      check_source ;;
		
		# confirmation and actual install
		confirm) confirmation ;;
		prepare_part) prepare_partition ;;
		copy_files) copy_files ;;
		success) success_msg ;;
		failed) failed_msg ;;
		
		# done
		exit) break ;;
		cancel) warn_cancel ;;
		*) 
			Xdialog --title "$APPTITLE" --infobox "Program bug. step=$step. Please report." 0 0 10000
			break
			;;
	esac
done

