#! /bin/sh
set -e

ESP_LIST="/etc/kernel/proxmox-boot-uuids"
ESPTYPE='c12a7328-f81f-11d2-ba4b-00a0c93ec93b'

MANUAL_KERNEL_LIST="/etc/kernel/proxmox-boot-manual-kernels"
PINNED_KERNEL_CONF="/etc/kernel/proxmox-boot-pin"
NEXT_BOOT_PIN="/etc/kernel/next-boot-pin"

MOUNTROOT="${TMPDIR:-/var/tmp}/espmounts"
# relative to the ESP mountpoint
PMX_ESP_DIR="EFI/proxmox"
PMX_LOADER_CONF="loader/loader.conf"
GRUB_PIN_SNIPPET="/etc/default/grub.d/proxmox-kernel-pin.cfg"

# `uki` mode: ESPs carrying the marker file get signed Unified Kernel
# Images (kernel + initrd + cmdline + signed PCR 11 prediction) in
# EFI/Linux/, auto-discovered by systemd-boot. The .pcrsig section lets
# systemd-cryptsetup unseal TPM2 LUKS keys without re-enrolling on
# kernel updates.
UKI_MODE_MARKER="${PMX_ESP_DIR}/.uki-mode" # relative to ESP mountpoint
PMX_UKI_DIR="EFI/Linux"                    # relative to ESP mountpoint

# Host-side mirror of UKI-mode ESPs by UUID. Survives `pbt format` (which
# wipes the on-ESP marker) so the mode intent is preserved across an
# operator-driven reformat.
UKI_ESP_LIST="/var/lib/proxmox-boot-tool/proxmox-boot-uki-uuids"

# Per-host PCR-policy signing keypair. Private signs each UKI's predicted
# PCR 11 measurement; public is embedded in the LUKS2 header via
# `systemd-cryptenroll --tpm2-public-key`. Loss = TPM2 cannot be re-sealed
# for new UKIs, but the recovery passphrase still unlocks.
PCR_POLICY_KEY_DIR="/var/lib/proxmox-boot-tool"
PCR_POLICY_KEY="${PCR_POLICY_KEY_DIR}/pcr-policy.key"
PCR_POLICY_PEM="${PCR_POLICY_KEY_DIR}/pcr-policy.pem"

# Secure Boot signing keypair for UKIs. Auto-generated per-host by
# ensure_sb_signing_key if missing (UKIs carry a per-host cmdline+initrd,
# so a centrally-held Proxmox key cannot pre-sign them); the operator
# must enroll the cert via mokutil on first boot. Override via env to
# point at an enterprise PKI cert that is already enrolled.
SB_KEY_DIR="/var/lib/proxmox-secure-boot-support"
SB_KEY="${PMX_SB_KEY:-$SB_KEY_DIR/db.key}"
SB_CERT="${PMX_SB_CERT:-$SB_KEY_DIR/db.crt}"

# 0 if the given ESP mountpoint is configured for uki mode. Honors either
# the on-ESP marker OR the host-persistent UUID list (so `pbt format`
# alone doesn't silently regress to default systemd-boot mode).
is_uki_esp() {
	[ -e "$1/$UKI_MODE_MARKER" ] && return 0
	[ -e "$UKI_ESP_LIST" ] || return 1
	uuid=""
	if command -V findmnt >/dev/null 2>&1; then
		uuid="$(findmnt -no UUID "$1" 2>/dev/null)"
	fi
	[ -n "$uuid" ] || return 1
	grep -qFx "$uuid" "$UKI_ESP_LIST"
}

