#!/bin/bash

# omarchy:summary=Set a new password for an encrypted drive
# omarchy:requires-sudo=true

# Changes a drive's LUKS key and nothing else. For the disk holding /, the new
# key must open the disk and the old one must not before the change counts. The
# journal records only the disk's UUID, key slots and phase, so a change that
# stops part way is finished by the next run. A platform whose boot checks prove
# the disk's key slots records the owner's new one before the run finishes.

# Passwords live in variables; never trace them.
set +x

REKEY_STATE=${XDG_STATE_HOME:-$HOME/.local/state}/omarchy/drive-password.state
source "$OMARCHY_PATH/install/provisioning/luks-rekey.sh"

# The LUKS partition beneath /, or nothing when / is not encrypted. Fails when
# it cannot tell, including a crypt layer whose LUKS partition lsblk cannot name.
root_luks_device() {
  local root_source ancestry device
  root_source=$(findmnt -no SOURCE /) && [[ -n $root_source ]] || return 1
  ancestry=$(lsblk -nsrpo NAME,TYPE,FSTYPE "${root_source%%[*}") || return 1
  device=$(awk '$3 == "crypto_LUKS" { print $1; exit }' <<<"$ancestry")
  if [[ -n $device ]]; then
    readlink -f "$device"
  elif awk '$2 == "crypt" { found = 1 } END { exit !found }' <<<"$ancestry"; then
    return 1
  fi
}

# Every LUKS device, from udev's probe as root_luks_device reads it. An
# unprivileged blkid only reads /run/blkid, empty until root runs blkid. lsblk
# repeats a device under each parent (RAID members), so each is listed once.
luks_drives() {
  lsblk -nrpo NAME,FSTYPE | awk '$2 == "crypto_LUKS" && !seen[$1]++ { print $1 }'
}

# Every cryptsetup call holds the disk lock, so one left running by a run that
# was killed finishes before the next run looks at the drive.

# Print the key slot the password opens. Returns 2 when it opens none, and
# another failure when the drive cannot be checked. cryptsetup open tries
# enrolled tokens (TPM2, FIDO2, keyring) before the key and reports a token's
# slot whatever key it was given, so allow only a token type no token has.
drive_key_slot() {
  local drive=$1 password=$2 out status=0
  out=$(printf '%s' "$password" | sudo flock "$disk_lock" env LC_ALL=C cryptsetup open --test-passphrase --verbose --token-type passphrase-only --key-file - "$drive" 2>&1) || status=$?
  (( status == 0 )) || return "$status"
  grep -o 'Key slot [0-9]* unlocked' <<<"$out" | grep -o '[0-9]*' | head -1
}

# The enabled key slots, comma-separated. A luks2-keyring token under "Tokens:"
# looks like a keyslot, so read only the keyslot section.
drive_slots() {
  local dump
  dump=$(sudo flock "$disk_lock" cryptsetup luksDump "$1") || return 1
  awk '
    /^[^ \t]/ { keyslots = ($0 == "Keyslots:") }
    keyslots && /^ +[0-9]+: luks2/ { sub(":", "", $1); print $1 }
    /^Key Slot [0-9]+: ENABLED/ { sub(":", "", $3); print $3 }' <<<"$dump" | paste -sd, -
}

