# shellcheck disable=SC2034
enroll_title=$"Full Disk Encryption Enrollment"
# shellcheck disable=SC2034
enroll_description=$"Enroll a device using TPM2 or FIDO2 key"

crypt_keyid=""
crypt_pw=""
crypt_tpm_pin=""

with_fido2=
with_tpm2=
with_recovery_key=

luks2_devices=()


have_luks2() {
	lsblk --noheadings -o FSTYPE | grep -q crypto_LUKS
}

# exit early without defining any helper functions if there are no luks devices
have_luks2 || return 0

enroll_systemd_firstboot() {
	[ -e /usr/bin/systemd-cryptenroll ] || return 0
	crypt_keyid="$(keyctl id %user:cryptenroll 2> /dev/null)" || return 0
	[ -n "$crypt_keyid" ] || return 0

	welcome_screen_with_console_switch

	local has_fido2=${JEOS_HAS_FIDO2:-}
	local has_tpm2=

	[ -z "$(systemd-cryptenroll --fido2-device=list 2>/dev/null)" ] || has_fido2=1
	[ ! -e '/sys/class/tpm/tpm0' ] || has_tpm2=lock

	while true; do
		local list=()

		if [ -z "$with_recovery_key" ]; then
			list+=('recovery-key' $'Enroll recovery key')
		fi
		if [ -z "$with_fido2" ] && [ -z "$with_tpm2" ] && [ -n "$has_fido2" ]; then
			list+=('FIDO2' $'Enroll FIDO2 token')
		fi
		if [ -z "$with_tpm2" ] && [ -z "$with_fido2" ] && [ -n "$has_tpm2" ]; then
			list+=('TPM2' $'Enroll TPM2 based token' 'TPM2_interactive' 'Enroll TPM2 based token with PIN')
		fi
		if [ -z "$crypt_pw" ]; then
			if [ -n "$password" ]; then
				list+=('root' $'Enroll root password')
			fi
			list+=('password' $'Enroll extra password')
		fi
		[ -n "$list" ] || break

		list+=('done' $'Done')

		d --no-tags --default-item "${list[0]}" --menu $"Disk Encryption" 0 0 "$(menuheight ${#list[@]})" "${list[@]}"
		if [ "$result" = 'done' ]; then
			if [ -z "$with_recovery_key" ] && [ -z "$crypt_pw" ] && [ -z "$with_fido2" ] && [ -z "$with_tpm2" ] && [ -z "$is_jeos_config" ]; then
				d_styled --msgbox $"Can not continue without selecting an enrollment" 5 52
				continue
			fi
			break;
		elif [ "$result" = 'FIDO2' ]; then
			with_fido2=1
		elif [ "$result" = 'TPM2' ]; then
			with_tpm2="$has_tpm2"
		elif [ "$result" = 'TPM2_interactive' ]; then
			while true; do
				d --insecure --passwordbox  $"Enter new PIN (actually just passphrase)" 0 0
				if [ -z "$result" ]; then
					d_styled --yesno $"Retry?" 0 0 || break
					continue
				fi
				crypt_tpm_pin="$result"
				d --insecure --passwordbox  $"Confirm PIN" 0 0
				[ "$crypt_tpm_pin" != "$result" ] || { with_tpm2="$has_tpm2"; break; }
				d --msgbox $"PINs don't match. Try again" 0 0
			done
		elif [ "$result" = 'recovery-key' ]; then
			with_recovery_key=1
		elif [ "$result" = 'root' ]; then
			crypt_pw="$password"
		elif [ "$result" = 'password' ]; then
			while true; do
				d --insecure --passwordbox  $"Enter encryption password" 0 0
				if [ -z "$result" ]; then
					d --aspect 29 --msgbox $"No encryption password set. You can add more keys manually using systemd-cryptenroll." 0 0
					break
				fi
				crypt_pw="$result"
				d --insecure --passwordbox  $"Confirm encryption password" 0 0
				[ "$crypt_pw" != "$result" ] || break
				d --msgbox $"Passwords don't match. Try again" 0 0
			done
		else
			d --msgbox "Error: $result" 0 0
		fi
	done

	return 0
}