# Generate the per-host PCR-policy keypair if missing. Idempotent.
# Called from `pbt init <esp> uki` and as a defensive measure from the
# kernel postinst hook (in case /etc was restored from backup without it).
ensure_pcr_policy_key() {
	if [ -f "$PCR_POLICY_KEY" ] && [ -f "$PCR_POLICY_PEM" ]; then
		# Verify the on-disk halves form a matching keypair: a mismatch
		# (e.g. /etc restored from backup against a different /var)
		# would make every existing UKI's .pcrsig unverifiable.
		if command -V openssl >/dev/null 2>&1; then
			if ! priv_pub="$(openssl rsa -in "$PCR_POLICY_KEY" -pubout 2>/dev/null)"; then
				warn "E: cannot read $PCR_POLICY_KEY (corrupted or wrong format)"
				warn "   - regenerate (and re-enroll TPM2)."
				return 1
			fi
			disk_pub="$(cat "$PCR_POLICY_PEM" 2>/dev/null)"
			if [ "$priv_pub" != "$disk_pub" ]; then
				warn "E: $PCR_POLICY_KEY and $PCR_POLICY_PEM do not form a"
				warn "   matching keypair - existing UKIs will fail TPM2 unseal."
				warn "   Restore from backup or regenerate (and re-enroll TPM2)."
				return 1
			fi
		fi
		return 0
	fi
	if ! command -V openssl >/dev/null 2>&1; then
		warn "E: openssl is required to generate the PCR-policy signing key"
		return 1
	fi
	mkdir -p "$PCR_POLICY_KEY_DIR"
	chmod 0700 "$PCR_POLICY_KEY_DIR"
	# RSA-2048 matches what systemd-measure expects; ECDSA support
	# across our Debian release matrix is uneven.
	old_umask="$(umask)"
	umask 0077
	# Roll back on any failure: a half-staged keypair would let a later
	# `mv` install a private key whose public half doesn't match the
	# pubkey embedded in the LUKS header.
	if ! openssl genpkey -algorithm RSA -pkeyopt rsa_keygen_bits:2048 \
		-out "${PCR_POLICY_KEY}.tmp"; then
		rm -f "${PCR_POLICY_KEY}.tmp" "${PCR_POLICY_PEM}.tmp"
		umask "$old_umask"
		warn "E: openssl genpkey failed - cannot generate PCR-policy key"
		return 1
	fi
	if ! openssl rsa -in "${PCR_POLICY_KEY}.tmp" -pubout \
		-out "${PCR_POLICY_PEM}.tmp"; then
		rm -f "${PCR_POLICY_KEY}.tmp" "${PCR_POLICY_PEM}.tmp"
		umask "$old_umask"
		warn "E: openssl rsa -pubout failed - cannot derive PCR-policy public key"
		return 1
	fi
	umask "$old_umask"
	chmod 0400 "${PCR_POLICY_KEY}.tmp"
	chmod 0444 "${PCR_POLICY_PEM}.tmp"
	mv "${PCR_POLICY_KEY}.tmp" "$PCR_POLICY_KEY"
	mv "${PCR_POLICY_PEM}.tmp" "$PCR_POLICY_PEM"
	echo "Generated per-host PCR-policy keypair at $PCR_POLICY_KEY"
}

# Generate a per-host Secure Boot signing keypair (RSA-2048 + self-signed
# X.509 cert) if missing. The operator MUST enroll the cert via mokutil
# before reboot or shim refuses the UKIs we sign; the function prints the
# enrollment hint after generation. Override via $PMX_SB_KEY/$PMX_SB_CERT
# to point at an enterprise PKI cert that is already enrolled.
ensure_sb_signing_key() {
	if [ -f "$SB_KEY" ] && [ -f "$SB_CERT" ]; then
		return 0
	fi
	if ! command -V openssl >/dev/null 2>&1; then
		warn "E: openssl is required to generate the Secure Boot signing keypair"
		return 1
	fi
	mkdir -p "$SB_KEY_DIR"
	chmod 0700 "$SB_KEY_DIR"

	# Per-host subject so MOK-enrolled certs are tellable apart in
	# firmware UI. Sanitize hostname / machine-id: Linux accepts chars
	# that openssl's `-subj /key=value/` parser silently mangles.
	machine_id="$(cat /etc/machine-id 2>/dev/null \
		| tr -c 'a-f0-9' '_' | head -c 32)"
	host_short="$(hostname -s 2>/dev/null \
		| tr -c 'A-Za-z0-9-' '_' | head -c 32)"
	[ -z "$host_short" ] && host_short="unknown"
	[ -z "$machine_id" ] && machine_id="unknown"
	subject="/CN=Proxmox UKI signing key ($host_short)/OU=$machine_id/O=Proxmox/"

	# Refuse to issue a cert with a clock-broken notBefore: dead-RTC
	# hosts boot at epoch and would produce a born-expired cert.
	year="$(date -u +%Y 2>/dev/null)"
	if [ -z "$year" ] || [ "$year" -lt 2024 ]; then
		warn "E: system clock reports year '$year'; refusing to issue a"
		warn "   long-lived cert with a wrong notBefore date. Set the"
		warn "   clock (NTP / hwclock --set) and retry."
		return 1
	fi

	old_umask="$(umask)"
	umask 0077
	# extendedKeyUsage=codeSigning is what shim/MOK looks for; CA:FALSE
	# is explicit so this leaf can't sign sub-certs.
	if ! openssl req -x509 -newkey rsa:2048 -nodes \
		-keyout "${SB_KEY}.tmp" \
		-out "${SB_CERT}.tmp" \
		-subj "$subject" \
		-addext "extendedKeyUsage = codeSigning" \
		-addext "basicConstraints = CA:FALSE" \
		-days 7300 2>/dev/null; then
		rm -f "${SB_KEY}.tmp" "${SB_CERT}.tmp"
		umask "$old_umask"
		warn "E: openssl req failed - cannot generate SB signing keypair"
		return 1
	fi
	umask "$old_umask"
	chmod 0400 "${SB_KEY}.tmp"
	chmod 0444 "${SB_CERT}.tmp"
	mv "${SB_KEY}.tmp" "$SB_KEY"
	mv "${SB_CERT}.tmp" "$SB_CERT"

	# Loud, non-skippable enrollment hint - missing this leaves the
	# host non-bootable on the next reboot.
	cat >&2 <<-EOM
	-------------------------------------------------------------------
	Generated per-host UKI signing keypair at $SB_KEY / $SB_CERT.

	You MUST enroll the public certificate via mokutil before reboot
	(or shim will refuse to load the UKIs we sign with this key):

	    mokutil --import $SB_CERT
	    # mokutil prompts for a one-time password (you choose this -
	    # it is NOT the LUKS passphrase, NOT a system account password,
	    # just an arbitrary string you'll need to retype during the
	    # next boot's MokManager dialog. Pick something memorable
	    # for the next 5 minutes; it's discarded after enrollment.)
	    reboot
	    # in MokManager: Enroll MOK -> Continue -> Yes ->
	    #   re-type the same one-time password from mokutil --import.

	If you have an enterprise PKI cert already enrolled in shim's MOK
	store or in firmware db, set PMX_SB_KEY / PMX_SB_CERT env vars to
	point at it and re-run this command - PBT will use those instead.
	-------------------------------------------------------------------
	EOM
}

