#!/usr/bin/env bash
# Installs MongoDB Community as a replica set with authentication: a new set, or one more server of a set.
# On Ubuntu from the MongoDB apt repo, or with DOCKER=1 in a container (any Linux with Docker).
# https://bsonjet.com/blog/install-mongodb-ubuntu-replica-set/
#
#   curl -fsSL https://bsonjet.com/blog/install-mongodb-ubuntu-replica-set/install-mongodb.sh | sudo bash
#
# Options (environment variables, e.g. "| sudo MONGO_HOST=db1.example.com bash"):
#   MONGO_HOST      the name clients and the other servers use to reach this server (default: this hostname).
#                   It must resolve to this machine; mongod listens on 127.0.0.1 and on the addresses it resolves to
#                   (never a public one, unless ALLOW_PUBLIC=1 - then only behind a firewall).
#   MONGO_VERSION   e.g. 8.0 (default: the newest)
#   RS_NAME         rs0 (default)
#   ADMIN_USER      admin (default); ADMIN_PASSWORD is generated when not given
#   DOCKER=1        run MongoDB in a container (Docker must be installed), managed in /opt/mongodb
#
# One more server of an existing set (the version and the set name come from the set):
#   JOIN            the connection string of the set, with a root user (MONGO_URI in /root/mongodb-credentials.txt)
#   KEYFILE         the keyfile of the set: its content (on the first server: sudo cat /etc/mongodb/keyfile),
#                   or the path to a copy of the file
#   ARBITER=1       an arbiter: votes in elections, keeps no data
set -euo pipefail
export DEBIAN_FRONTEND=noninteractive
# A fresh server often runs unattended-upgrades first: wait for its dpkg lock instead of failing.
apt-get() { command apt-get -o DPkg::Lock::Timeout=600 "$@"; }

MONGO_VERSION="${MONGO_VERSION:-}"
MONGO_HOST="${MONGO_HOST:-$(hostname)}"
RS_NAME="${RS_NAME:-rs0}"
ADMIN_USER="${ADMIN_USER:-admin}"
ADMIN_PASSWORD="${ADMIN_PASSWORD:-$(openssl rand -base64 24 | tr -d '/+=' | cut -c1-24)}"
JOIN="${JOIN:-}"
KEYFILE="${KEYFILE:-}"
[ -n "$KEYFILE" ] && [ -f "$KEYFILE" ] && KEYFILE=$(cat "$KEYFILE")     # a path to a copy of the file works too
ARBITER="${ARBITER:-}"
DOCKER="${DOCKER:-}"
PORT=27017
CREDENTIALS=/root/mongodb-credentials.txt
REPO="https://repo.mongodb.org/apt/ubuntu"
KEYRING=/usr/share/keyrings/mongodb-server.gpg
DOCKER_DIR=/opt/mongodb

say()  { printf '\033[1;34m==>\033[0m %s\n' "$*"; }
fail() { printf '\033[1;31mError:\033[0m %s\n' "$*" >&2; exit 1; }

# mongosh of the installed server: the package's, or the one in the container (it shares the server's network).
msh() {
    # </dev/null: run as "curl ... | bash", stdin is the rest of this script, and mongosh must not eat it.
    if [ "$DOCKER" = 1 ]; then docker exec -e ADMIN_USER -e ADMIN_PASSWORD -e NEW_HOST mongodb mongosh "$@" </dev/null
    else mongosh "$@" </dev/null; fi
}

# --- Checks: everything that would otherwise fail halfway --------------------------------------------------------
[ "$(id -u)" -eq 0 ] || fail "run it as root: curl ... | sudo bash"
. /etc/os-release
[ -z "$JOIN" ] || [ -n "$KEYFILE" ] || fail "JOIN needs the keyfile of the set: KEYFILE='...' (on the first server: sudo cat /etc/mongodb/keyfile)."
if [ "$DOCKER" = 1 ]; then
    docker compose version >/dev/null 2>&1 || fail "DOCKER=1 needs Docker with the compose plugin: https://docs.docker.com/engine/install/"
    [ -f "$DOCKER_DIR/compose.yaml" ] && fail "MongoDB already runs from $DOCKER_DIR. Use update-mongodb.sh."
else
    [ "${ID:-}" = ubuntu ] || fail "this script is for Ubuntu (found: ${PRETTY_NAME:-unknown}). With Docker, add DOCKER=1."
    command -v mongod >/dev/null && fail "MongoDB is already installed ($(mongod --version | head -1)). Use update-mongodb.sh."
fi
case "$(uname -m)" in
    x86_64) grep -qw avx /proc/cpuinfo || fail "this CPU has no AVX. MongoDB 5.0 and newer need it (old or emulated CPUs, some cheap VPS).";;
    aarch64) ;;
    *) fail "unsupported architecture: $(uname -m)";;
