#!/usr/bin/env bash
#
# Artix Linux service manager (runit)
# Forked and further developed, by: nikolar <nikolar@artixlinux.org>
# Fork Date: December 24, 2022
#
# Forked and further developed, by: linuxer <linuxer@artixlinux.org>
# Fork Date: August 15, 2020 as rsv
# Renamed and further developed as rsm from: January 10, 2021
#
# Original Author (vsv): Dave Eddy <dave@daveeddy.com>
# Date: August 29, 2018
# License: MIT

RSM_VERSION='v1.4.9'

export RUNSVDIR=${RUNSVDIR:-/etc/runit/runsvdir}
export SVDIR=${SVDIR:-/run/runit/service/}
export LOGDIR="/var/log/"
export RSVDIR="/run/runit/service/"
export FLDIR=${FLDIR:-/etc/runit/sv/}


charup='✔'
chardown='X'
charunknown='?'
progname=${0##*/}
num_re='^-?[0-9]+$'
svc_re='^[a-zA-Z0-9_\.-]+$'

shopt -s nullglob

# Get time in a human format, like 1 hour ago, 7 minutes ago, etc.
human() {
	local seconds=$1
	if ((seconds < 0)); then
		((seconds *= -1))
	fi

	local times=(
		$((seconds / 60 / 60 / 24 / 365)) # years
		$((seconds / 60 / 60 / 24 / 30))  # months
		$((seconds / 60 / 60 / 24 / 7))   # weeks
		$((seconds / 60 / 60 / 24))       # days
		$((seconds / 60 / 60))            # hours
		$((seconds / 60))                 # minutes
		$((seconds))                      # seconds
	)
	local names=(year month week day hour minute second)

	local i
	for ((i = 0; i < ${#names[@]}; i++)); do
		if ((${times[$i]} > 1)); then
			echo "${times[$i]} ${names[$i]}s"
			return
		elif ((${times[$i]} == 1)); then
			echo "${times[$i]} ${names[$i]}"
			return
		fi
	done
	echo '0 seconds'
}

# enable or disable colors based on the argument given, i.e.:
# setcolors on   # colors on
# setcolors off  # colors off
# setcolors auto # colors on or off depending on environment
setcolors() {
	local opt=$1
	local colors

	if [[ $opt == auto ]]; then
		# no colors if stdout is not a TTY
		if [[ ! -t 1 ]]; then
			opt='off'
		else
			# stdout is a tty, check tput capability for colors
			colors=$(tput colors 2>/dev/null || echo -1)
			if ! [[ $colors =~ $num_re ]]; then
				fatal "failed to parse output of \`tput colors\` ($colors)"
			fi

			if ((colors >= 8)); then
				opt='on'
			else
				opt='off'
			fi
		fi
	fi

	case "$opt" in
		on|yes|true)
			colorred=$(tput setaf 1)
			colorgreen=$(tput setaf 2)
			coloryellow=$(tput setaf 3)
			colorblue=$(tput setaf 4)
			colormagenta=$(tput setaf 5)
			colorcyan=$(tput setaf 6)
			#colorwhite=$(tput setaf 7)
			colorgray=$(tput setaf 8)
			colorbold=$(tput bold)
			colorreset=$(tput sgr0)
			;;
		off|no|false)
			colorred=
			colorgreen=
			coloryellow=
			colormagenta=
			colorcyan=
			#colorwhite=
			colorgray=
			colorbold=
			colorreset=
			;;
		*)
			echo "unknown color option: '$opt'" >&2
			exit 1
			;;
	esac
}

usage() {
	local cr=$colorreset
	local logo

	logo=$(getlogo)

	cat <<EOF
$colormagenta ______   ____  __  __
$colormagenta | | \ \ / ___\ | \/ |  $colorgreen Runit Service Manager for Artix Linux ($RSM_VERSION)
$colormagenta | |_/ / \____  | || |  $colorgreen Source: https://gitea.artixlinux.org/nikolar/Runit-Service-Manager
$colormagenta | | \ \ \____/ | || |  $colorgreen MIT License

$logo   $colorblue Manage and view runit services
$logo   $colorblue Made specifically for Void Linux but should work anywhere
$logo   $colorblue Author: Dave Eddy <dave@daveeddy.com> (bahamas10)
$logo   $colorblue Forked specifically for Artix Linux and renamed to rsm, for further development
$logo   $colorblue Maitained by nikolar <nikolar@artixlinux.org> (nikolar in #artix)

${coloryellow}USAGE:${colorgreen}
$progname [OPTIONS] [SUBCOMMAND] [<ARGS>]
$progname [-u] [-d <dir>] [-h] [-t] [SUBCOMMAND] [...]

${coloryellow}OPTIONS:${colorgreen}
${colorgreen}-c <yes|no|auto>         $cr Enable/disable color output, defaults to auto
${colorgreen}-d <dir>                 $cr Directory to look into, defaults to env SVDIR or /run/runit/service is unset
${colorgreen}-h                       $cr Print this message and exit
${colorgreen}-l                       $cr Show log processes, this is a shortcut for 'status -l'
${colorgreen}-t                       $cr Tree view, this is a shortcut for 'status -t'
${colorgreen}-u                       $cr User mode, this is a shortcut for '-d ~/runit/service'
${colorgreen}-v                       $cr Increase verbosity
${colorgreen}-V                       $cr Print the version number and exit

${coloryellow}ENV:${colorgreen}
SVDIR                    $cr The directory to use, passed to the 'sv' command, can
                          be overridden with '-d <dir>'

${coloryellow}SUBCOMMANDS:${colorgreen}
status [-lt] [filter]    $cr Default subcommand, show process status
			 $cr '-t' enables tree mode (process tree)
			 $cr '-l' enables log mode (show log processes)
			 $cr 'filter' is an optional string to match service names against
${colorgreen}
enable <svc> [...]       $cr Enable the service(s) (remove the "down" file, does not start service)
${colorgreen}
disable <svc> [...]      $cr Disable the service(s) (create the "down" file, does not stop service)

Any other subcommand gets passed directly to the 'sv' command, see sv(1) for the
full list of subcommands and information about what each does specifically.
Common subcommands:

${colorgreen}start <service>          $cr Start the service
${colorgreen}stop <service>           $cr Stop the service
${colorgreen}restart <service>        $cr Restart the service
${colorgreen}reload <service>         $cr Reload the service (send SIGHUP)
${colorgreen}logs <service>           $cr Outputs the service's logfilenames and their access & error logs from /var/log/<serice>/
${colorgreen}alllogs <service>        $cr The same like logs <service>
${colorgreen}errorlogs <service>      $cr Outputs the service's logfilenames and their errorlogs from /var/log/<serice>/


${coloryellow}EXAMPLES:${colorgreen}
${colorgreen}$progname                      $cr Show service status in /var/service
${colorgreen}$progname status               $cr Same as above
${colorgreen}$progname -t                   $cr Show service status + pstree output
${colorgreen}$progname status -t            $cr Same as above
${colorgreen}$progname status tty           $cr Show service status for any service that matches tty*
${colorgreen}$progname check uuidd          $cr Check the uuidd svc, wrapper for 'sv check uuidd'
${colorgreen}$progname restart sshd         $cr Restart sshd, wrapper for 'sv restart sshd'
${colorgreen}$progname -u                   $cr Show service status in ~/runit/service
${colorgreen}$progname -u restart ssh-agent $cr Restart ssh-agent in ~/runit/service/ssh-agent
EOF
}

verbose() {
	if ((verbosity > 0)); then
		echo '>' "$colorgray" "$@" "$colorreset" >&2
	fi
}

# print the logo with brackets colorized
getlogo() {
	printf '%s[%s%s%s]%s' \
	    "$colorcyan" \
	    "$colormagenta" "$progname" \
	    "$colorcyan" \
	    "$colorreset"
}

# prints a message
msg() {
	local logo

	logo=$(getlogo)

	echo "$logo" "$colorblue$*$colorreset"
}

# prints a fatal message and exists
fatal() {
	echo "${colorred}FATAL:" "$@" "$colorreset"
	exit 2
}

# rmsg - same as msg but colorized based on return status passed via $1
rmsg() {
	local code=$1
	local logo
	local statuscolor

	shift

	logo=$(getlogo)

	if ((code == 0)); then
		statuscolor=$colorgreen
	else
		statuscolor=$colorred
	fi

	echo "$logo" "$statuscolor$*$colorreset"
}

disable_message() {
	local svc=$1

	echo "service $svc disabled by user $USER on $(date)"
}

# process the status of a single service
process_service() {
	local svc=$1
	local now=$2
	local pid when up msg char err enabled msgcolor enabledcolor \
	    statecolor code islog prog down

	islog=false
	if [[ $svc == */log ]]; then
		islog=true
	fi

	# check permissions effectively
	err=$(sv status "$svc" 2>&1)
	code=$?

	# get service state, or set error messages
	if ((code == 0)); then
		err=
		state=$(< "$svc/supervise/stat")
		pid=$(< "$svc/supervise/pid")
		when=$(stat --printf %Y "$svc/supervise/pid")
		up=$((now - when))
		msg=$(human "$up")

		# if the service is enabled (wants up or wants down)
		if down=$(cat "$svc/down" 2>&1); then
			verbose "service disabled: $down"
			enabled='false'
		else
			enabled='true'
		fi

		# color time if the service hasn't been up long
		if ((up < 5)); then
			msgcolor=$colorred
		elif ((up < 30)); then
			msgcolor=$coloryellow
		else
			msgcolor=$colorgray
		fi
	else
		pid=
		enabled=
		state='n/a'
		msg=${err##*: }
		msgcolor=$colorred
	fi

	# sanity check pid
	if [[ -n $pid ]] && ! [[ $pid =~ $num_re ]]; then
		fatal "invalid pid: '$pid'"
	fi

	# figure out character and color to use for status
	if [[ -n $err ]]; then
		char=$charunknown
		statecolor=$coloryellow
	elif [[ $state == 'run' ]]; then
		char=$charup
		statecolor=$colorgreen
	else
		char=$chardown
		statecolor=$colorred
	fi

	# figure out enabled color
	if [[ $enabled == 'true' ]]; then
		enabledcolor=$colorgreen
	elif [[ $enabled == 'false' ]]; then
		enabledcolor=$colorred
	else
		enabledcolor=$coloryellow
	fi

	# msg color
	if [[ -z $msgcolor ]]; then
		msgcolor=$colorgray
	fi

	# format service name if it is a log service
	if $log && $islog && ! $tree; then
		svc='- log'
	fi

	# figure out program name
	if [[ -n $pid ]]; then
		IFS= read -d $'\0' -r prog _ < /proc/$pid/cmdline
		prog=${prog##*/}
		prog=${prog:0:17}
	fi

	# print service line
	printf ' %s%s%s %s%-20s%s %s%-7s%s %s%-9s%s %s%-8s%s %s%-17s%s %s%s%s\n' \
	    "$statecolor"   "$char"           "$colorreset" \
	    "$colorreset"   "$svc"            "$colorreset" \
	    "$statecolor"   "$state"          "$colorreset" \
	    "$enabledcolor" "${enabled:----}" "$colorreset" \
	    "$colormagenta" "${pid:----}"     "$colorreset" \
	    "$colorgreen"   "${prog:----}"    "$colorreset" \
	    "$msgcolor"     "$msg"            "$colorreset"

	# optionally print the pstree
	if $tree; then
		if [[ -n $pid ]] && ((pid > 0)); then
			echo "$colorgray"
			pstree -ac "$pid"
		fi
		echo "$colorreset"
	fi
}

# handle the status subcommand
do_status() {
	local OPTIND option
	local d
	local logo
	local now
	local svcs=()

	logo=$(getlogo)
	while getopts 'hlt' option; do
		case "$option" in
			h) usage; exit 0;;
			l) log=true;;
			t) tree=true;;
			*) usage; exit 1;;
		esac
	done
	shift "$((OPTIND - 1))"

	local filter=$1

	# loop service directories
	for d in ./*/; do
		d=${d%/}
		svc=${d##*/}

		# this is us being a little over-protective here
		if ! [[ $svc =~ $svc_re ]]; then
			rmsg -1 "unexpected characters in name: '$svc'" >&2
			continue
		fi

		if [[ ! -d $svc/supervise ]]; then
			verbose "skipping $svc - '$svc/supervise' not found"
			continue
		fi

		if [[ -n $filter && $svc != *"$filter"* ]]; then
			if [[ $verbosity != 1 ]]; then
				verbose "filtering out '$svc' because it does not match '$filter'"
			fi
			continue
		fi

		svcs+=("$svc")
		$log && svcs+=("$svc/log")
	done


	# print title if verbose
	if ((verbosity > 0)); then
		# service count
		local services
		local count=${#svcs[@]}

		if ((count == 1)); then
			services='service'
		else
			services='services'
		fi

		printf '%s> %s %s-%s %s%s %s(%s%s%s)%s %s-%s %s%d %s%s\n' \
		    "$colorgray" "$logo" \
		    "$colorgray" "$colorreset" \
		    "$colorblue" "$HOSTNAME" \
		    "$colorcyan" "$colorgreen" "${PWD/#$HOME/\~}" \
		    "$colorcyan" "$colorreset" \
		    "$colorgray" "$colorreset" \
		    "$colorblue" "$count" "$services" "$colorreset"
	fi
	echo

	# table header
	printf '%s   %-20s %-7s %-9s %-8s %-17s %s%s\n' \
	    "$colorbold" \
	    'SERVICE' 'STATE' 'ENABLED' 'PID' 'COMMAND' 'TIME' \
	    "$colorreset"

	# loop services
	printf -v now '%(%s)T' -1
	for svc in "${svcs[@]}"; do
		process_service "$svc" "$now"
	done
	echo
}

