#!/usr/bin/env bash

# -------------------------------------------------------------------------
#                                                                         -
#  Create bootable device from Ultimate Boot CD ISO image                 -
#                                                                         -
#  Created by Fonic <https://github.com/fonic>                            -
#  Date: 09/23/20 - 09/23/20                                              -
#                                                                         -
#  Based on:                                                              -
#  https://gist.github.com/fonic/647011012b48d6dcd65e8725fed48b0c         -
#  <ubcd.iso>/ubcd/tools/linux/ubcd2usb/readme.txt                        -
#                                                                         -
# -------------------------------------------------------------------------


# ------------------------------------
#                                    -
#  Functions                         -
#                                    -
# ------------------------------------

# Print normal message [$*: message]
function print_normal() {
	echo -e "$*"
}

# Print hilite message [$*: message]
function print_hilite() {
	echo -e "\e[1m$*\e[0m"
}

# Print good message [$*: message]
function print_good() {
	echo -e "\e[1;32m$*\e[0m"
}

# Print warn message [$*: message]
function print_warn() {
	echo -e "\e[1;33m$*\e[0m"
}

# Print error message [$*: message]
function print_error() {
	echo -e "\e[1;31m$*\e[0m"
}

# Check if command is available [$1: command]
function is_cmd_avail() {
	command -v "$1" &>/dev/null
	return $?
}

# Extend code of existing trap [$1: trap name, $2: insert mode ('start'/'end'), $3: code to insert]
# (https://stackoverflow.com/questions/3338030/multiple-bash-traps-for-the-same-signal)
# NOTE: gets the job done, but most likely not handling every possible scenario
function extend_trap() {
	local trap_name="$1" ins_mode="$2" ins_code="$3"
	local output re_code old_code new_code

	# Nothing to insert
	[[ -z "${ins_code}" ]] && return 0

	# Determine code of existing trap
	output="$(trap -p "${trap_name}")"
	if [[ -n "${output}" ]]; then
		re_code="^trap -- '(.*)' ${trap_name}$"
		if [[ "${output}" =~ ${re_code} ]]; then
			old_code="${BASH_REMATCH[1]//\\\'\'/}"
		else
			print_error "Unable to extend ${trap_name} trap: failed to match regex: '${re_code}' -> '${output}'"
			return 1
		fi
	else
		old_code=""
	fi

	# Determine new trap code
	if [[ -z "${old_code}" ]]; then
		new_code="${ins_code}"
	elif [[ "${ins_mode}" == "start" ]]; then
		new_code="${ins_code}; ${old_code}"
	elif [[ "${ins_mode}" == "end" ]]; then
		new_code="${old_code}; ${ins_code}"
	else
		print_error "Unable to extend ${trap_name} trap: invalid insert mode: '${ins_mode}'"
		return 1
	fi

	# Update trap
	trap -- "${new_code}" "${trap_name}"
	return $?
}; declare -f -t extend_trap

# Handler for error trap [no arguments]
function error_trap() {
	echo
	print_error "An error occurred, aborting."
	echo
	exit 1
}

# Handler for interrupt trap [no arguments]
function int_trap() {
	echo
	echo
	print_warn "Aborting on user request."
	echo
	exit 130
}


# ------------------------------------
#                                    -
#  Main Program                      -
#                                    -
# ------------------------------------

