Configure MongoDB 8.0 replica set with automatic failover for high availability

#!/usr/bin/env bash

set -euo pipefail

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

# Default values
REPLICA_SET_NAME="rs0"
MONGODB_PORT="27017"
ADMIN_PASSWORD=""
NODE_IPS=()
KEYFILE_PATH="/opt/mongodb/mongodb-keyfile"

# Usage function
usage() {
    cat << EOF
Usage: $0 --admin-password PASSWORD --nodes IP1,IP2,IP3 [OPTIONS]

Required:
  --admin-password PASSWORD    MongoDB admin password
  --nodes IP1,IP2,IP3         Comma-separated list of node IPs (minimum 3)

Optional:
  --replica-set-name NAME     Replica set name (default: rs0)
  --port PORT                 MongoDB port (default: 27017)
  --help                      Show this help message

Example:
  $0 --admin-password "SecurePass123!" --nodes "10.0.1.10,10.0.1.11,10.0.1.12"
EOF
    exit 1
}

# Logging functions
log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# Cleanup function
cleanup() {
    if [[ $? -ne 0 ]]; then
        log_error "Script failed. Cleaning up..."
        systemctl stop mongod 2>/dev/null || true
        rm -f /tmp/mongodb-keyfile 2>/dev/null || true
    fi
}

trap cleanup ERR

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --admin-password)
            ADMIN_PASSWORD="$2"
            shift 2
            ;;
        --nodes)
            IFS=',' read -ra NODE_IPS <<< "$2"
            shift 2
            ;;
        --replica-set-name)
            REPLICA_SET_NAME="$2"
            shift 2
            ;;
        --port)
            MONGODB_PORT="$2"
            shift 2
            ;;
        --help)
            usage
            ;;
        *)
            log_error "Unknown option: $1"
            usage
            ;;
    esac
done

# Validate arguments
if [[ -z "$ADMIN_PASSWORD" ]]; then
    log_error "Admin password is required"
    usage
fi

if [[ ${#NODE_IPS[@]} -lt 3 ]]; then
    log_error "At least 3 nodes are required for a replica set"
    usage
fi

# Check prerequisites
echo "[1/9] Checking prerequisites..."
if [[ $EUID -ne 0 ]]; then
    log_error "This script must be run as root"
    exit 1
fi

# Detect distribution
if [[ -f /etc/os-release ]]; then
    . /etc/os-release
    case "$ID" in
        ubuntu|debian)
            PKG_MGR="apt"
            PKG_INSTALL="apt install -y"
            PKG_UPDATE="apt update"
            FIREWALL_CMD="ufw"
            ;;
        almalinux|rocky|centos|rhel|ol)
            PKG_MGR="dnf"
            PKG_INSTALL="dnf install -y"
            PKG_UPDATE="dnf check-update || true"
            FIREWALL_CMD="firewall-cmd"
            ;;
        fedora)
            PKG_MGR="dnf"
            PKG_INSTALL="dnf install -y"
            PKG_UPDATE="dnf check-update || true"
            FIREWALL_CMD="firewall-cmd"
            ;;
        amzn)
            PKG_MGR="yum"
            PKG_INSTALL="yum install -y"
            PKG_UPDATE="yum check-update || true"
            FIREWALL_CMD="firewall-cmd"
            ;;
        *)
            log_error "Unsupported distribution: $ID"
            exit 1
            ;;
    esac
else
    log_error "Cannot detect distribution"
    exit 1
fi

log_info "Detected distribution: $ID"

# Install MongoDB 8.0
echo "[2/9] Installing MongoDB 8.0..."
if [[ "$PKG_MGR" == "apt" ]]; then
    curl -fsSL https://www.mongodb.org/static/pgp/server-8.0.asc | gpg --dearmor -o /usr/share/keyrings/mongodb-server-8.0.gpg
    echo "deb [arch=amd64,arm64 signed-by=/usr/share/keyrings/mongodb-server-8.0.gpg] https://repo.mongodb.org/apt/ubuntu jammy/mongodb-org/8.0 multiverse" > /etc/apt/sources.list.d/mongodb-org-8.0.list
    $PKG_UPDATE
    $PKG_INSTALL mongodb-org
else
    cat > /etc/yum.repos.d/mongodb-org-8.0.repo << 'EOF'
[mongodb-org-8.0]
name=MongoDB Repository
baseurl=https://repo.mongodb.org/yum/redhat/8/mongodb-org/8.0/x86_64/
gpgcheck=1
enabled=1
gpgkey=https://www.mongodb.org/static/pgp/server-8.0.asc
EOF
    $PKG_UPDATE
    $PKG_INSTALL mongodb-org
fi

