#!/bin/bash
# =============================================================================
# UFW Firewall Setup + Docker DOCKER-USER Chain
# =============================================================================
# wg0 = öffentliches Interface (getunnelte öffentliche IP)
#   → Nur Port 80 + 443 für Docker Container erlaubt
# Internes Interface (eth0/ens*) = lokales Netz
#   → Alles erlaubt (SSH, WireGuard, etc.)
#
# Ansatz: DOCKER-USER Chain statt "iptables: false"
#   → Docker Networking bleibt intakt (Container können kommunizieren)
#   → UFW kontrolliert welche Ports extern erreichbar sind
# =============================================================================

set -e

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

log()    { echo -e "${GREEN}[✓]${NC} $1"; }
warn()   { echo -e "${YELLOW}[!]${NC} $1"; }
error()  { echo -e "${RED}[✗]${NC} $1"; exit 1; }
info()   { echo -e "${BLUE}[i]${NC} $1"; }

# Root check
[ "$EUID" -ne 0 ] && error "Bitte als root ausführen: sudo $0"

# =============================================================================
# Interfaces ermitteln
# =============================================================================
INTERNAL_IF=$(ip route show default 2>/dev/null | grep -oP 'dev \K\S+' | head -1)
[ -z "$INTERNAL_IF" ] && error "Konnte internes Interface nicht ermitteln"
PUBLIC_IF="wg0"

info "Öffentliches Interface (WireGuard): ${PUBLIC_IF}"
info "Internes Interface:                 ${INTERNAL_IF}"

# Prüfen ob wg0 existiert
if ! ip link show wg0 &>/dev/null; then
    warn "wg0 existiert noch nicht - Regeln werden trotzdem gesetzt"
fi

echo ""
warn "Dieses Script wird UFW neu konfigurieren und Docker anpassen."
warn "Bestehende UFW Regeln werden GELÖSCHT."
read -p "Fortfahren? (j/N): " confirm
[[ "$confirm" != "j" && "$confirm" != "J" ]] && { info "Abgebrochen."; exit 0; }

# =============================================================================
# TEIL 1: Docker iptables reaktivieren
# =============================================================================
echo ""
info "=== TEIL 1: Docker iptables reaktivieren ==="

if systemctl is-active --quiet docker; then
    warn "Stoppe Docker kurz..."
    systemctl stop docker
    DOCKER_WAS_RUNNING=true
fi

DOCKER_DAEMON=/etc/docker/daemon.json

if [ -f "$DOCKER_DAEMON" ] && grep -q '"iptables".*false' "$DOCKER_DAEMON"; then
    cp "$DOCKER_DAEMON" "${DOCKER_DAEMON}.bak.$(date +%Y%m%d_%H%M%S)"
    python3 -c "
import json
with open('$DOCKER_DAEMON', 'r') as f:
    config = json.load(f)
config.pop('iptables', None)
with open('$DOCKER_DAEMON', 'w') as f:
    json.dump(config, f, indent=2)
"
    log "iptables: false aus daemon.json entfernt"
else
    log "Docker iptables bereits aktiv"
fi

# =============================================================================
# TEIL 2: UFW Grundkonfiguration
# =============================================================================
echo ""
info "=== TEIL 2: UFW Konfiguration ==="

ufw --force reset
log "UFW zurückgesetzt"

ufw default deny incoming
ufw default allow outgoing
log "Default Policy: deny incoming, allow outgoing"

# =============================================================================
# TEIL 3: Regeln für wg0 (öffentliches Interface)
# =============================================================================
echo ""
info "=== TEIL 3: Regeln für ${PUBLIC_IF} (öffentlich) ==="

# Nur HTTP + HTTPS auf wg0
ufw allow in on "$PUBLIC_IF" to any port 80 proto tcp comment 'HTTP public (wg0)'
ufw allow in on "$PUBLIC_IF" to any port 443 proto tcp comment 'HTTPS public (wg0)'
log "Port 80+443 (HTTP/HTTPS) auf ${PUBLIC_IF} geöffnet"

# =============================================================================
# TEIL 4: Regeln für internes Interface
# =============================================================================
echo ""
info "=== TEIL 4: Regeln für ${INTERNAL_IF} (intern) ==="

ufw allow in on "$INTERNAL_IF" comment 'Internes Netz - alles erlaubt'
log "Alles auf ${INTERNAL_IF} erlaubt (SSH, WireGuard, etc.)"

# =============================================================================
# TEIL 5: DOCKER-USER Chain (Docker Port-Kontrolle)
# =============================================================================
echo ""
info "=== TEIL 5: DOCKER-USER Chain ==="

