#!/bin/bash
# Save these options to be the persistent default
set -e -o pipefail
. /tmp/config
. /etc/functions.sh

TRACE_FUNC

while getopts "b:d:p:i:" arg; do
	case $arg in
		b) bootdir="$OPTARG" ;;
		d) paramsdev="$OPTARG" ;;
		p) paramsdir="$OPTARG" ;;
		i) index="$OPTARG" ;;
	esac
done

if [ -z "$bootdir" -o -z "$index" ]; then
	DIE "Usage: $0 -b /boot -i menu_option "
fi

if [ -z "$paramsdev" ]; then
	paramsdev="$bootdir"
fi

if [ -z "$paramsdir" ]; then
	paramsdir="$bootdir"
fi

bootdir="${bootdir%%/}"
paramsdev="${paramsdev%%/}"
paramsdir="${paramsdir%%/}"

TMP_MENU_FILE="/tmp/kexec/kexec_menu.txt"
ENTRY_FILE="$paramsdir/kexec_default.$index.txt"
HASH_FILE="$paramsdir/kexec_default_hashes.txt"
PRIMHASH_FILE="$paramsdir/kexec_primhdl_hash.txt"
KEY_DEVICES="$paramsdir/kexec_key_devices.txt"
KEY_LVM="$paramsdir/kexec_key_lvm.txt"

lvm_suggest=$(run_lvm vgscan 2>/dev/null | awk -F '"' {'print $1'} | tail -n +2)
num_lvm=$(echo "$lvm_suggest" | wc -l)
if [ "$num_lvm" -eq 1 ] && [ -n "$lvm_suggest" ]; then
	lvm_volume_group="$lvm_suggest"
elif [ -z "$lvm_suggest" ]; then
	num_lvm=0
fi
# $lvm_suggest is a multiline string, we need to convert it to a space separated string
lvm_suggest=$(echo $lvm_suggest | tr '\n' ' ')
DEBUG "LVM num_lvm: $num_lvm, lvm_suggest: $lvm_suggest"

# get all LUKS container devices
devices_suggest=$(blkid | cut -d ':' -f 1 | while read device; do
	if cryptsetup isLuks "$device"; then echo "$device"; fi
done | sort)
num_devices=$(echo "$devices_suggest" | wc -l)

if [ "$num_devices" -eq 1 ] && [ -s "$devices_suggest" ]; then
	key_devices=$devices_suggest
elif [ -z "$devices_suggest" ]; then
	num_devices=0
fi
# $devices_suggest is a multiline string, we need to convert it to a space separated string
devices_suggest=$(echo $devices_suggest | tr '\n' ' ')
DEBUG "LUKS num_devices: $num_devices, devices_suggest: $devices_suggest"

if [ "$num_lvm" -eq 0 ] && [ "$num_devices" -eq 0 ]; then
	#No encrypted partition found.
	no_encrypted_partition=1
fi

