#!/bin/bash
# Dynamic firewall setup script
set -euo pipefail

# Function to parse WireGuard endpoint
parse_wg_endpoint() {
    local wg_config="${1:-/etc/wireguard/wg0.conf}"
    
    if [[ ! -f "$wg_config" ]]; then
        echo "Error: WireGuard config not found at $wg_config"
        return 1
    fi
    
    grep -oP 'Endpoint = \K[0-9.]+:[0-9]+' "$wg_config" || {
        echo "Error: Could not parse endpoint from WireGuard config"
        return 1
    }
}

# Function to generate nftables rules
generate_nftables_rules() {
    local endpoint="$1"
    local ip="${endpoint%:*}"
    local port="${endpoint#*:}"
    
    cat << EOF
#!/usr/sbin/nft -f
# Secure firewall rules for WireGuard-only access
flush ruleset

table inet filter {
    chain input {
        type filter hook input priority 0; policy drop
        iif lo accept comment "Accept loopback"
        icmp type echo-request accept comment "Accept ping"
    }
    
    chain forward {
        type filter hook forward priority 0; policy drop
    }
    
    chain output {
        type filter hook output priority 0; policy drop
        oif lo accept comment "Accept loopback"
        udp dport "$port" ip daddr "$ip" accept comment "Allow WireGuard traffic"
        icmp type echo-request accept comment "Allow ping"
    }
}
EOF
}

# Function to apply firewall configuration
apply_firewall() {
    local wg_config="${1:-/etc/wireguard/wg0.conf}"
    
    if [[ -f "$wg_config" ]]; then
        endpoint=$(parse_wg_endpoint "$wg_config")
        if [[ -n "$endpoint" ]]; then
            generate_nftables_rules "$endpoint" > /etc/nftables.conf
            systemctl enable nftables
            systemctl restart nftables
            echo "Firewall configured for endpoint: $endpoint"
        else
            echo "Warning: Could not parse WireGuard endpoint, using default deny policy"
        fi
    else
        echo "Warning: WireGuard config not found, using default deny policy"
    fi
}

# Main setup
main() {
    echo "Setting up dynamic firewall..."
    apply_firewall
    echo "Firewall setup completed."
}

# Run main if script is executed directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi