#!/bin/bash
#
# NODE1 Firewall Hardening Script
# Version: 1.0
# Last Updated: 2026-01-26
#
# Usage: ./apply-node1-firewall.sh [--apply|--dry-run|--rollback]
#   --dry-run   Show what would be done (default)
#   --apply     Apply firewall rules
#   --rollback  Restore previous rules
#

set -e

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

# Admin IPs that should have full access (add your IPs here)
ADMIN_IPS=(
    # "YOUR_OFFICE_IP/32"
    # "YOUR_VPN_IP/32"
)

# Ports to DENY from public (will only be accessible locally)
DENY_PORTS=(
    "9102"   # Router
    "9300"   # Gateway (will be proxied via nginx)
    "6333"   # Qdrant
    "30633"  # Qdrant NodePort
    "9090"   # Prometheus
    "3030"   # Grafana
    "8890"   # Swapper
    "8000"   # Memory Service
    "9500"   # RAG Service
    "8001"   # Vision Encoder
    "8101"   # Parser Pipeline
)

# Ports to ALLOW from public
ALLOW_PORTS=(
    "22"     # SSH
    "80"     # HTTP (redirect to HTTPS)
    "443"    # HTTPS (nginx proxy)
)

# Parse arguments
MODE="dry-run"
for arg in "$@"; do
    case $arg in
        --apply) MODE="apply" ;;
        --dry-run) MODE="dry-run" ;;
        --rollback) MODE="rollback" ;;
        --help|-h)
            echo "Usage: $0 [--apply|--dry-run|--rollback]"
            exit 0
            ;;
    esac
done

echo "========================================"
echo "  NODE1 Firewall Hardening"
echo "  Mode: $MODE"
echo "========================================"
echo ""

# Backup current rules
backup_rules() {
    echo "Backing up current UFW rules..."
    sudo cp /etc/ufw/user.rules /etc/ufw/user.rules.backup.$(date +%Y%m%d_%H%M%S) 2>/dev/null || true
    sudo cp /etc/ufw/user6.rules /etc/ufw/user6.rules.backup.$(date +%Y%m%d_%H%M%S) 2>/dev/null || true
    echo "Backup saved to /etc/ufw/user.rules.backup.*"
}

# Apply deny rules
apply_deny_rules() {
    for port in "${DENY_PORTS[@]}"; do
        if [ "$MODE" = "apply" ]; then
            echo -e "${YELLOW}Denying${NC} port $port from public..."
            sudo ufw deny $port/tcp comment "Hardening: internal only" 2>/dev/null || true
        else
            echo "[DRY-RUN] Would deny port $port/tcp"
        fi
    done
}

# Apply allow rules for admin IPs
apply_admin_allowlist() {
    if [ ${#ADMIN_IPS[@]} -eq 0 ]; then
        echo -e "${YELLOW}Warning:${NC} No admin IPs configured in ADMIN_IPS array"
        echo "Add your IPs to enable remote admin access to internal ports"
        return
    fi
    
    for ip in "${ADMIN_IPS[@]}"; do
        for port in "${DENY_PORTS[@]}"; do
            if [ "$MODE" = "apply" ]; then
                echo -e "${GREEN}Allowing${NC} $ip to port $port..."
                sudo ufw allow from $ip to any port $port proto tcp comment "Admin access" 2>/dev/null || true
            else
                echo "[DRY-RUN] Would allow $ip to port $port/tcp"
            fi
        done
    done
}

# Ensure public ports are allowed
apply_allow_rules() {
    for port in "${ALLOW_PORTS[@]}"; do
        if [ "$MODE" = "apply" ]; then
            echo -e "${GREEN}Ensuring${NC} port $port is allowed..."
            sudo ufw allow $port/tcp 2>/dev/null || true
        else
            echo "[DRY-RUN] Would ensure port $port/tcp is allowed"
        fi
    done
}

# Rollback to previous rules
rollback_rules() {
    echo "Looking for backup files..."
    LATEST_BACKUP=$(ls -t /etc/ufw/user.rules.backup.* 2>/dev/null | head -1)
    
    if [ -z "$LATEST_BACKUP" ]; then
        echo -e "${RED}No backup files found!${NC}"
        exit 1
    fi
    
    echo "Restoring from: $LATEST_BACKUP"
    sudo cp "$LATEST_BACKUP" /etc/ufw/user.rules
    
    LATEST_BACKUP6=$(ls -t /etc/ufw/user6.rules.backup.* 2>/dev/null | head -1)
    if [ -n "$LATEST_BACKUP6" ]; then
        sudo cp "$LATEST_BACKUP6" /etc/ufw/user6.rules
    fi
    
    sudo ufw reload
    echo -e "${GREEN}Rollback complete${NC}"
}

# Main execution
case $MODE in
    "apply")
        echo "=== Applying firewall hardening ==="
        backup_rules
        echo ""
        apply_deny_rules
        echo ""
        apply_admin_allowlist
        echo ""
        apply_allow_rules
        echo ""
        echo "Reloading UFW..."
        sudo ufw reload
        echo ""
        echo -e "${GREEN}Hardening applied!${NC}"
        echo ""
        echo "=== Current UFW Status ==="
        sudo ufw status numbered | head -30
        ;;
    "rollback")
        rollback_rules
        ;;
    "dry-run")
        echo "=== DRY RUN - No changes will be made ==="
        echo ""
        echo "Would backup current rules..."
        echo ""
        echo "Ports to DENY from public:"
        for port in "${DENY_PORTS[@]}"; do
            echo "  - $port/tcp"
        done
        echo ""
        echo "Ports to ALLOW from public:"
        for port in "${ALLOW_PORTS[@]}"; do
            echo "  - $port/tcp"
        done
        echo ""
        if [ ${#ADMIN_IPS[@]} -gt 0 ]; then
            echo "Admin IPs with full access:"
            for ip in "${ADMIN_IPS[@]}"; do
                echo "  - $ip"
            done
        else
            echo -e "${YELLOW}Note: No admin IPs configured${NC}"
        fi
        echo ""
        echo "Run with --apply to execute these changes"
        ;;
esac

echo ""
echo "========================================"