#!/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 "${1:-}"
  echo "Firewall setup completed."
}

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