esac
if command -v ss >/dev/null && ss -ltn "sport = :$PORT" | grep -q LISTEN; then fail "something already listens on port $PORT."; fi

if ! command -v curl >/dev/null || ! command -v openssl >/dev/null || { [ "$DOCKER" != 1 ] && ! command -v gpg >/dev/null; }; then
    apt-get update -qq
    apt-get install -y -qq curl gnupg openssl >/dev/null
fi

# The name in the replica set config is what every client gets back - it must resolve, here and for the clients.
mapfile -t HOST_IPS < <(getent ahostsv4 "$MONGO_HOST" | awk '{print $1}' | sort -u)
[ "${#HOST_IPS[@]}" -gt 0 ] || fail "'$MONGO_HOST' does not resolve. Set MONGO_HOST to a name that points to this server."
# mongod listens on these addresses: never on a public one by accident (a VPS hostname often resolves to its public IP).
for ip in "${HOST_IPS[@]}"; do
    case "$ip" in
        127.*|10.*|192.168.*|172.1[6-9].*|172.2[0-9].*|172.3[01].*|100.6[4-9].*|100.[7-9][0-9].*|100.1[01][0-9].*|100.12[0-7].*) ;;
        *) [ "${ALLOW_PUBLIC:-}" = 1 ] || fail "'$MONGO_HOST' resolves to the public address $ip, so MongoDB would listen on the internet.
       Add '127.0.1.1 $MONGO_HOST' to /etc/hosts (only this server connects), or set MONGO_HOST to a private name,
       or run with ALLOW_PUBLIC=1 behind a firewall that allows only your servers.";;
    esac
done
BIND_IP=$(printf '%s\n' 127.0.0.1 "${HOST_IPS[@]}" | sort -u | paste -sd, -)

# The newest MongoDB, from the list MongoDB publishes, so the script needs no update for a new release.
# (The apt repo has a folder per version and no "latest".)
[ -n "$MONGO_VERSION" ] || MONGO_VERSION=$(curl -fsSL https://downloads.mongodb.org/current.json 2>/dev/null \
    | grep -oE '"version": *"[0-9]+\.[0-9]+\.[0-9]+"' | grep -oE '[0-9]+\.[0-9]+' | sort -uV | tail -1)
[ -n "$MONGO_VERSION" ] || fail "cannot read the newest MongoDB version from downloads.mongodb.org. Set MONGO_VERSION."

has_server() {
    curl -fsSL "$REPO/dists/$VERSION_CODENAME/mongodb-org/$1/multiverse/binary-$(dpkg --print-architecture)/Packages.gz" 2>/dev/null \
        | gunzip | grep '^Package: mongodb-org-server$' >/dev/null
}
# Ubuntu newer than the MongoDB repo is the classic first failure: the repo folder may even exist, without the server.
check_repo() {
    has_server "$1" && return 0
    fail "MongoDB $1 has no server package for Ubuntu $VERSION_ID ($VERSION_CODENAME) yet. The newest Ubuntu is often not there for months; 24.04 LTS is. Or add DOCKER=1."
}
add_repo() {
    local major="${1%%.*}"
    # The key is named server-9.asc from 9.0 on, server-8.0.asc before that.
    curl -fsSL "https://pgp.mongodb.com/server-$major.asc" -o /tmp/mongodb.asc 2>/dev/null \
        || curl -fsSL "https://pgp.mongodb.com/server-$major.0.asc" -o /tmp/mongodb.asc
    gpg --batch --yes --dearmor -o "$KEYRING" /tmp/mongodb.asc
    echo "deb [ arch=$(dpkg --print-architecture) signed-by=$KEYRING ] $REPO $VERSION_CODENAME/mongodb-org/$1 multiverse" \
        > /etc/apt/sources.list.d/mongodb-org.list
    apt-get update -qq
}

# The unit of the package and the official image set GLIBC_TUNABLES=glibc.pthread.rseq=0, and with it mongod 8.0+
# refuses to start on Linux 6.19 and newer (SERVER-121912). Check the newest installed kernel, not only the running
# one: the next kernel update and reboot would stop the database otherwise.
NEWEST_KERNEL=$(ls /boot/vmlinuz-* 2>/dev/null | sed 's|.*/vmlinuz-||' | sort -V | tail -1)
[ -n "$NEWEST_KERNEL" ] || NEWEST_KERNEL=$(uname -r)
KERNEL_FIX=""
[ "$(printf '%s\n' 6.19 "$NEWEST_KERNEL" | sort -V | head -1)" = 6.19 ] && KERNEL_FIX=1

