#!/bin/bash
# =============================================================================
# KALI USB COMPLETO: Formateo + Persistencia + Teclado ES + Driver 8812AU
# Ejecutar como root desde Kali live: sudo bash kali-todo.sh
# =============================================================================
set -e

echo "=========================================="
echo "  KALI USB — Setup completo"
echo "=========================================="
echo ""

# Detectar USB (el más grande que no es el ISO booteado)
echo "[1/6] Detectando USB..."
BOOT_DEV=$(lsblk -dn -o PKNAME /dev/sda1 2>/dev/null || echo "sda")
for dev in sdb sdc sdd sde; do
    if [ -b "/dev/$dev" ]; then
        SIZE=$(lsblk -dn -o SIZE "/dev/$dev" 2>/dev/null | tr -d ' ')
        SIZE_GB=$(echo "$SIZE" | awk '{gsub(/G/,"",$1); print $1}')
        if [ $(echo "$SIZE_GB > 10" | bc 2>/dev/null || echo 0) -eq 1 ] || echo "$SIZE" | grep -qE "[0-9]{2,}G"; then
            USB="/dev/$dev"
            echo "  USB encontrado: $USB ($SIZE)"
            break
        fi
    fi
done

if [ -z "$USB" ]; then
    echo "  ERROR: No se encontró USB grande. Listando todos:"
    lsblk -dn -o NAME,SIZE,MODEL
    echo ""
    read -p "  Escribe el dispositivo (ej: sdb): " USB_NAME
    USB="/dev/$USB_NAME"
fi

echo ""
echo "AVISO: Se va a BORRAR TODO el USB: $USB"
echo "  Los datos se perderán permanentemente."
echo ""
read -p "  ¿Continuar? (escribe BORRAR): " CONFIRM
if [ "$CONFIRM" != "BORRAR" ]; then
    echo "  Cancelado."
    exit 0
fi

# -----------------------------------------------------------
echo ""
echo "[2/6] Limpiando y particionando $USB..."
swapoff -a 2>/dev/null || true

# Limpiar disco
dd if=/dev/zero of="$USB" bs=1M count=10 2>/dev/null
sync

# Crear tabla GPT + 2 particiones
parted -s "$USB" mklabel gpt
parted -s "$USB" mkpart primary fat32 1MiB 4GiB
parted -s "$USB" set 1 boot on
parted -s "$USB" mkpart primary ext4 4GiB 100%
parted -s "$USB" name 2 persistence
sync
sleep 2

# Detectar nombres de particiones
BOOT_PART="${USB}1"
PERSIST_PART="${USB}2"

# Si el kernel usa números en lugar de letras
if [ ! -b "$BOOT_PART" ]; then
    BOOT_PART="${USB}p1"
    PERSIST_PART="${USB}p2"
fi

echo "  Boot: $BOOT_PART"
echo "  Persistence: $PERSIST_PART"

# -----------------------------------------------------------
echo ""
echo "[3/6] Formateando particiones..."
mkfs.vfat -F 32 -n KALI "$BOOT_PART"
mkfs.ext4 -F -L persistence "$PERSIST_PART"

# -----------------------------------------------------------
echo ""
echo "[4/6] Montando ISO de Kali y copiando archivos..."
mkdir -p /mnt/iso /mnt/boot /mnt/persist

# Encontrar ISO de Kali
ISO=""
for f in /home/kali/Downloads/*.iso /home/kali/*.iso /tmp/*.iso /run/live/medium/**/*.iso; do
    if [ -f "$f" ] && echo "$f" | grep -qi kali; then
        ISO="$f"
        break
    fi
done