# The key slots other than FOUND that the password also opens, comma-separated,
# or nothing. Fails when a slot cannot be checked.
other_key_slots() {
  local drive=$1 password=$2 found=$3 slots slot status
  local -a others=()
  slots=$(drive_slots "$drive") || return 1
  for slot in ${slots//,/ }; do
    [[ $slot != "$found" ]] || continue
    status=0
    printf '%s' "$password" | sudo flock "$disk_lock" cryptsetup open --test-passphrase --token-type passphrase-only --key-slot "$slot" --key-file - "$drive" >/dev/null 2>&1 || status=$?
    if (( status == 0 )); then
      others+=("$slot")
    elif (( status != 2 )); then
      return 1
    fi
  done
  local IFS=,
  printf '%s' "${others[*]}"
}

kill_slot() {
  local drive=$1 password=$2 slot=$3
  printf '%s' "$password" | sudo flock "$disk_lock" cryptsetup luksKillSlot -q --key-file - "$drive" "$slot"
}

# luksChangeKey writes the new key before it retires the old one (LUKS2 swaps
# them in one header update), so a failure that leaves the new key unable to
# open the drive changed nothing: that returns 2. One that stops in between
# leaves both keys, so the old key is retired here until it opens nothing.
# Sets new_slot.
change_luks_key() {
  local drive=$1 current=$2 new=$3 old_slot=$4 remains attempt status
  local -a pbkdf=(--iter-time 2000)

  echo "Changing full-disk encryption password for $drive"
  if sudo flock "$disk_lock" cryptsetup isLuks --type luks2 "$drive"; then
    pbkdf+=(--pbkdf argon2id)
  fi
  # Both keys travel over stdin, the current one on the first line.
  printf '%s\n%s' "$current" "$new" |
    sudo flock "$disk_lock" bash -c 'IFS= read -r current; drive=$1; shift; exec cryptsetup luksChangeKey "$@" --key-file <(printf "%s" "$current") "$drive" <(cat)' bash "$drive" "${pbkdf[@]}"

  new_slot=$(drive_key_slot "$drive" "$new") || return
  for (( attempt = 0; attempt < 8; attempt++ )); do
    status=0
    remains=$(drive_key_slot "$drive" "$current") || status=$?
    (( status == 2 )) && return 0
    (( status == 0 )) && [[ $remains != "$new_slot" ]] || return 1
    kill_slot "$drive" "$new" "$remains" || return 1
  done
  return 1
}

# The entrypoint that records the system disk's key slots for the platform's
# boot checks, or nothing where none does. A machine without its platform's
# boot package at all (dispatch exit 3) predates it, and nothing checks its
# slots. Fails with the dispatcher's message otherwise, such as a boot package
# too old to ship luks-slots.
slot_recorder() {
  local out status=0
  out=$("$OMARCHY_PATH/bin/omarchy-lifecycle-dispatch" --resolve luks-slots 2>&1) || status=$?
  if (( status == 0 )); then
    printf '%s' "$out"
  elif (( status != 3 )); then
    printf '%s\n' "$out" >&2
    return 1
  fi
}

# Refuse a current password that opens any slot but the owner's where the
# platform records the slots: on a disk an earlier setup gave a second key (a
# recovery key), changing that one would leave the owner's password as it was.
# luks-slots --owner prints the recorded slot (0), or exits 4 when no owner slot
# is recorded, as on a Mac set up before its boot package recorded slots; any
# other status is a failure the entrypoint explains, and stops the change.
# Where nothing is recorded, a password that opens the disk's only key slot is
# the owner's, and the change records its new slot (unrecorded_owner).
owner_slot_matches() {
  local slot=$1 recorder owner errors reason status=0
  recorder=$(slot_recorder) || return 1
  [[ -n $recorder ]] || return 0
  errors=$(mktemp) || return 1
  owner=$(sudo "$OMARCHY_PATH/bin/omarchy-lifecycle-dispatch" luks-slots --owner 2>"$errors") || status=$?
  reason=$(<"$errors")
  rm -f "$errors"
  if (( status == 4 )); then
    only_key_slot "$slot" || return 1
    unrecorded_owner=1
    return 0
  elif (( status != 0 )); then
    echo "Could not read which key slot holds the system disk password: ${reason:-luks-slots --owner exited $status}"
    return 1
  elif [[ ! $owner =~ ^[0-9]+$ ]]; then
    echo "Could not read which key slot holds the system disk password: the platform's boot package named '$owner', which is not a key slot."
    return 1
  fi
  if [[ $owner != "$slot" ]]; then
    echo "That password opens another key slot on the system disk, such as a recovery key from an earlier setup. Enter the system disk password."
    return 1
  fi
}

# With no owner slot recorded, the password must open the system disk's one key
# slot: with more keys there is no telling which is the owner's, and recording
# only the owner's would fail the platform's boot check, which wants every kept
# slot recorded. Says how to settle it.
only_key_slot() {
  local slot=$1 slots others
  slots=$(drive_slots "$root_drive") || { echo "Could not read the key slots of $root_drive."; return 1; }
  [[ $slots != "$slot" ]] || return 0
  others=$(tr ',' '\n' <<<"$slots" | grep -vx -- "$slot" | paste -sd, -)
  echo "No key slot is recorded as the system disk password's, and $root_drive has keys in slots $slots, so which one is yours can't be told. The password you entered opens slot $slot."
  if [[ $others =~ ^[0-9]+$ ]]; then
    echo "If that password is yours and slot $others holds a recovery key you keep, record both: sudo omarchy-lifecycle-dispatch luks-slots owner=$slot recovery=$others"
    echo "Otherwise remove the key you don't use: sudo cryptsetup luksKillSlot $root_drive <slot>"
  else
    echo "Remove the keys you don't use: sudo cryptsetup luksKillSlot $root_drive <slot>"
  fi
  echo "Then run omarchy-drive-password again."
  return 1
}

# Record the slot the owner's password opens now; a no-op where nothing does.
record_owner_slot() {
  local recorder
  recorder=$(slot_recorder) || return 1
  [[ -z $recorder ]] || sudo "$OMARCHY_PATH/bin/omarchy-lifecycle-dispatch" luks-slots "owner=$1"
}

# Phase luks: the key change may or may not have landed. Phase record: the new
# key alone opens the recorded slot, which the boot checks still need.
change_system_password() {
  local current=$1 new=$2 old_slot=$3 slots status=0

  slots=$(drive_slots "$root_drive")
  if [[ ,$slots, != *",$old_slot,"* ]]; then
    echo "Could not read the key slots of $root_drive."
    return 1
  fi
  # A disk taken as the owner's for holding one key must still hold only it.
  if (( unrecorded_owner )) && [[ $slots != "$old_slot" ]]; then
    echo "The key slots of $root_drive changed while you entered the new password. Run omarchy-drive-password again."
    return 1
  fi
  mkdir -p "${REKEY_STATE%/*}" &&
    rekey_state_put uuid "$root_uuid" old_slot "$old_slot" slots "$slots" checked 1 phase luks || return 1

  change_luks_key "$root_drive" "$current" "$new" "$old_slot" || status=$?
  if (( status == 2 )); then
    rm -f "$REKEY_STATE"
    echo "The disk password did not change."
    return 1
  elif (( status != 0 )); then
    echo "Could not confirm the disk password change. Run omarchy-drive-password again to finish."
    return 1
  fi
  rekey_state_put slot "$new_slot" phase record || return 1

  if ! record_owner_slot "$new_slot"; then
    echo "Could not record the new key slot for the boot checks. Run omarchy-drive-password again to finish."
    return 1
  fi
  rm -f "$REKEY_STATE"
}

# A journal left on this system's disk by a change that stopped part way.
change_pending() {
  local phase uuid
  phase=$(rekey_state_get phase) || return 1
  uuid=$(rekey_state_get uuid)
  if [[ $phase == "luks" || $phase == "record" || $phase == "accounts" ]] && [[ -n $uuid && $uuid == "$root_uuid" ]]; then
    return 0
  fi
  rm -f "$REKEY_STATE"
  return 1
}

# Settle the disk on the password it opens with now, then record its slot.
# Before the change was confirmed that is the old key (removing the slot the
# change added, if any) or the new key in the one slot the change added
# (retiring the old slot); after, only the confirmed slot. Any other key on the
# disk, such as a recovery passphrase, is refused, and so is a disk that gained
# more than one slot since.
finish_interrupted_change() {
  local phase old_slot slots slot header password opened now added extra count=0 status=0

  phase=$(rekey_state_get phase)
  # Earlier versions named the confirmed phase "accounts".
  [[ $phase != "accounts" ]] || phase=record
  old_slot=$(rekey_state_get old_slot)
  slots=$(rekey_state_get slots)
  slot=$(rekey_state_get slot)

  if [[ $phase == "luks" && ! ( $old_slot =~ ^[0-9]+$ && ,$slots, == *",$old_slot,"* ) ]] ||
    [[ $phase == "record" && ! $slot =~ ^[0-9]+$ ]]; then
    echo "$REKEY_STATE is unreadable. Remove it once the disk opens with only the password you want."
    return 1
  fi

  # Earlier versions journaled the change without checking that the old
  # password opened one slot alone (checked=1), so with other slots on the disk
  # a copy of it may have outlived the change. No password entered now can show
  # which slot was whose, so the owner settles that by hand.
  if [[ -z $(rekey_state_get checked) ]]; then
    now=$(drive_slots "$root_drive") || { echo "Could not read the key slots of $root_drive."; return 1; }
    if [[ $now == *,* ]]; then
      echo "The last password change for $root_drive was started by an earlier version of omarchy-drive-password, and the disk has key slots $now, so it may still hold a copy of the old password."
      echo "Settle it by hand, in this order:"
      echo "1. Find the slot the password you want to keep opens, and keep it: sudo cryptsetup open --test-passphrase --token-type passphrase-only --verbose $root_drive names it. If only the old password opens the disk, that is the one you keep."
      echo "2. If the password you keep is not the old one, check each other slot with the old password (sudo cryptsetup open --test-passphrase --token-type passphrase-only --key-slot <slot> $root_drive) and remove each one it opens (sudo cryptsetup luksKillSlot $root_drive <slot>). Never remove the slot from step 1."
      if [[ -n $(slot_recorder 2>/dev/null) ]]; then
        echo "3. Record the slots you kept for the boot checks: sudo omarchy-lifecycle-dispatch luks-slots owner=<slot> recovery=<slot>, with an empty recovery= when no recovery key is left."
      fi
      echo "Then remove $REKEY_STATE, and run omarchy-drive-password again if the disk should still take a new password."
      return 1
    fi
  fi

  echo "The last password change for $root_drive stopped part way."
  if [[ $phase == "record" ]]; then
    header="New encryption password"
  else
    header="Encryption password that opens $root_drive now (the new one if the change went through)"
  fi
  password=$(gum input --password --header "$header") || return 1

  opened=$(drive_key_slot "$root_drive" "$password") || status=$?
  if (( status != 0 && status != 2 )); then
    echo "Could not check the password against $root_drive."
    return 1
  fi

  if [[ $phase == "luks" && -n $opened ]]; then
    now=$(drive_slots "$root_drive") || { echo "Could not read the key slots of $root_drive."; return 1; }
    for extra in ${now//,/ }; do
      if [[ ,$slots, != *",$extra,"* ]]; then
        added=$extra
        count=$(( count + 1 ))
      fi
    done
    if (( count > 1 )); then
      echo "The key slots of $root_drive changed since the interrupted change. Finish it by hand."
      return 1
    fi

    if [[ $opened == "$old_slot" ]]; then
      [[ -z $added ]] || kill_slot "$root_drive" "$password" "$added" || return 1
    elif [[ -n $added && $opened == "$added" ]]; then
      if [[ ,$now, == *",$old_slot,"* ]]; then
        kill_slot "$root_drive" "$password" "$old_slot" || return 1
      fi
    else
      opened=""
    fi
    [[ -z $opened ]] || rekey_state_put slot "$opened" phase record || return 1
  elif [[ $phase == "record" && $opened != "$slot" ]]; then
    opened=""
  fi

  if [[ -z $opened ]]; then
    echo "That password does not finish the change on $root_drive. Run omarchy-drive-password again to finish."
    return 1
  fi

  if ! record_owner_slot "$opened"; then
    echo "Could not record the key slot for the boot checks. Run omarchy-drive-password again to finish."
    return 1
  fi
  rm -f "$REKEY_STATE"
  echo "$root_drive now opens with the password you entered, and the change is settled."
}

change_selected_drive() {
  local encrypted_drives drive current old_slot others new_password confirmation status=0 system_drive=false

  encrypted_drives=$(luks_drives)
  if [[ -z $encrypted_drives ]]; then
    echo "No encrypted drives available."
    return 1
  fi

  if (( $(wc -l <<<"$encrypted_drives") == 1 )); then
    drive="$encrypted_drives"
  else
    drive="$(omarchy-drive-select "$encrypted_drives")"
  fi

  if [[ -z $drive ]]; then
    echo "No drive selected."
    return 0
  fi

  if [[ -n $root_drive && $(readlink -f "$drive") == "$root_drive" ]]; then
    system_drive=true
    drive=$root_drive
  fi

  current=$(gum input --password --header "Current encryption password") || return 1
  old_slot=$(drive_key_slot "$drive" "$current") || status=$?
  if (( status == 2 )); then
    echo "That password does not open $drive."
    return 1
  elif (( status != 0 )); then
    echo "Could not check the password against $drive."
    return 1
  fi
  if [[ $system_drive == "true" ]]; then
    owner_slot_matches "$old_slot" || { echo "The system disk password did not change."; return 1; }
    # An interrupted change is finished from the one slot the password opened;
    # another copy of it would outlive the change.
    others=$(other_key_slots "$drive" "$current" "$old_slot") || { echo "Could not check the key slots of $drive."; return 1; }
    if [[ -n $others ]]; then
      echo "That password also opens key slot $others on $drive. Remove the extra copy first (sudo cryptsetup luksKillSlot $drive <slot>), then run omarchy-drive-password again."
      if [[ -n $(slot_recorder 2>/dev/null) ]]; then
        echo "Then record the slots you keep for the boot checks: sudo omarchy-lifecycle-dispatch luks-slots owner=<slot> recovery=<slot>, with an empty recovery= when no recovery key is left."
      fi
      echo "The system disk password did not change."
      return 1
    fi
  fi

  new_password=$(gum input --password --header "New encryption password") || return 1
  [[ -n $new_password ]] || { echo "Password cannot be empty."; return 1; }

  confirmation=$(gum input --password --header "Confirm new encryption password") || return 1
  [[ $new_password == "$confirmation" ]] || { echo "Passwords do not match."; return 1; }
  [[ $new_password != "$current" ]] || { echo "The new password is the current one."; return 1; }
  # Its slot is how the change is followed, so it must not open one already.
  status=0
  drive_key_slot "$drive" "$new_password" >/dev/null || status=$?
  if (( status == 0 )); then
    echo "That password already opens a key slot on $drive. Choose another."
    return 1
  elif (( status != 2 )); then
    echo "Could not check the password against $drive."
    return 1
  fi
  status=0

  if [[ $system_drive == "true" ]]; then
    change_system_password "$current" "$new_password" "$old_slot"
  else
    change_luks_key "$drive" "$current" "$new_password" "$old_slot" || status=$?
    if (( status == 2 )); then
      echo "The password of $drive did not change."
      return 1
    elif (( status != 0 )); then
      echo "Could not confirm the password change on $drive."
      return 1
    fi
  fi
}

disk_lock=${REKEY_STATE%/*}/drive-password.disk.lock
mkdir -p "${REKEY_STATE%/*}" && touch "$disk_lock" && exec {lock_fd}>"${REKEY_STATE%/*}/drive-password.lock" && flock -n "$lock_fd" ||
  { echo "Another drive password change is running."; exit 1; }

unrecorded_owner=0
root_drive=$(root_luks_device) || { echo "Could not tell which drive holds the system."; exit 1; }
root_uuid=""
if [[ -n $root_drive ]]; then
  root_uuid=$(lsblk -dno UUID "$root_drive") && [[ -n $root_uuid ]] || { echo "Could not read the UUID of $root_drive."; exit 1; }
fi

if change_pending; then
  finish_interrupted_change
else
  change_selected_drive
fi