# --- Joining a set: its name and version come from the set itself -------------------------------------------------
if [ -n "$JOIN" ]; then
    say "Reading the replica set to join"
    if [ "$DOCKER" = 1 ]; then
        SET_INFO=$(docker run --rm --network host "mongo:$MONGO_VERSION" mongosh "$JOIN" --quiet \
            --eval 'print(rs.conf()._id + " " + db.version())' 2>/dev/null) || SET_INFO=""
    else
        check_repo "$MONGO_VERSION"
        add_repo "$MONGO_VERSION"
        apt-get install -y -qq mongodb-mongosh >/dev/null
        SET_INFO=$(mongosh "$JOIN" --quiet --eval 'print(rs.conf()._id + " " + db.version())' 2>/dev/null) || SET_INFO=""
    fi
    [ -n "$SET_INFO" ] || fail "cannot log in with JOIN. It needs a user with the root role, and this server must reach the set by its names."
    RS_NAME="${SET_INFO%% *}"
    MONGO_VERSION=$(echo "${SET_INFO##* }" | cut -d. -f1,2)
fi

# --- Install and start -------------------------------------------------------------------------------------------
# The keyfile: one line, so it is easy to copy to the next server (mongod ignores whitespace in a keyfile anyway).
write_keyfile() {
    install -d -m 700 "$(dirname "$1")"
    if [ -n "$KEYFILE" ]; then printf '%s\n' "$KEYFILE" > "$1"; else openssl rand -base64 756 | tr -d '\n' > "$1"; fi
    chown "$2" "$1"
    chmod 400 "$1"          # mongod refuses a keyfile that others can read
}

if [ "$DOCKER" = 1 ]; then
    say "Starting MongoDB $MONGO_VERSION in Docker ($DOCKER_DIR)"
    write_keyfile "$DOCKER_DIR/keyfile" 999:999          # 999 = the mongodb user of the official image
    echo "MONGO_VERSION=$MONGO_VERSION" > "$DOCKER_DIR/.env"
    # The server's own network: reachable under MONGO_HOST like without Docker, and the firewall applies as usual
    # (a published Docker port would go around ufw). The update script only changes MONGO_VERSION in .env.
    cat > "$DOCKER_DIR/compose.yaml" <<EOF
# Written by install-mongodb.sh. Update with update-mongodb.sh, which changes MONGO_VERSION in .env.
services:
  mongodb:
    image: mongo:\${MONGO_VERSION}
    container_name: mongodb
    restart: unless-stopped
    network_mode: host
    command: [mongod, --replSet, $RS_NAME, --keyFile, /etc/mongodb/keyfile, --bind_ip, "$BIND_IP", --port, "$PORT"]
    environment:
      # The image sets glibc.pthread.rseq=0, and with it mongod refuses to start on Linux 6.19+ (SERVER-121912).
      GLIBC_TUNABLES: ""
    volumes:
      - data:/data/db
      - ./keyfile:/etc/mongodb/keyfile:ro
volumes:
  data:
EOF
    say "Pulling mongo:$MONGO_VERSION (a few hundred MB)"
    docker compose -f "$DOCKER_DIR/compose.yaml" pull --quiet || fail "cannot pull the image mongo:$MONGO_VERSION, see the error above."
    docker compose -f "$DOCKER_DIR/compose.yaml" up -d --quiet-pull >/dev/null 2>&1 || fail "the container did not start: docker compose -f $DOCKER_DIR/compose.yaml up"
else
    check_repo "$MONGO_VERSION"          # when joining: the version of the set, maybe not the newest
    add_repo "$MONGO_VERSION"
    say "Installing MongoDB $MONGO_VERSION"
    apt-get install -y -qq mongodb-org >/dev/null
    if [ -n "$KERNEL_FIX" ]; then
        install -d /etc/systemd/system/mongod.service.d
        cat > /etc/systemd/system/mongod.service.d/kernel-6.19.conf <<'EOF'
[Service]
# Written by install-mongodb.sh: mongod refuses to start on Linux 6.19+ with this tunable (SERVER-121912)
UnsetEnvironment=GLIBC_TUNABLES
EOF
        systemctl daemon-reload
    fi
    write_keyfile /etc/mongodb/keyfile mongodb:mongodb
    chown mongodb:mongodb /etc/mongodb
    [ -f /etc/mongod.conf.orig ] || cp /etc/mongod.conf /etc/mongod.conf.orig
    cat > /etc/mongod.conf <<EOF
# Written by install-mongodb.sh (the original is /etc/mongod.conf.orig)
storage:
  dbPath: /var/lib/mongodb
systemLog:
  destination: file
  logAppend: true
  path: /var/log/mongodb/mongod.log
net:
  port: $PORT
  bindIp: $BIND_IP
processManagement:
  timeZoneInfo: /usr/share/zoneinfo
replication:
  replSetName: $RS_NAME
security:
  authorization: enabled
  keyFile: /etc/mongodb/keyfile
EOF
    systemctl enable --now mongod >/dev/null 2>&1
fi

