#!/bin/sh
# initramfs-framework module for LUKS encryption with fTPM support

# Configuration
BOOT_DEV="/dev/mmcblk1p1"           # Boot partition (FAT, unencrypted)
ROOT_DEV="/dev/mmcblk1p2"           # Root partition (will be encrypted)
CRYPT_NAME="root_crypt"
CRYPT_DEV="/dev/mapper/${CRYPT_NAME}"
BOOT_MNT="/boot_part"
TPM_PRIMARY_CTX="/tmp/tpm_primary.ctx"
TPM_KEY_PRIV="/tmp/tpm_key.priv"
TPM_KEY_PUB="/tmp/tpm_key.pub"
TPM_KEY_CTX="/tmp/tpm_key.ctx"
TPM2_HANDLE="0x81080001"            # TPM persistent handle for LUKS key
ENCRYPTION_MARKER="${BOOT_MNT}/.encryption_in_progress"

# Wait for MMC device to appear
wait_for_device() {
    local device="$1"
    local timeout="${2:-10}"

    msg "Waiting for storage device ${device}..."
    for i in $(seq 1 ${timeout}); do
        if [ -b "${device}" ]; then
            return 0
        fi
        sleep 1
    done
    return 1
}

# Initialize fTPM and check availability
init_ftpm() {
    msg "Initializing secure hardware (fTPM)..."

    # Start TEE supplicant (required for fTPM TA to work)
    if [ -x /usr/sbin/tee-supplicant ]; then
        /usr/sbin/tee-supplicant -d &
        TEE_SUPPLICANT_PID=$!
        sleep 5
    else
        info "Warning: Trusted execution environment not available"
        return 1
    fi

    # Load fTPM kernel module
    if ! /sbin/modprobe tpm_ftpm_tee; then
        info "Warning: TPM module failed to load"
        return 1
    fi

    # Wait for TPM device
    for i in $(seq 1 10); do
        if [ -c /dev/tpmrm0 ]; then
            export TPM2TOOLS_TCTI="device:/dev/tpmrm0"
            return 0
        fi
        sleep 1
    done

    info "Warning: fTPM not available - encryption will be skipped"
    return 1
}

# Generate 32-byte random key using TPM RNG
generate_random_key() {
    /usr/bin/tpm2_getrandom --hex 32
}

# Seal data with TPM and store in persistent handle
tpm_seal_key() {
    local KEY_DATA="$1"

    # Create primary key in owner hierarchy
    /usr/bin/tpm2_createprimary -C o -c "${TPM_PRIMARY_CTX}" -Q || return 1

    # Create sealed object
    echo -n "${KEY_DATA}" | \
        /usr/bin/tpm2_create -C "${TPM_PRIMARY_CTX}" \
        -u "${TPM_KEY_PUB}" -r "${TPM_KEY_PRIV}" \
        -i- -Q || return 1

    # Load sealed object into TPM
    /usr/bin/tpm2_load -C "${TPM_PRIMARY_CTX}" \
        -u "${TPM_KEY_PUB}" -r "${TPM_KEY_PRIV}" \
        -c "${TPM_KEY_CTX}" -Q || return 1

    # Make key persistent at handle (stored in TPM NV RAM - RPMB)
    /usr/bin/tpm2_evictcontrol -C o -c "${TPM_KEY_CTX}" "${TPM2_HANDLE}" || return 1

    return 0
}

# Unseal data from TPM persistent handle
tpm_unseal_key() {
    # Check if persistent handle exists
    if ! /usr/bin/tpm2_getcap handles-persistent | grep -q "${TPM2_HANDLE}"; then
        debug "ERROR: TPM persistent handle not found"
        return 1
    fi

    # Unseal key directly from persistent handle
    /usr/bin/tpm2_unseal -c "${TPM2_HANDLE}" || return 1

    return 0
}