AFTER_RULES=/etc/ufw/after.rules
cp "$AFTER_RULES" "${AFTER_RULES}.bak.$(date +%Y%m%d_%H%M%S)"

# Alte Docker-Fixes aus before.rules entfernen (falls vorhanden)
BEFORE_RULES=/etc/ufw/before.rules
if grep -q "DOCKER-UFW-FIX" "$BEFORE_RULES"; then
    cp "$BEFORE_RULES" "${BEFORE_RULES}.bak.$(date +%Y%m%d_%H%M%S)"
    TMPFILE=$(mktemp)
    awk '/# DOCKER-UFW-FIX/{skip=1} skip && /^COMMIT/{skip=0; next} !skip' "$BEFORE_RULES" > "$TMPFILE"
    mv "$TMPFILE" "$BEFORE_RULES"
    warn "Alte Docker-Fixes aus before.rules entfernt"
fi

# DOCKER-USER Regeln in after.rules einfügen
if grep -q "DOCKER-USER-FIX" "$AFTER_RULES"; then
    warn "DOCKER-USER Regeln bereits vorhanden - überspringe"
else
    cat >> "$AFTER_RULES" << EOF

# DOCKER-USER-FIX - Docker Container Port-Kontrolle
# Kontrolliert welcher externe Traffic Docker Container erreicht
# Docker Networking bleibt intakt, nur der Zugang von aussen wird gefiltert
*filter
:DOCKER-USER - [0:0]
-A DOCKER-USER -j RETURN -s 172.16.0.0/12
-A DOCKER-USER -j RETURN -m conntrack --ctstate RELATED,ESTABLISHED
-A DOCKER-USER -p tcp -m tcp --dport 80 -i ${PUBLIC_IF} -j RETURN
-A DOCKER-USER -p tcp -m tcp --dport 443 -i ${PUBLIC_IF} -j RETURN
-A DOCKER-USER -i ${INTERNAL_IF} -j RETURN
-A DOCKER-USER -j DROP
COMMIT
EOF
    log "DOCKER-USER Regeln in after.rules eingefügt"
fi

# =============================================================================
# TEIL 6: IP Forwarding aktivieren
# =============================================================================
echo ""
info "=== TEIL 6: IP Forwarding ==="

SYSCTL_CONF=/etc/ufw/sysctl.conf
if grep -q "^net/ipv4/ip_forward=1" "$SYSCTL_CONF" 2>/dev/null; then
    log "IP Forwarding bereits aktiv"
else
    sed -i 's/#net\/ipv4\/ip_forward=1/net\/ipv4\/ip_forward=1/' "$SYSCTL_CONF"
    grep -q "net/ipv4/ip_forward=1" "$SYSCTL_CONF" || echo "net/ipv4/ip_forward=1" >> "$SYSCTL_CONF"
    log "IP Forwarding aktiviert"
fi

if ! grep -q "net.ipv4.ip_forward = 1" /etc/sysctl.conf; then
    echo "net.ipv4.ip_forward = 1" >> /etc/sysctl.conf
fi
sysctl -p /etc/sysctl.conf > /dev/null 2>&1
log "sysctl neu geladen"

# =============================================================================
# TEIL 7: UFW aktivieren und Docker neu starten
# =============================================================================
echo ""
info "=== TEIL 7: Aktivierung ==="

ufw --force enable
log "UFW aktiviert"

if [ "$DOCKER_WAS_RUNNING" = true ]; then
    systemctl start docker
    sleep 3
    log "Docker neu gestartet"
fi

# =============================================================================
# ZUSAMMENFASSUNG
# =============================================================================
echo ""
echo -e "${GREEN}============================================${NC}"
echo -e "${GREEN}  UFW Setup abgeschlossen!${NC}"
echo -e "${GREEN}============================================${NC}"
echo ""
info "Ansatz: DOCKER-USER Chain (Docker Networking intakt)"
echo ""
info "Interface-Zuordnung:"
info "  ${PUBLIC_IF} (öffentlich) → nur Port 80 + 443"
info "  ${INTERNAL_IF} (intern)   → alles erlaubt"
echo ""
info "Docker Container:"
info "  Extern (${PUBLIC_IF}) → nur Port 80 + 443 erreichbar"
info "  Intern (${INTERNAL_IF}) → alle Ports erreichbar"
echo ""
info "Aktive Regeln:"
ufw status verbose
echo ""
warn "WICHTIG: Teste SSH Verbindung in einem NEUEN Terminal bevor du dieses schließt!"
warn "Falls Docker Container kein Internet haben: 'docker restart <container>'"