#!/bin/ksh93

typeset -r VERSION='1.0' FPROG=${.sh.file} PROG=${FPROG##*/} SDIR=${FPROG%/*} \
	SLURMCONF=/etc/slurm/slurm.conf

function showUsage {
	[[ -n $1 ]] && X='-?' ||  X='--man'
	getopts -a ${PROG} "${ print ${USAGE} ; }" OPT $X
}

function asOctalMode {
	typeset DEFAULT="$1" MOD="$2"
	typeset -i8 O

	[[ -z ${MOD} ]] && print "${DEFAULT}" && return 0

	if [[ ${MOD} =~ ^[0-9]+$ && ${MOD} < 6 ]]; then
		if [[ ${MOD} =~ # ]]; then
			O=${MOD}
		elif [[ ${MOD:0:1} == '0' ]]; then
			O="8#${MOD}"
		else
			O=${MOD}
		fi
		S=${ print $O ; }
		print "0$S"
		return 0
	fi

	print -u2 "Invalid Mode '${MOD}' - exiting."
	return 1
}

function initDefaults {
	typeset SVC=${OPT[SVC]##*/} LOG=${OPT[LOG]} OWNER=${OPT[OWNER]} S T CFG
	integer N

	[[ -z ${SVC} && -n ${RUNTIME_DIRECTORY} && -e ${RUNTIME_DIRECTORY} ]] && \
		SVC=${RUNTIME_DIRECTORY##*/}
	[[ ${SVC} =~ ^slurm(ctl|db|rest)?d$ ]] && OPT[SVC]=${SVC} || \
		{ SVC= ; print -u2 "Unknown service '${SVC}' - exiting."; return 1; }

	if [[ -n ${OPT[CFG]} ]]; then
		if [[ ! -r ${OPT[CFG]} ]]; then
			print -u2 "Unable to read slurm config '${OPT[CFG]}' - exiting."
			return 1
		fi
		CFG=${OPT[CFG]}
	else
		CFG=${SLURMCONF%/*}/${SVC}.conf
		OPT[CFG]="${CFG}"
	fi

	if [[ -z ${LOG} ]]; then
		[[ ${SVC} == 'slurmdbd' ]] && S='LogFile=' || S="S${SVC:1}LogFile="
		N=${#S}
		while read L T ; do
			[[ ${L:0:N} == $S ]] && LOG=${L:N} && break
		done < ${CFG}
		[[ -z ${LOG} ]] && LOG=/var/log/${SVC}/main.log
		[[ ${LOG:0:1} == '/' ]] && OPT[LOG]="${LOG}" || \
			OPT[LOG]="/var/log/${LOG}"
	fi
	[[ -z ${OPT[LOG]} ]] && return 2

	[[ -z ${OWNER} ]] && OWNER="${LOGNAME:-root}"
	S=${ getent passwd ${OWNER} ; }
	T=( ${S//:/ } )
	[[ -n ${T[2]} ]] && OPT[OWNER]="$T" || \
		{ print -u2 "Invalid user '${OWNER}' - exiting."; return 3; } 

	if [[ -n ${OPT[GROUP]} ]]; then
		S=${ getent group ${OPT[GROUP]} ; }
		T=( ${S//:/ } )
		[[ -n ${T[2]} ]] && OPT[GROUP]="$T" || \
			{ print -u2 "Invalid group '${OT[GROUP]}' - exiting."; return 4; } 
	fi

	OPT[DMOD]=${ asOctalMode '0750' "${OPT[DMOD]}" ; }
	OPT[FMOD]=${ asOctalMode '0644' "${OPT[FMOD]}" ; }
	[[ -z ${OPT[DMOD]} || -z ${OPT[FMOD]} ]] && return 5

	return 0
}

function adjustPermissions {
	typeset LOG=${OPT[LOG]} LDIR=${LOG%/*} OWNER=${OPT[OWNER]} S T X
	integer PRE=1

	[[ -z ${OPT[GROUP]} ]] || OWNER+=":${OPT[GROUP]}"
	if [[ ! -d ${LDIR} ]]; then
		mkdir "${LDIR}" || return 1
		chown "${OWNER}" "${LDIR}"
		chmod 0750 "${LDIR}"
	fi
	[[ ${OPT[START]} == 'post' ]] && PRE=0

	if (( PRE )); then
		print 'Running in PRE mode.'
	else
		print 'Running in POST mode.'
	fi


	# i.e. pre exec
	S=( ${ stat -c '%a %U %G' "${LDIR}" ; } )
	[[ ${S[0]} != ${OPT[DMOD]:1} ]] && chmod ${OPT[DMOD]} "${LDIR}" && \
		print -u2 "Changed dir mode: '0$S' => '${OPT[DMOD]}'."
	[[ ${S[1]} != ${OPT[OWNER]} ]] && chown ${OPT[OWNER]} "${LDIR}" && \
		print -u2 "Changed dir owner: '${S[1]}' => '${OPT[OWNER]}'."
	[[ -n ${OPT[GROUP]} && ${S[2]} != ${OPT[GROUP]} ]] && \
		chgrp ${OPT[GROUP]} "${LDIR}" && \
		print -u2 "Changed dir group: '${S[2]}' => '${OPT[GROUP]}'."

	if (( ! PRE )); then
		SECONDS=0
		while [[ ! -f ${LOG} ]]; do
			(( SECONDS > 5 )) && break
			sleep 1
		done
		if [[ -n ${PIDFILE} ]]; then
			# systemd is so dumb. Need to check by our own, whether the daemon
			# succeeded to start.
			SECONDS=0
			while [[ ! -f ${PIDFILE} ]]; do
				(( SECONDS > 30 )) && return 1
				sleep 1
			done
			X=$(<${PIDFILE})
			[[ -z $X ]] && return 2		# already died
			sleep 1						# time to die
			kill -0 $X || return 3		# died
		fi
	fi
	if [[ -f ${LOG} ]]; then
		S=( ${ stat -c '%a %U %G' "${LOG}" ; } )
		[[ ${S[0]} != ${OPT[FMOD]:1} ]] && chmod ${OPT[FMOD]} "${LOG}" && \
			print -u2 "Changed file mode: '0$S' => '${OPT[FMOD]}'."
		[[ ${S[1]} != ${OPT[OWNER]} ]] && chown ${OPT[OWNER]} "${LOG}" && \
			print -u2 "Changed file owner: '${S[1]}' => '${OPT[OWNER]}'."
		[[ -n ${OPT[GROUP]} && ${S[2]} != ${OPT[GROUP]} ]] && \
			chgrp ${OPT[GROUP]} "${LOG}" && \
			print -u2 "Changed file group: '${S[2]}' => '${OPT[GROUP]}'."
	fi
	return 0
}

function printEnv {
	[[ ${OPT[ENV]} == '1' ]] || return

	typeset F="/tmp/env-slurm.${INVOCATION_ID:-1}" \
		LINE='############################################'
	if [[ ${OPT[START]} != 'post' ]]; then
		print "${LINE}\n# PRE\n${LINE}" >$F
	else
		print "\n${LINE}\n# POST\n${LINE}" >>$F
	fi
	env >>$F
}

function doMain {
	printEnv
	initDefaults || return 99
	adjustPermissions || return 95
	return 0
}

USAGE="[-?${VERSION}"' ]
[-copyright?Copyright (c) 2023 Jens Elkner. All rights reserved.]
[-license?CDDL 1.0]
[+NAME?'"${PROG}"' - little helper used in systemd slurm services to adjust permissions.]
[+DESCRIPTION?This script gets used in systemd slurm services to adjust permissions of log directories and log files, because slurm services are too paranoid and incomplete directory related implementation of systemd really sucks. Options can be specified multiple times on the command line. The last one takes precedences over previous options.]
[h:help?Print this help and exit.]
[F:functions?Print a list of all functions available.]
[T:trace]:[functionList?A comma separated list of functions of this script to trace (convinience for troubleshooting).] 
[+?]
[N:noop?Do nothing but exit with an exit code of 0.]
[c:config]:[path?Use the given path as slurm config file instead of '"${SLURMCONF}"' to deduce required settings.]
[d:dirmode]:[num?An octal number, which describes the mode to set to the directory containing the logfile. Default: 0750]
[e:env?Log the environment and script args to \b/tmp/env-slurm.$INVOCATION_ID\b .]
[f:filemode]:[num?An octal number, which describes the mode to set to the logfile, after the service has been started. For whatever reason slurmctld, etc. sets it always to 0600, which really sucks. Default: 0644]
[g:group]:[name?The name or ID of the group, which the logfile and its parent directory should have. Usually slurmctld, etc. does on start a chmod using the SlurmUser'"'"'s UID:GID. Default: none, i.e. group permissions are not enforced if not specified.]
[L:log]:[path?The absolute path of the logfile. Default: /var/log/\aSVC\a\b/main.log\b - with SVC == service name.]
[o:owner]:[name?The name or ID of the owner, which the logfile and its parent directory should have. Default: The user executing this script.]
[p:post?Whether this is run as a ExecStartPost script. Per default it is assumed that it runs as an ExecStartPre script.]
[s:service]:[name?The name of the related service: slurmd, slurmctld, slurmdbd, or slurmrestd. Default: deduced from the RUNTIME_DIRECTORY env variable.]
'

unset OPT; typeset -A OPT
X="${ print ${USAGE} ; }"
while getopts "${X}" OPT ; do
	case ${OPT} in
		h) showUsage ; exit 0 ;;
		T)	if [[ ${OPTARG} == 'ALL' ]]; then
				typeset -ft ${ typeset +f ; }
			else
				typeset -ft ${OPTARG//,/ }
			fi
			;;
		F) typeset +f && exit 0 ;;
		N) exit 0 ;;
		c) OPT[CFG]=${OPTARG} ;;
		d) OPT[DMOD]=${OPTARG} ;;
		e) OPT[ENV]=1 ;;
		f) OPT[FMOD]=${OPTARG} ;;
		g) OPT[GROUP]=${OPTARG} ;;
		L) OPT[LOG]=${OPTARG} ;;
		o) OPT[OWNER]=${OPTARG} ;;
		p) OPT[START]='post' ;;
		s) OPT[SVC]=${OPTARG} ;;
		*) showUsage 1 ; exit 1 ;;
	esac
done

X=$((OPTIND-1))
shift $X && OPTIND=1
unset X

doMain "$@"