# Process command line
if (( $# != 2 )); then
	print_normal "Usage: $(basename "$0") <iso> <device>"
	exit 2
fi
isoimg="$1"
device="$2"
if [[ ! -f "${isoimg}" ]]; then
	print_error "Error: file '${isoimg}' does not exist, aborting."
	exit 2
fi
if [[ ! -b "${device}" ]]; then
	print_error "Error: device '${device}' does not exist, aborting."
	exit 2
fi

# Check command availability
result=0
#for cmd in "fdisk" "wipefs" "sync" "partprobe" "mkfs.fat" "mktemp" "mount" "cp" "umount" "rmdir" "dd" "uname"; do
for cmd in "fdisk" "dd" "sync" "partprobe" "mkfs.fat" "mktemp" "mount" "cp" "umount" "rmdir" "uname"; do
	if ! is_cmd_avail "${cmd}"; then
		print_error "Error: command '${cmd}' is not available"
		result=1
	fi
done
if (( ${result} != 0 )); then
	print_error "Error: required command(s) unavailable, aborting."
	exit 1
fi

# Check root privileges
if (( ${EUID} != 0 )); then
	print_error "Error: root privileges required, aborting."
	exit 1
fi

# Set up error handling / traps
set -e
trap "error_trap" ERR
trap "int_trap" INT
trap - EXIT

# Print warning message
echo
lines=()
lines+=("")
lines+=("DEVICE '${device}' WILL BE COMPLETELY ERASED/WIPED")
lines+=("DOUBLE-CHECK THAT YOU SPECIFIED THE CORRECT DEVICE")
lines+=("")
maxlen=0
for line in "${lines[@]}"; do (( ${#line} > ${maxlen} )) && maxlen=${#line}; done
for line in "${lines[@]}"; do printf "\e[1;33m!!!  %-${maxlen}s  !!!\e[0m\n" "${line}"; done
echo
fdisk -l "${device}"
echo

# Ask for user confirmation
echo -en "\e[1mType 'OK' to continue: \e[0m"
read input
if [[ "${input}" != "OK" ]]; then
	echo
	exit 130
fi
echo

# Erase/wipe device
print_hilite "Erasing/wiping device..."
#wipefs --all "${device}" # not working as expected; it would seem this wipes only the partition table, but not filesystem signatures; using dd instead
dd if=/dev/zero of=${device} bs=1M count=16 oflag=direct # overwrite a couple of megs to clean out partition table and fs signature
echo

# Sync and partprobe
print_hilite "Syncing and partprobing..."
sync
partprobe "${device}"
echo

# Create partition table
print_hilite "Creating partition table..."
{
	echo "o"   # Create new empty DOS partition table
	echo "n"   # Add new partition
	echo ""    # Partition type (default: 'primary' if less than 4 partitions present, else 'extended')
	echo ""    # Partition number (default: 1)
	echo ""    # First sector (default: 2048)
	echo ""    # Last sector (default: last sector of device)
	echo "t"   # Change partition type (automatically selects partition 1 if only one partition present)
	#echo "0c" # Type '0c W95 FAT32 (LBA)'
	echo "0b"  # Type '0b W95 FAT32'
	echo "a"   # Toggle bootable flag (automatically selects partition 1 if only one partition present)
	echo "w"   # Write changes to disk and quit
} | fdisk "${device}"

# Sync and partprobe
print_hilite "Syncing and partprobing..."
sync
partprobe "${device}"
echo

# Create filesystem
print_hilite "Creating filesystem..."
mkfs.fat -F 32 -n "UBCD" "${device}1"
echo

# Sync
print_hilite "Syncing..."
sync
echo

# Mount ISO image
print_hilite "Mounting ISO image..."
isomnt="$(mktemp -d)"
mount -o loop,ro "${isoimg}" "${isomnt}"
extend_trap EXIT start "mountpoint -q \"${isomnt}\" && { print_hilite \"Unmounting ISO image...\"; umount \"${isomnt}\" || umount -f \"${isomnt}\" || umount -l \"${isomnt}\"; rmdir \"${isomnt}\"; echo; }"
echo

# Mount device
print_hilite "Mounting device..."
devmnt="$(mktemp -d)"
mount "${device}1" "${devmnt}"
extend_trap EXIT start "mountpoint -q \"${devmnt}\" && { print_hilite \"Unmounting device...\"; umount \"${devmnt}\" || umount -f \"${devmnt}\" || umount -l \"${devmnt}\"; rmdir \"${devmnt}\"; echo; }"
echo

# Copy ISO image contents to device
print_hilite "Copying ISO image contents to device..."
cp -r "${isomnt}"/* "${devmnt}"
echo

# Sync
print_hilite "Syncing..."
sync
echo

# Unmount device
print_hilite "Unmounting device..."
umount "${devmnt}" || umount -f "${devmnt}" || umount -l "${devmnt}"
rmdir "${devmnt}"
echo

# Sync
print_hilite "Syncing..."
sync
echo

# Write Syslinux MBR to device
print_hilite "Writing Syslinux MBR to device..."
dd if="${isomnt}/ubcd/tools/linux/ubcd2usb/mbr.bin" of="${device}" oflag=direct
echo

# Sync and partprobe
print_hilite "Syncing and partprobing..."
sync
partprobe "${device}"
echo

# Install Syslinux to device
print_hilite "Installing Syslinux to device..."
machine_type="$(uname -m)"
if [[ "${MACHINE_TYPE}" == "x86_64" ]]; then
	"${isomnt}/ubcd/tools/linux/ubcd2usb/syslinux64" -i -s -d /boot/syslinux "${device}1"
else
	"${isomnt}/ubcd/tools/linux/ubcd2usb/syslinux" -i -s -d /boot/syslinux "${device}1"
fi
echo

# Sync
print_hilite "Syncing..."
sync
echo

# Unmount ISO image
print_hilite "Unmounting ISO image..."
umount "${isomnt}" || umount -f "${isomnt}" || umount -l "${isomnt}"
rmdir "${isomnt}"
echo

# Print results
print_good "Success!"
echo
trap - EXIT