# Build a signed UKI for $kver at absolute path $output: bundle of
# kernel + initrd + cmdline, signed with $SB_KEY (shim/systemd-boot
# accepts it) and carrying a .pcrsig section signed with $PCR_POLICY_KEY
# (systemd-cryptsetup can unseal TPM2 keys against this UKI's predicted
# PCR 11).
build_signed_uki() {
	kver="$1"
	cmdline="$2"
	output="$3"

	if ! command -V ukify >/dev/null 2>&1; then
		warn "E: ukify is not available - install systemd-ukify"
		return 1
	fi
	if [ ! -f "$SB_KEY" ] || [ ! -f "$SB_CERT" ]; then
		warn "E: Secure Boot signing keypair missing ($SB_KEY / $SB_CERT) -"
		warn "   install/configure proxmox-secure-boot-support"
		return 1
	fi
	if [ ! -f "$PCR_POLICY_KEY" ]; then
		warn "E: PCR-policy signing key $PCR_POLICY_KEY missing - cannot"
		warn "   produce a TPM2-unlockable UKI"
		return 1
	fi

	# `--pcr-private-key=` alone is enough to embed the .pcrsig section;
	# `--measure` is omitted because we don't need the predicted PCR
	# values printed on every kernel install.
	if ! ukify build \
		--linux "/boot/vmlinuz-$kver" \
		--initrd "/boot/initrd.img-$kver" \
		--cmdline "$cmdline" \
		--uname "$kver" \
		--secureboot-private-key "$SB_KEY" \
		--secureboot-certificate "$SB_CERT" \
		--pcr-private-key "$PCR_POLICY_KEY" \
		--pcr-public-key "$PCR_POLICY_PEM" \
		--pcr-banks sha256 \
		--output "$output.tmp"; then
		rm -f "$output.tmp" # don't leave half-written .tmp around
		return 1
	fi
	mv "$output.tmp" "$output"
}

# adapted from /etc/kernel/postinst.d/apt-auto-removal as present in
# debian's apt package:
#
# Mark as not-for-autoremoval those kernel packages that are:
#  - the currently booted version, if still installed
#  - the kernel version we've been called for
#  - the latest kernel version (as determined by debian version number)
#  - the second-latest kernel version
#  - the latest kernel version of each series (e.g. 4.13, 4.15, 5.0) by
#    marking the meta-packages
#  - the currently pinned kernel if any