#Reusable function when user wants to define new TPM DUK for lvms/disks
prompt_for_existing_encrypted_lvms_or_disks() {
	TRACE_FUNC
	DEBUG "num_lvm: $num_lvm, lvm_suggest: $lvm_suggest, num_devices: $num_devices, devices_suggest: $devices_suggest"

	# Create an associative array to store the suggested LVMs and their paths
	declare -A lvms_array
	# Loop through the suggested LVMs and add them to the array
	for lvm in $lvm_suggest; do
		lvms_array[$lvm]=$lvm
	done

	# Get the number of suggested LVMs
	num_lvms=${#lvms_array[@]}

	if [ "$num_lvms" -gt 1 ]; then
		DEBUG "Multiple LVMs found: $lvm_suggest"
		selected_lvms_not_existing=1
		# Create an array to store the selected LVMs
		declare -a key_lvms_array

		while [ $selected_lvms_not_existing -ne 0 ]; do
			# Read the user input and store it in a variable
			INPUT "Encrypted LVMs? (type 'all' to select all, or space-separated subset of: $lvm_suggest):" -r key_lvms

			# 'all' expands to every discovered LVM
			if [ "$key_lvms" = "all" ]; then
				key_lvms="$lvm_suggest"
				DEBUG "User chose 'all' LVMs: $key_lvms"
				selected_lvms_not_existing=0
				continue
			fi

			# Split the user input by spaces and validate each element
			IFS=' ' read -r -a key_lvms_array <<<"$key_lvms"

			if [ ${#key_lvms_array[@]} -eq 0 ]; then
				continue
			fi

			valid=1
			for lvm in "${key_lvms_array[@]}"; do
				if [[ ! ${lvms_array[$lvm]+_} ]]; then
					valid=0
					break
				fi
			done

			if [[ $valid -eq 1 ]]; then
				selected_lvms_not_existing=0
			fi
		done
	elif [ "$num_lvms" -eq 1 ]; then
		INFO "Single Encrypted LVM found at $lvm_suggest"
		key_lvms=$lvm_suggest
	else
		DEBUG "No encrypted LVMs found"
	fi

	# Create an associative array to store the suggested devices and their paths
	declare -A devices_array
	# Loop through the suggested devices and add them to the array
	for device in $devices_suggest; do
		devices_array[$device]=$device
	done

	# Get the number of suggested devices
	num_devices=${#devices_array[@]}

	if [ "$num_devices" -gt 1 ]; then
		DEBUG "Multiple LUKS devices found: $devices_suggest"
		selected_luksdevs_not_existing=1
		# Create an array to store the selected devices
		declare -a key_devices_array

		while [ $selected_luksdevs_not_existing -ne 0 ]; do
			# Read the user input and store it in a variable
			INPUT "Encrypted devices? (type 'all' to select all, or space-separated subset of: $devices_suggest):" -r key_devices

			# 'all' expands to every discovered LUKS device
			if [ "$key_devices" = "all" ]; then
				key_devices="$devices_suggest"
				DEBUG "User chose 'all' LUKS devices: $key_devices"
				selected_luksdevs_not_existing=0
				continue
			fi

			# Split the user input by spaces and validate each element
			IFS=' ' read -r -a key_devices_array <<<"$key_devices"

			if [ ${#key_devices_array[@]} -eq 0 ]; then
				continue
			fi

			valid=1
			for device in "${key_devices_array[@]}"; do
				if [[ ! ${devices_array[$device]+_} ]]; then
					valid=0
					break
				fi
			done

			if [[ $valid -eq 1 ]]; then
				selected_luksdevs_not_existing=0
			fi
		done
	elif [ "$num_devices" -eq 1 ]; then
		INFO "Single Encrypted Disk found at $devices_suggest"
		key_devices=$devices_suggest
	else
		DEBUG "No encrypted devices found"
	fi

	DEBUG "Multiple LUKS devices selected: $key_devices"

}

if [ ! -r "$TMP_MENU_FILE" ]; then
	DIE "No menu options available, please run kexec-select-boot.sh"
fi

entry=$(head -n $index $TMP_MENU_FILE | tail -1)
if [ -z "$entry" ]; then
	DIE "Invalid menu index $index"
fi

save_key="n"

if [ "$CONFIG_TPM" = "y" ] && [ "$CONFIG_TPM_NO_LUKS_DISK_UNLOCK" != "y" ] && [ "$CONFIG_BASIC" != y ]; then
	DEBUG "TPM is enabled and TPM_NO_LUKS_DISK_UNLOCK is not set"
	DEBUG "Checking if a a LUKS TPM Disk Unlock Key was previously set up from $KEY_DEVICES"
	#check if $KEY_DEVICES file exists and is not empty
	if [ -r "$KEY_DEVICES" ] && [ -s "$KEY_DEVICES" ]; then
		DEBUG "LUKS TPM Disk Unlock Key was previously set up from $KEY_DEVICES"
		INPUT "Do you want to reseal a Disk Unlock Key in the TPM [y/N]:" -n 1 change_key_confirm

		if [ "$change_key_confirm" = "y" \
			-o "$change_key_confirm" = "Y" ]; then
			old_lvm_volume_group=""
			if [ -r "$KEY_LVM" ]; then
				old_lvm_volume_group=$(cat $KEY_LVM) || true
				old_key_devices=$(cat $KEY_DEVICES |
					cut -d\  -f1 |
					grep -v "$old_lvm_volume_group" |
					xargs) || true
			else
				old_key_devices=$(cat $KEY_DEVICES |
					cut -d\  -f1 | xargs) || true
			fi

			lvm_suggest="$old_lvm_volume_group"
			devices_suggest="$old_key_devices"
			save_key="y"
		fi
	else
		DEBUG "No previous LUKS TPM Disk Unlock Key was set up, confirming to add a Disk Unlock Key (DUK) to the TPM"
		INPUT "Do you wish to add a LUKS TPM Disk Unlock Key (DUK) to the TPM [y/N]:" -n 1 add_key_confirm

		if [ "$add_key_confirm" = "y" \
			-o "$add_key_confirm" = "Y" ]; then
			DEBUG "User confirmed desire to add a Disk Unlock Key (DUK) to the TPM"
			save_key="y"
		fi
	fi

	if [ "$save_key" = "y" ]; then
		if [ -n "$old_key_devices" ] || [ -n "$old_lvm_volume_group" ]; then
			DEBUG "Previous LUKS TPM Disk Unlock Key was set up for $old_key_devices $old_lvm_volume_group"
			INPUT "Do you want to reuse configured Encrypted LVM groups/Block devices? (Y/n):" -n 1 reuse_past_devices
			if [ "$reuse_past_devices" = "y" ] || [ "$reuse_past_devices" = "Y" ] || [ -z "$reuse_past_devices" ]; then
				if [ -z "$key_devices" ] && [ -n "$old_key_devices" ]; then
					key_devices="$old_key_devices"
				fi
				if [ -z "$lvm_volume_group" ] && [ -n "$old_lvm_volume_group" ]; then
					lvm_volume_group="$old_lvm_volume_group"
				fi
			#User doesn't want to reuse past devices, so we need to prompt him from devices_suggest and lvm_suggest
			else
				prompt_for_existing_encrypted_lvms_or_disks
			fi
		else
			DEBUG "No previous LUKS TPM Disk Unlock Key was set up, setting up"
			prompt_for_existing_encrypted_lvms_or_disks
		fi

		save_key_params="-s -p $paramsdev"
		if [ -n "$lvm_volume_group" ]; then
			save_key_params="$save_key_params -l $lvm_volume_group $key_devices"
		else
			save_key_params="$save_key_params $key_devices"
		fi
		kexec-save-key.sh $save_key_params ||
			DIE "Failed to save the LUKS TPM Disk Unlock Key"
	fi
fi

# try to switch to rw mode
mount -o rw,remount $paramsdev

if [ ! -d $paramsdir ]; then
	mkdir -p $paramsdir ||
		DIE "Failed to create params directory"
fi

# All writes go to a staging directory; files are moved to their final
# locations only after kexec-sign-config.sh succeeds.  This prevents a failed
# signing attempt from leaving /boot with updated config but no matching
# signature.
stagedir=$(mktemp -d /tmp/kexec-default-XXXXXX)
cleanup_stagedir() { rm -rf "$stagedir"; }
trap cleanup_stagedir EXIT
DEBUG "Staging directory created: $stagedir"

# Seed staging with existing kexec*.txt so sign-config signs everything
for f in "$paramsdir"/kexec*.txt; do
	[ -e "$f" ] && cp "$f" "$stagedir/"
done
# Seed staging with existing crypttab overrides if present
[ -e "$bootdir/kexec_initrd_crypttab_overrides.txt" ] && \
	cp "$bootdir/kexec_initrd_crypttab_overrides.txt" "$stagedir/"
DEBUG "Seeded $stagedir with existing config files from $paramsdir"

if [ "$CONFIG_TPM2_TOOLS" = "y" ]; then
	if [ -f /tmp/secret/primary.handle ]; then
		DEBUG "Hashing TPM2 primary key handle..."
		sha256sum /tmp/secret/primary.handle > "$stagedir/kexec_primhdl_hash.txt" ||
			DIE "ERROR: Failed to Hash TPM2 primary key handle!"
		DEBUG "TPM2 primary key handle hash written to $stagedir/kexec_primhdl_hash.txt"
	else
		DIE "ERROR: TPM2 primary key handle file does not exist!"
	fi
fi

# Remove old kexec_default.*.txt from staging; new entry written below
rm "$stagedir"/kexec_default.*.txt 2>/dev/null || true
echo "$entry" >"$stagedir/kexec_default.$index.txt"
(
	hash_pipeline_exit=0
	cd $bootdir || exit 1
	kexec-boot.sh -b "$bootdir" -e "$entry" -f |
		xargs sha256sum >"$stagedir/kexec_default_hashes.txt" ||
		hash_pipeline_exit=$?
	exit $hash_pipeline_exit
) || DIE "Failed to create hashes of boot files"
if [ ! -r "$stagedir/kexec_default.$index.txt" -o ! -r "$stagedir/kexec_default_hashes.txt" ]; then
	DIE "Failed to write default config to staging"
fi

if [ "$save_key" = "y" ]; then
	# logic to parse OS initrd to extract crypttab, its filepaths and its OS defined options
	initrd_decompressed="/tmp/initrd_extract"
	mkdir -p "$initrd_decompressed"
	# Get initrd filename selected to be default initrd that OS could be using to configure LUKS on boot by deploying crypttab files
	current_default_initrd=$(cat "$stagedir/kexec_default_hashes.txt" | grep initr | awk -F " " {'print $NF'} | sed 's/\.\//\/boot\//g')

	DEBUG "Extracting $current_default_initrd to find crypttab files"
	unpack_initramfs.sh "$current_default_initrd" "$initrd_decompressed"
	crypttab_files=$(find "$initrd_decompressed" -type f | grep crypttab 2>/dev/null) || true

	if [ ! -z "$crypttab_files" ]; then
		DEBUG "Found crypttab files in $current_default_initrd"
		rm -f "$stagedir/kexec_initrd_crypttab_overrides.txt" || true

		#Parsing each crypttab file found
		echo "$crypttab_files" | while read crypttab_file; do
			# Change crypttab file path to be relative to initrd for string manipulation
			final_initrd_filepath=${crypttab_file#/tmp/initrd_extract}
			DEBUG "Final initramfs crypttab path:$final_initrd_filepath"
			# Keep only non-commented lines for crypttab entries
			current_crypttab_entries=$(cat "$crypttab_file" | grep -v "^#")
			DEBUG "Found initrd crypttab entries $final_initrd_filepath:$current_crypttab_entries"
			# Modify each retained crypttab line for /secret.key under intramfs to be considered as a keyfile
			modified_crypttab_entries=$(echo "$current_crypttab_entries" | sed 's/none/\/secret.key/g')
			DEBUG "Modified crypttab entries $final_initrd_filepath:$modified_crypttab_entries"
			echo "$modified_crypttab_entries" | while read modified_crypttab_entry; do
				echo "$final_initrd_filepath:$modified_crypttab_entry" >>"$stagedir/kexec_initrd_crypttab_overrides.txt"
			done
		done

		#insert current default boot's initrd crypttab locations into tracking file to be overwritten into initramfs at kexec-inject-key
		STATUS "The following OS crypttab entries were modified from default boot's initrd:"
		cat "$stagedir/kexec_initrd_crypttab_overrides.txt"
		STATUS_OK "Heads added /secret.key to those entries and saved overrides for signing"
		DEBUG "Crypttab overrides will be included in signed digests and injected via cpio at kexec"
	else
		INFO "No crypttab found in initrd; a generic crypttab will be generated"
		if [ -e "$stagedir/kexec_initrd_crypttab_overrides.txt" ]; then
			DEBUG "Removing stale $stagedir/kexec_initrd_crypttab_overrides.txt (no crypttab found in initrd)"
			rm -f "$stagedir/kexec_initrd_crypttab_overrides.txt"
		fi
	fi

	# Cleanup
	cd /
	rm -rf /tmp/initrd_extract || true
fi

# sign and auto-roll config counter
extparam=
if [ "$CONFIG_TPM" = "y" ]; then
	if [ "$CONFIG_IGNORE_ROLLBACK" != "y" ]; then
		extparam=-r
	fi
fi
if [ "$CONFIG_BASIC" != "y" ]; then
	DO_WITH_DEBUG kexec-sign-config.sh -p "$stagedir" $extparam ||
		DIE "Failed to sign default config"
fi
# Signing succeeded; move all staged files to their final locations.
# kexec-sign-config.sh may have remounted $paramsdev ro; remount rw for the moves.
mount -o rw,remount $paramsdev
rm "$paramsdir"/kexec_default.*.txt 2>/dev/null || true
for f in "$stagedir"/*; do
	[ -e "$f" ] || continue
	fname="$(basename "$f")"
	if [ "$fname" = "kexec_initrd_crypttab_overrides.txt" ]; then
		mv "$f" "$bootdir/$fname"
	else
		mv "$f" "$paramsdir/$fname"
	fi
done
DEBUG "Staged files from $stagedir moved to $paramsdir (crypttab overrides to $bootdir)"
# switch back to ro mode
mount -o ro,remount $paramsdev
