#!/usr/bin/bash
## setup auto-unlock LUKS2 encrypted root on Fedora/Silverblue/maybe others
set -eou pipefail

[ "$UID" -eq 0 ] || { echo "This script must be run as root."; exit 1;}

echo "WARNING: Do NOT use this if your CPU is vulnerable to faulTPM!"
echo "All AMD Zen1, Zen2, and Zen3 Processors are known to be affected!"
echo "Zen 4 Processors are affected by a similar exploit which can be mitigated with a BIOS update!"
echo "If you have an AMD CPU, you likely shouldn't use this!"
echo "----------------------------------------------------------------------------"
echo "This script uses systemd-cryptenroll to enable TPM2 auto-unlock."
echo "You can review systemd-cryptenroll's manpage for more information."
echo "This script will modify your system."
echo "It will enable TPM2 auto-unlock of your LUKS partition for your root device!"
echo "It will bind to PCR 7 and 14 which is tied to your secureboot and moklist state."
read -p "Are you sure are good with this and want to enable TPM2 auto-unlock? (y/N): " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
  [[ "$0" = "${BASH_SOURCE[0]}" ]] && exit 1 || return 1 # handle exits from shell or function but don't exit interactive shell
fi

## Inspect Kernel Cmdline for rd.luks.uuid
RD_LUKS_UUID="$(xargs -n1 -a /proc/cmdline | grep rd.luks.uuid | cut -d = -f 2)"

# Check to make sure cmdline rd.luks.uuid exists
if [[ -z ${RD_LUKS_UUID:-} ]]; then
  printf "LUKS device not defined on Kernel Commandline.\n"
  printf "This is not supported by this script.\n"
  printf "Exiting...\n"
  exit 1
fi

# Check to make sure that the specified cmdline uuid exists.
if ! grep -q "${RD_LUKS_UUID}" <<< "$(lsblk)" ; then
  printf "LUKS device not listed in block devices.\n"
  printf "Exiting...\n"
  exit 1
fi

# Cut off the luks-
LUKS_PREFIX="luks-"
if grep -q ^${LUKS_PREFIX} <<< "${RD_LUKS_UUID}"; then
  DISK_UUID=${RD_LUKS_UUID#"$LUKS_PREFIX"}
else
  echo "LUKS UUID format mismatch."
  echo "Exiting..."
  exit 1
fi

SET_PIN_ARG=""
read -p "Would you like to set a PIN? (y/N): " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
  SET_PIN_ARG=" --tpm2-with-pin=yes "
fi

# Specify Crypt Disk by-uuid
CRYPT_DISK="/dev/disk/by-uuid/$DISK_UUID"

# Check to make sure crypt disk exists
if [[ ! -L "$CRYPT_DISK" ]]; then
  printf "LUKS device not listed in block devices.\n"
  printf "Exiting...\n"
  exit 1
fi

if cryptsetup luksDump "$CRYPT_DISK" | grep systemd-tpm2 > /dev/null; then
  KEYSLOT=$(cryptsetup luksDump "$CRYPT_DISK" | sed -n '/systemd-tpm2$/,/Keyslot:/p' | grep Keyslot|awk '{print $2}')
  echo "TPM2 already present in LUKS keyslot $KEYSLOT of $CRYPT_DISK."
  read -p "Wipe it and re-enroll? (y/N): " -n 1 -r
  echo
  if [[ $REPLY =~ ^[Yy]$ ]]; then
    systemd-cryptenroll --wipe-slot=tpm2 "$CRYPT_DISK"
  else
    echo
    echo "Either clear the existing TPM2 keyslot before retrying, else choose 'y' next time."
    echo "Exiting..."
    [[ "$0" = "${BASH_SOURCE[0]}" ]] && exit 1 || return 1
  fi
fi

## Run crypt enroll
echo "Enrolling TPM2 unlock requires your existing LUKS2 unlock password"
systemd-cryptenroll --tpm2-device=auto --tpm2-pcrs=7+14 $SET_PIN_ARG "$CRYPT_DISK"


if lsinitrd 2>&1 | grep -q tpm2-tss > /dev/null; then
  ## add tpm2-tss to initramfs
  if rpm-ostree initramfs | grep tpm2 > /dev/null; then
    echo "TPM2 already present in rpm-ostree initramfs config."
    rpm-ostree initramfs
    echo "Re-running initramfs to pickup changes above."
  fi
  rpm-ostree initramfs --enable --arg=--force-add --arg=tpm2-tss
else
  ## initramfs already containts tpm2-tss
  echo "TPM2 already present in initramfs."
fi

## Now reboot
echo
echo "TPM2 LUKS auto-unlock configured. Reboot now."


# References:
#  https://www.reddit.com/r/Fedora/comments/uo4ufq/any_way_to_get_systemdcryptenroll_working_on/
#  https://0pointer.net/blog/unlocking-luks2-volumes-with-tpm2-fido2-pkcs11-security-hardware-on-systemd-248.html