kernel_keep_versions() {
	eval "$(apt-config shell DPKG Dir::bin::dpkg/f)"
	test -n "$DPKG" || DPKG="/usr/bin/dpkg"

	list="$("${DPKG}" -l | awk '/^[ih][^nc][ ]+(proxmox|pve)-kernel-[0-9]+\./ && $2 !~ /-dbg(:.*)?$/ && $2 !~ /-dbgsym(:.*)?$/ { print $2; }' \
	   | sed -e 's#^pve-kernel-##' -e 's#^proxmox-kernel-##' -e 's#-signed$##' -e 's#:[^:]\+ # #')"

	sorted_list="$(echo "$list" | sort --unique --reverse --version-sort)"

	[ -n "$1" ] && install_version="$1"

	running_version="$(uname -r | tr 'A-Z' 'a-z')"

	# ignore the currently running version if attempting a reproducible build
	if [ -n "${SOURCE_DATE_EPOCH}" ]; then
		running_version=""
	elif [ ! -e "/boot/vmlinuz-$running_version" ]; then
		# ignore the current version if it got removed, the "auto-remove" logic
		# will not be affected, because either it is installed and thus we keep
		# it in the list, or it's already removed anyway
		running_version=""
	fi

	latest_2_versions="$(echo "$sorted_list" | grep -E '^[^ ]+-pve' | head -n2 )"

	series_metapackages="$(echo "$sorted_list" | grep -Ev '^[^ ]+-pve' | head -n2)"

	oldseries="$(echo "$series_metapackages" | tail -n1)"
	oldseries_latest_kernel="$(echo "$sorted_list" | grep -E "^${oldseries}\.[^ ]+-pve" | head -n1 )"

	if [ -e "$MANUAL_KERNEL_LIST" ]; then
		manual_kernels="$(cat "$MANUAL_KERNEL_LIST")"
	fi

	pinned_kernel="$(get_first_line "$PINNED_KERNEL_CONF")"
	nextboot_kernel="$(get_first_line "$NEXT_BOOT_PIN")"

	kernels="$(cat <<-EOF
		$running_version
		$install_version
		$manual_kernels
		$latest_2_versions
		$series_metapackages
		$oldseries_latest_kernel
		$pinned_kernel
		$nextboot_kernel
		EOF
	)"

	echo "$kernels" | sort -u | sed -e '/^$/ d'
}

#bootable kernels are the same as the no_autoremove ones without the meta-package
boot_kernel_list() {
	list="$(kernel_keep_versions "$@")"

	echo "$list" | grep -vE '^[0-9]+\.[0-9]+$' || true
}

warn() {
	echo "$@" 1>&2
}

reexec_in_mountns() {
	if [ -z "$PVE_EFIBOOT_UNSHARED" ]; then
		export PVE_EFIBOOT_UNSHARED=1
		echo "Re-executing '$0' in new private mount namespace.."
		unshare --mount --propagation private "$0" "$@"
		exit 0
	fi
}

loop_esp_list() {
	if [ ! -e ${ESP_LIST} ]; then
		return 2
	fi

	# Input redirection, NOT a pipe: a `cat ... | while` puts the body
	# in a subshell, swallowing any_esp_failed=1 set by update_esp_func.
	while IFS= read -r curr_uuid; do
		if [ -z "$curr_uuid" ]; then
			continue
		fi
		"$@"
	done < "${ESP_LIST}"
}

get_first_line() {
	file="$1"
	if [ ! -e  "$file" ]; then
	    echo ""
	    return
	fi

	while IFS= read -r line || [ -n "$line" ]; do
		break
	done < "${file}"
	echo "$line"
}

set_grub_default() {
	kver="$1"

	if [ -z "${kver}" ]; then
		rm -f "${GRUB_PIN_SNIPPET}"
	else
		# grub menu entry ids contain the internal root-device id (e.g. for zfs the GUID of
		# the pool printed in hex) as this is independent of the ESP (or grub location)
		# take it from /boot/grub/grub.cfg
		root_devid=$(sed -rn "s/.*gnulinux-advanced-(.+)['] \{$/\1/p" \
			/boot/grub/grub.cfg)
		entry="gnulinux-advanced-${root_devid}>gnulinux-${kver}-advanced-${root_devid}"
		echo "GRUB_DEFAULT=\"${entry}\"" > "${GRUB_PIN_SNIPPET}"
	fi
}

set_systemd_boot_default() {
	mountpoint="$1"
	kver="$2"
	if [ -z "${kver}" ]; then
		entry="proxmox-*"
	elif is_uki_esp "${mountpoint}"; then
		# Type-2 entries are auto-discovered; default selector matches
		# the entry id (filename without .efi), not a .conf name.
		entry="proxmox-${kver}"
	else
		entry="proxmox-${kver}.conf"
	fi

	# replaces the current default entry, if one exists else append it at the end of the file
	sed -ri "/^default /{h;s/ .*\$/ ${entry}/};\${x;/^$/{s//default ${entry}/;H};x}" \
		"${mountpoint}/$PMX_LOADER_CONF"

}
