#!/usr/bin/env bash

####################################################################################################
#
#	FIX GRUB FOR ENCRYPTED INSTALLATION
#
# This script is based on the Troubleshooting part of the Manual Full System Encryption for Ubuntu.
# See https://help.ubuntu.com/community/ManualFullSystemEncryption/Troubleshooting
#
# This is NOT AN OFFICIAL Ubuntu script, it can and will break your system. Use at own risk.
#
####################################################################################################


#===================================================================================================
#
#	Section: Initialise
#
#===================================================================================================


#---------------------------------------------------------------------------------------------------
#	Set up the script.
#---------------------------------------------------------------------------------------------------


function initialise ()
{
	trap ctrlC SIGINT				# Trap Ctrl+C.

	# Miscellaneous items for display.
	declare -gr SPACER='----------------------------------------------------------------------------------------------------'
	declare -gr E=$'\e[1;31;103m'			# E for Error: highlighted text.
	declare -gr W=$'\e[1;31;103m'			# W for Warning: highlighted text.
	declare -gr B=$'\e[1m'				# B for Bold.
	declare -gr R=$'\e[0m'				# R for Reset.

	# Display a warning to the user.
	clear

	# Give the user some instruction.
	cat <<-END


		${SPACER}


		In these instructions, when I ask you to press ${B}Y${R} (for Yes) or ${B}N${R} (for No),
		you may type these in either uppercase or lowercase.

		${B}Everything else is case sensitive!${R}


	END

	read -rp 'Press Enter to continue reading. '

	# Show a warning.
	cat <<-END


		${SPACER}


		${B}WARNING${R}

		This script is based on the Troubleshooting part of the Manual Full System Encryption for Ubuntu.
		See https://help.ubuntu.com/community/ManualFullSystemEncryption/Troubleshooting

		This is NOT AN OFFICIAL Ubuntu script, it can and will break your system. Use at own risk.

		Please confirm that you have read and followed the instructions before proceeding.


	END

	# Ask for confirmation.
	local ANSWER
	read -rp "Type ${B}Y${R} to proceed, or anything else to cancel, and press Enter: ${B}" ANSWER
	echo "${R}"

	# Terminate if required.
	if [[ "${ANSWER,}" != 'y' ]]
	then
		echo
		echo 'Terminated. I did nothing.'
		echo
		exit 1
	fi

} # initialise


#---------------------------------------------------------------------------------------------------
#	Trap for Ctrl+C.
#---------------------------------------------------------------------------------------------------


function ctrlC ()
{
	cat <<-END
		${SPACER}

		        ********        Terminated with Ctrl+C.        ********

		${SPACER}
	END

	exit 2

} # ctrlC


#---------------------------------------------------------------------------------------------------
#	Display an error messsage.
#
# Parameter
#	1	The error message.
#	2	The exit code (numeric); if absent, won't exit.
#---------------------------------------------------------------------------------------------------


function error ()
{
	local -r ERROR="${1}"
	local -r EXIT_CODE=${2}

	echo -e "\n${E}${ERROR}${R}\n" >&2		# Display the message, highlighted.

	[[ -n ${EXIT_CODE} ]] && exit ${EXIT_CODE}	# Exit with the code if given.

} # error


#===================================================================================================
#
#	Section: Gather Data
#
#===================================================================================================

#---------------------------------------------------------------------------------------------------
#	Gather all of the required data.
#
# Return
#	PARTITION_ESP		/dev/[partition] e.g. /dev/sda2, /dev/nvme0n1p2
#	PARTITION_SYSTEM	/dev/[partition] e.g. /dev/sda5, /dev/nvme0n1p5
#	PASSPHRASE_SYSTEM	Passphrase for the system partition
#---------------------------------------------------------------------------------------------------


function gatherData ()
{
	findPartitions					# The various partitions.
	findSystemPassphrase			# The system partition passphrase.
}


