ssm/ssm

#!/usr/bin/env bash
shopt -s nullglob

# Utility functions
## Make setting default values a bit less awkward
default() {
	declare -n _p=$1; shift

	[[ "$_p" ]] || {
		for v in "$@"; do
			_p+=( "$v" )
		done
	}
}; readonly -f default

## Die. Why not?
die() {
	declare code=${1:-0}

	[[ "$2" ]] && printf '%s\n' "$2"
	exit "$code"
}; readonly -f die

## Run the command and wait for it to die
svc() {
	declare job_pid

	svc::cleanup() {
		kill -n "$service_stop_signal" "$job_pid"

		pid_wait "$job_pid" && {
			rm -f "$svc_pidfile" "$service_ready_flag"
		}
	}; trap 'svc::cleanup' TERM

	"$@" & job_pid=$!

	printf '%s' "$job_pid" > "$svc_pidfile"
	wait "$job_pid"

	svc::cleanup
}; readonly -f svc

## Respawn
respawn() {
	declare jobs job_pid

	respawn::cleanup() {
		jobs=( $(jobs -p) )

		if [[ "$jobs" ]]; then
			kill -n 15 "${jobs[@]}"
			wait "${jobs[@]}"
		fi

		rm -f "$svc_pidfile" "$service_ready_flag"

		exit 0
	}; trap 'respawn::cleanup' TERM

	respawn::sigpass() {
		declare sig=$1 pid=$2
		kill -n "$sig" "$pid"
	}

	respawn::set_traps() {
		for s in "${service_signals[@]}"; do
			trap "respawn::sigpass $s \$job_pid" "$s"
		done
	}; respawn::set_traps

	while true; do
		exec "$@" & job_pid=$!

		while nullexec kill -n 0 "$job_pid"; do
			wait "$job_pid"
		done
	done
}; readonly -f respawn

## Run a command with its output discarded
nullexec() { "$@" &>/dev/null; }
readonly -f nullexec

## Wait for a pid to die
pid_wait() {
	declare cnt=0

	while nullexec kill -0 "$1"; do
		(( cnt >= (service_stop_timeout*10) )) && return 1
		sleep 0.1
		(( cnt++ ))
	done

	return 0
}; readonly -f pid_wait

## See if NAME is a function
is_function() {
	declare name=$1 name_type

	name_type=$( type -t "$name" )

	if [[ $name_type == 'function' ]]; then
		return 0
	fi

	return 1
}; readonly -f is_function

## Simple timer
timer() {
	declare cnt=0 timeout=$1
	shift

	while ! "$@"; do
		(( cnt >= (timeout*10) )) && return 1
		sleep 0.1
		(( cnt++ ))
	done

	return 0
}; readonly -f timer

## Is a service ready?
is_ready() [[ -f "$service_ready_flag" ]]
readonly -f is_ready

## Wait for this service to get ready
wait_ready() { timer "$service_ready_timeout" is_ready; }
readonly -f wait_ready

## Depend on other services to be started
depend() {
	declare s

	for s in "$@"; do
		if ! "$_self" "$s" qstatus; then
			nullexec "$_self" "$s" start || {
				failed_deps+=( "$s" )
				return 1
			}
		fi
	done
}; readonly -f depend

## Create tmpfiles
mktmpfiles() {
	declare f

	for f in "${service_tmpfiles[@]}"; do
		IFS=':' read -r f_path f_type f_args <<< "$f"

		if ! [[ -e $f_path ]]; then
			case "$f_type" in
				symlink) ln -s "$f_args" "$f_path";;
				file|dir) IFS=':' read -r f_perms f_owner f_group <<< "$f_args"
					if [[ $f_type == 'file' ]]; then
						> "$f_path"
						chmod "${f_perms:-644}" "$f_path"
					elif [[ "$f_type" == 'dir' ]]; then
						mkdir -p -m "${f_perms:-755}" "$f_path"
					fi

					if [[ "$f_owner" || "$f_group" ]]; then
						chown "${f_owner:-root}:${f_group:-root}" "$f_path"
					fi;;
			esac
		fi
	done
}; readonly -f mktmpfiles

## Depend on other services to be ready
depend_ready() {
	declare s

	depend "$@" || return 1

	for s in "$@"; do
		"$_self" "$s" wait_ready || {
			failed_deps+=( "$s" )
			return 1
		}
	done
}; readonly -f depend_ready