say "Configuring $MONGO_HOST:$PORT for the replica set $RS_NAME"
for _ in $(seq 90); do msh --quiet --port $PORT --eval 1 >/dev/null 2>&1 && break; sleep 1; done
msh --quiet --port $PORT --eval 1 >/dev/null 2>&1 \
    || fail "mongod did not start, see: $( [ "$DOCKER" = 1 ] && echo "docker logs mongodb" || echo "journalctl -u mongod and /var/log/mongodb/mongod.log")"

save_credentials() {
    umask 077
    printf '# MongoDB %s, replica set %s - written by install-mongodb.sh on %s\nMONGO_URI=%q\n' \
        "$MONGO_VERSION" "$RS_NAME" "$(date -I)" "$1" > "$CREDENTIALS"
}

# --- One more server of a set: added from the primary ------------------------------------------------------------
if [ -n "$JOIN" ]; then
    export NEW_HOST="$MONGO_HOST:$PORT"
    j() { msh "$JOIN" --quiet --eval "$1"; }
    state() { j 'const m = rs.status().members.find(x => x.name === process.env.NEW_HOST); print(m ? m.stateStr : "missing")'; }
    if [ "$ARBITER" = 1 ]; then
        # Adding an arbiter changes the implicit default write concern, so MongoDB wants it set explicitly first.
        # w:1 keeps the set writable while the other data server is down.
        j 'if (db.adminCommand({ getDefaultRWConcern: 1 }).defaultWriteConcernSource === "implicit")
               db.adminCommand({ setDefaultRWConcern: 1, defaultWriteConcern: { w: 1 } })' >/dev/null
        j 'rs.addArb(process.env.NEW_HOST)' >/dev/null
        WANT=ARBITER
    else
        # Without a vote until it has the data: the set keeps its majority during the initial sync.
        j 'rs.add({ host: process.env.NEW_HOST, priority: 0, votes: 0 })' >/dev/null
        WANT=SECONDARY
    fi
    say "Added to $RS_NAME, waiting until it is $WANT"
    unreachable=0
    while :; do
        s=$(state)
        [ "$s" = "$WANT" ] && break
        case "$s" in
            *"not reachable"*) unreachable=$((unreachable + 1))
                [ $unreachable -lt 24 ] || fail "the primary cannot reach $MONGO_HOST:$PORT. Check that the name resolves there and the firewall allows it.";;
            *) unreachable=0;;
        esac
        echo "    $s"; sleep 5
    done
    if [ "$ARBITER" != 1 ]; then
        j 'const c = rs.conf(); const m = c.members.find(x => x.host === process.env.NEW_HOST);
           m.votes = 1; m.priority = 1; rs.reconfig(c)' >/dev/null
    fi
    save_credentials "$JOIN"
    say "Done. $MONGO_HOST is $WANT of $RS_NAME:"
    j 'rs.status().members.forEach(m => print("    " + m.name + "  " + m.stateStr))'
    exit 0
fi

# --- A new set: no user exists yet, so the localhost exception lets us initiate it and create the first user -----
msh --quiet --port $PORT --eval "rs.initiate({ _id: '$RS_NAME', members: [{ _id: 0, host: '$MONGO_HOST:$PORT' }] })" >/dev/null
for _ in $(seq 60); do
    [ "$(msh --quiet --port $PORT --eval 'db.hello().isWritablePrimary' 2>/dev/null)" = true ] && break; sleep 1
done
[ "$(msh --quiet --port $PORT --eval 'db.hello().isWritablePrimary')" = true ] || fail "the replica set did not elect a primary"

# The password goes through the environment, not the command line (ps would show it).
export ADMIN_USER ADMIN_PASSWORD
msh --quiet --port $PORT --eval \
    'db.getSiblingDB("admin").createUser({ user: process.env.ADMIN_USER, pwd: process.env.ADMIN_PASSWORD, roles: ["root"] })' >/dev/null

URI="mongodb://$ADMIN_USER:$ADMIN_PASSWORD@$MONGO_HOST:$PORT/?replicaSet=$RS_NAME&authSource=admin"
save_credentials "$URI"

# (mongod --version: db.version() needs a login now that the user exists)
if [ "$DOCKER" = 1 ]; then VERSION_RUNNING=$(docker exec mongodb mongod --version | awk 'NR==1{print $3}'); WHERE=" in Docker"
else VERSION_RUNNING=$(mongod --version | awk 'NR==1{print $3}'); WHERE=""; fi
say "Done. MongoDB ${VERSION_RUNNING#v} runs as the replica set $RS_NAME$WHERE."
echo
echo "  Connection string (also in $CREDENTIALS, readable by root only):"
echo "  $URI"
echo
echo "  mongod listens on $BIND_IP only. To reach it from another machine, make '$MONGO_HOST' resolve"
echo "  to an address of this server there, and never open port $PORT to the whole internet."