# 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
}

# A dialog is the only report this module has, and once it is dismissed
# nothing is left: jeos-firstboot.service has "StandardOutput=tty" and
# lets StandardError inherit it, so not even stderr reaches the journal.
# The captured output of sdbootutil is the whole diagnostic of a failed
# enrollment, and this is the last place where it still exists.
#
# `sdbootutil-enroll` needs none of this: it is a service, so its output
# is already journaled
#
# Never fails the caller: a missing logger must not turn a reportable
# problem into an unreported one
log_enroll() {
	logger -t jeos-firstboot-enroll -p "user.${1:?}" -- "$2" 2> /dev/null || :
}

# The same text in the dialog and in the journal
report_error() {
	log_enroll err "$1"
	d --msgbox "$1" 0 0
}

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

enroll_systemd_firstboot() {
	welcome_screen_with_console_switch

	# There is a LUKS2 device, so returning quietly here means a
	# machine that keeps asking for the password that the installer
	# set, with nothing to say why.  The dialog times out because
	# firstboot also runs unattended, when every other answer comes
	# from a credential.  `d` is not used for it: dialog exits 255
	# when the time is over, and that is the same code as ESC, so it
	# would ask "Do you really want to quit?"
	local msg
	[ -e /usr/bin/systemd-cryptenroll ] || {
		msg=$"systemd-cryptenroll is not installed, so the encrypted devices cannot be enrolled."
		log_enroll err "$msg"
		d_styled --timeout 60 --msgbox "$msg" 0 0 || :
		return 0
	}
	# The key that disk-encryption-tool adds lives only in the memory
	# of the boot that created it, so a firstboot that is delayed to
	# a later boot finds nothing here
	crypt_keyid="$(keyctl id %user:cryptenroll 2> /dev/null)" || true
	[ -n "$crypt_keyid" ] || {
		msg=$"The password of the encrypted devices is not available, so they cannot be enrolled. Use 'jeos-config enroll' to do it later."
		log_enroll err "$msg"
		d_styled --timeout 60 --msgbox "$msg" 0 0 || :
		return 0
	}

	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
			report_error "Error: $result"
		fi
	done

	return 0
}

write_issue_file() {
	local content="$1"
	# The recovery key and the recovery PIN are different secrets
	# that can both be produced by the same run, so they cannot
	# share a file: the second call would truncate the first
	local name="${2:-90-recovery-key}"
	local issuefile="/run/issue.d/$name.issue"

	if [ -x '/usr/sbin/issue-generator' ]; then
		issuefile="/run/issue.d/$name.conf"
	fi
	mkdir -p "/run/issue.d"
	echo "$content" > "$issuefile"
	# Not a trailing `[ -x ... ] && run issue-generator`: as the
	# last command of the function that returns 1 when the generator
	# is not installed
	if [ -x '/usr/sbin/issue-generator' ]; then
		run issue-generator
	fi
}

# `systemd-pcrlock` generates a recovery PIN for the NV index when
# nothing supplies one, and sdbootutil only prints it, so without this
# the credential needed to re-authorize the policy later is lost with
# the captured output.  When a recovery key was enrolled first it is
# reused as the PIN, and then there is nothing new to show
write_recovery_pin() {
	grep -q "^Recovery PIN:" <<< "$1" || return 0
	write_issue_file "$1" 91-recovery-pin
}

# Remove the keyslot that disk-encryption-tool left behind, which it
# marks with a token of its own, of type "enrollment-key", and only
# when it generated the key itself: with a user-supplied key there is
# no token, the keyslot is an ordinary passphrase and it is left alone
#
# The type comes from the LUKS2 JSON and not from the "SLOT TYPE" table
# that systemd-cryptenroll prints, where every token it does not know
# is "other".  A Clevis or a tang binding also prints "other", and
# wiping that removes somebody else's way into the device
#
# The slot to remove is the one that carries the token, and not the
# slot 0: they are the same only when disk-encryption-tool was the first
# to write to the header.  When it was not, wiping the slot 0 removes
# the password of whoever owns the machine and keeps the transient key,
# which is exactly backwards
#
# And it is removed only once a slot that somebody can actually use is
# left.  systemd-cryptenroll does refuse to empty a header, but counting
# slots is not enough here: the key of this one is random, never shown,
# and gone when the boot ends, so a device left with only that is a
# device nobody can open
#
# Keep this identical to the copy in sdbootutil-enroll.  Nothing is
# reported from here so that both can share it, as this one has a dialog
# instead of a terminal.  Returns 1 when the slot was kept because
# nothing else opens the device, and 2 when the wipe itself failed
wipe_enrollment_key() {
	local dev="${1:?}"
	local slot header usable
	local transient=()
	local rc=0

	# luksDump reads the metadata, so there is no password to give
	header="$(cryptsetup luksDump --dump-json-metadata "$dev" 2> /dev/null)" || return 2

	# The keyslots the tokens point at ($t) and the keyslots that
	# exist ($s).  Only the intersection is wiped, as wiping a slot
	# that is not there is just an error.  cryptsetup on its own does
	# not produce that: it refuses to import a token naming a missing
	# slot, and empties the "keyslots" of a token whose slot is
	# killed.  But this header is written by several tools
	# shellcheck disable=SC2016 # $t and $s are jq variables
	local q='[.tokens[]? | select(.type == "enrollment-key") | .keyslots[]?] as $t
		| [.keyslots | keys[]] as $s | '

	mapfile -t transient < <(jq -r "$q"'($s - ($s - $t))[]' <<< "$header")

	[ "${#transient[@]}" -ne 0 ] || return 0

	# Everything else opens the device: a keyslot with no token is a
	# passphrase, and one with a token of another type belongs to
	# whoever enrolled it
	usable="$(jq -r "$q"'($s - $t) | length' <<< "$header")" || return 2

	[ "$usable" -gt 0 ] || return 1

	for slot in "${transient[@]}"; do
		systemd-cryptenroll --wipe-slot="$slot" "$dev" || rc=2
	done

	return "$rc"
}

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

	do_enroll
}

