ssm/ssm

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

# Utility functions
is_function() [[ $(type -t $1 2>/dev/null) == 'function' ]]

var() {
	declare varname=$1; shift

	if ! is_function "$varname"; then
		eval "
			${varname}() {
				declare mode=set
				declare -n _var=\"${varname}\"

				if (( \$# )); then
					case \"\$1\" in
						('=') mode=set;;
						('+=') mode=append;;
						('_=') mode=prepend;;
						(':=') mode=default;;

						('==') mode=compare;;
						('=~') mode=regex;;

						('u') mode=includes;;

						(*) die 71 \"Syntax error in \$FUNCNAME!\";;
					esac
					shift
				else
					mode='bool'
				fi

				case \"\$mode\" in
					(set) _var=( \"\$@\" );;
					(append) _var+=( \"\$@\" );;
					(prepend) _var=( \"\$@\" \"\${var[@]}\" );;
					(default) [[ \"\$_var\" ]] || _var=( \"\$@\" );;

					(compare) [[ \"\$_var\" == \"\$*\" ]];;
					(regex) [[ \"\$_var\" =~ \$@ ]];;

					(bool)
						case \"\${_var,,}\" in
							(''|'false'|'0') return 1;;
							(*) return 0;;
						esac
					;;

					(includes)
						for i in \"\${_var[@]}\"; do
							[[ \"\$i\" == \"\$*\" ]] && return 0
						done

						return 1
					;;
				esac
			}; readonly -f \"${varname}\"

			${varname}++() {
				declare -n _var=\"${varname}\"
				(( ${varname}++ ))
			}

			${varname}--() {
				declare -n _var=\"${varname}\"
				(( ${varname}-- ))
			}
		"
	fi

	if (( $# )); then
		case "$1" in
			('='|'=='|'=~'|'+='|'_='|':=')
				"$varname" "$@"
			;;

			(*)
				for v in "$@"; do
					var "$v"
				done
			;;
		esac
	fi
}

## 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
	var job_pid

	# Cgroups
	if cgroups; then
		mkdir -p "$cgroup_home/$service_cgroup_name"
		echo "$BASHPID" > "$cgroup_home/$service_cgroup_name/cgroup.procs"
	fi

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

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

	svc::reload() {
		kill -n "$service_reload_signal" "$job_pid"
	}; trap 'svc::reload' HUP

	"$@" 1>"$service_logfile_out" 2>"$service_logfile_err" & job_pid = "$!"

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

	svc::cleanup
}; readonly -f svc

## Respawn
respawn() {
	declare job_pid
	var job_pid

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

		rm -f "$svc_pidfile" "$service_ready_flag"

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

	respawn::sigpass() {
		declare sig pid
		var sig = "$1"
		var 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
	var 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

## Simple timer
timer() {
	declare cnt timeout
	var cnt = 0
	var 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; var 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 line
	var key value line

	while read -r line; do
		line =~ '^\[(.+)\]' && var 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 "stop() { $value; }";;
					(execreload) eval "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

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

	if cgroups; then
		if service_cgroup_exclusive; then
			if ! service_cgroup_empty; then
				return 15
			fi
		fi
	fi

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

	mktmpfiles || return 13

	rm -f "$service_stopped_flag"

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

		if timer "$service_ready_timeout" ready; then
			printf '1' > "$service_ready_flag"
		else
			return 5
		fi
	elif service_oneshot; then
		"${service_command[@]}" 1>"$service_logfile_out" 2>"$service_logfile_err"; res=$?
		(( res )) && return "$res"
		printf '1' > "$service_enabled_flag"
	else
		# Put ourselves into the cgroup, so that even when we die, whatever we started stays in it
		if cgroups; then
			mkdir -p "$cgroup_home/$service_cgroup_name"
			echo "$BASHPID" > "$cgroup_home/$service_cgroup_name/cgroup.procs"
		fi

		exec "${service_command[@]}" 1>"$service_logfile_out" 2>"$service_logfile_err" &
	fi

	return 0
}

## Reload the service
## Usually just sends HUP
reload() {
	service_running || return 3

	kill -n "$service_reload_signal" "$service_pid"
}

## 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

		nullexec kill -n "$service_stop_signal" "$service_pid" || return 1

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

		# Cgroup stuff
		if cgroups; then
			if service_cgroup_wait; then
				for p in "${service_cgroup_procs[@]}"; do
					pid_wait "$p" &
					wait || return 5
				done
			fi
		fi

		return 0
	fi
}

info() {
	declare _status_label _status _type _info_items
		var _status_label = 'Running'
		var _status = 'no'
		var _type = 'daemon'
		var _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[*]}" \
		"Respawn"         "${service_respawn:-false}" \
		"Config path"     "$service_config" \
		"Output log"      "$service_logfile_out" \
		"Error log"       "$service_logfile_err"

	if _status == 'yes'; then
		_info_items += \
			"PIDfile"  "${service_pidfile:-none}" \
			"PID"      "${service_pid:-none}"
	fi

	# Show the cgroup
	if cgroups; then
		_info_items += "Cgroup" "$cgroup_home/$service_cgroup_name" \
		               "Cgroup procs" "${#service_cgroup_procs[@]}"
	fi

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

}

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

edit() { $EDITOR "$service_config"; }
logs() { printf '%s\n' "$service_logfile_out" "$service_logfile_err"; }

## 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
## Empty declarations
var service_pid \
	service_pidfile \
	service_type \
	service_depends_ready \
	service_command \
	service_config \
	service_path \
	service_name \
	service_args \
	service_logfile_out \
	service_logfile_err \
	service_ready_flag \
	service_enabled_flag \
	service_stopped_flag \
	service_cgroup_name \
	service_cgroup_procs \
	cgroup_home \
	failed_deps \
	svc_pidfile \
	cfg_path \
	cfg_file \
	cfg_dir \
	rundir \
	logdir \
	_self

## check for some environment stuff
var EDITOR := 'vim'
var PAGER := 'less'
var XDG_CONFIG_HOME := "$HOME/.config"
var XDG_RUNTIME_DIR := "/run/user/$UID"

## Let's set some defaults
# These are meaningful to reconfigure.
var service_respawn = 0 # Respawn the service if it exits
var service_workdir = '/'
var service_stop_timeout = 30
var service_ready_timeout = 15
var service_signals = 1 10 12
var service_reload_signal = 1
var service_stop_signal = 15
var service_cgroup_exclusive = 0 # Refuse to start the service if its cgroup is not empty
var service_cgroup_wait = 0 # Wait on all the members of the cgroup to exit when stopping the service.

# Global config
var systemd = 0 # Enable systemd-related functions.
var systemd_service_path = /etc/systemd/system /run/systemd/system /lib/systemd/system
var cgroups = 0 # Enable cgroup-related functions
var cgroups_check_pid = 1 # Enable checking if the main service PID is in the correct cgroup before doing anythin with the service
var usrdir = '/usr/share/ssm'

# These are not
var service_managed = 1
var service_oneshot = 0
var service_running = 0
var service_enabled = 0
var service_stopped = 0
var service_systemd = 0
var service_nologs = 0
var service_cgroup_empty = 1

# These depend on who we are
if (( $UID )); then
	rundir = "$XDG_RUNTIME_DIR/ssm"
	logdir = "$HOME/log/ssm"
	cgroup_home = "/sys/fs/cgroup/user/$UID/ssm"
	cfg_file = "$XDG_CONFIG_HOME/ssm/ssm.conf"
else
	rundir = '/run/ssm'
	logdir = '/var/log/ssm'
	cgroup_home = "/sys/fs/cgroup/ssm"
	cfg_file = '/etc/ssm/ssm.conf'
fi

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

# Warn the user of deprecated stuff.
for p in "/etc/ssm/init.d" "$XDG_CONFIG_HOME/ssm/init.d"; do
	if [[ -d "$p" ]]; then
		printf 'WARNING: `%s` was renamed to `%s`! Please move your scripts accordingly!\n' "$p" "${p%init.d}services" >&2
		service_path += "$p"
	fi
done

# Source the config
if [[ -f "$cfg_file" ]]; then
	source "$cfg_file" || die 37 "Failed to load config: $cfg_file"
fi

# Common config path
cfg_path = "$XDG_CONFIG_HOME/ssm" '/etc/ssm'

# Load custom config and 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

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

# 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

	service_config || {
		# Search for a systemd service too
		systemd && {
			for i in "${systemd_service_path[@]/%//$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_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
service_pidfile := "$svc_pidfile"
service_logfile_out := "$logdir/${service_name}.out.log"
service_logfile_err := "$logdir/${service_name}.err.log"
service_ready_flag := "$rundir/$service_name.ready"
service_enabled_flag := "$rundir/$service_name.enabled"
service_stopped_flag := "$rundir/$service_name.stopped"
service_cgroup_name := "$service_name"

# A shortcut for disabling logging
if service_nologs; then
	service_logfile_out = '/dev/null'
	service_logfile_err = '/dev/null'
fi

# Let's see if there's a PID
if [[ -f "$service_pidfile" ]]; then
	read -r service_pid < "$service_pidfile"

	# Let's see if it's running
	if nullexec kill -0 "$service_pid"; then
		service_running = 1
	fi
fi

# 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 cgroups, if enabled
if cgroups; then
	if [[ -d "$cgroup_home/$service_cgroup_name" ]]; then
		while read -r line; do
			service_cgroup_procs += "$line"
		done < "$cgroup_home/$service_cgroup_name/cgroup.procs"

		if (( ${#service_cgroup_procs[@]} )); then
			service_cgroup_empty = 0
		fi
	fi

	# If there's a service PID, check if it's in the service's cgroup
	if cgroups_check_pid; then
		if service_pid; then
			service_cgroup_procs u "$service_pid" || die 29 "Recorded service PID is not in the service's cgroup, bailing!"
		fi
	fi
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" \
			15 "Refusing to start $service_name: the service cgroup is not empty and \$service_cgroup_exclusive is set"
	;;

	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 )) && exit "$res"

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