write_issue_file() {
	local recovery_key="$1"
	local issuefile="/run/issue.d/90-recovery-key.issue"

	[ -x '/usr/sbin/issue-generator' ] && issuefile="/run/issue.d/90-recovery-key.conf"
	mkdir -p "/run/issue.d"
	echo "$recovery_key" > "$issuefile"
	[ -x '/usr/sbin/issue-generator' ] && run issue-generator
}

enroll_post() {
	[ -e /usr/bin/systemd-cryptenroll ] || return 0
	[ -n "$crypt_keyid" ] || return 0

	do_enroll
}

do_enroll() {
	local out r error=0
	[ -z "$with_recovery_key" ] || {
		# Note that if --no-reuse-initrd is used, then a new
		# initrd will be created and will break the
		# measurement of the initial components if later the
		# TPM2 enrollment is called
		extra=
		if [ -z "$with_tpm2" ] && [ -z "$with_fido2" ]; then
			extra="--no-reuse-initrd"
		fi
		d --infobox "Enrolling recovery-key ..." 3 40
		out="$(run sdbootutil enroll --method=recovery-key "$extra" 2>&1)"
		r="$?"
		if [ $r -ne 0 ]; then
			d --msgbox "Error (recovery-key): $out" 0 0
			error=1
		else
			write_issue_file "$out"
		fi
	}

	[ -z "$crypt_pw" ] || {
		# Note that if --no-reuse-initrd is used, then a new
		# initrd will be created and will break the
		# measurement of the initial components if later the
		# TPM2 enrollment is called
		extra=
		if [ -z "$with_tpm2" ] && [ -z "$with_fido2" ]; then
			extra="--no-reuse-initrd"
		fi
		d --infobox "Enrolling password ..." 3 40
		out="$(PW="$crypt_pw" run sdbootutil enroll --method=password "$extra" 2>&1)"
		r="$?"
		[ $r -eq 0 ] || {
			d --msgbox "Error (password): $out" 0 0
			error=1
		}
	}

	if [ -n "$with_tpm2" ]; then
		if [ -n "$crypt_tpm_pin" ]; then
			d --infobox "Enrolling TPM2+PIN ..." 3 40
			out="$(SDB_ADD_INITIAL_COMPONENT=1 PIN="$crypt_tpm_pin" run sdbootutil enroll --method=tpm2+pin 2>&1)"
			r="$?"
			[ $r -eq 0 ] || {
				d --msgbox "Error (TPM2+PIN): $out" 0 0
				error=1
			}

		else
			d --infobox "Enrolling TPM2 ..." 3 40
			out="$(SDB_ADD_INITIAL_COMPONENT=1 run sdbootutil enroll --method=tpm2 2>&1)"
			r="$?"
			[ $r -eq 0 ] || {
				d --msgbox "Error (TPM2): $out" 0 0
				error=1
			}
		fi
	fi

	[ -z "$with_fido2" ] || {
		d --infobox "Enrolling FIDO2 ..." 3 40
		out="$(run sdbootutil enroll --method=fido2 2>&1)"
		r="$?"
		[ $r -eq 0 ] || {
			d --msgbox "Error (FIDO2): $out" 0 0
			error=1
		}
	}

	if [ "$error" -eq 1 ]; then
		d --msgbox "One or more enrollment methods failed. Resolve the issue and re-try with 'jeos-config enroll'" 0 0
	else
		# Clean the enrollment key.  disk-encryption-tool
		# creates it in the keyslot 0 with the name
		# "enrollment-key", that is showed by
		# systemd-cryptenroll as "other"
		local slots
		while read -r dev; do
			slots=$(systemd-cryptenroll "$dev")
			if grep -q "other" <<<"$slots"; then
				systemd-cryptenroll --wipe-slot=0 "$dev"
			fi
		done < <(sdbootutil list-devices)
	fi
}

enroll_jeos_config() {
	is_jeos_config=1
	d --insecure --passwordbox  $"Enter decryption password" 0 0
	[ -n "$result" ] || return 0
	echo -n "$result" | keyctl padd user cryptenroll @u &> /dev/null

	enroll_systemd_firstboot
	do_enroll
}
