#!/bin/bash

# Script to set up LUKS encryption for home partition on internal drive
# This will encrypt the home partition in-place

set -euo pipefail

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

log_info() {
    echo -e "${GREEN}[INFO]${NC} $(date '+%Y-%m-%d %H:%M:%S'): $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $(date '+%Y-%m-%d %H:%M:%S'): $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $(date '+%Y-%m-%d %H:%M:%S'): $1"
}

log_step() {
    echo -e "${BLUE}[STEP]${NC} $(date '+%Y-%m-%d %H:%M:%S'): $1"
}

# Check if running as root
if [[ $EUID -ne 0 ]]; then
    log_error "This script must be run as root (use sudo)"
    exit 1
fi

echo -e "${BLUE}=== LUKS Encryption Setup for Internal Drive ===${NC}"
echo
echo "This will set up LUKS encryption for your home partition on the internal drive."
echo "The process will:"
echo "1. Create a backup image of the current home data"
echo "2. Recreate the home LV with LUKS encryption"
echo "3. Restore the data to the encrypted volume"
echo "4. Update system configuration"
echo
echo -e "${RED}WARNING: This process requires sufficient free space for backup!${NC}"

# Check available space
free_space=$(vgs --noheadings --units g --nosuffix -o vg_free internal-vg | tr -d ' ' | tr ',' '.')
home_size=$(lvs --noheadings --units g --nosuffix -o lv_size internal-vg/home | tr -d ' ' | tr ',' '.')

echo "Home partition size: ${home_size}GB"
echo "Available free space: ${free_space}GB"

if (( $(echo "$free_space < $home_size" | bc -l) )); then
    log_error "Not enough free space for backup. Need ${home_size}GB free space."
    exit 1
fi

read -p "Continue with LUKS encryption setup? (yes/no): " confirm
if [[ "$confirm" != "yes" ]]; then
    log_info "Operation cancelled"
    exit 0
fi

BACKUP_LV="home_backup_temp"
MOUNT_POINT="/mnt/luks_setup"

log_step "Creating backup of home data..."

# Create backup LV
lvcreate -L "${home_size}G" -n "$BACKUP_LV" internal-vg

# Copy home data to backup
log_info "Copying home data to backup volume..."
dd if=/dev/internal-vg/home of="/dev/internal-vg/$BACKUP_LV" bs=1M status=progress

log_step "Removing and recreating home LV..."

# Remove the current home LV
lvremove -f internal-vg/home

# Create new home LV
lvcreate -L "${home_size}G" -n home internal-vg

log_step "Setting up LUKS encryption..."

# Setup LUKS on the new LV
echo "Please enter your desired LUKS passphrase:"
cryptsetup luksFormat /dev/internal-vg/home

echo "Please enter your LUKS passphrase again to open the volume:"
cryptsetup open /dev/internal-vg/home luks-home-internal

# Format the encrypted volume
mkfs.ext4 -L home /dev/mapper/luks-home-internal

log_step "Restoring home data..."

# Mount backup and encrypted volumes
mkdir -p "$MOUNT_POINT/backup" "$MOUNT_POINT/encrypted"
mount "/dev/internal-vg/$BACKUP_LV" "$MOUNT_POINT/backup"
mount /dev/mapper/luks-home-internal "$MOUNT_POINT/encrypted"

# Copy data back
log_info "Copying data from backup to encrypted volume..."
rsync -avHAXS --progress "$MOUNT_POINT/backup/" "$MOUNT_POINT/encrypted/"

# Clean up mounts
umount "$MOUNT_POINT/backup" "$MOUNT_POINT/encrypted"
cryptsetup close luks-home-internal

# Remove backup LV
lvremove -f "internal-vg/$BACKUP_LV"

log_step "Updating system configuration..."

# Get the UUID of the LUKS device
LUKS_UUID=$(cryptsetup luksUUID /dev/internal-vg/home)

# Mount the internal root to update configuration
mount /dev/internal-vg/root "$MOUNT_POINT"

# Update /etc/crypttab
echo "luks-home-internal UUID=$LUKS_UUID none luks" >> "$MOUNT_POINT/etc/crypttab"

# Update /etc/fstab
cat > "$MOUNT_POINT/etc/fstab" << EOF
# Internal LVM Configuration with LUKS
/dev/internal-vg/root     /           ext4    defaults                    0 1
/dev/internal-vg/boot     /boot       ext4    defaults                    0 2
/dev/mapper/luks-home-internal /home   ext4    defaults                    0 2
/dev/internal-vg/swap     none        swap    sw                          0 0
/dev/nvme0n1p1            /boot/efi   vfat    umask=0077                  0 1
EOF

# Update initramfs to include LUKS support
mount --bind /dev "$MOUNT_POINT/dev"
mount --bind /proc "$MOUNT_POINT/proc"
mount --bind /sys "$MOUNT_POINT/sys"
mount --bind /run "$MOUNT_POINT/run"

chroot "$MOUNT_POINT" /bin/bash -c "update-initramfs -u -k all"
chroot "$MOUNT_POINT" /bin/bash -c "update-grub"

# Clean up
umount "$MOUNT_POINT/dev" "$MOUNT_POINT/proc" "$MOUNT_POINT/sys" "$MOUNT_POINT/run"
umount "$MOUNT_POINT"

log_info "LUKS encryption setup completed successfully!"
echo
echo -e "${GREEN}SUCCESS!${NC} Home partition is now encrypted with LUKS"
echo "Next steps:"
echo "1. Reboot from the internal drive"
echo "2. You will be prompted for the LUKS passphrase during boot"
echo "3. Verify that everything works correctly"
echo
echo -e "${YELLOW}Important:${NC} Remember your LUKS passphrase! Without it, your home data will be inaccessible."