manage-swarm

#! /usr/bin/env bash

# Check for basic commands which are required before this script
# can be useful.
for checkCommand in docker-machine docker jq curl; do
	if [ ! -x "$(command -v "${checkCommand}")" ]; then
		echo "\"${checkCommand}\" is required, but missing" >&2

		exit 1
	fi
done

# Providers
## DigitalOcean
### Prompt for a parameter which can only be answered by querying the DigitalOcean API
function create_digitalocean_ask() {
	local what
	local apiKey
	local values value

	what="$1"
	apiKey="$2"

	echo 'Loading...' >&3

	case "${what}" in
		size)
			values=($(curl -sSL -X GET -H "Content-Type: application/json" -H "Authorization: Bearer ${apiKey}" "https://api.digitalocean.com/v2/sizes"  | jq -crM '.sizes[] | .slug' | sort))
			;;
		region)
			values=($(curl -sSL -X GET -H "Content-Type: application/json" -H "Authorization: Bearer ${apiKey}" "https://api.digitalocean.com/v2/regions"  | jq -crM '.regions[] | .slug' | sort))
			;;
		*)
			echo "internal error: unable to ask about ${what}" >&2
			return 1
			;;
	esac

	echo "Please select one of the following for the ${what}:" >&3
	select value in "${values[@]}"; do
		if [ -n "${value}" ]; then
			echo "${value}"
			return 0
		fi
	done

	return 1
}

### Create a new remote Docker Swarm by creating new nodes on DigitalOcean
function create_digitalocean() {
	local name token count
	local idx startIdx
	local pids
	local node1 ip
	local arg nextArg nextArgName docker_machine_args
	local machineName

	name="$1"
	token="$2"
	count='3'
	shift; shift;

	if [ -z "${name}" -o -z "${token}" ]; then
		echo "Usage: manage-swarm create digitalocean <swarmName> <apiToken> [--count <n>]" >&2
		echo "                    [--size <n>] [--region <name>]" >&2

		return 1
	fi

	docker_machine_args=()

	startIdx='1'
	nextArg=''
	docker_machine_args=(
		--digitalocean-tags "docker-swarm,${name}"
	)
	for arg in "$@"; do
		if [ -n "${nextArg}" ]; then
			if [ "${nextArg}" = '+' ]; then
				if [ "${arg}" = 'ask' ]; then
					exec 3>&1
					arg="$(create_digitalocean_ask "${nextArgName}" "${token}")" || return 1
					exec 3>&-
				fi
				docker_machine_args+=("${arg}")
			else
				eval "${nextArg}="'"${arg}"'
			fi
			nextArg=''
			continue
		fi
		case "${arg}" in
			--size|--region)
				docker_machine_args+=("--digitalocean-${arg:2}")
				nextArg='+'
				nextArgName="${arg:2}"
				;;
			--count)
				nextArg='count'
				;;
			--startIndex)
				nextArg='startIdx'
				;;
		esac
	done

	pids=()
	for ((idx = 1; idx <= count; idx++)) {
		machineName="${name}-$[${startIdx} + ${idx} - 1]"
		docker-machine create --driver digitalocean --digitalocean-image=ubuntu-18-04-x64 \
                   --digitalocean-monitoring=true	\
 --digitalocean-access-token "${token}" "${docker_machine_args[@]}" "${machineName}" &
		pids+=("$!")
	}
	wait "${pids[@]}"

	node1="${name}-1"
	ip="$(docker-machine inspect "${node1}" | jq -crM .Driver.IPAddress)" || return 1
	if [ "${startIdx}" = '1' ]; then
		docker-machine ssh "${node1}" docker swarm init --advertise-addr "${ip}" || return 1
	fi

	joinCommand="$(docker-machine ssh "${node1}" docker swarm join-token worker | grep ' --token ')"
	if [ -z "${joinCommand}" ]; then
		echo "Failed to create join command" >&2
		return 1
	fi
	for ((idx = 1; idx <= count; idx++)) {
		machineName="${name}-$[${startIdx} + ${idx} - 1]"
		if [ "${machineName}" = "${node1}" ]; then
			continue
		fi
		docker-machine ssh "${name}-$[${startIdx} + ${idx} - 1]" "${joinCommand}" || return 1
	}

	if [ "${startIdx}" = '1' ]; then
		mkdir -p "${swarmEnvDir}"
		docker-machine env "${name}-1" > "${swarmEnvDir}/${name}"
	fi

	echo "Complete !"

	return 0
}

### Expand a Docker Swarm by spawning additional droplets on DigitalOcean and adding them
### to the Docker Swarm.
function expand_digitalocean() {
	local name info startIdx count
	local token size region

	name="$1"
	info="$2"
	startIdx="$3"
	count="$4"
	shift; shift; shift; shift

	token="$(echo "${info}" | jq -crM '.Driver.AccessToken')"
	region="$(echo "${info}" | jq -crM '.Driver.Region')"
	size="$(echo "${info}" | jq -crM '.Driver.Size')"

	create_digitalocean "${name}" "${token}" --count "${count}" --size "${size}" --region "${region}" --startIndex "${startIdx}" "$@" || return $?

	return 0
}

