#!/usr/sbin/nft -f
#
# Example nftables config for:
# - Public web ports (80/443) via Traefik
# - SSH + AI services accessible ONLY over ZeroTier (zt+) or Wireguard
# - Default deny for everything else

flush ruleset

#
# Adjust these to match your system
#
define wan_if      = "eno1"   # WAN interface (optional use; mainly for clarity)
# define wan_if = "$(bash /path/to/detect-wan.sh)"
define vpn_if_pref = "zt+"    # ZeroTier interfaces start with "zt"
define wg_if       = "wg0"
# define lan_if    = "br0"    # If you later want LAN-specific rules


table inet filter {

  #
  # Named sets for ports
  #

  # Ports allowed *only* over ZeroTier|Wireguard (VPN) interface(s)
  set vpn_tcp_ports {
    type inet_service
    comment "SSH + AI services via ZeroTier only"
    elements = {
      22,          # SSH
      7860,        # Stable Diffusion
      8000,        # vLLM or similar
      8080-8089,   # Llamafile, Sandbox UIs, etc. (range example)
      8501,        # Sandbox Fusion / Streamlit-style
      11434,       # Ollama
      21114-21119  # RustDesk server TCP ports
    }
  }

  set vpn_udp_ports {
    type inet_service
    flags interval
    comment "RustDesk UDP via ZeroTier|Wireguard only"
    elements = {
      21116
    }
  }

  # Ports allowed on WAN (public Internet)
  #   Add mail etc. to 'elements' if you ever expose them:
  #   25, 465, 587, 993 # SMTP
  set wan_tcp_ports {
    type inet_service
    comment "Public-facing services (Traefik, etc.)"
    elements = {
      80,          # HTTP (redirect to HTTPS)
      443          # HTTPS
    }
  }

  chain input {
    type filter hook input priority 0;

    #
    # Baseline allow rules
    #

    # Allow loopback
    iif lo accept

    # Allow established/related
    ct state established,related accept

    # ICMP (ping, path MTU, etc.)
    ip protocol icmp accept
    ip6 nexthdr icmpv6 accept

    #
    # Interface/port-specific rules
    #

    # --- ZeroTier-only services (SSH + AI stack) ---
    # Any TCP port in @zt_tcp_ports is allowed only when coming
    # from an interface whose name starts with "zt".
    iifname $vpn_if_pref tcp dport @vpn_tcp_ports accept
    iifname $vpn_if_pref udp dport @vpn_udp_ports accept

    # WireGuard listener on WAN
    iifname $wan_if udp dport $wg_port accept

    # Private services via WireGuard OR ZeroTier
    iifname $wg_if   tcp dport @vpn_tcp_ports accept
    iifname $zt_if_pref tcp dport @vpn_tcp_ports accept
    iifname $wg_if   udp dport @vpn_udp_ports accept
    iifname $zt_if_pref udp dport @vpn_udp_ports accept

    # --- Public web (Traefik) ---
    # Allow HTTP/HTTPS on any interface (typically WAN),
    # but you could restrict to $wan_if if you want:
    # iifname $wan_if tcp dport @wan_tcp_ports accept
    tcp dport @wan_tcp_ports accept

    #
    # Anything not matched above is dropped.
    #
    counter drop
  }

  chain forward {
    type filter hook forward priority 0;

    # By default, block forwarding. If you later do routing/NAT,
    # you can add more specific rules here.
    drop
  }

  chain output {
    type filter hook output priority 0;

    # Start with allow-all outbound.
    # You can tighten this later if you want strict egress control.
    accept
  }
}