#!/bin/bash
# Mount a USB device
. /etc/functions.sh
. /etc/gui_functions.sh
. /etc/luks-functions.sh

TRACE_FUNC

function usage() {
	cat <<USAGE_END
usage: $0 [options...] <--mode [ro|rw]> <--device device> <--mountpoint mountpoint> <--pass passphrase>
       $0 --help
       
parameters: 
  --mode: ro or rw (default ro)
  --device: device to mount (default: first USB device found)
  --mountpoint: where to mount the device (default: /media)
  --pass: passphrase for LUKS device (default: none)
  --help: Show this help
USAGE_END
}

MODE="ro"
DEVICE=""
MOUNTPOINT="/media"
PASS=""

#Only assign --mode, --device, --mountpoint and --pass parameters only if variables following them are not empty
while [ $# -gt 0 ]; do
	case "$1" in
	--mode)
		if [ -n "$2" ]; then
			MODE="$2"
			shift
			shift
		fi
		;;
	--device)
		if [ -n "$2" ]; then
			DEVICE="$2"
			shift
			shift
		fi
		;;
	--mountpoint)
		if [ -n "$2" ]; then
			MOUNTPOINT="$2"
			shift
			shift
		fi
		;;
	--pass)
		if [ -n "$2" ]; then
			PASS="$2"
			shift
			shift
		fi
		;;
	*)
		usage
		exit 1
		;;
	esac
done

#Show parameters content but not LUKS passphrase: if empty, show "empty", if provided, show "provided"
DEBUG "Parameters: --mode=$MODE, --device=${DEVICE:-empty}, --mountpoint=$MOUNTPOINT, --pass=${PASS:+provided}"

enable_usb
enable_usb_storage

if [ ! -d "$MOUNTPOINT" ]; then
	DEBUG "Creating $MOUNTPOINT directory"
	mkdir -p "$MOUNTPOINT" >/dev/null 2>&1
else
	DEBUG "Cleaning $MOUNTPOINT directory"
	umount "$MOUNTPOINT" >/dev/null 2>&1 || true
fi

list_usb_storage >/tmp/usb_block_devices
if [ -z "$(cat /tmp/usb_block_devices)" ]; then
	if [ -x /bin/whiptail ]; then
		whiptail_warning --title 'USB Drive Missing' \
			--msgbox "Insert your USB drive and press Enter to continue." 0 80
	else
		INPUT "USB Drive Missing! Insert your USB drive and press Enter to continue."
	fi
	sleep 1
	list_usb_storage >/tmp/usb_block_devices
	if [ -z "$(cat /tmp/usb_block_devices)" ]; then
		if [ -x /bin/whiptail ]; then
			whiptail_error --title 'ERROR: USB Drive Missing' \
				--msgbox "USB Drive Missing! Aborting mount attempt.\n\nPress Enter to continue." 0 80
		else
			DIE "USB Drive Missing! Aborting mount."
		fi
		exit 1
	fi
fi

USB_MOUNT_DEVICE=""
# Check if the user has specified a USB device
if [ -n "$DEVICE" ]; then
	DEBUG "Checking if "$DEVICE" is a USB detected block device"
	if grep -q "$DEVICE" /tmp/usb_block_devices; then
		DEBUG "Selected device is a USB block device"
		USB_MOUNT_DEVICE="$DEVICE"
	else
		DIE "ERROR: Selected $DEVICE is not a USB block device"
	fi
