#!/bin/bash
#
# Copyright (c) 2017, Philip Kovacs
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

prog=`basename $0`
user=root
group=root
yesno=no

red='\033[0;31m'
green='\033[0;32m'
nc='\033[0m'
pass="${green}[*]${nc}"
fail="${red}[!]${nc}"
ask="[?]"

# the rpm spec configures these directory variables
dir_conf="%{_sysconfdir}/%{name}"
dir_log="%{_var}/log/%{name}"
dir_run="%{_rundir}/%{name}"
dir_spool="%{_var}/spool/%{name}"
dir_tmpfiles_d="%{_tmpfilesdir}"

file_slurm_conf="${dir_conf}/slurm.conf"
file_slurmdbd_conf="${dir_conf}/slurmdbd.conf"
file_tmpfiles_d_conf="${dir_tmpfiles_d}/slurm.conf"

function usage ()
{
	echo  -e "Sets the slurm installation on this node to run as the specified user and group\n"
	echo  "Usage: $prog [-u <user>] [-g <group>] [-h] [-y]"
	echo  "   -u user>    : the slurm file owner and SlurmUser ($user)"
	echo  "   -g <group>  : the slurm file group ($group)"
	echo  "   -y          : answer yes to all questions ($yesno)"
	echo  "   -h          : print help"
}

function answer ()
{
	while true
	do
		if [ "${yesno}" = "yes" ]; then
			echo -e "${ask} $1 [${yesno}]" 
			__answer="yes"
			break;
		fi

		echo -e -n "${ask} "
		read -e -p "$1 [${yesno}] " yn
		case $yn in
			"")     __answer=$yesno
				break;;
			[yY]*)  __answer="yes" 
				break;;
			[nN]*)  __answer="no"
				break;;
			*) 
				echo "Please answer yes or no";;
		esac
	done
}

#
# Parse options
#
while getopts "u:g:yh" opt
do
	case "$opt" in
	h) 
		usage
		exit 2;;
	u)
		user=$OPTARG;;
	g)
		group=$OPTARG;;
	y)
		yesno=yes;;
	*) 
		echo ""
		usage
		exit 2;;
	esac
done

#
# Run this script only as root
#
if [ $UID -eq 0 ]; then
	echo -e "${pass} running as root... good"
else
	echo -e "${fail} $prog must be run as root."
	exit 1
fi

#
# Validate the user and group
#
valid_user="`getent passwd $user || :`"
if [ -n "$valid_user" ]; then
	echo -e "${pass} found user $user... good"
else
	echo -e "${fail} the specified user was not found... ${user}"
	exit 2
fi
valid_group="`getent group $group || :`"
if [ -n "$valid_group" ]; then
	echo -e "${pass} found group $group... good"
else
	echo -e "${fail} the specified group was not found... ${group}"
	exit 2
fi

#
# Slurm services must not be running
#
slurmctld_running_pid="`ps -e | grep slurmctld | grep -v grep | awk '{print $1}'`"
if [ -z "$slurmctld_running_pid" ]; then
	echo -e "${pass} slurmctld is not running... good"
else
	echo -e "${fail} slurmctld is running... stop it with [systemctl stop slurmctld]"
	exit 2
fi
slurmd_running_pid="`ps -e | grep slurmd | grep -v grep | awk '{print $1}'`"
if [ -z "$slurmd_running_pid" ]; then
	echo -e "${pass} slurmd is not running... good"
else
	echo -e "${fail} slurmd is running... stop it with [systemctl stop slurmd]"
	exit 2
fi
slurmdbd_running_pid="`ps -e | grep slurmdbd | grep -v grep | awk '{print $1}'`"
if [ -z "$slurmdbd_running_pid" ]; then
	echo -e "${pass} slurmdbd is not running... good"
else
	echo -e "${fail} slurmdbd is running... stop it with [systemctl stop slurmdbd]"
	exit 2
fi

#
# Update SlurmUser in the slurm configs
#
for file in "$file_slurm_conf" "$file_slurmdbd_conf"
do
	if [ -f "$file" ]; then
		answer "update SlurmUser in $file ?"
		if [ "$__answer" = "yes" ]; then
			sed -i "s|^SlurmUser=.*|SlurmUser=${user}|g" $file
			if [ $? -eq 0 ]; then
				echo -e "${pass} $file updated successfully"
			else
				echo -e "${fail} error updating $file"
				exit 1
			fi
		fi
	fi
done

#
# Update ownership of slurm directories that must be owned by the slurm user
#
for dir in "$dir_log" "$dir_run" "$dir_spool" "$dir_spool/ctld"
do
	if [ -d "$dir" ]; then
		answer "update ownership of $dir ?"
		if [ "$__answer" = "yes" ]; then
			chown ${user}:${group} $dir
			if [ $? -eq 0 ]; then
				echo -e "${pass} $dir updated successfully"
			else
				echo -e "${fail} error updating $dir"
				exit 1
			fi
		fi
	fi
done

#
# Update ownership of slurmctld spool files which are owned by the slurm user
#
if [ -d "${dir_spool}/ctld" ]; then
	answer "update ownership of files in ${dir_spool}/ctld ?"
	if [ "$__answer" = "yes" ]; then
		if [ $? -eq 0 ]; then
			find ${dir_spool}/ctld -mindepth 1 -exec chown ${user}:${group} {} \; 
			echo -e "${pass} ${dir_spool}/ctld files updated successfully"
		else
			echo -e "${fail} error updating ${dir_spool}/ctld files"
			exit 1
		fi
	fi
fi

#
# Update the tmpfiles.d config file to the new slurm user
#
if [ -f "$file_tmpfiles_d_conf" ]; then
	answer "update tmpfiles.d config $file_tmpfiles_d_conf ?"
	if [ "$__answer" = "yes" ]; then
		sed -i "s|0755 \(.*\) -$|0755 ${user} ${group} -|g" $file_tmpfiles_d_conf
		if [ $? -eq 0 ]; then
			echo -e "${pass} $file_tmpfiles_d_conf updated successfully"
		else
			echo -e "${fail} error updating $file_tmpfiles_d_conf"
			exit 1
		fi
	fi
fi

exit 0