#!/bin/bash

#############################################################
#                    scripts/killswitch.sh                 #
#############################################################

cat > scripts/killswitch.sh << 'EOFKILLSWITCH'
#!/bin/bash

# VPN Gateway Killswitch Script
# CRITICAL: This script ensures NO traffic leaks without VPN
# Version: 1.0.0

set -e

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

# Get network configuration from install
if [ -f /opt/vpn-gateway/network.conf ]; then
    source /opt/vpn-gateway/network.conf
else
    # Fallback to auto-detection
    LAN_IF=$(ip route | grep default | awk '{print $5}' | head -n1)
    LAN_NET=$(ip route | grep "$LAN_IF" | grep -v default | awk '{print $1}' | head -n1)
fi

# Logging
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" >> /var/log/vpn-killswitch.log
    echo -e "${GREEN}[+]${NC} $1"
}

error() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ERROR: $1" >> /var/log/vpn-killswitch.log
    echo -e "${RED}[!]${NC} $1"
}

warning() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] WARNING: $1" >> /var/log/vpn-killswitch.log
    echo -e "${YELLOW}[*]${NC} $1"
}

enable_killswitch() {
    log "ENABLING PERMANENT KILLSWITCH - NO INTERNET WITHOUT VPN"
    
    # Backup current rules (just in case)
    iptables-save > /tmp/iptables.backup.$(date +%s) 2>/dev/null || true
    
    # Reset all rules
    iptables -F
    iptables -X
    iptables -t nat -F
    iptables -t nat -X
    iptables -t mangle -F
    iptables -t mangle -X
    iptables -t raw -F
    iptables -t raw -X
    
    # DEFAULT POLICIES: DROP EVERYTHING
    iptables -P INPUT DROP
    iptables -P FORWARD DROP
    iptables -P OUTPUT DROP
    
    log "Default policies set to DROP"
    
    # CRITICAL: Allow loopback (required for system operation)
    iptables -A INPUT -i lo -j ACCEPT
    iptables -A OUTPUT -o lo -j ACCEPT
    
    # Allow LAN communication (for WebUI and local clients)
    if [ -n "$LAN_IF" ] && [ -n "$LAN_NET" ]; then
        iptables -A INPUT -i $LAN_IF -s $LAN_NET -j ACCEPT
        iptables -A OUTPUT -o $LAN_IF -d $LAN_NET -j ACCEPT
        log "LAN communication allowed: $LAN_IF ($LAN_NET)"
    fi
    
    # Allow DNS for initial VPN connection (root only, limited)
    iptables -A OUTPUT -p udp --dport 53 -m owner --uid-owner root -m limit --limit 10/min -j ACCEPT
    iptables -A OUTPUT -p tcp --dport 53 -m owner --uid-owner root -m limit --limit 10/min -j ACCEPT
    
    # Allow DHCP client (if needed)
    iptables -A OUTPUT -p udp --sport 68 --dport 67 -j ACCEPT
    iptables -A INPUT -p udp --sport 67 --dport 68 -j ACCEPT
    
    # Allow established/related connections
    iptables -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    iptables -A OUTPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    iptables -A FORWARD -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    
    # Allow forwarding from LAN (will only work when VPN is up)
    if [ -n "$LAN_IF" ] && [ -n "$LAN_NET" ]; then
        iptables -A FORWARD -i $LAN_IF -s $LAN_NET -j ACCEPT
    fi
    
    # Log dropped packets (optional, can be verbose)
    # iptables -A INPUT -j LOG --log-prefix "INPUT-DROP: " --log-level 4
    # iptables -A OUTPUT -j LOG --log-prefix "OUTPUT-DROP: " --log-level 4
    
    # IPv6: Complete blocking (unless you use IPv6 VPN)
    ip6tables -P INPUT DROP
    ip6tables -P FORWARD DROP
    ip6tables -P OUTPUT DROP
    ip6tables -A INPUT -i lo -j ACCEPT
    ip6tables -A OUTPUT -o lo -j ACCEPT
    
    # Save rules for persistence
    mkdir -p /etc/iptables
    iptables-save > /etc/iptables/rules.v4
    ip6tables-save > /etc/iptables/rules.v6
    
    # Ensure persistence across reboots
    if ! systemctl is-enabled netfilter-persistent >/dev/null 2>&1; then
        systemctl enable netfilter-persistent 2>/dev/null || true
    fi
    
    log "KILLSWITCH ACTIVE - System is now protected"
    log "Rules saved to /etc/iptables/"
    
    # Verify killswitch is working
    verify_killswitch
}