#---------------------------------------------------------------------------------------------------
#	Find the partitions.
#
# Return
#	PARTITION_ESP		/dev/[partition] e.g. /dev/sda2, /dev/nvme0n1p2
#	PARTITION_SYSTEM	/dev/[partition] e.g. /dev/sda5, /dev/nvme0n1p5
#---------------------------------------------------------------------------------------------------


function findPartitions ()
{
	local ANSWER=''					# For user responses.

	# The ESP partition.
	while true
	do
		cat <<-END

			${SPACER}

			Which is your EFI System Partition (ESP)?
			Type just the part that goes after /dev/, no spaces, and press Enter.

		END
		read -rp "/dev/${B}" ANSWER
		echo "${R}"

		# Validate the partition.
		validatePartition "${ANSWER}" fat32 'EFI System Partition' false && break
	done
	declare -gr PARTITION_ESP=/dev/${ANSWER}

	# The system partition.
	while true
	do
		cat <<-END

			${SPACER}
			Which is your System partition?
			Type just the part that goes after /dev/, no spaces, and press Enter.

		END
		read -rp "/dev/${B}" ANSWER
		echo "${R}"

		# Validate the partition.
		validatePartition "${ANSWER}" cleared system true && break
	done
	declare -gr PARTITION_SYSTEM=/dev/${ANSWER}

} # findPartitions


#---------------------------------------------------------------------------------------------------
#	Validate a partition.
#
# Parameter
#	1	The partition, e.g. sda2, nvme0n1p2
#	2	The required file system, either fat32 or cleared
#	3	The expected GPT label for the partition.
#	4	Whether or not to enforce a GPT label. true or false
#
# Return
#	0 if valid, non-zero otherwise.
#
# Output
#	An error message to &2 if invalid.
#---------------------------------------------------------------------------------------------------