# Perform in-place LUKS encryption (first boot)
encrypt_root_filesystem() {
    msg "=========================================="
    msg "First boot: Encrypting root filesystem"
    msg "=========================================="

    # Set marker to track encryption progress
    touch "${ENCRYPTION_MARKER}"
    sync

    # Generate random encryption key using TPM RNG
    msg "Generating encryption key..."
    LUKS_KEY=$(generate_random_key)

    if [ -z "${LUKS_KEY}" ]; then
        msg "ERROR: Failed to generate encryption key"
        rm -f "${ENCRYPTION_MARKER}"
        return 1
    fi

    # Seal key with TPM before encryption starts
    msg "Securing key with TPM..."
    if ! tpm_seal_key "${LUKS_KEY}"; then
        msg "ERROR: Failed to secure key"
        rm -f "${ENCRYPTION_MARKER}"
        return 1
    fi

    # Filesystem check before encryption
    msg "Checking filesystem integrity..."
    /usr/sbin/e2fsck -f -y "${ROOT_DEV}"
    E2FSCK_RET=$?
    if [ ${E2FSCK_RET} -ge 4 ]; then
        msg "ERROR: Filesystem check failed"
        rm -f "${ENCRYPTION_MARKER}"
        return 1
    fi

    # Shrink filesystem before encryption to leave room for LUKS header
    msg "Preparing filesystem for encryption..."
    /usr/sbin/resize2fs -M "${ROOT_DEV}" || {
        msg "ERROR: Failed to prepare filesystem"
        rm -f "${ENCRYPTION_MARKER}"
        return 1
    }

    # Verify partition has sufficient space for LUKS header
    msg "Verifying space for encryption..."
    MIN_BLOCKS=$(/usr/sbin/resize2fs -P "${ROOT_DEV}" 2>&1 | awk '/[Mm]inimum.*:/ {print $NF}')

    # Get filesystem block size and device size
    BLOCK_SIZE=$(/usr/sbin/tune2fs -l "${ROOT_DEV}" 2>/dev/null | awk '/^Block size:/ {print $NF}')
    DEV_NAME=$(basename "${ROOT_DEV}")
    PART_SECTORS=$(cat /sys/class/block/"${DEV_NAME}"/size 2>/dev/null)

    if [ -z "${MIN_BLOCKS}" ] || [ -z "${BLOCK_SIZE}" ] || [ -z "${PART_SECTORS}" ]; then
        msg "ERROR: Unable to determine partition geometry"
        rm -f "${ENCRYPTION_MARKER}"
        return 1
    fi

    # Convert filesystem blocks to 512-byte sectors
    MIN_SECTORS=$((MIN_BLOCKS * BLOCK_SIZE / 512))
    LUKS_SECTORS=65536  # 32MB in 512-byte sectors

    if [ $((PART_SECTORS - MIN_SECTORS)) -lt ${LUKS_SECTORS} ]; then
        msg "ERROR: Insufficient space for LUKS header (need 32MB free)"
        rm -f "${ENCRYPTION_MARKER}"
        return 1
    fi

    # Perform in-place encryption
    msg "=========================================="
    msg "Encrypting filesystem..."
    msg "This will take several minutes."
    msg "DO NOT POWER OFF THE DEVICE!"
    msg "=========================================="

    echo -n "${LUKS_KEY}" | \
        /usr/sbin/cryptsetup reencrypt --encrypt \
        --type luks2 \
        --cipher aes-xts-plain64 \
        --key-size 256 \
        --hash sha256 \
        --reduce-device-size 32M \
        --key-file - \
        "${ROOT_DEV}" || {
        msg "ERROR: Encryption failed"
        rm -f "${ENCRYPTION_MARKER}"
        return 1
    }

    msg "=========================================="
    msg "Encryption completed successfully!"
    msg "=========================================="

    # Remove encryption marker
    rm -f "${ENCRYPTION_MARKER}"
    sync

    # Unlock the newly encrypted device
    msg "Activating encrypted filesystem..."
    echo -n "${LUKS_KEY}" | \
        /usr/sbin/cryptsetup luksOpen "${ROOT_DEV}" "${CRYPT_NAME}" --key-file - || {
        msg "ERROR: Failed to activate encrypted filesystem"
        return 1
    }

    # Resize filesystem to fit the encrypted device
    msg "Optimizing filesystem..."
    /usr/sbin/resize2fs -f "${CRYPT_DEV}" || {
        msg "ERROR: Failed to optimize filesystem"
        return 1
    }

    # Verify filesystem after resize
    /usr/sbin/e2fsck -f -y "${CRYPT_DEV}" || {
        info "WARNING: Filesystem verification had issues, but continuing"
    }

    return 0
}

