setup.sh

#!/bin/bash

echo ""
echo "🚀 Starting setup...v2.8"
echo ""

set -euo pipefail

##########################################
## Begin: Global definitions
##########################################

# Define color codes and checklist variables
TOTAL_STEPS=14
CURRENT_STEP=0

# Function to show progress
show_progress() {
    CURRENT_STEP=$((CURRENT_STEP + 1))
    echo "🔄 Progress: Step $CURRENT_STEP of $TOTAL_STEPS"
    echo "📋 $1"
    echo ""  # Add extra spacing after each step
}

##########################################
## End: Global definitions
##########################################

# Set up logging
setup_logging() {
    LOG_FILE="/var/log/setup-$(date +%Y%m%d-%H%M%S).log"
    exec 1> >(tee -a "$LOG_FILE")
    exec 2> >(tee -a "$LOG_FILE" >&2)
    echo "📝 Logging setup started at $(date)"
    echo ""  # Spacing after logging setup
}

# Handle errors
handle_error() {
    local line_no=$1
    local error_code=$2
    echo "Error occurred at line ${line_no}, exit code ${error_code}"
    exit "${error_code}"
}
trap 'handle_error ${LINENO} $?' ERR

# --------------------------------------
# Detect physical interface (MOVED HERE, ABOVE MAIN SCRIPT LOGIC)
detect_physical_interface() {
    # Try to find interface with the router ID first
    local iface=$(ip -br addr show | grep "$ROUTER_ID" | awk '{print $1}')
    
    if [ -z "$iface" ]; then
        # If not found, try to find the main physical interface
        iface=$(ip -o link show | grep -v "lo\|docker\|veth\|br-\|vxlan\|ipvlan\|dummy" | head -n1 | awk -F': ' '{print $2}')
    fi
    
    if [ -z "$iface" ]; then
        echo "❌ Error: Could not detect physical network interface"
        exit 1
    fi
    
    echo "Found physical interface: $iface"
    echo ""
    PHYSICAL_INTERFACE="$iface"
}

# --------------------------------------
# Usage Documentation (Point 11)
# --------------------------------------

show_help() {
    cat << EOF
Usage: $0 -<region><node_number>

Regions available:
    ny - New York     (203.153.52.0/24)
    sg - Singapore    (203.153.54.0/24)
    bm - Bombay       (203.153.53.0/24)
    fr - Frankfurt    (203.153.55.0/24)
    mn - Manilla      (103.165.98.0/24)

Example:
    $0 -ny38
    $0 -sg42
EOF
}

# --------------------------------------
# Network Connectivity Check (Point 3)
# --------------------------------------

check_connectivity() {
    show_progress "🌐 Checking network connectivity..."
    if ! ping -c 1 8.8.8.8 &>/dev/null; then
        echo "No internet connectivity. Please check your network connection."
        exit 1
    fi
}

# --------------------------------------
# Configuration Validation (Point 5)
# --------------------------------------

validate_configs() {
    echo "Validating configurations..."
    # Validate nftables configuration
    if ! nft -c -f /etc/nftables.conf; then
        echo "Invalid nftables configuration"
        exit 1
    fi

    # Validate Docker network configuration
    if ! docker network inspect space &>/dev/null; then
        echo "Docker network 'space' was not created properly"
        exit 1
    fi
}

# --------------------------------------
# Health Checks (Point 9)
# --------------------------------------

perform_health_checks() {
    echo "Performing health checks..."
    # Check Docker service
    if ! systemctl is-active --quiet docker; then
        echo "Docker service is not running"
        exit 1
    fi

    # Check nftables service
    if ! systemctl is-active --quiet nftables; then
        echo "nftables service is not running"
        exit 1
    fi

    # Check DockNet service
    if ! systemctl is-active --quiet docknet; then
        echo "DockNet service is not running"
        exit 1
    fi
}