function validatePartition ()
{
	local -r PARTITION="${1}"			# The partition.
	local -r REQUIRED_FS=${2}			# The required file system.
	local -r GPT_LABEL="${3}"			# The expected GPT label for the partition.
	local -r ENFORCE_LABEL=${4}			# Whether or not to enforce a label match.

	# Check the syntax.
	if ! [[ "${PARTITION}" =~ ^[a-z0-9]+$ ]]
	then
		error 'The syntax looks wrong. Please try again.'
		return 1
	fi

	#  Extract the drive.
	local DRIVE=$( readlink /sys/class/block/${PARTITION} )
	[[ -n ${DRIVE} ]] && DRIVE=${DRIVE%/*}
	[[ -n ${DRIVE} ]] && DRIVE=${DRIVE##*/}

	if [[ -z "${DRIVE}" ]]
	then
		error 'That drive or partition appears not to exist. Please try again. (code 1)'
		return 2
	fi

	# Find the specific partition. Needs sudo for some partitions only.
	local -r PARTITION_DETAILS=$( sudo blkid /dev/${PARTITION} 2>/dev/null )

	if [[ -z "${PARTITION_DETAILS}" ]]
	then
		error 'That partition appears not to exist. Please try again.'
		return 4
	fi

	# Extract the file system.
	local FILE_SYSTEM=$( grep -E --only-matching ' TYPE="[a-z0-9]+" ' <<<"${PARTITION_DETAILS}" | cut --delimiter=\" --field=2 )

	if [[ -z "${FILE_SYSTEM}" ]]
	then
		FILE_SYSTEM=cleared			# A cleared file system returns nothing.
	elif [[ ${FILE_SYSTEM} == vfat ]]
	then
		FILE_SYSTEM=fat32			# For now, assume FAT32, but we check later.
	fi

	# Check that the file system is correct.
	if [[ ${FILE_SYSTEM} != ${REQUIRED_FS} ]]
	then
		error "The file system should be ${REQUIRED_FS} but is instead ${FILE_SYSTEM}."
		return 5
	fi

	# For vfat, check that it really is FAT32 and not FAT16 or something else.
	if [[ ${REQUIRED_FS} == fat32 ]]
	then
		FAT32=$( sudo file --special-files /dev/${PARTITION} 2>/dev/null | grep -F --only-matching 'FAT (32 bit)' )
		if [[ -z ${FAT32} ]]
		then
			error "The file system should be ${REQUIRED_FS} but is instead a different vfat."
			return 6
		fi
	fi

	# Find the partition label.
	local -r PARTITION_LABEL="$( grep -E --only-matching ' PARTLABEL="[^\"]+" ' <<<"${PARTITION_DETAILS}" | cut --delimiter=\" --field=2 )"

	# If we don't have to enforce the label, do a case-insensitive check; otherwise check case.
	if ${ENFORCE_LABEL}
	then
		if [[ "${PARTITION_LABEL}" == "${GPT_LABEL}" ]]
		then
			local -r LABEL_MATCH=true
		else
			local -r LABEL_MATCH=false
		fi
	elif [[ "${PARTITION_LABEL,,}" == "${GPT_LABEL,,}" ]]
	then
		local -r LABEL_MATCH=true
	else
		local -r LABEL_MATCH=false
	fi

	# Check for a name discrepancy.
	if ! ${LABEL_MATCH}
	then
		# Tell the user of the discrepancy.
		if [[ -z "${PARTITION_LABEL}" ]]
		then
			error "The partition label should be \"${GPT_LABEL}\" but is instead blank."
		else
			error "The partition label should be \"${GPT_LABEL}\" but is instead \"${PARTITION_LABEL}\"."
		fi
		${ENFORCE_LABEL} && return 7		# Try again if we must enforce the label.

		# Ask the user whether or not to continue.
		local ANSWER
		read -rp "If this is anyway correct, press ${B}Y${R} to continue or anything else to try again. " ANSWER
		[[ "${ANSWER,}" != 'y' ]] && return 7	# Try again.
	fi

	return 0					# Happy to continue.

} # validatePartition


#---------------------------------------------------------------------------------------------------
#	Find the system partition passphrase.
#
# Return
#	PASSPHRASE_SYSTEM
#---------------------------------------------------------------------------------------------------


function findSystemPassphrase ()
{
	# Read the system passphrase.
	local SYSTEM=''
	while true
	do
		cat <<-END

			${SPACER}

			Please enter the passphrase that you chose for your system partition.
			Check carefully before you press Enter.

		END
		read -rp "System passphrase: ${B}" SYSTEM
		echo "${R}"

		(( ${#SYSTEM} > 0 )) && break
	done
	declare -gr PASSPHRASE_SYSTEM="${SYSTEM}"
} # findSystemPassphrase


#===================================================================================================
#
#	Section: Fix GRUB Section
#
#===================================================================================================


#---------------------------------------------------------------------------------------------------
#	Fix GRUB process
#---------------------------------------------------------------------------------------------------


function fixGrubProcess ()
{
	unlockPartition System system ${PARTITION_SYSTEM} "${PASSPHRASE_SYSTEM}"	# Unlock the system partition.
	mountPartitions																# Mount the partitions.
	validateFileExistence "/mnt/root/usr/local/sbin/refreshgrub"				# Check if refreshgrub script exists.
	chrootRefreshGrub															# Enter chroot and run refreshgrub.
} # fixGrubProcess


#---------------------------------------------------------------------------------------------------
#	Unlock a partition
#
# Parameters
#	1	Human-readable name for the partition
#	2	Partition label
#	3	The partition, e.g. /dev/sda2, /dev/nvme0n1p2
#	4	The passphrase
#---------------------------------------------------------------------------------------------------


function unlockPartition ()
{
	local -r HUMAN_NAME=${1}
	local -r LABEL=${2}
	local -r PARTITION=${3}
	local -r PASSPHRASE="${4}"

	echo
	echo "Unlocking the ${HUMAN_NAME} partition..."

	# Unlock the partition.
	echo -n "${PASSPHRASE}" | sudo cryptsetup open --type=luks --key-file=- ${PARTITION} ${LABEL}

	local -ir RET=${?}				# Catch the return code.

	(( RET )) && error "There was an error unlocking the ${HUMAN_NAME} partition." ${RET}

} # unlockPartition


#---------------------------------------------------------------------------------------------------
#	Mount the partitions to prepare for fixing grub.
#---------------------------------------------------------------------------------------------------


function mountPartitions ()
{
	local -i RET					# To catch errors.

	sudo vgchange -ay				# Activate logical volumes.
	RET=${?}
	(( RET )) && error 'Error activating logical volumes.' ${RET}

	sleep 1
	sudo mkdir /mnt/root				# Create a mount point for root.
	RET=${?}
	(( RET )) && error 'Error creating a mount point for root.' ${RET}

	sudo mount /dev/mapper/system-root /mnt/root	# Mount root.
	RET=${?}
	(( RET )) && error 'Error mounting root.' ${RET}

	# Mount /boot.
	sudo mount /dev/mapper/system-boot /mnt/root/boot
	RET=${?}
	(( RET )) && error 'Error mounting /boot.' ${RET}

	sudo mount ${PARTITION_ESP} /mnt/root/boot/efi	# Mount the EFI.
	RET=${?}
	(( RET )) && error 'Error mounting the EFI System Partition (ESP).' ${RET}

	sudo mount --bind /dev /mnt/root/dev		# Mount dev in preparation for chroot.
	RET=${?}
	(( RET )) && error 'Error mounting dev.' ${RET}

	sudo mount --bind /run /mnt/root/run		# Mount run in preparation for chroot.
	RET=${?}
	(( RET )) && error 'Error mounting run.' ${RET}

	sudo mount --bind /proc /mnt/root/proc		# Mount proc in preparation for chroot.
	RET=${?}
	(( RET )) && error 'Error mounting proc.' ${RET}

	sudo mount --bind /sys /mnt/root/sys		# Mount sys in preparation for chroot.
	RET=${?}
	(( RET )) && error 'Error mounting sys.' ${RET}
} # mountPartitions


#---------------------------------------------------------------------------------------------------
#	Validate existence of a file.
#
# Parameter
#	1	The path to the file
#
# Output
#	An error message to &2 if invalid.
#---------------------------------------------------------------------------------------------------


function validateFileExistence ()
{
	local -r FILE_PATH="${1}"

	if [ ! -f "${FILE_PATH}" ]
	then
		error "File ${FILE_PATH} not found." 1
	fi

} # validateFileExistence


#---------------------------------------------------------------------------------------------------
#	Enter chroot and continue within it.
#---------------------------------------------------------------------------------------------------


function chrootRefreshGrub ()
{
	local -i RET					# To catch errors.

	# Begin chroot and continue within it; wait until it has finished before continuing.
	sudo chroot /mnt/root /usr/local/sbin/refreshgrub
	RET=${?}
	(( RET )) && error 'chroot returned an error.' ${RET}

} # chrootRefreshGrub


#===================================================================================================
#
#	Section: Finish up the script.
#
#===================================================================================================


#---------------------------------------------------------------------------------------------------
#	Let the user know what to do next.
#---------------------------------------------------------------------------------------------------


function finalise ()
{
	cat <<-END


		${SPACER}


		${W}The installation process is complete!${R}

		Thank you for your patience.

		${B}Please reboot the computer.${R}

		You may close this terminal window at any time.
	END


} # finalise


#===================================================================================================
#
#	Section: Main script control.
#
#===================================================================================================


####################################################################################################
#	MAIN SCRIPT CONTROL
####################################################################################################


initialise						# Set up the script.
gatherData						# Gather all the required data.
fixGrubProcess					# Fix GRUB.
finalise						# Finish up the script.
