#!/bin/bash
source ~/.bashrc
# Configure Firewall Rules for Proxmox Hosts

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"

# Load environment variables
if [ -f "$PROJECT_ROOT/.env" ]; then
    set -a
    source <(grep -v '^#' "$PROJECT_ROOT/.env" | grep -v '^$' | sed 's/#.*$//' | grep '=')
    set +a
fi

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

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

SSH_KEY="${SSH_KEY:-$HOME/.ssh/id_ed25519_proxmox}"
PROXMOX_HOSTS=("192.168.1.206" "192.168.1.49")  # ML110 and R630

main() {
    log_info "Configuring Firewall Rules for Proxmox Hosts"
    echo ""

    for host in "${PROXMOX_HOSTS[@]}"; do
        log_info "Configuring firewall on $host..."
        
        # Check if we can connect
        if ! ssh -i "$SSH_KEY" -o StrictHostKeyChecking=no "root@${host}" "pveversion" &>/dev/null; then
            log_warn "Cannot connect to $host. Skipping..."
            continue
        fi

        # Enable firewall if not already enabled
        log_info "Enabling firewall..."
        ssh -i "$SSH_KEY" -o StrictHostKeyChecking=no "root@${host}" <<'EOF'
set -e

# Enable firewall
pve-firewall enable || true

# Create security group for cluster communication
pve-firewall security-group add cluster-comm --comment "Cluster communication"
pve-firewall security-group rule add cluster-comm --action ACCEPT --proto tcp --dport 8006 --comment "Proxmox Web UI"
pve-firewall security-group rule add cluster-comm --action ACCEPT --proto tcp --dport 22 --comment "SSH"
pve-firewall security-group rule add cluster-comm --action ACCEPT --proto udp --dport 5404:5412 --comment "Corosync cluster"
pve-firewall security-group rule add cluster-comm --action ACCEPT --proto tcp --dport 3128 --comment "SPICE proxy"
pve-firewall security-group rule add cluster-comm --action ACCEPT --proto tcp --dport 111 --comment "RPC"
pve-firewall security-group rule add cluster-comm --action ACCEPT --proto tcp --dport 2049 --comment "NFS"
pve-firewall security-group rule add cluster-comm --action ACCEPT --proto tcp --dport 5900:5999 --comment "VNC"
pve-firewall security-group rule add cluster-comm --action ACCEPT --proto tcp --dport 60000:60050 --comment "Migration"

# Create security group for VM services
pve-firewall security-group add vm-services --comment "VM service ports"
pve-firewall security-group rule add vm-services --action ACCEPT --proto tcp --dport 3000 --comment "Gitea/Grafana"
pve-firewall security-group rule add vm-services --action ACCEPT --proto tcp --dport 9090 --comment "Prometheus"
pve-firewall security-group rule add vm-services --action ACCEPT --proto tcp --dport 6443 --comment "K3s API"
pve-firewall security-group rule add vm-services --action ACCEPT --proto tcp --dport 10250 --comment "Kubelet"

# Configure datacenter firewall options
pve-firewall options set enable 1
pve-firewall options set log_level_in 6  # Log dropped packets
pve-firewall options set log_level_out 6

# Allow cluster communication between nodes
pve-firewall cluster add-rule cluster-comm --action ACCEPT --source 192.168.1.0/24 --comment "Allow cluster subnet"

echo "Firewall configured successfully"
EOF

        log_info "✓ Firewall configured on $host"
        echo ""
    done

    log_info "Firewall configuration complete!"
    echo ""
    log_warn "Review firewall rules:"
    log_info "  - Check rules: pve-firewall status"
    log_info "  - View security groups: pve-firewall security-group list"
    log_info "  - Test connectivity after applying rules"
    echo ""
    log_info "Default rules allow:"
    log_info "  - Cluster communication (ports 5404-5412 UDP)"
    log_info "  - Proxmox Web UI (port 8006)"
    log_info "  - SSH (port 22)"
    log_info "  - VM services (ports 3000, 9090, 6443, 10250)"
    log_info "  - Migration ports (60000-60050)"
}

main "$@"