result() {
	declare rc=$1; shift
	declare -A msgs

	while (( $# )); do
		[[ "$2" ]] || return 1

		msgs["$1"]="$2"
		shift 2
	done

	[[ "${msgs[$rc]}" ]] || msgs["$rc"]="Failed!"

	printf '%s\n' "${msgs[$rc]}"
}; readonly -f result

read_systemd_service() {
	declare section key value

	while read -r line; do
		[[ $line =~ ^\[(.+)\] ]] && section=${BASH_REMATCH[1],,}
		[[ $line =~ ^([^=]+)=(.+) ]] && {
			key=${BASH_REMATCH[1],,}
			value=${BASH_REMATCH[2]}
		}

		case $section in
			(service)
				case $key in
					(pidfile) service_pidfile=$value;;
					(execstart) eval "service_command=( $value )";;
					(execstop) eval "service_command_stop=( $value )";;
					(execreload) eval "service_command_reload=( $value )";;
					(restart) [[ $value == 'always' ]] && service_respawn=1;;
				esac
			;;
		esac
	done < "$1"
}

# Overloadable functions
## Start the service, write down the svc pid
start() {
	(( service_running )) && return 3

	rm -f "$service_stopped_flag"

	[[ -f "${service_command[0]}" ]] || return 9

	depend "${service_depends[@]}" || return 7
	depend_ready "${service_depends_ready[@]}" || return 7

	mktmpfiles || return 13

	if (( service_managed )); then
		if (( service_respawn )); then
			svc respawn "${service_command[@]}" &>"$service_logfile" &
		else
			svc "${service_command[@]}" &>"$service_logfile" &
		fi

		if timer "$service_ready_timeout" ready; then
			printf '1' > "$service_ready_flag"
		else
			return 5
		fi
	elif (( service_oneshot )); then
		"${service_command[@]}" &>"$service_logfile"; res=$?
		(( $res )) && return "$res"
		printf '1' > "$service_enabled_flag"
	else
		exec "${service_command[@]}" &
	fi

	return 0
}

## Reload the service
## Usually just sends HUP
reload() {
	(( service_running )) || return 3
	"${service_command_reload[@]}"
}

## Stop the service
## Returns:
##   3: Service is not running.
stop() {
	if (( service_oneshot )); then
		(( service_enabled )) || return 3

		rm -f "$service_enabled_flag"

		return 0
	else
		(( service_running )) || return 3

		"${service_command_stop[@]}" || return 1

		pid_wait "$service_pid" || return 5
		> "$service_stopped_flag"

		return 0
	fi
}

info() {
	declare \
		_status_label='Running' \
		_status='no'
		_type='daemon'
		_info_items=()

	(( service_oneshot )) && {
		_status_label='Enabled'
		_type='oneshot'
	}

	status && _status='yes'

	_info_items=(
		"Name"            "$service_name"
		"Type"            "$_type"
		"$_status_label"  "$_status"
		"Exec"            "${service_command[*]} ${service_args[*]}"
		"Respawn"         "${service_respawn:-false}"
		"Config path"     "${service_config}"
	)

	[[ "$_status" == 'yes' ]] && {
		_info_items+=(
			"PIDfile"  "${service_pidfile:-none}"
			"PID"      "${service_pid:-none}"
		)
	}

	printf "%12s: %s\n" "${_info_items[@]}"
}

# Restart just calls the script twice by default
restart() {
	"$_self" "$service_name" stop
	"$_self" "$service_name" start
}

edit() { "${EDITOR:-vim}" "$service_config"; }
logs() { ${PAGER:-less} "$service_logfile"; }

## Status is a bit of a special case. It's talkative.
status() {
	(( service_running )) && return 0
	(( service_enabled )) && return 0
	(( service_stopped )) && return 7

	return 1
}

## For use in scripts
qstatus() { nullexec status; }

## By default there is no ready check
ready() { :; }

# Main code
main() {
	# Figure out our full path
	case "$0" in
		(/*) _self=$0;;
		(*) _self="$PWD/$0";;
	esac

	# Needs to be global
	declare -g service_pid

	# Let's set some defaults
	service_managed=1

	usrdir='/usr/share/ssm'

	if (( $UID )); then
		# XDG stuff
		default XDG_CONFIG_HOME "$HOME/.config"
		default XDG_RUNTIME_DIR "/run/user/$UID"

		service_path=( "$XDG_CONFIG_HOME/ssm/services" )
		cfg_path=( "$XDG_CONFIG_HOME/ssm" )

		# Warn the user of deprecated stuff.
		if [[ -d "$XDG_CONFIG_HOME/ssm/init.d" ]]; then
			printf 'WARNING: `%s` was renamed to `%s`! Please move your scripts accordingly!\n' \
				"$XDG_CONFIG_HOME/ssm/init.d" \
				"$XDG_CONFIG_HOME/ssm/services" >&2

			service_path+=( "$XDG_CONFIG_HOME/ssm/init.d" )
		fi

		rundir="$XDG_RUNTIME_DIR/ssm"
		logdir="$HOME/log/ssm"
	else
		rundir='/run/ssm'
		logdir='/var/log/ssm'
	fi

	# Warn the user of deprecated stuff.
	if [[ -d "/etc/ssm/init.d" ]]; then
		printf 'WARNING: `/etc/ssm/init.d` was renamed to `/etc/ssm/services`! Please move your scripts accordingly!\n' >&2
		service_path+=( "/etc/ssm/init.d" )
	fi

	# Common service path
	service_path+=( '/etc/ssm/services' "$rundir/services" "$usrdir/services" )

	# Common config path
	cfg_path+=( '/etc/ssm' )

	# Load custom functions, reversing the PATH order
	for (( idx=${#cfg_path[@]}-1; idx>=0; idx-- )); do
		cfg_dir="${cfg_path[idx]}"

		for f in "$cfg_dir/functions"/*; do
			source "$f" || die 9 "Failed to source functions from $f"
		done
	done
	
	# Now create the needed runtime stuff
	for d in "$rundir" "$logdir"; do
		mkdir -p "$d" || die 3 "Failed to create runtime dir: $d"
	done

	# If $1 is a full path, source it.
	# If not, search for it in the service path.
	if [[ $1 == /* ]]; then
		service_config=$1
	else
		for i in "${service_path[@]/%//$1}"; do
			[[ -f "$i" ]] && {
				service_config=$i
				break
			}
		done

		# Search for a systemd service too
		for i in {/etc/systemd,/run/systemd,/lib/systemd,/usr/lib/systemd}/system/$1.service; do
			[[ -f "$i" ]] && {
				service_name=$1
				service_systemd=1
				service_config=$i
				read_systemd_service "$i"
				break
			}
		done
	fi

	# Die if there is no service config file
	[[ "$service_config" ]] || die 19 "Service not found: $1"
	
	# We can handle other people's service configs, poorly
	if (( service_systemd )); then
		# I'm pretty sure we'll need this
		:
	else
		# Service name is the basename
		service_name="${1##*/}"

		# Get the service defaults
		for p in "${cfg_path[@]/%//$service_name}"; do
			[[ -f "$p" ]] && {
				source "$p" || die 5 "Failed to read service defaults: $p"
			}
		done

		# Get the service config
		source -- "$service_config" "${@:3}" || die 7 "Failed to read the service config: $service_config"
	fi
	
	# Legacy
	[[ "$service_args" ]] && service_command=( "${service_command[@]}" "${service_args[@]}" )
	[[ "$service_respawn" == 'true' ]] && service_respawn=1
	[[ "$service_type" == 'oneshot' ]] && service_oneshot=1

	(( service_oneshot )) && service_managed=0
	[[ "$service_pidfile" ]] && service_managed=0

	if ! (( service_managed )); then
		(( service_respawn )) && die 21 "Refusing to respawn a service that manages itself."
	fi

	# Semi-hardcoded stuff
	svc_pidfile="$rundir/$service_name.pid"

	# Service-level defaults
	default service_pidfile         "$svc_pidfile"
	default service_logfile         "$logdir/$service_name.log"
	default service_ready_flag      "$rundir/$service_name.ready"
	default service_enabled_flag    "$rundir/$service_name.enabled"
	default service_stopped_flag    "$rundir/$service_name.stopped"
	default service_workdir         '/'
	default service_stop_timeout    30
	default service_ready_timeout   15
	default service_stop_signal     15
	default service_reload_signal   1
	default service_signals         1 10 12

	# Let's see if there's a PID
	if [[ -f "$service_pidfile" ]]; then
		service_pid=$(<$service_pidfile)
	
		# Let's see if it's running
		if nullexec kill -0 "$service_pid"; then
			service_running=1
		fi
	fi

	# The stop and reload commands need the service pid
	default service_command_stop    kill -n "$service_stop_signal" "$service_pid"
	default service_command_reload  kill -n "$service_reload_signal" "$service_pid"

	# Maybe the service is enabled?
	if [[ -f "$service_enabled_flag" ]]; then
		# Yay, it is!
		service_enabled=1
	fi

	# Let's see if the service was deliberately stopped
	if [[ -f "$service_stopped_flag" ]]; then
		# Ooh, it was.
		service_stopped=1
	fi

	# Check if action is even defined
	is_function "$2" || die 17 "Function $2 is not defined for $service_name."

	# cd into the workdir, if defined.
	[[ "$service_workdir" ]] && {
		cd "$service_workdir" || die $?
	}

	# Run pre_$action function
	if is_function "pre_$2"; then
		"pre_$2" || {
			printf 'pre_%s failed!\n' "$2"
			die 13
		}
	fi

	# Run the function
	"$2"; res=$?

	case "$2" in
		stop)
			result "$res" \
				0 "Stopped $service_name" \
				3 "$service_name is not running" \
				5 "Operation timed out"
		;;

		start)
			result "$res" \
				0 "Started $service_name" \
				3 "$service_name is already running" \
				5 "Readyness check for $service_name timed out" \
				7 "Failed to start dependencies for $service_name: ${failed_deps[@]}" \
				9 "service_command does not exist: ${service_command[0]}" \
				13 "Failed to create temporary files for $service_name"
		;;

		reload)
			result "$res" \
				0 "Reloaded $service_name"
		;;

		status)
			if (( service_oneshot )); then
				result "$res" \
					0 "$service_name is enabled" \
					1 "$service_name is not enabled"
			else
				result "$res" \
					0 "$service_name is running" \
					1 "$service_name is not running" \
					7 "$service_name was stopped"
			fi
		;;
	esac

	(( res )) && return "$res"

	# Run post_$action function
	if is_function "post_$2"; then
		"post_$2" || {
			printf 'post_%s failed!\n' "$2"
			die 15
		}
	fi
}; readonly -f main

main "$@"