# CLI function to create a new remote Docker Swarm by calling the "create" command
# for the provider.
function create() {
	local provider
	local swarmFile

	provider="$1"
	name="$2"
	shift

	if [ -z "${provider}" ]; then
		echo "Usage: manage-swarm create <provider> <swarmName> ..." >&2

		return 1
	fi

	swarmFile="${swarmEnvDir}/${name}"

	if [ -n "${name}" -a -f "${swarmFile}" ]; then
		echo "Swarm already exists: ${swarmFile}" >&2

		return 1
	fi

	create_${provider} "$@" || return "$?"

	exportConfigToSwarm "${name}"

	return 0
}

# CLI function to expand an existing remote Docker Swarm by calling the "expand" command
# for the provider.
function expand() {
	local name count
	local provider
	local info
	local highestID nextID

	name="$1"
	count="$2"
	shift; shift

	if [ -z "${count}" ]; then
		count='1'
	fi

	if [ -z "${name}" ]; then
		echo "Usage: manage-swarm expand <swarmName> [<count> [<args>...]]" >&2

		return 1
	fi

	swarmFile="${swarmEnvDir}/${name}"
	if [ ! -f "${swarmFile}" ]; then
		echo "Swarm does not exist: ${swarmFile}" >&2

		return 1
	fi

	info="$(docker-machine inspect "${name}-1")"
	highestID="$(docker-machine ls --filter NAME="${name}-[0-9]" -q | sed 's@^.*-@@' | sort -n | tail -n 1)"
	nextID=$[${highestID} + 1]

	provider="$(echo "${info}" | jq -crM '.DriverName')"
	expand_${provider} "${name}" "${info}" "${nextID}" "${count}" "$@" || return $?

	exportConfigToSwarm "${name}"

	return 0
}

# CLI function to destroy a remote Docker Swarm, this just calls "docker-machine"
# so it is agnostic of all providers, as long as they are supported by
# "docker-machine"
function destroy() {
	local name
	local swarmFile

	name="$1"

	if [ -z "${name}" ]; then
		echo "Usage: manage-swarm destroy <swarmName>" >&2

		return 1
	fi

	swarmFile="${swarmEnvDir}/${name}"
	if [ ! -f "${swarmFile}" ]; then
		echo "Swarm does not exist: ${swarmFile}" >&2

		return 1
	fi

	(
		. "${swarmFile}"
		toDelete=( $(
			docker node ls -q | while IFS='' read -r id; do
				machine="$(docker node inspect "${id}" | jq -crM '.[0].Description.Hostname')"
				echo "${machine}"
			done
		) )

		retval='0'
		for machine in "${toDelete[@]}"; do
			docker-machine rm --force "${machine}" || retval='1'
		done

		exit "${retval}"
	) || return 1

	rm -f "${swarmFile}"

	return 0
}

# CLI function to print the status of a given Docker Swarm, it will
# list the stacks running on that swarm.
function status() {
	local name
	local swarmFile

	name="$1"

	if [ -z "${name}" ]; then
		echo "Usage: manage-swarm status <swarmName>" >&2

		return 1
	fi

	swarmFile="${swarmEnvDir}/${name}"
	if [ ! -f "${swarmFile}" ]; then
		echo "Swarm does not exist: ${swarmFile}" >&2

		return 1
	fi

	(
		. "${swarmFile}"
		docker node ls
		docker stack ls
	) || return 1

	return 0
}

# CLI function to print a list of known remote swarms
function list() {
	local file

	cd "${swarmEnvDir}" 2>/dev/null || return 0
	for file in *; do
		if [ ! -s "${file}" ]; then
			continue
		fi

		echo "${file}"
	done

	return 0
}