# Add this new function after perform_health_checks()
backup_existing_configs() {
    echo "Backing up existing configurations..."
    
    # Backup nftables config if exists
    if [ -f "/etc/nftables.conf" ]; then
        cp /etc/nftables.conf /etc/nftables.conf.backup.$(date +%Y%m%d-%H%M%S)
        echo "Backed up nftables configuration"
    fi
    
    # Backup Docker daemon config if exists
    if [ -f "/etc/docker/daemon.json" ]; then
        cp /etc/docker/daemon.json /etc/docker/daemon.json.backup.$(date +%Y%m%d-%H%M%S)
        echo "Backed up Docker daemon configuration"
    fi
}

# Add this new function
check_docknet_version() {
    echo "Checking DockNet version..."
    local version
    
    # Try to get version from running service
    if systemctl is-active --quiet docknet; then
        version=$(journalctl -u docknet -n 50 --no-pager | grep "Starting docknet v" | tail -n1 | grep -oP 'v\d+\.\d+\.\d+' || echo "unknown")
        if [ "$version" != "unknown" ]; then
            echo "Current running DockNet version: $version"
        else
            echo "Could not determine current DockNet version from logs"
        fi
    else
        echo "DockNet service is not running"
    fi
}

# Modify the cleanup_existing_docknet function
cleanup_existing_docknet() {
    echo "Checking for existing docknet installation..."
    if [ -f "/usr/local/bin/docknet" ]; then
        echo "Found existing docknet binary, checking version..."
        check_docknet_version
        
        echo "Stopping DockNet service..."
        systemctl stop docknet || true
        
        echo "Backing up existing docknet binary..."
        mv /usr/local/bin/docknet /usr/local/bin/docknet.backup.$(date +%Y%m%d-%H%M%S)
        
        echo "Existing docknet binary backed up and removed"
    fi
}

# Modify the update_docker_network function
update_docker_network() {
    echo "Checking Docker network configuration..."
    
    # Check if network exists
    if docker network inspect space >/dev/null 2>&1; then
        echo "Existing Docker network 'space' found"
        
        # Compare all existing configs with desired config
        CURRENT_PARENT=$(docker network inspect space -f '{{index .Options "parent"}}')
        CURRENT_SUBNETS=$(docker network inspect space -f '{{range .IPAM.Config}}{{.Subnet}} {{end}}')
        CURRENT_DRIVER=$(docker network inspect space -f '{{.Driver}}')
        CURRENT_IPVLAN_MODE=$(docker network inspect space -f '{{index .Options "ipvlan_mode"}}')
        
        # Build expected subnets string for comparison
        EXPECTED_SUBNETS="10.0.0.0/20 ${DOCKER_SUBNET} 103.165.99.0/24 "
        
        # Check if any configuration differs
        if [ "$CURRENT_PARENT" != "$PHYSICAL_INTERFACE" ] || \
           [ "$CURRENT_SUBNETS" != "$EXPECTED_SUBNETS" ] || \
           [ "$CURRENT_DRIVER" != "ipvlan" ] || \
           [ "$CURRENT_IPVLAN_MODE" != "l3s" ]; then
            
            echo "Network configuration changes detected:"
            [ "$CURRENT_PARENT" != "$PHYSICAL_INTERFACE" ] && echo "- Interface changed: $CURRENT_PARENT -> $PHYSICAL_INTERFACE"
            [ "$CURRENT_SUBNETS" != "$EXPECTED_SUBNETS" ] && echo "- Subnets changed: $CURRENT_SUBNETS -> $EXPECTED_SUBNETS"
            [ "$CURRENT_DRIVER" != "ipvlan" ] && echo "- Driver changed: $CURRENT_DRIVER -> ipvlan"
            [ "$CURRENT_IPVLAN_MODE" != "l3s" ] && echo "- IPVLAN mode changed: $CURRENT_IPVLAN_MODE -> l3s"
            
            # Check if network has any containers
            if docker network inspect space -f '{{len .Containers}}' | grep -q '^0$'; then
                echo "No containers using network, safe to recreate"
                echo "Removing existing Docker network..."
                docker network rm space || true
                
                echo "Creating Docker network with new configuration..."
                docker network create -d ipvlan \
                    --subnet=10.0.0.0/20 \
                    --subnet=${DOCKER_SUBNET} \
                    --subnet=103.165.99.0/24 \
                    -o parent=${PHYSICAL_INTERFACE} \
                    -o ipvlan_mode=l3s space
            else
                echo "⚠️ WARNING: Network configuration differs but containers are running."
                echo "❌ Cannot recreate network while containers are using it."
                echo "👉 Please manually migrate containers before updating network configuration."
                return 1
            fi
        else
            echo "✅ Network configuration matches desired state, no changes needed"
        fi
    else
        echo "Docker network 'space' does not exist"
        echo "Creating new Docker network..."
        docker network create -d ipvlan \
            --subnet=10.0.0.0/20 \
            --subnet=${DOCKER_SUBNET} \
            --subnet=103.165.99.0/24 \
            -o parent=${PHYSICAL_INTERFACE} \
            -o ipvlan_mode=l3s space
    fi
}

