#!/bin/bash

# omarchy:summary=Set up FIDO2 authentication for sudo and polkit
# omarchy:requires-sudo=true

set -e
set -o pipefail


check_fido2_hardware() {
  tokens=$(fido2-token -L 2>/dev/null)
  if [[ -z $tokens ]]; then
    echo -e "\e[31m\nNo FIDO2 device detected. Please plug it in (you may need to unlock it as well).\e[0m"
    return 1
  fi
  return 0
}

setup_pam_config() {
  # Configure sudo
  if ! grep -q pam_u2f.so /etc/pam.d/sudo; then
    echo "Configuring sudo for FIDO2 authentication..."
    sudo sed -i '1i auth    sufficient pam_u2f.so cue authfile=/etc/fido2/fido2' /etc/pam.d/sudo
  fi

  # Configure polkit
  if [[ -f /etc/pam.d/polkit-1 ]] && ! grep -q 'pam_u2f.so' /etc/pam.d/polkit-1; then
    echo "Configuring polkit for FIDO2 authentication..."
    sudo sed -i '1i auth      sufficient pam_u2f.so cue authfile=/etc/fido2/fido2' /etc/pam.d/polkit-1
  elif [[ ! -f /etc/pam.d/polkit-1 ]]; then
    echo "Creating polkit configuration with FIDO2 authentication..."
    sudo tee /etc/pam.d/polkit-1 >/dev/null <<'EOF'
auth      sufficient pam_u2f.so cue authfile=/etc/fido2/fido2
auth      required pam_unix.so

account   required pam_unix.so
password  required pam_unix.so
session   required pam_unix.so
EOF
  fi
}

echo -e "\e[32mSetting up FIDO2 device for authentication.\n\e[0m"

# Install required packages
echo "Installing required packages..."
omarchy-pkg-add libfido2 pam-u2f

if ! check_fido2_hardware; then
  exit 1
fi

# Create the pamu2fcfg file
authdir=/etc/fido2
authfile=/etc/fido2/fido2

# install -d follows a symlink here and applies the mode and ownership to
# whatever it points at, so the credential would be staged and published inside
# the link target and that directory reopened to root:root 755. This is the
# threat omarchy-remove-security-fido2 already names on its side.
if [[ -L $authdir || ( -e $authdir && ! -d $authdir ) ]]; then
  echo -e "\e[31m\n$authdir is not a FIDO2 configuration directory.\e[0m"
  echo "Run omarchy-remove-security-fido2 first, then set FIDO2 up again."
  exit 1
fi

# -f follows symlinks, so the already-registered check below reads a symlinked
# authfile as a registration and leaves it in place, and is false for a
# directory, so it tries to register over one. Only a regular file is a valid
# pam_u2f authfile.
if [[ -L $authfile || ( -e $authfile && ! -f $authfile ) ]]; then
  echo -e "\e[31m\n$authfile is not a FIDO2 registration file.\e[0m"
  echo "Run omarchy-remove-security-fido2 first, then set FIDO2 up again."
  exit 1
fi

if [[ ! -f $authfile ]]; then
  sudo install -d -m 755 -o root -g root "$authdir"
  echo -e "\e[32m\nLet's setup your device by confirming on the device now.\e[0m"
  echo -e "Touch your FIDO2 key when it lights up...\n"

  # A unique sibling created by root cannot be replaced by another process
  # running as this user. Stream pamu2fcfg into it instead of asking root to
  # reopen a caller-owned path: an observed temporary name could otherwise be
  # replaced with a symlink before the privileged copy. The final rename is
  # atomic, and -T refuses a directory at the destination. Mode 644 keeps the
  # root-owned global authfile readable when pam_u2f uses openasuser; only root
  # can still rewrite it.
  stage=""

  # mktemp's output is an operand for four privileged commands below, one of
  # them an rm. Take only the name this script asked for rather than whatever
  # came back on stdout.
  safe_stage_path() {
    local candidate=$1
    local prefix="$authfile.new."
    local suffix

    [[ $candidate == "$prefix"* ]] || return 1
    suffix=${candidate#"$prefix"}
    [[ $suffix =~ ^[[:alnum:]]{6}$ ]]
  }

  cleanup_stage() {
    local status=$?

    if safe_stage_path "$stage"; then
      sudo rm -f -- "$stage" || true
    fi

    return "$status"
  }

  trap cleanup_stage EXIT
  stage=$(sudo mktemp "$authfile.new.XXXXXX")

  if ! safe_stage_path "$stage" || [[ ! -f $stage || -L $stage ]]; then
    echo -e "\e[31m\nCould not create a safe staging file beside $authfile.\e[0m"
    exit 1
  fi

  if pamu2fcfg | sudo tee "$stage" >/dev/null && [[ -s $stage ]]; then
    sudo chmod 644 "$stage"
    sudo mv -Tf "$stage" "$authfile"
    stage=""
    trap - EXIT
    echo -e "\e[32mFIDO2 device registered successfully!\e[0m"
  else
    echo -e "\e[31m\nFIDO2 registration failed. Please try again.\e[0m"
    exit 1
  fi
else
  echo "FIDO2 device already registered."
fi

# Configure PAM
setup_pam_config

# Test with sudo
echo -e "\nTesting FIDO2 authentication with sudo..."
echo -e "Touch your FIDO2 key when prompted.\n"

if sudo echo "FIDO2 authentication test successful"; then
  echo -e "\e[32m\nPerfect! FIDO2 authentication is now configured.\e[0m"
  echo "You can use your FIDO2 key for sudo and polkit authentication."
else
  echo -e "\e[31m\nVerification failed. You may want to check your configuration.\e[0m"
fi