do_enroll() {
	local out r error=0

	# `sdbootutil` writes the secrets it generates (the recovery key,
	# the recovery PIN) in stdout, and every diagnostic in stderr, so
	# the two are captured apart.  Merging them with "2>&1" is what
	# would hand a secret to `report_error` the moment an enrollment
	# fails after having printed one, and half of that report is the
	# `logger` call in `log_enroll`, which reaches the journal on disk
	# and stays there.  The dialog is fine, it dies with the tty, but
	# both get the same text.  What has to be reported is the
	# diagnostic, and that is stderr
	local errfile
	errfile="$(mktemp -t jeos-firstboot-enroll.XXXXXX)"

	# What was chosen, so that the journal can answer "was anything
	# enrolled at all, and what" without the person who ran firstboot
	# having to remember the menu
	log_enroll info "Enrolling:${with_recovery_key:+ recovery-key}${crypt_pw:+ password}${with_tpm2:+ tpm2}${crypt_tpm_pin:+ (with PIN)}${with_fido2:+ fido2}"

	[ -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> "$errfile")"
		r="$?"
		if [ $r -ne 0 ]; then
			report_error "Error (recovery-key): $(cat "$errfile")"
			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> "$errfile")"
		r="$?"
		[ $r -eq 0 ] || {
			report_error "Error (password): $(cat "$errfile")"
			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> "$errfile")"
			r="$?"
			if [ $r -ne 0 ]; then
				report_error "Error (TPM2+PIN): $(cat "$errfile")"
				error=1
			else
				write_recovery_pin "$out"
			fi
		else
			d --infobox "Enrolling TPM2 ..." 3 40
			out="$(SDB_ADD_INITIAL_COMPONENT=1 run sdbootutil enroll --method=tpm2 2> "$errfile")"
			r="$?"
			if [ $r -ne 0 ]; then
				report_error "Error (TPM2): $(cat "$errfile")"
				error=1
			else
				write_recovery_pin "$out"
			fi
		fi
	fi

	[ -z "$with_fido2" ] || {
		# systemd-cryptenroll requires the presence of the user
		# on the security token, but it prints the request in
		# the output that is captured here, so the user needs
		# to be warned in advance
		d --infobox "Enrolling FIDO2 ...\n\nConfirm the presence on the security token when it blinks" 6 60
		out="$(run sdbootutil enroll --method=fido2 2> "$errfile")"
		r="$?"
		[ $r -eq 0 ] || {
			report_error "Error (FIDO2): $(cat "$errfile")"
			error=1
		}
	}

	if [ "$error" -eq 1 ]; then
		report_error "One or more enrollment methods failed. Resolve the issue and re-try with 'jeos-config enroll'"
	elif [ -n "$with_recovery_key$crypt_pw$with_tpm2$with_fido2" ]; then
		# Clean the enrollment key that disk-encryption-tool
		# created
		#
		# Only when something else was enrolled: "Done" can leave
		# the menu with nothing selected under jeos-config, and
		# then that keyslot is still the only one that opens the
		# device
		local dev
		while read -r dev; do
			wipe_enrollment_key "$dev"
			case "$?" in
				1) report_error "Keeping the enrollment key of $dev, as no other method can open it" ;;
				2) report_error "Failed to remove the enrollment key of $dev" ;;
			esac
		done < <(sdbootutil list-devices)
		log_enroll info "Enrollment finished"
	fi

	rm -f "$errfile"
}

# A password that does not open the device is only noticed by the
# enrollment itself, one initrd rebuild later, and is reported as an
# enrollment failure
valid_password() {
	local pw="$1" dev

	[ -n "$pw" ] || return 1
	while read -r dev; do
		# Without --disable-external-tokens an enrolled TPM2 or
		# FIDO2 token opens the device and any password is
		# reported as valid
		cryptsetup luksOpen --test-passphrase --disable-external-tokens "$dev" <<<"$pw" &> /dev/null && return 0
	done < <(sdbootutil list-devices)
	return 1
}

enroll_jeos_config() {
	is_jeos_config=1
	while true; do
		d --insecure --passwordbox  $"Enter decryption password" 0 0
		[ -n "$result" ] || return 0
		valid_password "$result" && break
		d_styled --yesno $"The password does not open any encrypted device. Retry?" 0 0 || return 0
	done
	# Everything that comes next takes the password from here, so
	# without it the enrollment can only fail, once per method
	echo -n "$result" | keyctl padd user cryptenroll @u &> /dev/null || {
		report_error $"Failed to store the password in the kernel keyring"
		return 0
	}

	enroll_systemd_firstboot
	do_enroll
}