# Modify the verify_docknet_running function
verify_docknet_running() {
    echo "Verifying DockNet service status..."
    echo ""
    local max_attempts=10
    local attempt=1
    local service_ok=false
    
    # Initial setup
    chmod 755 /usr/local/bin/docknet
    chown root:root /usr/local/bin/docknet
    systemctl daemon-reload
    
    # Always do a clean restart
    echo "Performing clean restart of DockNet service..."
    systemctl stop docknet
    sleep 3
    systemctl start docknet
    sleep 5  # Give it time to initialize
    
    while [ $attempt -le $max_attempts ]; do
        echo "Checking DockNet service (attempt $attempt/$max_attempts)..."
        
        if systemctl is-active --quiet docknet; then
            # Check logs for success
            local logs=$(journalctl -u docknet -n 50 --no-pager)
            
            # Look for BGP connection specifically
            if echo "$logs" | grep -q "Peer Up.*State=BGP_FSM_OPENCONFIRM"; then
                service_ok=true
                echo "✅ BGP peer connection established successfully"
                break
            else
                echo "Service started but waiting for BGP connection. Current status:"
                echo "$logs" | grep -E "Starting docknet|Peer.*State=|Add a peer" | tail -n3
                
                # If we see errors, restart the service
                if echo "$logs" | grep -q "error\|failed" || ! echo "$logs" | grep -q "Add a peer"; then
                    echo "Detected issues, restarting DockNet..."
                    systemctl restart docknet
                    sleep 5
                fi
            fi
        else
            echo "Service stopped unexpectedly, restarting..."
            systemctl restart docknet
            sleep 5
        fi
        
        attempt=$((attempt + 1))
        echo "Waiting 5 seconds before next check..."
        echo ""
        sleep 5
    done
    
    if [ "$service_ok" = true ]; then
        echo "✅ DockNet service is running and BGP peer is connected"
        echo ""
        check_docknet_version
        echo ""
        return 0
    else
        # Check if service is at least running
        if systemctl is-active --quiet docknet; then
            # Get the last few relevant log lines for debugging
            echo "Current service status:"
            journalctl -u docknet -n 20 --no-pager | grep -E "Starting docknet|Advertised route|NFT element|Add a peer|Peer Up" || true
            echo ""
            echo "ℹ️ DockNet service is running but BGP peer connection not detected yet"
            echo "👉 This is normal during initial startup, BGP may take a moment to establish"
            check_docknet_version
            echo ""
            return 0
        else
            echo "❌ Error: DockNet service failed to start"
            echo "Last few log entries:"
            journalctl -u docknet -n 20 --no-pager
            echo ""
            return 1
        fi
    fi
}

# --------------------------------------
# Variable Declarations
# --------------------------------------

# Define global variables for regions
declare -A REGION_SUBNETS=(
    [ny]="203.153.52.0/24"
    [sg]="203.153.54.0/24"
    [bm]="203.153.53.0/24"
    [fr]="203.153.55.0/24"
    [mn]="103.165.98.0/24"
)

declare -A REGION_AS=(
    [ny]="64512"
    [sg]="64514"
    [bm]="64513"
    [fr]="64515"
    [mn]="64598"
)