# handle the enable and disable subcommands
do_enable_disable() {
	local cmd=$1
	local svc
	local ret=0
	local file
	shift

	if (($# < 1)); then
		rmsg -1 "Argument expected for '$cmd'"
		return 1
	fi

	# Loop all arguments as services
	msg "Running $progname $cmd $*"

	for svc in "$@"; do
		# Validate service name
		if ! [[ $svc =~ $svc_re ]]; then
			rmsg -1 "unexpected characters in name: '$svc'" >&2
			continue
		fi

		# ensure service exists
		if ! [[ -d ${FLDIR}${svc} ]]; then
			rmsg -1 "service directory '$svc' does not exist" >&2
			continue
		fi

		file=$svc/down

		case "$cmd" in
			enable)
				verbose "enable '$svc': rm -f '$file'"
				if ! rm -f "$file"; then
					rmsg 1 "failed to enable '$svc'"
					ret=1
				else
					if [ -L "${SVDIR}" ]
					then
						ln -s "${FLDIR}${svc}" "${SVDIR}"
					else
						ln -s "${FLDIR}${svc}" "${RUNSVDIR}/default" # default runlevel
					fi
					do_status
				fi
				;;
			disable)
				verbose "disabling '$svc': creating '$file'"
				if ! disable_message "$svc" > "$file"; then
					rmsg 1 "failed to disable '$svc'"
					ret=1
				else
					if [ -L "${SVDIR}" ]
					then
						unlink "${SVDIR}${svc}"
					else
						unlink "${RUNSVDIR}/default/${svc}" # default runlevel
					fi
					do_status
				fi
				;;
			*)
				rmsg 2 "unknown command: $cmd" >&2
				return 2
				;;
		esac
	done

	rmsg "$ret" "[$progname $cmd $*], exit code: $ret"
	return "$ret"
}

