set -Eeuo pipefail

sudo apt install clevis clevis-tpm2 clevis-luks clevis-initramfs tss2

# Define colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1" >&2
}


# 3. Detect LUKS partition and prompt for password
log_info "=== LUKS-Verschlüsselung (TPM2) vorbereiten ==="

LUKS_PARTITION=""
LUKS_PASSPHRASE=""

# Get all crypto_LUKS block devices
LUKS_DEVS=($(lsblk -p -ro NAME,FSTYPE | awk '$2 == "crypto_LUKS" {print $1}'))

if [ ${#LUKS_DEVS[@]} -eq 0 ]; then
    log_warn "Keine LUKS-Partitionen (crypto_LUKS) auf dem System gefunden. TPM2-Schritt wird übersprungen."
elif [ ${#LUKS_DEVS[@]} -eq 1 ]; then
    LUKS_PARTITION="${LUKS_DEVS[0]}"
    log_info "Automatisch erkannte LUKS-Partition: $LUKS_PARTITION"
else
    log_info "Mehrere LUKS-Partitionen gefunden. Suche nach der Root-Systempartition..."
    ROOT_DEV=$(findmnt -n -o SOURCE /)
    
    # Traverse dependencies upward to find the underlying physical partition
    if [[ "$ROOT_DEV" == /dev/mapper/* ]] || [[ "$ROOT_DEV" == /dev/dm-* ]]; then
        DEPS=($(lsblk -p -s -no NAME,FSTYPE "$ROOT_DEV" | awk '$2 == "crypto_LUKS" {print $1}'))
        if [ ${#DEPS[@]} -gt 0 ]; then
            LUKS_PARTITION="${DEPS[0]}"
            log_info "Automatisch ermittelte LUKS-Partition für das Root-System: $LUKS_PARTITION"
        fi
    fi
    
    # Selection fallback if auto-detection failed
    if [ -z "$LUKS_PARTITION" ]; then
        log_info "Die LUKS-Partition für das Root-System konnte nicht eindeutig bestimmt werden."
        echo "Bitte wähle die zu verschlüsselnde Partition aus:"
        select opt in "${LUKS_DEVS[@]}" "Überspringen"; do
            if [ "$opt" = "Überspringen" ]; then
                log_info "LUKS-Konfiguration wird übersprungen."
                break
            elif [ -n "$opt" ]; then
                LUKS_PARTITION="$opt"
                log_info "Ausgewählte Partition: $LUKS_PARTITION"
                break
            else
                echo "Ungültige Auswahl."
            fi
        done </dev/tty
    fi
fi

# Ask for LUKS passphrase if partition is determined and not already bound
if [ -n "$LUKS_PARTITION" ]; then
    ALREADY_BOUND=false
    # Check if clevis/luksmeta binding exists (handles idempotency if run multiple times)
    if command -v clevis &>/dev/null; then
        if clevis luks list -d "$LUKS_PARTITION" 2>/dev/null | grep -q "tpm2"; then
            ALREADY_BOUND=true
        fi
    elif cryptsetup luksDump "$LUKS_PARTITION" 2>/dev/null | grep -qi -E "clevis|luksmeta"; then
        ALREADY_BOUND=true
    fi
    
    if [ "$ALREADY_BOUND" = true ]; then
        log_info "Die Partition $LUKS_PARTITION ist bereits mit TPM2 verknüpft. Überspringe Passworteingabe."
    else
        while true; do
            echo -n "Gib das aktuelle LUKS-Passwort für die Partition $LUKS_PARTITION ein (für die TPM2-Bindung): "
            read -rs LUKS_PASSPHRASE </dev/tty
            echo
            if [ -z "$LUKS_PASSPHRASE" ]; then
                log_warn "Das LUKS-Passwort darf nicht leer sein!"
                continue
            fi
            
            log_info "Überprüfe LUKS-Passwort..."
            if echo -n "$LUKS_PASSPHRASE" | cryptsetup luksOpen --test-passphrase "$LUKS_PARTITION" &>/dev/null; then
                log_success "LUKS-Passwort ist korrekt!"
                break
            else
                log_error "Ungültiges LUKS-Passwort! Bitte versuche es erneut."
                LUKS_PASSPHRASE=""
            fi
        done
    fi
fi


# 8. LUKS TPM2 Binding
log_info "===  LUKS-TPM2-Bindung durchführen ==="
if [ -n "$LUKS_PARTITION" ] && [ -n "$LUKS_PASSPHRASE" ]; then
    # Check if TPM2 chip is available
    if [ ! -e "/dev/tpm0" ] && [ ! -e "/dev/tpmrm0" ]; then
        log_warn "Kein TPM 2.0 Chip unter /dev/tpm0 oder /dev/tpmrm0 gefunden! Bindung wird übersprungen."
    else
        # Detect UEFI / Legacy BIOS
        if [ ! -d "/sys/firmware/efi" ]; then
            log_warn "Das System wurde im Legacy-BIOS-Modus gestartet (kein UEFI)."
            log_warn "TPM2-Entsperrung beim Booten erfordert in der Regel UEFI!"
        fi
        
        # Check Secure Boot status
        secure_boot=false
        if command -v mokutil &>/dev/null; then
            if mokutil --sb-state 2>/dev/null | grep -qi "enabled"; then
                secure_boot=true
            fi
        fi
        
        clevis_config='{"pcr_ids":"7"}'
        if [ "$secure_boot" = false ]; then
            log_warn "Secure Boot ist deaktiviert oder wird nicht unterstützt."
            log_info "Verwende leere PCR-Konfiguration '{}' für die Bindung (empfohlen für VMs/deaktiviertes Secure Boot)."
            clevis_config='{}'
        fi
        
        # Enable Clevis services for systemd integration (only if available)
        if systemctl list-unit-files | grep -q "clevis-luks-askpass.path"; then
            log_info "Aktiviere Clevis-Services..."
            systemctl enable clevis-luks-askpass.path || log_warn "Konnte clevis-luks-askpass.path nicht aktivieren."
        fi
        
        log_info "Binde LUKS-Partition $LUKS_PARTITION an TPM 2.0 (Konfig: $clevis_config)..."
        # We wrap the command in a conditional block to handle errors gracefully
        if echo -n "$LUKS_PASSPHRASE" | clevis luks bind -y -d "$LUKS_PARTITION" -k - tpm2 "$clevis_config"; then
            log_success "LUKS-Partition wurde erfolgreich an TPM 2.0 gebunden."
            log_info "Konfiguriere TPM-Kernel-Module für das Initramfs..."
            initramfs_modules="/etc/initramfs-tools/modules"
            if [ -f "$initramfs_modules" ]; then
                for mod in tpm tpm_tis tpm_tis_core tpm_crb; do
                    if ! grep -q "^$mod" "$initramfs_modules"; then
                        echo "$mod" >> "$initramfs_modules"
                        log_info "Modul '$mod' zu $initramfs_modules hinzugefügt."
                    fi
                done
            fi
            
            log_info "Aktualisiere Initramfs..."
            update-initramfs -u -k all
            log_success "Initramfs wurde aktualisiert."
        else
            log_error "Bindung an TPM 2.0 fehlgeschlagen! Bitte überprüfe das LUKS-Passwort."
        fi
    fi
else
    log_info "Keine LUKS-Bindung auszuführen (bereits gebunden, übersprungen oder keine Partition gefunden)."
fi