declare -A DOCKER_SUBNETS=(
    [ny]="203.153.52.0/24"
    [sg]="203.153.54.0/24"
    [bm]="203.153.53.0/24"
    [fr]="203.153.55.0/24"
    [mn]="103.165.98.0/24"
)

# --------------------------------------
# Functions
# --------------------------------------

# Parse argument to extract region and node number
parse_argument() {
    local arg=$1
    # Remove any leading dash
    arg=${arg#-}
    # Extract region (first two characters)
    REGION=${arg:0:2}
    # Extract node number (remaining characters)
    NODE_NUM=${arg:2}

    # Convert region to lowercase for consistency
    REGION=$(echo "$REGION" | tr '[:upper:]' '[:lower:]')
}

# Validate input arguments
validate_input() {
    if [ -z "$REGION" ] || [ -z "$NODE_NUM" ]; then
        echo "Error: Invalid input format."
        show_help
        exit 1
    fi

    if [[ ! ${REGION_SUBNETS[$REGION]+_} ]]; then
        echo "Error: Invalid region."
        show_help
        exit 1
    fi

    if ! [[ "$NODE_NUM" =~ ^[0-9]+$ ]]; then
        echo "Error: Node number must be numeric."
        exit 1
    fi
}

# Set derived variables based on region and node number
set_derived_variables() {
    MAIN_SUBNET=${REGION_SUBNETS[$REGION]}
    LOCAL_AS=${REGION_AS[$REGION]}
    DOCKER_SUBNET=${DOCKER_SUBNETS[$REGION]}

    # Extract base IP from subnet (remove CIDR notation)
    BASE_IP=${MAIN_SUBNET%/*}

    # Calculate router IP (always .1) and router ID (based on node number)
    ROUTER_IP="${BASE_IP%.*}.1"
    ROUTER_ID="${BASE_IP%.*}.$NODE_NUM"
    
    # Get the physical interface for this node
    get_physical_interface "$ROUTER_ID"

    # Add this after set_derived_variables
    backup_existing_configs
}

# Add this function after set_derived_variables()
get_physical_interface() {
    local node_ip="$1"
    # Get interface name that has the node IP
    PHYSICAL_INTERFACE=$(ip -br addr show | grep "$node_ip" | awk '{print $1}')
    
    if [ -z "$PHYSICAL_INTERFACE" ]; then
        echo "❌ Error: Could not find physical interface for IP $node_ip"
        exit 1
    fi
    
    echo "Found physical interface: $PHYSICAL_INTERFACE"
}

# --------------------------------------
# Main Execution
# --------------------------------------

# Main script execution
setup_logging

# Parse command line arguments first
if [[ $# -ne 1 || ! $1 =~ ^-[a-z]{2}[0-9]{1,2}$ ]]; then
    show_help
    exit 1
fi

# Extract region and node number
REGION=${1:1:2}
NODE_NUMBER=${1:3}

# Validate region and set variables
if [[ ! ${REGION_SUBNETS[$REGION]+_} ]]; then
    echo "Invalid region code: $REGION"
    show_help
    exit 1
fi

# Set network variables based on region
MAIN_SUBNET=${REGION_SUBNETS[$REGION]}
LOCAL_AS=${REGION_AS[$REGION]}
ROUTER_IP=$(echo $MAIN_SUBNET | sed 's|0/24|1|')
ROUTER_ID=$(echo $MAIN_SUBNET | sed 's|0/24|'$NODE_NUMBER'|')
DOCKER_SUBNET=${DOCKER_SUBNETS[$REGION]}

# Detect physical interface
show_progress "🔍 Detecting network interface..."
detect_physical_interface

# First, check and install required packages
show_progress "📦 Updating system packages..."
if command -v apt &> /dev/null; then
    # Debian/Ubuntu
    DEBIAN_FRONTEND=noninteractive apt-get -qq update && DEBIAN_FRONTEND=noninteractive apt-get -qq upgrade -y
elif command -v yum &> /dev/null; then
    # CentOS/RHEL
    yum update -y
fi

# Install Docker if not present
show_progress "🐳 Installing Docker..."
if ! command -v docker &> /dev/null; then
    echo "Installing Docker..."
    curl -fsSL https://get.docker.com -o get-docker.sh
    sh get-docker.sh
    systemctl start docker
    systemctl enable docker
else
    echo "Docker is already installed, skipping installation"
fi

# Now proceed with network checks and other setup
check_connectivity

# Show configuration parameters
show_progress "⚙️ Configuring with the following parameters:"
echo "Region: $REGION"
echo "Node Number: $NODE_NUMBER"
echo "Main Subnet: $MAIN_SUBNET"
echo "Local AS: $LOCAL_AS"
echo "Router IP: $ROUTER_IP"
echo "Router ID: $ROUTER_ID"
echo "Docker Subnet: $DOCKER_SUBNET"

# Modify the nftables installation section
show_progress "🔄 Checking nftables installation..."
if command -v nft &> /dev/null; then
    echo "nftables is already installed, skipping installation"
else
    echo "Installing nftables..."
    if command -v apt &> /dev/null; then
        # Debian/Ubuntu
        DEBIAN_FRONTEND=noninteractive apt-get -qq install -y nftables
    elif command -v yum &> /dev/null; then
        # CentOS/RHEL
        yum install -y nftables
    fi
fi

# Define update_nftables_service function BEFORE it is called
update_nftables_service() {
    echo "Checking nftables service configuration..."
    # Simple nftables service check
    if ! systemctl is-active --quiet nftables; then
        echo "nftables service is not running"
        return 1
    fi
    echo "nftables service is running"
    return 0
}

# Enable and start nftables service
echo "Verifying nftables service status..."
if systemctl is-enabled nftables &>/dev/null; then
    echo "nftables is already enabled, skipping enable."
else
    # Check both possible symlink locations
    if [ -e "/etc/systemd/system/sysinit.target.wants/nftables.service" ] || \
       [ -e "/usr/lib/systemd/system/sysinit.target.wants/nftables.service" ]; then
        echo "nftables service symlink exists but service shows as disabled."
        echo "Attempting to fix service state..."
        
        # Remove existing symlink
        rm -f /etc/systemd/system/sysinit.target.wants/nftables.service
        rm -f /usr/lib/systemd/system/sysinit.target.wants/nftables.service
        
        # Reload systemd to clear cached state
        systemctl daemon-reload
        
        # Try enabling again
        echo "Re-enabling nftables service..."
        systemctl enable nftables
    else
        echo "Enabling nftables service..."
        systemctl enable nftables
    fi
fi

# Ensure service is started
systemctl start nftables || {
    echo "Failed to start nftables service. Attempting to fix..."
    systemctl stop nftables
    sleep 2
    systemctl daemon-reload
    sleep 2
    systemctl start nftables
}

# Configure Docker Daemon
show_progress "🔧 Configuring Docker daemon..."
mkdir -p /etc/docker || {
    echo "❌ Error: Failed to create /etc/docker directory"
    exit 1
}
if [ -f /etc/docker/daemon.json ]; then
    cp /etc/docker/daemon.json /etc/docker/daemon.json.bak
fi

cat > /etc/docker/daemon.json << EOF
{
    "iptables": false,
    "bridge": "none"
}
EOF

# Backup existing configurations
backup_existing_configs

# Create NFT rules
show_progress "🔒 Creating NFT rules..."
cat > /etc/nftables.conf << EOF
#!/usr/sbin/nft -f

flush ruleset

table inet firewall {
	set media {
		type ipv4_addr
		flags dynamic
	}

	set sip {
		type ipv4_addr
		flags dynamic
	}

	set gw {
		type ipv4_addr
		flags dynamic
	}

	chain input {
		type filter hook input priority filter; policy drop;
		ct state { established, related } accept
		iif "lo" accept
		ip saddr { 10.0.0.0/8, 103.165.98.0/23, 203.153.52.0/22 } accept
		ip daddr @media udp dport 55000-60000 accept
		ip daddr @sip udp dport 5060 accept
		ip daddr @sip tcp dport 5061 accept
		ip daddr @gw tcp dport { 80, 443 } accept
		ip daddr @gw udp dport 443 accept
		ip protocol { icmp, igmp } limit rate 5/second burst 5 packets accept
	}

	chain output {
		type filter hook output priority filter; policy accept;
	}

	chain prerouting {
		type nat hook prerouting priority dstnat; policy accept;
	}

	chain postrouting {
		type nat hook postrouting priority srcnat; policy accept;
        oifname "${PHYSICAL_INTERFACE}" ip saddr 10.0.0.0/8 ip daddr != { 10.0.0.0/8, ${MAIN_SUBNET}, 103.165.99.0/24 } masquerade
	}
}

EOF

# Make nftables.conf executable
chmod +x /etc/nftables.conf

# Apply NFT rules
echo "Applying NFT rules..."
nft -f /etc/nftables.conf

# Ensure nftables rules persist across reboots
if [ -f /etc/sysconfig/nftables.conf ]; then
    # RHEL/CentOS
    cp /etc/nftables.conf /etc/sysconfig/nftables.conf
fi

# Add this function to handle config file migration
handle_docknet_config() {
    echo "Setting up DockNet configuration..."
    
    # Create directory for new config location
    mkdir -p /etc/docker
    
    # Backup and remove old config if it exists
    if [ -f "/usr/local/bin/config.yml" ]; then
        echo "Found old config file, backing it up..."
        cp /usr/local/bin/config.yml /usr/local/bin/config.yml.backup.$(date +%Y%m%d-%H%M%S)
        rm -f /usr/local/bin/config.yml
        echo "Old config file backed up and removed"
    fi
    
    # Backup existing config in new location if it exists
    if [ -f "/etc/docker/docknet.yml" ]; then
        echo "Backing up existing docknet config..."
        cp /etc/docker/docknet.yml /etc/docker/docknet.yml.backup.$(date +%Y%m%d-%H%M%S)
    fi
    
    # Create new config file
    echo "Creating DockNet configuration..."
    cat > /etc/docker/docknet.yml << EOF
nfTable: "firewall"
bgp:
  neighborIp: ${ROUTER_IP}
  neighborAs: 45513
  localAs: ${LOCAL_AS}
  routerId: ${ROUTER_ID}
networks:
  - 203.153.52.0/22
  - 103.165.98.0/23
EOF
}

# Set up Docker login
show_progress "🔑 Configuring Docker login..."
mkdir -p /root/.docker
cat > /root/.docker/config.json << EOF
{
	"auths": {
		"cd.ss.dev:5000": {
			"auth": "c2FtZXNwYWNlOmlMbVU2Q1lzYUNZZkI0anY="
		},
		"deploy.ss.dev:5000": {
			"auth": "c2FtZXNwYWNlOmlMbVU2Q1lzYUNZZkI0anY="
		},
		"https://index.docker.io/v1/": {
			"auth": "c2FtZXNwYWNlOnZlWHBhai10aXNzaXEtNWZ5bXBp"
		},
		"registry.samespace.com": {
			"auth": "c2FtZXNwYWNlOnBZcmU4UXFiek1IWFduX2FfcHpi"
		}
	}
}
EOF

# Restart Docker service
systemctl daemon-reload
systemctl restart docker


# Create systemd service file
echo "Configuring systemd service..."
cat > /etc/systemd/system/docknet.service << EOF
[Unit]
Description=DockNet Service
After=network.target nftables.service
PartOf=nftables.service

[Service]
User=root
ExecStart=/usr/local/bin/docknet
Restart=always
WorkingDirectory=/usr/local/bin

[Install]
WantedBy=multi-user.target
EOF

# Add this before downloading DockNet binary
cleanup_existing_docknet

# Download DockNet binary
show_progress "⬇️ Downloading docknet binary..."
DOCKNET_URL="https://raw.githubusercontent.com/samespace/Public/master/docknet"
DOWNLOAD_PATH="/usr/local/bin/docknet"

# Check if curl is installed
if ! command -v curl &> /dev/null; then
    echo "Installing curl..."
    if command -v apt &> /dev/null; then
        DEBIAN_FRONTEND=noninteractive apt-get -qq install -y curl
    elif command -v yum &> /dev/null; then
        yum install -y curl
    fi
fi

# Try to download with curl
if ! curl --max-time 30 -L -f -s -o "$DOWNLOAD_PATH" "$DOCKNET_URL"; then
    echo "Failed to download docknet binary from $DOCKNET_URL"
    exit 1
fi

# Make it executable
chmod +x "$DOWNLOAD_PATH"

# Verify the file exists and is executable
if [ ! -x "/usr/local/bin/docknet" ]; then
    echo "Error: docknet binary is not executable or doesn't exist"
    exit 1
fi

# Handle DockNet configuration
handle_docknet_config

# Start the service
show_progress "🚀 Starting docknet service..."
systemctl daemon-reload
systemctl enable docknet.service

# Verify service is running correctly
if ! verify_docknet_running; then
    echo "ERROR: DockNet service failed to start properly"
    exit 1
fi

# Add version check
check_docknet_version

# Function to handle sysctl settings
configure_sysctl() {
    echo "Configuring system settings..."
    echo ""
    
    # Check and set IP forwarding
    local ip_forward_val=$(sysctl -n net.ipv4.ip_forward 2>/dev/null || echo "0")
    if [ "$ip_forward_val" != "1" ]; then
        echo "Enabling IP forwarding..."
        sed -i '/^net.ipv4.ip_forward[[:space:]]*=/d' /etc/sysctl.conf
        echo "net.ipv4.ip_forward = 1" >> /etc/sysctl.conf
        sysctl -w net.ipv4.ip_forward=1 >/dev/null
        echo "✅ IP forwarding enabled"
        echo ""
    else
        echo "✅ IP forwarding already enabled"
        echo ""
    fi
    
    # Check and set FS AIO
    local aio_val=$(sysctl -n fs.aio-max-nr 2>/dev/null || echo "0")
    if [ "$aio_val" != "2097152" ]; then
        echo "Setting Scylla FS AIO..."
        sed -i '/^fs.aio-max-nr[[:space:]]*=/d' /etc/sysctl.conf
        echo "fs.aio-max-nr = 2097152" >> /etc/sysctl.conf
        sysctl -w fs.aio-max-nr=2097152 >/dev/null
        echo "✅ FS AIO configured"
        echo ""
    else
        echo "✅ FS AIO already configured"
        echo ""
    fi
}

# Replace the old sysctl calls with the new function
configure_sysctl

# Ensure Docker network 'space' is created before validating
update_docker_network || {
    echo "❌ Failed to configure Docker network properly"
    echo ""
    exit 1
}
validate_configs
perform_health_checks

# Add this new function to handle dps script installation
install_dps_script() {
    echo "Setting up dps script..."
    local DPS_URL="https://raw.githubusercontent.com/samespace/Public/master/dps.sh"
    local DPS_PATH="/usr/local/bin/dps"
    
    # Backup existing dps if it exists
    if [ -f "$DPS_PATH" ]; then
        echo "Found existing dps script, backing it up..."
        mv "$DPS_PATH" "${DPS_PATH}.backup.$(date +%Y%m%d-%H%M%S)"
        echo "Existing dps script backed up"
    fi
    
    # Download the script
    echo "Downloading dps script..."
    if ! curl -L -f -s -o "$DPS_PATH" "$DPS_URL"; then
        echo "Failed to download dps script from $DPS_URL"
        return 1
    fi
    
    # Make it executable
    chmod +x "$DPS_PATH"
    
    # Verify the file exists and is executable
    if [ ! -x "$DPS_PATH" ]; then
        echo "Error: dps script is not executable or doesn't exist"
        return 1
    fi
    
    echo "dps script installed successfully"
    return 0
}

# Install dps script
install_dps_script || {
    echo "⚠️ WARNING: Failed to install dps script"
    echo ""
}

echo ""
echo "✅ Setup completed successfully! All $TOTAL_STEPS steps completed."
echo ""