# show logs
do_show_logs() {
	local cmd=$1
	local svc
	local ret=0
	local file
	shift

	# Loop all arguments as services
	msg "Running $progname $cmd $*"

	for svc in "$@"; do
		# Validate service name
		if ! [[ $svc =~ $svc_re ]]; then
			rmsg -1 "unexpected characters in name: '$svc'" >&2
			continue
		fi
	done


	if [ -h "${RSVDIR}${cmd}" ]
	then
		if [ -d "${RSVDIR}${cmd}" ]
		then
			# TODO check for logs before globing non-existent directories
			printf "\n%20s\n" "${colorgreen}The following log files found:"
			local logs_files_array=($(ls /var/log/$cmd/{*.*,current}))
			printf "\n$colorblue"
			printf '%s\n' "${logs_files_array[@]}"

			printf "\n$colormagenta"
			read -n 1 -s -r -p "Press any key to continue"

			printf "\n"
			printf "\n${colorcyan}Log files full output:\n"
			printf "${coloryellow}\n"
			printf "${coloryellow}${opt}"

			tail -n 10  ${LOGDIR}${cmd}/{*.*,current} | less
			printf "\n"
		fi
	else
			printf "\n%20s\n" "${colorred}The service's ${cmd} log files have not been found or do not exist"
	fi

	rmsg "$ret" "[$progname $cmd $*], exit code: $ret"
	return "$ret"
}

