#!/bin/bash
# =============================================================================
# nftables Firewall for Security Scanner VM
# Run this inside the VM as root
# Defense in depth: works alongside PVE-level firewall
# =============================================================================
set -euo pipefail

echo "============================================"
echo " nftables Firewall - Security Scanner VM"
echo "============================================"

# --- Configuration (MUST be set via environment variables) ---
ADMIN_IPS="${SCANNER_ADMIN_IPS:-192.168.68.0/24}"
INTERNAL_NETS="${SCANNER_INTERNAL_NETS:-10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16}"
DNS_SERVERS="${SCANNER_DNS_SERVERS:-192.168.68.1}"

# --- Guard: require ADMIN_IPS ---
if [[ -z "${ADMIN_IPS}" ]]; then
    echo "[ERROR] SCANNER_ADMIN_IPS environment variable is not set." >&2
    echo "  Example: export SCANNER_ADMIN_IPS='192.168.68.100, 192.168.68.101'" >&2
    exit 1
fi

# --- Install nftables ---
apt install -y nftables
systemctl enable nftables

# --- Write nftables config ---
cat > /etc/nftables.conf << EOF
#!/usr/sbin/nft -f

flush ruleset

# =============================================================================
# Definitions
# =============================================================================
define ADMIN_IPS = { ${ADMIN_IPS} }
define INTERNAL_NETS = { ${INTERNAL_NETS} }
define DNS_SERVERS = { ${DNS_SERVERS} }

# Ports for scanner services
define SCANNER_WEB_PORTS = { 443, 9392 }
define SCANNER_SSH_PORT = 22

# =============================================================================
# Main table
# =============================================================================
table inet firewall {

    # -------------------------------------------------------------------------
    # Sets for dynamic blocking
    # -------------------------------------------------------------------------
    set blocked_ips {
        type ipv4_addr
        flags timeout
        comment "Dynamically blocked IPs"
    }

    set ssh_bruteforce {
        type ipv4_addr
        flags dynamic, timeout
        timeout 15m
        comment "SSH brute force tracking"
    }

    # -------------------------------------------------------------------------
    # INPUT chain - What can reach the scanner
    # -------------------------------------------------------------------------
    chain input {
        type filter hook input priority 0; policy drop;

        # Drop blocked IPs immediately
        ip saddr @blocked_ips drop

        # Allow loopback
        iif "lo" accept

        # Allow established/related connections
        ct state established,related accept

        # Drop invalid packets
        ct state invalid counter drop

        # --- ICMP (rate limited) ---
        ip protocol icmp icmp type {
            echo-request,
            echo-reply,
            destination-unreachable,
            time-exceeded
        } limit rate 10/second accept

        # --- SSH from admin IPs only (with brute-force protection) ---
        tcp dport \$SCANNER_SSH_PORT ip saddr \$ADMIN_IPS \
            ct state new \
            limit rate 5/minute \
            accept

        # Track SSH brute force attempts
        tcp dport \$SCANNER_SSH_PORT ct state new \
            add @ssh_bruteforce { ip saddr limit rate 3/minute } \
            counter drop

        # --- OpenVAS Web UI from admin IPs only ---
        tcp dport \$SCANNER_WEB_PORTS ip saddr \$ADMIN_IPS accept

        # --- Docker internal communication (use iifname for wildcard) ---
        iifname "docker*" accept
        iifname "br-*" accept

        # --- Log dropped packets (rate limited) ---
        limit rate 10/second \
            log prefix "[NFT-INPUT-DROP] " level warn

        # Default: drop (set by policy)
    }

    # -------------------------------------------------------------------------
    # FORWARD chain - Docker container traffic
    # -------------------------------------------------------------------------
    chain forward {
        type filter hook forward priority 0; policy accept;

        # Allow Docker container communication
        iifname "docker*" accept
        oifname "docker*" accept
        iifname "br-*" accept
        oifname "br-*" accept
    }

    # -------------------------------------------------------------------------
    # OUTPUT chain - What the scanner can reach
    # -------------------------------------------------------------------------
    chain output {
        type filter hook output priority 0; policy drop;

        # Allow loopback
        oif "lo" accept

        # Allow established/related
        ct state established,related accept

        # --- DNS ---
        udp dport 53 accept
        tcp dport 53 accept

        # --- NTP ---
        udp dport 123 accept

        # --- Scanning: allow ALL traffic to internal networks ---
        # This is required for port scanning, service detection, vuln scanning
        ip daddr \$INTERNAL_NETS accept

        # --- Package updates and vuln DB sync (HTTPS/HTTP) ---
        # Outbound HTTP/HTTPS is allowed for vuln DB updates, package management,
        # and Docker image pulls. Restricted at PVE layer to update_targets IPSET.
        tcp dport { 80, 443 } accept

        # --- ICMP out ---
        ip protocol icmp accept

        # --- Log dropped outbound (rate limited) ---
        limit rate 10/second \
            log prefix "[NFT-OUTPUT-DROP] " level warn

        # Default: drop (set by policy)
    }
}

# =============================================================================
# Separate table for connection tracking and logging
# =============================================================================
table inet logging {

    chain prerouting {
        type filter hook prerouting priority -200; policy accept;

        # Log new inbound connections (rate limited, for audit trail)
        ct state new \
            limit rate 5/second \
            log prefix "[NFT-NEW-CONN] " level info
    }
}
EOF

# --- Apply rules ---
echo "[+] Applying nftables rules..."
nft -f /etc/nftables.conf

# --- Verify ---
echo "[+] Current ruleset:"
nft list ruleset

echo ""
echo "[+] nftables firewall configured and active."
echo ""
echo "Useful commands:"
echo "  nft list ruleset              # Show all rules"
echo "  nft list set inet firewall blocked_ips   # Show blocked IPs"
echo "  nft add element inet firewall blocked_ips { 1.2.3.4 timeout 1h }  # Block IP"
echo "  systemctl restart nftables    # Reload rules"