#!/bin/bash
# =============================================================================
# KALI LINUX — Setup completo: Persistencia + Drivers Realtek 8812AU
# Ejecutar como root desde el USB live de Kali
# =============================================================================
set -e

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

echo -e "${CYAN}============================================${NC}"
echo -e "${CYAN}  KALI USB SETUP — Persistencia + 8812AU${NC}"
echo -e "${CYAN}============================================${NC}"
echo ""

# -----------------------------------------------------------
# PASO 1: Detectar el USB
# -----------------------------------------------------------
echo -e "${YELLOW}[1/5] Detectando USB...${NC}"

USB_DEVICE=""
for dev in /dev/sd{b,c,d,e}; do
    if [ -b "$dev" ]; then
        SIZE=$(lsblk -dn -o SIZE "$dev" 2>/dev/null | tr -d ' ')
        FSTYPE=$(lsblk -dn -o FSTYPE "$dev" 2>/dev/null | tr -d ' ')
        LABEL=$(lsblk -dn -o LABEL "$dev" 2>/dev/null | tr -d ' ')
        echo "  $dev — $SIZE — $FSTYPE — $LABEL"
        # El USB más grande probablemente es nuestro
        USB_DEVICE="$dev"
    fi
done

if [ -z "$USB_DEVICE" ]; then
    echo -e "${RED}  ERROR: No se encontró ningún USB${NC}"
    exit 1
fi

echo -e "${GREEN}  USB detectado: $USB_DEVICE${NC}"
echo ""

# -----------------------------------------------------------
# PASO 2: Verificar si ya hay persistencia
# -----------------------------------------------------------
echo -e "${YELLOW}[2/5] Verificando persistencia...${NC}"

# Buscar partición etiquetada "persistence"
PERSIST_PART=$(lsblk -dn -o NAME,LABEL | grep -i persistence | awk '{print $1}' | head -1)

if [ -n "$PERSIST_PART" ]; then
    echo -e "${GREEN}  Persistencia ya configurada en /dev/$PERSIST_PART${NC}"
    PERSIST_DEVICE="/dev/$PERSIST_PART"
else
    echo "  No se encontró persistencia. Creando..."
    
    # Encontrar la última partición del USB
    LAST_PART=$(lsblk -dn -o NAME "$USB_DEVICE" | tail -1 | sed 's/.*\([a-z][a-z0-9]*\)$/\1/')
    echo "  Última partición detectada: $LAST_PART"
    
    echo ""
    echo -e "${YELLOW}  IMPORTANTE: Voy a redimensionar la última partición del USB${NC}"
    echo -e "${YELLOW}  para crear espacio para persistencia.${NC}"
    echo ""
    echo "  Opciones:"
    echo "    1) Usar GParted (recomendado — gráfico)"
    echo "    2) Usar fdisk (terminal)"
    echo ""
    read -p "  ¿Qué prefieres? [1/2]: " CHOICE
    
    case $CHOICE in
        2)
            echo ""
            echo "  Abriendo fdisk en $USB_DEVICE..."
            echo "  Instrucciones:"
            echo "    - 'd' para borrar la última partición"
            echo "    - 'n' para crear nueva partición (más pequeña)"
            echo "    - 'n' de nuevo para crear partición de persistencia con el espacio restante"
            echo "    - 'w' para guardar"
            echo ""
            fdisk "$USB_DEVICE"
            
            # Después de fdisk, detectar la nueva partición
            sleep 2
            NEW_PART=$(lsblk -dn -o NAME "$USB_DEVICE" | tail -1 | sed 's/.*\([a-z][a-z0-9]*\)$/\1/')
            PERSIST_DEVICE="/dev/$NEW_PART"
            
            echo ""
            read -p "  Formatear $PERSIST_DEVICE como ext4 y etiquetar 'persistence'? [s/N]: " CONFIRM
            if [ "$CONFIRM" = "s" ] || [ "$CONFIRM" = "S" ]; then
                mkfs.ext4 -F -L "persistence" "$PERSIST_DEVICE"
            else
                echo "  Saltado. Formatea manualmente con: mkfs.ext4 -F -L persistence $PERSIST_DEVICE"
                exit 0
            fi
            ;;
        *)
            echo ""
            echo "  Abriendo GParted..."
            echo "  1. Selecciona $USB_DEVICE"
            echo "  2. Redimensiona la última partición (dejando ~4GB para persistencia)"
            echo "  3. Crea una partición ext4 nueva con label 'persistence'"
            echo "  4. Aplica cambios y cierra GParted"
            echo ""
            gparted "$USB_DEVICE" 2>/dev/null || gnome-disks "$USB_DEVICE" 2>/dev/null
            
            # Detectar después de GParted
            sleep 2
            PERSIST_PART=$(lsblk -dn -o NAME,LABEL | grep -i persistence | awk '{print $1}' | head -1)
            if [ -z "$PERSIST_PART" ]; then
                echo -e "${RED}  No se detectó la partición 'persistence'.${NC}"
                echo "  Asegúrate de que el label sea exactamente 'persistence'"
                exit 1
            fi
            PERSIST_DEVICE="/dev/$PERSIST_PART"
            ;;
    esac