# Unlock encrypted root filesystem (subsequent boots)
unlock_encrypted_root() {
    msg "Unlocking encrypted filesystem..."

    # Unseal key from TPM persistent handle
    LUKS_KEY=$(tpm_unseal_key)

    if [ -z "${LUKS_KEY}" ]; then
        msg "ERROR: Failed to retrieve encryption key from TPM"
        msg "Attempting passphrase fallback..."

        # Try to unlock with passphrase (interactive)
        /usr/sbin/cryptsetup luksOpen "${ROOT_DEV}" "${CRYPT_NAME}" || {
            fatal "ERROR: Failed to unlock encrypted filesystem"
        }
    else
        # Unlock with unsealed key
        echo -n "${LUKS_KEY}" | \
            /usr/sbin/cryptsetup luksOpen "${ROOT_DEV}" "${CRYPT_NAME}" --key-file - || {
            fatal "ERROR: Failed to unlock with TPM key"
        }
    fi

    msg "Encrypted filesystem unlocked"
}

# Module enabled check
luksftpm_enabled() {
    # Always run this module - it handles both encrypted and unencrypted cases
    return 0
}

# Module main function
luksftpm_run() {
    # Wait for storage device
    if ! wait_for_device "${ROOT_DEV}" 10; then
        info "Storage device not found, skipping encryption module"
        return 0
    fi

    # Mount boot partition
    msg "Mounting boot partition..."
    mkdir -p "${BOOT_MNT}"
    if ! mount "${BOOT_DEV}" "${BOOT_MNT}"; then
        info "ERROR: Failed to mount boot partition, attempting standard boot..."
        mkdir -p ${ROOTFS_DIR}
        mount "${ROOT_DEV}" ${ROOTFS_DIR}
        return 0
    fi

    # Initialize fTPM
    TPM_AVAILABLE=0
    if init_ftpm; then
        TPM_AVAILABLE=1
    fi

    # Check filesystem encryption status
    msg "Checking filesystem encryption status..."

    MOUNT_DEV="${ROOT_DEV}"

    if /usr/sbin/cryptsetup isLuks "${ROOT_DEV}"; then
        msg "Filesystem is encrypted"
        unlock_encrypted_root
        MOUNT_DEV="${CRYPT_DEV}"
    else
        msg "Filesystem is not encrypted"

        # Check if encryption is enabled and TPM is available
        if [ $TPM_AVAILABLE -eq 1 ]; then
            # Check for encryption marker (resume interrupted encryption)
            if [ -f "${ENCRYPTION_MARKER}" ]; then
                msg "Resuming interrupted encryption..."
                if ! encrypt_root_filesystem; then
                    msg "ERROR: Failed to resume encryption"
                    msg "Booting without encryption..."
                    MOUNT_DEV="${ROOT_DEV}"
                else
                    MOUNT_DEV="${CRYPT_DEV}"
                fi
            else
                # First boot - perform encryption
                if encrypt_root_filesystem; then
                    MOUNT_DEV="${CRYPT_DEV}"
                else
                    msg "ERROR: Encryption failed - booting without encryption"
                    MOUNT_DEV="${ROOT_DEV}"
                fi
            fi
        else
            msg "TPM not available - skipping encryption"
            MOUNT_DEV="${ROOT_DEV}"
        fi
    fi

    # Unmount boot partition before switching root
    umount "${BOOT_MNT}"

    # Mount root filesystem to $ROOTFS_DIR (framework expects this)
    msg "Mounting root filesystem..."
    mkdir -p ${ROOTFS_DIR}
    mount "${MOUNT_DEV}" ${ROOTFS_DIR} || {
        fatal "ERROR: Failed to mount root filesystem!"
    }

    # Clean up tmpfs and sensitive variables
    rm -f "${TPM_PRIMARY_CTX}" "${TPM_KEY_PUB}" "${TPM_KEY_PRIV}" "${TPM_KEY_CTX}"
    unset LUKS_KEY TPM_AVAILABLE MOUNT_DEV TEE_SUPPLICANT_PID

    msg "Boot complete"
}
