#!/bin/bash

echo "[*] Checking initramfs contents for potential exposure..."

INITRD=$(ls /boot/initrd.img-* | sort -V | tail -n1)
echo "[*] Found initramfs: $INITRD"

# Try gzip first, fallback to cpio-only
echo "[*] Extracting initramfs for analysis..."
TMPDIR=$(mktemp -d)
cd "$TMPDIR"

if file "$INITRD" | grep -q 'gzip'; then
    gzip -cd "$INITRD" | cpio -idmu > /dev/null 2>&1
else
    echo "[!] Warning: fallback to uncompressed initrd..."
    cat "$INITRD" | cpio -idmu > /dev/null 2>&1
fi

# Check for risky binaries
echo "[*] Analyzing for sensitive binaries..."
FOUND=false
for bin in sh cryptsetup lvm busybox mount umount blkid; do
    if find . -name "$bin" | grep -q .; then
        echo "[!] ⚠ Found sensitive binary: $bin"
        FOUND=true
    fi
done

cd /
rm -rf "$TMPDIR"

if ! $FOUND; then
    echo "[*] No critical binaries found. GRUB password not required."
    exit 0
fi

# Check for existing GRUB password
if grep -q "GRUB2_PASSWORD" /etc/grub.d/40_custom; then
    echo "[*] GRUB password already set. Skipping..."
    exit 0
fi

# Prompt user for password
echo "[*] System is vulnerable. Setting GRUB password..."
read -s -p "Enter GRUB password: " PASSWORD
echo
read -s -p "Confirm GRUB password: " PASSWORD2
echo
if [ "$PASSWORD" != "$PASSWORD2" ]; then
    echo "[!] Passwords do not match. Aborting."
    exit 1
fi

# Hash password
HASH=$(echo "$PASSWORD" | grub-mkpasswd-pbkdf2 | awk '/grub.pbkdf2/{print $NF}')
unset PASSWORD PASSWORD2

if [ -z "$HASH" ]; then
    echo "[!] Failed to generate password hash."
    exit 1
fi

# Insert into /etc/grub.d/40_custom
echo "[*] Writing password to /etc/grub.d/40_custom..."
cat <<EOF >> /etc/grub.d/40_custom

set superusers="root"
password_pbkdf2 root $HASH
EOF

# Update GRUB
echo "[*] Updating GRUB config..."
update-grub

echo "[+] GRUB password is now active. Test by rebooting and pressing 'e' on boot menu."