# Internal function to create an importable configuration from "docker-machine" for
# a remote swarm and store it on every node within the swarm.  This is needed so that
# other users of the swarm can access each machine in the swarm when creating or
# managing stacks on the swarm.  It needs to be done after every create or expand
# of the swarm.  After either of those operations, all users will need to (re-)import
# the swarm.
function exportConfigToSwarm() {
	local name
	local swarmFile
	local tmpdir
	local retval

	name="$1"

	if [ -z "${name}" ]; then
		echo "Usage: exportConfigToSwarm <swarmName>" >&2

		return 1
	fi

	swarmFile="${swarmEnvDir}/${name}"
	if [ ! -f "${swarmFile}" ]; then
		echo "Swarm does not exist: ${swarmFile}" >&2

		return 1
	fi

	tmpdir="$(mktemp -d)" || return 1
	tarfile="$(mktemp)" || return 1

	(
		. "${swarmFile}"

		dockerMachineBaseDir="${HOME}/.docker/machine/machines"
		docker node ls --format '{{.Hostname}}' | while IFS='' read -r hostname; do
			dockerMachineDir="${dockerMachineBaseDir}/${hostname}"
			configDir="${tmpdir}/${hostname}"

			cp -rp "${dockerMachineDir}" "${configDir}"
			cp "${dockerMachineBaseDir}/../certs/key.pem" "${configDir}/key.pem"
			cp "${dockerMachineBaseDir}/../certs/cert.pem" "${configDir}/cert.pem"
			cp "${dockerMachineBaseDir}/../certs/ca.pem" "${configDir}/ca.pem"
			cp "${dockerMachineBaseDir}/../certs/ca-key.pem" "${configDir}/ca-key.pem"

			cat "${configDir}/config.json" | sed '
				s|"/.*/\.docker/machine/machines/'"${hostname}"'\(/*\)|"@DOCKER_MACHINE_DIR@\1|g
				s|"/.*/\.docker/machine/certs/key\.pem"|"@DOCKER_MACHINE_DIR@/key.pem"|g
				s|"/.*/\.docker/machine/certs/cert\.pem"|"@DOCKER_MACHINE_DIR@/cert.pem"|g
				s|"/.*/\.docker/machine/certs/ca\.pem"|"@DOCKER_MACHINE_DIR@/ca.pem"|g
				s|"/.*/\.docker/machine/certs/ca-key\.pem"|"@DOCKER_MACHINE_DIR@/ca-key.pem"|g
				s|"/.*/\.docker/machine/certs"|"@DOCKER_MACHINE_DIR@"|g
				s|"/.*/\.docker/machine"|"@DOCKER_MACHINE_DIR@/../.."|g
			' > "${configDir}/config.json.in"
			rm -f "${configDir}/config.json"
		done

		(
			set -e

			cd "${tmpdir}"
			tar -cf - * | gzip -9c > "${tarfile}"
			gzip -dc "${tarfile}" | tar -tf - >/dev/null
		) || exit 1

		docker node ls --format '{{.Hostname}}' | while IFS='' read -r hostname; do
			docker-machine scp "${tarfile}" "${hostname}:/etc/wavelet-swarm.tar.gz" </dev/null
		done
	)
	retval="$?"

	rm -rf "${tmpdir}"
	rm -f "${tarfile}"

	return "${retval}"
}

# CLI function to import the "docker-machine" configuration for a remote swarm.
# It operates by SSH'ing into the IP specified and pulling the stored configuration
# by exportConfigToSwarm.
function import() {
	local ip
	local swarmFile swarmName
	local tmpdir
	local retval

	ip="$1"

	if [ -z "${ip}" ]; then
		echo "Usage: manage-swarm import <ip>" >&2

		return 1
	fi

	tmpdir="$(mktemp -d)" || return 1

	(
		set -e

		dockerMachineBaseDir="${HOME}/.docker/machine/machines"
		mkdir -p "${dockerMachineBaseDir}"

		cd "${tmpdir}"

		scp "root@${ip}:/etc/wavelet-swarm.tar.gz" wavelet-swarm.tar.gz

		gzip -dc wavelet-swarm.tar.gz | tar -xf - >/dev/null
		rm -f wavelet-swarm.tar.gz

		for hostname in *; do
			dockerMachineDir="${dockerMachineBaseDir}/${hostname}"
			if [ -d "${dockerMachineDir}" ]; then
				continue
			fi

			cp -rp "${hostname}" "${dockerMachineDir}/"
			sed "s|@DOCKER_MACHINE_DIR@|${dockerMachineDir}|g" "${dockerMachineDir}/config.json.in" > "${dockerMachineDir}/config.json"
			rm -f "${dockerMachineDir}/config.json.in"
		done
	)
	retval="$?"

	if [ "${retval}" = '0' ]; then
		swarmName="$(cd "${tmpdir}" && echo *-1 | sed 's@-1$@@')"
		swarmFile="${swarmEnvDir}/${swarmName}"

		mkdir -p "${swarmEnvDir}"
		docker-machine env "${swarmName}-1" > "${swarmFile}"
	fi

	rm -rf "${tmpdir}"

	return "${retval}"
}

# CLI function to print usage information
function usage() {
	echo "Usage: manage-swarm create <provider> <swarmName> ..."
	echo "       manage-swarm import <ip>"
	echo "       manage-swarm expand <swarmName> [<count>]"
	echo "       manage-swarm destroy <swarmName>"
	echo "       manage-swarm status <swarmName>"
	echo "       manage-swarm list"
	echo "       manage-swarm help"
}

# To make operating easier, change directories to the directory this script is in
cd "$(dirname "${BASH_SOURCE[0]}")" || exit 1

action="$1"
shift

configDir="$(dirname "${BASH_SOURCE[0]}")/config"
globalEnvFile="${configDir}/default"
if [ -f "${globalEnvFile}" ]; then
	. "${globalEnvFile}"
fi

swarmEnvDir="${configDir}/swarms"

# Map the CLI commands to their functions
case "${action}" in
	create|destroy|expand|status|list|import|exportConfigToSwarm)
		"${action}" "$@"
		exit "$?"
		;;
	ls)
		list "$@"
		exit "$?"
		;;
	help)
		usage
		exit 0
		;;
	*)
		usage >&2
		exit 1
		;;
esac