fi

echo -e "${GREEN}  Partición de persistencia: $PERSIST_DEVICE${NC}"
echo ""

# -----------------------------------------------------------
# PASO 3: Montar y configurar persistencia
# -----------------------------------------------------------
echo -e "${YELLOW}[3/5] Configurando persistencia...${NC}"

mkdir -p /mnt/usb
mount "$PERSIST_DEVICE" /mnt/usb

# Verificar si ya tiene persistence.conf
if [ -f /mnt/usb/persistence.conf ]; then
    echo -e "${GREEN}  persistence.conf ya existe:$(cat /mnt/usb/persistence.conf)${NC}"
else
    echo "/ union" > /mnt/usb/persistence.conf
    echo -e "${GREEN}  persistence.conf creado: / union${NC}"
fi

umount /mnt/usb
echo -e "${GREEN}  Persistencia configurada OK${NC}"
echo ""

# -----------------------------------------------------------
# PASO 4: Instalar drivers Realtek 8812AU
# -----------------------------------------------------------
echo -e "${YELLOW}[4/5] Instalando drivers Realtek 8812AU...${NC}"

# Instalar dependencias
echo "  Instalando dependencias de compilación..."
apt update -qq
apt install -y -qq dkms linux-headers-$(uname -r) git build-essential airmon-ng 2>/dev/null

# Clonar driver
DRIVER_DIR="/opt/rtl8812au"
if [ -d "$DRIVER_DIR" ]; then
    echo -e "${GREEN}  Driver ya descargado en $DRIVER_DIR${NC}"
    cd "$DRIVER_DIR"
    git pull -qq 2>/dev/null || true
else
    echo "  Descargando driver 8812AU..."
    git clone https://github.com/aircrack-ng/rtl8812au.git "$DRIVER_DIR"
    cd "$DRIVER_DIR"
fi

# Compilar e instalar
echo "  Compilando driver..."
make clean 2>/dev/null || true
make
make install

echo -e "${GREEN}  Driver 8812AU instalado${NC}"
echo ""

# -----------------------------------------------------------
# PASO 5: Cargar módulo y verificar
# -----------------------------------------------------------
echo -e "${YELLOW}[5/5] Verificando...${NC}"

modprobe 8812au 2>/dev/null || modprobe rtl8812au 2>/dev/null || true

if ip link show | grep -q wl; then
    WIFI_IF=$(ip link show | grep wl | awk -F: '{print $2}' | tr -d ' ' | head -1)
    echo -e "${GREEN}  WiFi detectado: $WIFI_IF${NC}"
else
    echo -e "${YELLOW}  WiFi no detectado aún (puede necesitar reboot)${NC}"
fi

echo ""
echo -e "${CYAN}============================================${NC}"
echo -e "${GREEN}  SETUP COMPLETADO${NC}"
echo -e "${CYAN}============================================${NC}"
echo ""
echo "  Resumen:"
echo "    ✓ Persistencia configurada"
echo "    ✓ Drivers Realtek 8812AU instalados"
echo ""
echo "  Próximos pasos:"
echo "    1. Reinicia el USB (selecciona 'KDE + Persistent' o 'XFCE + Persistent')"
echo "    2. Al reiniciar, los cambios se guardarán"
echo "    3. Para usar el WiFi: airmon-ng start $WIFI_IF"
echo ""
echo "  NOTA: Si el módulo no carga al reiniciar, ejecuta:"
echo "    sudo modprobe 8812au"
echo "    # o"
echo "    sudo modprobe rtl8812au"
echo ""