#!/bin/sh
# bootsh - create a bootable removable device from an image file
# SPDX-FileCopyrightText: 2021 Sotiris Papatheodorou
# SPDX-License-Identifier: GPL-3.0-or-later
set -eu

# Usage: check_dependencies
# Check if the depndencies are present and exit with an appropriate error
# message otherwise.
check_dependencies() {
	printf 'column\nlsblk\nsync\n' | while IFS= read -r dep; do
		if ! command -v "$dep" > /dev/null 2>&1; then
			printf 'Error: dependency %s not found\n' "$dep"
			exit 1
		fi
	done
}

# Usage: trim_string STRING
# SPDX-FileCopyrightText: 2019 Dylan Araps
# SPDX-License-Identifier: MIT
# https://github.com/dylanaraps/pure-sh-bible#trim-leading-and-trailing-white-space-from-string
trim_string() {
    trim=${1#${1%%[![:space:]]*}}
    trim=${trim%${trim##*[![:space:]]}}
    printf '%s\n' "$trim"
}

# Usage: print_block_devices
# Pretty-print all block devices.
print_block_devices() {
	if command -v udevadm > /dev/null 2>&1; then
		udevadm settle > /dev/null 2>&1
	fi
	# The following command was copied from bootiso
	# SPDX-FileCopyrightText: 2018-2020 Jules Randolph
	# SPDX-License-Identifier: GPL-3.0-or-later
	lsblk --nodeps --list --output NAME,VENDOR,MODEL,SIZE,TRAN,HOTPLUG
}

# Usage: print_removable_devices
# Pretty-print all removable block devices.
print_removable_devices() {
	devices=$(print_block_devices)
	# Keep the header and all devices that are hotpluggable.
	printf '%s\n' "$devices" | head -n 1 | sed 's/ *HOTPLUG$//'
	printf '%s\n' "$devices" | grep ' *1$' | grep -v 'sata *1$' |  sed 's/ *1$//'
}

# Usage: removable_device_names
# Print the names of all removable devices, one per line.
removable_device_names() {
	print_removable_devices | tail -n +2 | awk '{ print $1 }'
}

# Usage: device_capacity NAME
# Print the capacity of device NAME in human-readable form.
device_capacity() {
	lsblk --nodeps --list --output NAME,SIZE | grep -F "$1" | awk '{ print $2 }'
}

# Usage: device_mount_points NAME
# Print the mounted partitions of device NAME and their respective mount
# points, separated by tabs, one per line.
device_mount_points() {
	mount | grep -F "$1" | sed 's|^\('"$1"'[0-9][0-9]*\) on \(/.*\) type .*$|\1\t\2|'
}

# Usage: valid_input INPUT VALID_INPUTS
# Return 0 when INPUT exactly matches one of the lines in VALID_INPUTS.
valid_input() {
	# Matching INPUT as a fixed pattern can match substrings, e.g. sd matches
	# sda. Matching INPUT as a regular expression can produce weird matches if
	# INPUT is a regular expression. Just do both matches to be more robust.
	printf '%s\n' "$2" | grep -qF "$1" && printf '%s\n' "$2" | grep -q "^$1"'$'
}



check_dependencies

# Check for valid input argument.
if [ "$#" -ne 1 ]; then
	printf 'Usage: %s IMAGE\n' "${0##*/}"
	exit 2
fi
if [ ! -s "$1" ]; then
	printf 'Error: %s must be a non-empty file\n' "$1"
	exit 1
fi

# Find removable devices.
names=$(removable_device_names)
if [ -z "$names" ]; then
	printf 'Error: no removable devices found\n'
	exit 1
fi
names_pretty=$(print_removable_devices)

# Ensure at least one removable device was found.
if [ "$(printf '%s\n' "$names_pretty" | wc -l)" -lt 2 ]; then
	printf 'No removable devices found\n'
	exit 1
fi

# Prompt the user to select the device to write the image to.
target=""
while [ -z "$target" ]; do
	printf '%s\n' "$names_pretty"
	printf 'Enter the NAME of one of the above removable devices or press Ctrl+C to exit:\n'
	read -r device_input
	device_input=$(trim_string "$device_input")
	if valid_input "$device_input" "$names"; then
		# Ask the user to confirm the device's capacity to avoid writing to the
		# wrong device.
		while true; do
			capacity=$(device_capacity "$device_input")
			printf '\nSelected removable device %s with capacity %s\n' \
				"$device_input" "$capacity"
			printf 'Enter the device capacity (%s) to continue or press Ctrl+C to exit:\n' \
				"$capacity"
			read -r capacity_input
			capacity_input=$(trim_string "$capacity_input")
			if [ "$capacity_input" = "$capacity" ]; then
				target=/dev/"$device_input"
				break
			fi
		done
	fi
done

# Unmount any mounted partitions of the target device.
device_mount_points "$target" | while IFS= read -r line; do
	partition=$(printf '%s\n' "$line" | awk 'BEGIN { FS="\t" }; { print $1 }')
	mount_point=$(printf '%s\n' "$line" | awk 'BEGIN { FS="\t" }; { print $2 }')
	sync
	if ! umount "$mount_point" > /dev/null 2>&1; then
		printf 'Error: could not unmount partition %s mounted on %s\n' \
			"$partition" "$mount_point"
		exit 1
	fi
done

# Test if privilege elevation is needed.
superuser_command=""
if [ "$(id -u)" != "0" ]; then
	if command -v sudo > /dev/null 2>&1; then
		superuser_command="sudo"
	elif command -v doas > /dev/null 2>&1; then
		superuser_command="doas"
	fi
fi

# Write the image to the device.
printf '\nWriting image %s to device %s\n' "$1" "$target"
"$superuser_command" dd bs=1024k if="$1" of="$target" > /dev/null
sync
printf 'Done! You may remove device %s\n' "$target"