if [ -z "$ISO" ]; then
    # Intentar montar el medio de booteo actual
    if [ -d /run/live/medium ]; then
        mount --bind /run/live/medium /mnt/iso
        ISO="/mnt/iso"
    else
        echo "  No se encontró ISO. Copia archivos manualmente."
        echo "  Monta $BOOT_PART en /mnt/boot y copia el contenido de la ISO."
        mount "$BOOT_PART" /mnt/boot
        mount "$PERSIST_PART" /mnt/persist
        echo "/ union" > /mnt/persist/persistence.conf
        echo ""
        echo "  Monta la ISO manualmente:"
        echo "    sudo mount -o loop kali-xxx.iso /mnt/iso"
        echo "    sudo cp -r /mnt/iso/* /mnt/boot/"
        echo "    sudo umount /mnt/iso"
        echo ""
        echo "  Persistence configurada. Sigue con el paso 5."
        SKIP_COPY=1
    fi
fi

if [ -z "$SKIP_COPY" ]; then
    mount "$BOOT_PART" /mnt/boot
    mount "$PERSIST_PART" /mnt/persist

    # Copiar contenido de ISO al USB
    echo "  Copiando archivos de boot..."
    rsync -ah --progress "$ISO/" /mnt/boot/ 2>/dev/null || cp -r "$ISO"/* /mnt/boot/

    # Configurar persistencia
    echo "/ union" > /mnt/persist/persistence.conf

    # Modificar GRUB para teclado español
    echo "  Configurando teclado español en GRUB..."
    if [ -f /mnt/boot/grub/grub.cfg ]; then
        sed -i 's/quiet/quiet keyboard-configuration\/layoutcode=es/g' /mnt/boot/grub/grub.cfg
    fi
    if [ -f /mnt/boot/grub.cfg ]; then
        sed -i 's/quiet/quiet keyboard-configuration\/layoutcode=es/g' /mnt/boot/grub.cfg
    fi
    # Syslinux
    if [ -f /mnt/boot/isolinux/isolinux.cfg ]; then
        sed -i 's/quiet/quiet keyboard-configuration\/layoutcode=es/g' /mnt/boot/isolinux/isolinux.cfg
    fi
    if [ -f /mnt/boot/syslinux/syslinux.cfg ]; then
        sed -i 's/quiet/quiet keyboard-configuration\/layoutcode=es/g' /mnt/boot/syslinux/syslinux.cfg
    fi
fi

# -----------------------------------------------------------
echo ""
echo "[5/6] Instalando driver Realtek 8812AU..."
apt update -qq
apt install -y -qq dkms linux-headers-$(uname -r) git build-essential bc 2>/dev/null

DRIVER_DIR="/opt/rtl8812au"
if [ -d "$DRIVER_DIR" ]; then
    cd "$DRIVER_DIR" && git pull -qq
else
    git clone https://github.com/aircrack-ng/rtl8812au.git "$DRIVER_DIR"
    cd "$DRIVER_DIR"
fi

make clean 2>/dev/null || true
make
make install
modprobe 8812au 2>/dev/null || modprobe rtl8812au 2>/dev/null || true
echo "  Driver 8812AU instalado"

# -----------------------------------------------------------
echo ""
echo "[6/6] Configurando teclado español persistente..."
cat > /mnt/persist/etc/keyboard-setup.sh << 'KBD'
#!/bin/bash
loadkeys es 2>/dev/null || true
setxkbmap es 2>/dev/null || true
KBD
chmod +x /mnt/persist/etc/keyboard-setup.sh 2>/dev/null || true

# Desmontar
umount /mnt/boot /mnt/persist 2>/dev/null || true
umount /mnt/iso 2>/dev/null || true

echo ""
echo "=========================================="
echo "  ¡TODO CONFIGURADO!"
echo "=========================================="
echo ""
echo "  Resumen:"
echo "    ✓ USB particionado (GPT: FAT32 boot + ext4 persistence)"
echo "    ✓ Kali Linux instalado en USB"
echo "    ✓ Persistencia configurada"
echo "    ✓ Teclado español en GRUB"
echo "    ✓ Driver Realtek 8812AU compilado e instalado"
echo ""
echo "  Para usar:"
echo "    1. Reinicia y arranca desde el USB"
echo "    2. En GRUB, selecciona 'Kali' y pulsa 'e'"
echo "    3. Verifica que 'keyboard-configuration/layoutcode=es' está"
echo "    4. F10 para arrancar"
echo "    5. Los cambios se guardan por la persistencia"
echo ""