else
	# Check for the common case: a single USB disk with one partition
	if [ $(cat /tmp/usb_block_devices | wc -l) -eq 1 ]; then
		USB_MOUNT_DEVICE="$(cat /tmp/usb_block_devices)"
	fi
	# When a passphrase is provided and multiple devices are present,
	# auto-select the LUKS partition (e.g. GPG backup drive: LUKS private + exFAT public).
	# This avoids burdening the user with selecting the right partition.
	if [ -z "$USB_MOUNT_DEVICE" ] && [ -n "$PASS" ]; then
		luks_dev=""
		luks_count=0
		while IFS= read -r dev; do
			if cryptsetup isLuks "$dev" 2>/dev/null; then
				luks_dev="$dev"
				luks_count=$((luks_count + 1))
			fi
		done </tmp/usb_block_devices
		if [ "$luks_count" -eq 1 ]; then
			DEBUG "Auto-selected LUKS partition: $luks_dev (passphrase provided)"
			USB_MOUNT_DEVICE="$luks_dev"
		fi
	fi
	# otherwise, let the user pick
	if [ -z ${USB_MOUNT_DEVICE} ]; then
		>/tmp/usb_disk_list
		for i in $(cat /tmp/usb_block_devices); do
			#appends label to the device name
			echo $i $(blkid | grep $i | grep -o 'LABEL=".*"' | cut -f2 -d '"') >>/tmp/usb_disk_list
		done

		if [ -x /bin/whiptail ]; then
			MENU_OPTIONS=()
			n=0
			while read option; do
				n=$(expr $n + 1)
				MENU_OPTIONS+=("$n" "$option")
			done </tmp/usb_disk_list

			MENU_OPTIONS+=("a" "Abort")
			whiptail_type $BG_COLOR_MAIN_MENU --title "Select your USB disk" \
				--menu "Choose your USB disk [1-$n, a to abort]:" 0 80 8 \
				-- "${MENU_OPTIONS[@]}" \
				2>/tmp/whiptail
			if [ $? -ne 0 ]; then
				DIE "ERROR: Selecting USB disk/partition aborted."
			fi
			option_index=$(cat /tmp/whiptail)
		else
			STATUS "Select your USB disk:"
			n=0
			while read option; do
				n=$(expr $n + 1)
				printf '%d. %s\n' "$n" "$option" >"${HEADS_TTY:-/dev/stderr}"
			done </tmp/usb_disk_list

			INPUT "Choose your USB disk [1-$n, a to abort]:" -r option_index
		fi

		if [ "$option_index" = "a" ]; then
			exit 5
		fi
		USB_MOUNT_DEVICE=$(head -n $option_index /tmp/usb_disk_list | tail -1 | sed 's/\ .*$//')
	fi
fi

DEBUG "Checking if $USB_MOUNT_DEVICE is a LUKS device/partition"
if cryptsetup isLuks "$USB_MOUNT_DEVICE"; then
	DEBUG "Selected USB partition is a LUKS device"
	#Selected USB partition is a LUKS device
	if [ -e /dev/mapper/"usb_mount_$(basename "$USB_MOUNT_DEVICE")" ]; then
		DEBUG "Closing currently mapped LUKS device"
		cryptsetup close "usb_mount_$(basename "$USB_MOUNT_DEVICE")"
	fi
	DEBUG "Opening LUKS device $USB_MOUNT_DEVICE"
	#Pass LUKS passphrase to cryptsetup only if we received one
	if [ -z "$PASS" ]; then
		#We haven't received a passphrase
		cryptsetup open "$USB_MOUNT_DEVICE" "usb_mount_$(basename "$USB_MOUNT_DEVICE")" ||
			DIE "ERROR: Failed to open ${USB_MOUNT_DEVICE} LUKS device"
	else
		#We received a pasphrase
		cryptsetup open "$USB_MOUNT_DEVICE" "usb_mount_$(basename "$USB_MOUNT_DEVICE")" --key-file <(echo -n "${PASS}") ||
			DIE "ERROR: Failed to open ${USB_MOUNT_DEVICE} LUKS device"
	fi

	WARN "Note that you cannot boot from a mounted encrypted device"
	DEBUG "Setting USB_MOUNT_DEVICE=/dev/mapper/"usb_mount_$(basename "$USB_MOUNT_DEVICE")""
	USB_MOUNT_DEVICE="/dev/mapper/"usb_mount_$(basename "$USB_MOUNT_DEVICE")""
else
	# Selected USB partition is not a LUKS device
	DEBUG "Selected USB partition is not a LUKS device, continuing..."
fi

# Mount the USB device
if [ "$MODE" = "rw" ]; then
	DEBUG "Mounting $USB_MOUNT_DEVICE as read-write"
	mount -o rw "$USB_MOUNT_DEVICE" "$MOUNTPOINT" || DIE "ERROR: Failed to mount ${USB_MOUNT_DEVICE} as read-write"
else
	DEBUG "Mounting $USB_MOUNT_DEVICE as read-only"
	mount -o ro "$USB_MOUNT_DEVICE" "$MOUNTPOINT" || DIE "ERROR: Failed to mount ${USB_MOUNT_DEVICE} as read-only"
fi