# show error logs
do_show_err_logs() {
	local cmd=$1
	local svc
	local ret=0
	local file
	shift

	# Loop all arguments as services
	msg "Running $progname $cmd $*"

	for svc in "$@"; do
		# Validate service name
		if ! [[ $svc =~ $svc_re ]]; then
			rmsg -1 "unexpected characters in name: '$svc'" >&2
			continue
		fi
	done

	if [ -h "${RSVDIR}${cmd}" ]
	then
		if ! [ -d ${RSVDIR}${cmd} ]
		then
			printf "\n%20s\n" "${colorred}The following error log files found:"
			local logs_files_array=($(ls /var/log/$cmd/*error*.*))
			printf "\n$colorblue"
			printf '%s\n' "${logs_files_array[@]}"

			printf "\n$colormagenta"
			read -n 1 -s -r -p "Press any key to continue"

			printf "\n"
			printf "\n${colorcyan}Log files full output:\n"
			printf "${coloryellow}\n"
			printf "${coloryellow}${opt}"

			tail -n 10  ${LOGDIR}${cmd}/*error*.* | more
			printf "\n"
		fi
	else
			printf "\n%20s\n" "${colorred}The service's ${cmd} error log files have not been found or do not exist"
	fi

	rmsg "$ret" "[$progname $cmd $*], exit code: $ret"
	return "$ret"
}

# handle any other subcommand
do_sv_cmd() {
	if (($# < 2)); then
		rmsg -1 "Argument expected for 'sv $1'"
		return 1
	fi

	msg "Running sv command (SVDIR=$SVDIR sv $*):"
	sv "$@"
	local ret=$?
	rmsg "$ret" "[sv $*], exit code: $ret"
	return "$ret"
}

color_arg=
tree=false
log=false
do_usage=-1
verbosity=1
while getopts 'c:d:hltuvV' option; do
	case "$option" in
		c) color_arg="$OPTARG";;
		d) SVDIR=$OPTARG;;
		h) do_usage=0;;
		l) log=true; cmd='status';;
		t) tree=true; cmd='status';;
		u) SVDIR=~/runit/service;;
		v) ((verbosity++));;
		V) echo "$RSM_VERSION"; exit 0;;
		*) do_usage=1;;
	esac
done
shift "$((OPTIND - 1))"

if [ -n "$NO_COLOR" ] && [ -z "$color_arg" ]
then
	setcolors off
elif [ -n "$color_arg" ]
then
	setcolors "$color_arg"
else
	setcolors auto
fi


# we wait until the colors are optionally set to output the usage message
if ((do_usage > -1)); then
	usage
	exit "$do_usage"
fi

# move to the service directory
if [ -L "$SVDIR" ]
then
	cd "$SVDIR" || fatal "failed to enter dir: $SVDIR"
else
	cd "$RUNSVDIR/default" || fatal "failed to enter dir: $FLDIR"
fi

# figure out 'cmd' command
if [[ -z $cmd ]]; then
	cmd=${1:-status}
	shift
fi

case "$cmd" in
	status)
		do_status "$@"
		;;
	enable|disable)
		do_enable_disable "$cmd" "$@"
		;;
	logs)
		do_show_logs "$@"
		;;
	alllogs)
		do_show_logs "$@"
		;;
	errorlogs)
		do_show_err_logs "$@"
		;;

	*)
		if $tree; then
			rmsg -1 "-t can only be specified with 'status'"
			exit 1
		fi
		if $log; then
			rmsg -1 "-l can only be specified with 'status'"
			exit 1
		fi
		do_sv_cmd "$cmd" "$@"
		;;
esac