disable_killswitch() {
    error "KILLSWITCH CANNOT BE DISABLED FOR SECURITY REASONS!"
    warning "This is a security feature. The killswitch must remain active."
    warning "If you really need to disable it, you must manually flush iptables rules."
    warning "This would leave your system vulnerable to leaks!"
    
    # Re-enable killswitch immediately
    enable_killswitch
    
    return 1
}

verify_killswitch() {
    log "Verifying killswitch status..."
    
    # Check DROP policies
    if iptables -L -n | grep -q "Chain INPUT (policy DROP)" && \
       iptables -L -n | grep -q "Chain OUTPUT (policy DROP)" && \
       iptables -L -n | grep -q "Chain FORWARD (policy DROP)"; then
        log "✓ Killswitch policies verified: DROP"
    else
        error "✗ Killswitch policies NOT active!"
        enable_killswitch
        return 1
    fi
    
    # Test that we cannot reach internet directly
    if ping -c 1 -W 1 8.8.8.8 >/dev/null 2>&1; then
        if wg show wg0 >/dev/null 2>&1; then
            log "✓ Internet accessible (VPN is connected)"
        else
            error "✗ LEAK DETECTED: Internet accessible without VPN!"
            enable_killswitch
            return 1
        fi
    else
        if wg show wg0 >/dev/null 2>&1; then
            warning "VPN appears connected but internet not reachable"
        else
            log "✓ Internet blocked (VPN not connected)"
        fi
    fi
    
    log "Killswitch verification complete"
}

status_killswitch() {
    echo -e "${GREEN}=== VPN Killswitch Status ===${NC}"
    echo ""
    
    # Check policies
    echo "Firewall Policies:"
    iptables -L -n | grep "Chain.*policy" | while read line; do
        if echo "$line" | grep -q "DROP"; then
            echo -e "  ${GREEN}✓${NC} $line"
        else
            echo -e "  ${RED}✗${NC} $line"
        fi
    done
    
    echo ""
    echo "Active Rules Summary:"
    echo "  INPUT:   $(iptables -L INPUT -n | grep -c ACCEPT) ACCEPT rules"
    echo "  OUTPUT:  $(iptables -L OUTPUT -n | grep -c ACCEPT) ACCEPT rules"
    echo "  FORWARD: $(iptables -L FORWARD -n | grep -c ACCEPT) ACCEPT rules"
    
    echo ""
    echo "VPN Status:"
    if wg show wg0 >/dev/null 2>&1; then
        echo -e "  ${GREEN}✓${NC} WireGuard interface active"
        echo "  Endpoint: $(wg show wg0 endpoints | awk '{print $2}')"
    else
        echo -e "  ${YELLOW}⚠${NC} WireGuard interface not active"
    fi
    
    echo ""
    echo "Leak Test:"
    if ping -c 1 -W 1 8.8.8.8 >/dev/null 2>&1; then
        if wg show wg0 >/dev/null 2>&1; then
            echo -e "  ${GREEN}✓${NC} Internet accessible via VPN"
        else
            echo -e "  ${RED}✗ LEAK: Internet accessible without VPN!${NC}"
        fi
    else
        echo -e "  ${GREEN}✓${NC} Internet blocked (killswitch working)"
    fi
}

# Main logic
case "${1:-enable}" in
    enable|start)
        enable_killswitch
        ;;
    disable|stop)
        disable_killswitch
        ;;
    verify|check)
        verify_killswitch
        ;;
    status)
        status_killswitch
        ;;
    restart|reload)
        enable_killswitch
        ;;
    *)
        echo "Usage: $0 {enable|disable|verify|status|restart}"
        echo "  enable  - Enable killswitch (default)"
        echo "  disable - Disable killswitch (blocked for security)"
        echo "  verify  - Verify killswitch is working"
        echo "  status  - Show detailed status"
        echo "  restart - Restart killswitch"
        exit 1
        ;;
esac