# Configure MongoDB
echo "[3/9] Configuring MongoDB for replica set..."
cat > /etc/mongod.conf << EOF
storage:
  dbPath: /var/lib/mongo
  journal:
    enabled: true

systemLog:
  destination: file
  logAppend: true
  path: /var/log/mongodb/mongod.log

net:
  port: ${MONGODB_PORT}
  bindIp: 0.0.0.0

processManagement:
  timeZoneInfo: /usr/share/zoneinfo

replication:
  replSetName: ${REPLICA_SET_NAME}

security:
  authorization: enabled
  keyFile: ${KEYFILE_PATH}
EOF

# Generate keyfile
echo "[4/9] Generating replica set keyfile..."
mkdir -p /opt/mongodb
openssl rand -base64 756 > "$KEYFILE_PATH"
chmod 400 "$KEYFILE_PATH"
chown mongodb:mongodb "$KEYFILE_PATH"

log_info "Keyfile generated at: $KEYFILE_PATH"
log_warn "Copy this keyfile to all other nodes with the same path and permissions"

# Configure firewall
echo "[5/9] Configuring firewall..."
if [[ "$FIREWALL_CMD" == "ufw" ]]; then
    if command -v ufw >/dev/null; then
        for ip in "${NODE_IPS[@]}"; do
            ufw allow from "$ip" to any port "$MONGODB_PORT" >/dev/null 2>&1 || true
        done
        ufw --force reload >/dev/null 2>&1 || true
    fi
else
    if command -v firewall-cmd >/dev/null; then
        for ip in "${NODE_IPS[@]}"; do
            firewall-cmd --permanent --add-rich-rule="rule family=\"ipv4\" source address=\"$ip\" port protocol=\"tcp\" port=\"$MONGODB_PORT\" accept" >/dev/null 2>&1 || true
        done
        firewall-cmd --reload >/dev/null 2>&1 || true
    fi
fi

# Start MongoDB
echo "[6/9] Starting MongoDB service..."
systemctl enable mongod
systemctl start mongod
sleep 5

# Wait for MongoDB to be ready
echo "[7/9] Waiting for MongoDB to be ready..."
for i in {1..30}; do
    if mongosh --quiet --eval "db.adminCommand('ismaster')" >/dev/null 2>&1; then
        break
    fi
    sleep 2
    if [[ $i -eq 30 ]]; then
        log_error "MongoDB failed to start properly"
        exit 1
    fi
done

# Initialize replica set (only on first node)
echo "[8/9] Initializing replica set..."
CURRENT_IP=$(hostname -I | awk '{print $1}')
if [[ "$CURRENT_IP" == "${NODE_IPS[0]}" ]]; then
    log_info "Initializing replica set on primary node"
    
    # Create replica set config
    cat > /tmp/rs_init.js << EOF
rs.initiate({
  _id: "${REPLICA_SET_NAME}",
  members: [
EOF
    
    for i in "${!NODE_IPS[@]}"; do
        priority=$((i == 0 ? 2 : 1))
        cat >> /tmp/rs_init.js << EOF
    { _id: $i, host: "${NODE_IPS[$i]}:${MONGODB_PORT}", priority: $priority }$([ $i -lt $((${#NODE_IPS[@]} - 1)) ] && echo "," || echo "")
EOF
    done
    
    cat >> /tmp/rs_init.js << EOF
  ]
})
EOF
    
    mongosh --quiet < /tmp/rs_init.js
    
    # Wait for replica set to initialize
    sleep 10
    
    # Create admin user
    log_info "Creating admin user"
    mongosh --quiet --eval "
    use admin;
    db.createUser({
      user: 'admin',
      pwd: '$ADMIN_PASSWORD',
      roles: [ 'root' ]
    })
    "
    
    rm -f /tmp/rs_init.js
    
    log_info "Replica set initialized successfully"
else
    log_info "This is a secondary node - replica set will be initialized from primary"
fi

# Verification
echo "[9/9] Verifying installation..."
if systemctl is-active --quiet mongod; then
    log_info "✓ MongoDB service is running"
else
    log_error "✗ MongoDB service is not running"
    exit 1
fi

if [[ -f "$KEYFILE_PATH" ]]; then
    log_info "✓ Keyfile exists and has correct permissions"
else
    log_error "✗ Keyfile missing"
    exit 1
fi

log_info "MongoDB 8.0 replica set installation completed!"
echo
log_warn "Next steps:"
echo "1. Copy the keyfile from $KEYFILE_PATH to all other nodes"
echo "2. Run this script on all other nodes with the same parameters"
echo "3. Wait for all nodes to join the replica set"
echo "4. Test failover using: mongosh --host ${NODE_IPS[0]} --port $MONGODB_PORT -u admin -p"