#!/bin/bash

#
# Nemea control script
#


# Copyright (C) 2013 CESNET
#
# LICENSE TERMS
#
# 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.
# 3. Neither the name of the Company nor the names of its contributors
#    may be used to endorse or promote products derived from this
#    software without specific prior written permission.
#
# ALTERNATIVELY, provided that this notice is retained in full, this
# product may be distributed under the terms of the GNU General Public
# License (GPL) version 2 or later, in which case the provisions
# of the GPL apply INSTEAD OF those given above.
#
# This software is provided ``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 company 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.

# Move to directory where the script is located so relative paths work
# as expected even if the script is called from elsewhere.
#cd $(dirname $0)
#SCRIPTHOMEDIR="$(cd "$(dirname "${0}")" && pwd )"

prefix="/usr/local"
exec_prefix="${prefix}"
bindir="${exec_prefix}/bin"
#includedir=${prefix}/include
#libdir=${exec_prefix}/lib
#datarootdir=${prefix}/share
#datadir=${datarootdir}
sessiondir="/usr/local/share/nemea"

SCRIPTHOMEDIR="${bindir}"
# Paths to find various programs
NEMEABASE=${SCRIPTHOMEDIR}"/.."
#NEMEAMODULES="${NEMEABASE}/modules"
#NEMEASESSIONS="${NEMEABASE}/sessions"
NEMEAMODULES="${bindir}"
NEMEASESSIONS="${sessiondir}"
NEMEADEFAULTCONFIG="nemeadefault.conf"
NEMEALOGSUFFIX=".log"
NEMEAERRLOGSUFFIX="-err.log"
TERMINATEDPROCDIRSUFFIX="-TERMINATED"
SESSFINEXTENSION="-FINISHED"
BACKUPSUFFIX="-last.backup"

SESSIONPROCESSES="running_processes"
TERMSESSEXTENSION="-TERMINATED"

#sessionmodes=("list" "status" "start" "stop" "restart" "refresh" "restore")
#nonsessionmodes=("check" "status-all")
sessionmodes=("start" "stop" "restart")
nonsessionmodes=("check")

declare -A sessrunningprocesses
declare -A sessrplines

################################################################################
## Auxilliary function implementation
###################################>>

# *** Check if element is in array ***
# - usage: par1=integer variable for result (0=NOT FOUND/1=FOUND) par2=element par...=array
# / example (TRUE): is_in_array s "status-all" "${nonsessionmodes[@]}"
# / example (FALSE): is_in_array s "status-all" "${sessionmodes[@]}"
is_in_array () {
   local __found=$1
	local __element=$2
	shift 2
	eval $__found='0'
  	for arg; do
  		if [ "${arg}" == "${__element}" ]; then
		   eval $__found='1'
		fi
   done
}

# *** Get module names from program parameter, check validity also ***
getparmodules () {
	parmodules=()
	while (( "$#" )); do
	   if [ ${nemearunmodules[$1]+_} ]; then
	      parmodules+=(${1})
		else
		   echo "Warning: Ignoring unknown module ${1}."
		fi
      shift
   done
   echo
}

# *** Check if process is realy running ***
# - usage: par1=integer variable for result (0=NOT RUNNING/1=RUNNING) par2=PID par3=process (module) name
proc_is_running () {
   local __found=$1
	local __pid=$2
	local __name=$3
	local __proc=$(ps ax | grep $__pid | grep -v "grep" | grep $__name)
	if [ -z "${__proc}" ]; then
		eval $__found='0'
 	else
		eval $__found='1'
	fi
}

# *** Get running processes (PID and name) for session, check if realy running also ***
getsessrunningprocesses () {
	unset sessrunningprocesses
	unset sessrpconfigurations
	unset sessrplines
	local _f=${SESSIONDIRECTORY}/${SESSIONPROCESSES}
	declare -A _terminatedprocesses
	if [ -f $_f ]; then
	 	while read line; do
	     	local _pid=`echo "${line}" | cut -f1`
	      local _name=`echo "${line}" | cut -f2`
	      local _config=`echo "${line}" | cut -f3`
	      local _running=
	      local _processdir=${SESSIONDIRECTORY}/${_config}
	      if [ "${_pid}" ]; then
				proc_is_running _running ${_pid} ${_name}
				if (( $_running == 1 )); then
	            sessrunningprocesses[$_pid]=${_name}
	            sessrpconfigurations[$_pid]=${_config}
	   			sessrplines[$_pid]=${line}
				else
				   echo "Warning: Process ${_name} (PID: $_pid) not running while it is should." ###EDIT THIS
      		 	_terminatedprocesses[$j]=${_config}
      		 	if [ -d ${_processdir} ]; then
         			movetofinished ${_processdir} ${TERMINATEDPROCDIRSUFFIX}
					fi
				fi
			fi
		done < ${_f}
	else
    	echo "Warning: Missing file ${SESSIONPROCESSES} for session ${sessionid}."
	fi
	if [ ${#_terminatedprocesses[@]} -ne 0 ]; then
	   updatesession
	fi
	unset _terminatedprocesses
}

# *** Start one process, expecting two parameters - first is variable for pid of started process, second is module configuration name
startprocess () {
 	local __pid=$1
 	local __procconfigname=${2}
	local __processdir=${SESSIONDIRECTORY}/${__procconfigname}
	mkdir ${SESSIONDIRECTORY}/${__procconfigname}
	sleep 1
	if [[ ! -d ${SESSIONDIRECTORY} ]]; then
		echo "Error: cannot create directory \"${__processdir}\" for process ${runmodulenames[$__procconfigname]}, configuration ${__procconfigname}."
		exit 4
	fi
 local __statuslog="${__processdir}/${__procconfigname}${NEMEALOGSUFFIX}"
	local __errorlog="${__processdir}/${__procconfigname}${NEMEAERRLOGSUFFIX}"
	cd ${NEMEAMODULES};
	echo "Starting module ${runmodulenames[$__procconfigname]}, configuration ${__procconfigname}."
	echo "`date` nohup ./${nemearunmodules[$__procconfigname]} >${__statuslog} 2>${__errorlog} < /dev/null &" >> ${NEMEASESSIONLOG}
	nohup ./${nemearunmodules[$__procconfigname]} >${__statuslog} 2>${__errorlog} < /dev/null & eval $__pid='$!'
	cd ${OLDPWD}
	sleep 1
}

# *** Move folder with finished session or terminated process to its "end-folder" ***
movetofinished () {
 	local _path=$1
 	local _ext=$2
 	local _complete=0
	local _i=0
	while (( $_complete == 0)); do
		local _dest=${_path}${_ext}
		if (( $_i > 0)); then
         _dest=${_dest}${_i}
		fi
		if [ ! -d $_dest ]; then
			mv ${_path} ${_dest}
			_complete=1
		else
			_i=$(expr $_i + 1)
		fi
	done
}

# *** Update running_processes for given session***
updatesession () {
   local _f=${SESSIONDIRECTORY}/${SESSIONPROCESSES}
   #create temporary backup
	mv $_f ${_f}${BACKUPSUFFIX}
	for i in ${!sessrpconfigurations[@]}; do
     	local _isin
     	is_in_array _isin "${sessrpconfigurations[$i]}" "$@"
    	if (( $_isin == 0 )); then
    		echo "${sessrplines[$i]}" >> $_f
		fi
	done
	rm ${_f}${BACKUPSUFFIX}
}
##########################################<<
## END OF Auxilliary function implementation
################################################################################
## Running modes implementation
#############################>>
# *** Check for libtrap ***
check() {
	echo "Checking for libtrap...."
	libtrap=`yum list installed | grep "libtrap"`
	if [ -z "${libtrap}" ]; then
		echo "ERROR: Libtrap packages not found."
	else
		echo "Libtrap packages found:"
		echo "${libtrap}"
	fi
}

# *** List available modules (for given session = configuration) ***
list () {
  	echo "-------------------------------"
  	echo "List of available Nemea modules"
  	echo "-------------------------------"
  	if (( $defaultflag == 1 )); then
  	   for i in ${!nemearunmodules[*]}; do
  	      local _name=$(getmodulenamefromkey $i)
  	      local _modulpath=`find ${NEMEAMODULES} -type d -path "*/${_name}"`
  	      echo "Name: ${_name}"
  	      echo "Configuration: ${i}"
      	echo "Cmd:  ${nemearunmodules[$i]}"
      	echo "Path: ${_modulpath}"
      	echo
		done
  	else
  	   for i in ${parmodules[*]}; do
  	      local _name=$(getmodulenamefromkey $i)
  	      local _modulpath=`find ${NEMEAMODULES} -type d -path "*/${_name}"`
  	      echo "Name: ${_name}"
  	      echo "Configuration: ${i}"
      	echo "Cmd:  ${nemearunmodules[$i]}"
      	echo "Path: ${_modulpath}"
      	echo
		done
  	fi
}

# *** Print running modules for session ***
status () {
   if [ ! -d ${SESSIONDIRECTORY} ]; then
      echo "No session ${sessionid} is running."
      exit 1
	fi
   local _f=${SESSIONDIRECTORY}/${SESSIONPROCESSES}
   if [ ! -f $_f ]; then
      echo "Error: Missing file ${SESSIONPROCESSES} for session ${sessionid}."
      exit 6
	fi
	echo "Information on TRAP in /var/log/messages (you must be root)."
	cat /var/log/messages | grep -i "trap" | tail -n 20
	echo

	echo "Listing of running Nemea processes for session ${sessionid}:"
	getsessrunningprocesses
	# ----- DEFAULT -------------------------------------------------------------
	if (( $defaultflag == 1 )); then
    	for i in ${!sessrplines[@]}; do
		   echo "${sessrplines[$i]}"
		done
   # ----- PARAM MODULES (CONFIGS) ---------------------------------------------
  	else
      for i in ${parmodules[@]}; do
			local _found
			is_in_array _found ${i} ${sessrpconfigurations[@]}
   		if (( $_found == 1 )); then
   		   for j in ${!sessrpconfigurations[@]}; do
					if [ "${sessrpconfigurations[$j]}" == "${i}" ]; then
						echo "${sessrplines[$j]}"
					fi
				done
			else
			   echo "Process ${i} is not running."
			fi
		done
  	fi
}

# *** Print running modules for all sessions ***
status_all () {
	echo "Information on TRAP in /var/log/messages (you must be root)"
	cat /var/log/messages | grep -i "trap" | tail -n 20
	echo

	echo "Listing of running Nemea processes (for all sessions)"
	if (( $defaultflag == 1 )); then
		echo
  	else
		echo
  	fi
}

# *** Start module(s) for given session ***
start () {
	# Check if session (folder exists)
	if [ -d ${SESSIONDIRECTORY} ]; then
		echo "Session $sessionid allready exists, updating existing session."
		getsessrunningprocesses
	else
		echo "Creating new session $sessionid."
		mkdir ${SESSIONDIRECTORY}
		sleep 1
		if [[ ! -d ${SESSIONDIRECTORY} ]]; then
			echo "Error: cannot create directory \"${SESSIONDIRECTORY}\" for session $sessionid."
			exit 4
		fi
	fi
	echo

   # ----- DEFAULT -------------------------------------------------------------
	if (( $defaultflag == 1 )); then
		local _configs=( ${!nemearunmodules[@]} )
	else
	# ----- PARAM MODULES (CONFIGS) ---------------------------------------------
	   local _configs=( ${parmodules[@]} )
	fi
	# ---------------------------------------------------------------------------

   local _f=${SESSIONDIRECTORY}/${SESSIONPROCESSES}
   echo "Starting Nemea processes `date`"
   echo "Removing old socket files" | tee -a ${NEMEASESSIONLOG}
   echo
   rm -f /tmp/trap*.sock;
   for i in ${_configs[@]}; do
		local _isin
		is_in_array _isin "${i}" "${sessrpconfigurations[@]}"
		if (( $_isin == 1 )); then
	 		echo "Warning: Process ${i} is allready running."
		else
			local _actualpid
			startprocess _actualpid ${i}
	   	local _running
	   	proc_is_running _running $_actualpid ${runmodulenames[$i]}
			if (( $_running == 1 )); then
	      	echo ${_actualpid}$'\t'${runmodulenames[$i]}$'\t'${i}$'\t'${nemearunmodules[$i]} >> $_f
			else
			   movetofinished ${_processdir} ${TERMINATEDPROCDIRSUFFIX}
		   	echo "`date` Process ${runmodulenames[$i]} failed on start (details in: ${_processdir}${TERMINATEDPROCDIRSUFFIX}/${i}${NEMEAERRLOGSUFFIX})." >> ${NEMEASESSIONLOG}
		   	echo "Warning: Failed to start module ${runmodulenames[$i]}, see ${_processdir}${TERMINATEDPROCDIRSUFFIX}/${i}${NEMEAERRLOGSUFFIX} for more details."
			fi
		fi
	done
	echo
  	status
}

# *** Stop module(s) running in given session ***
stop () {
	local _justrestarting=$1
	getsessrunningprocesses
 	declare -A _terminatedprocesses
 	# ----- DEFAULT -------------------------------------------------------------
	if (( $defaultflag == 1 )); then
      for i in ${!sessrpconfigurations[@]}; do
         local _processdir=${SESSIONDIRECTORY}/${sessrpconfigurations[$i]}
         if [ $_justrestarting ]; then
            parmodules+=(${sessrpconfigurations[$i]})
            defaultflag=0 #hack
         fi
			echo "Killing process ${sessrunningprocesses[$i]}, configuration ${sessrpconfigurations[$i]}(PID: ${i})."
			kill $i
			sleep 5
			local _running
	   	proc_is_running _running $i ${sessrunningprocesses[$i]}
			if (( $_running == 0 )); then
			   echo "`date` Process ${sessrunningprocesses[$i]}, configuration ${sessrpconfigurations[$i]}(PID: ${i}) killed." >> ${NEMEASESSIONLOG}
				_terminatedprocesses[$i]=${sessrpconfigurations[$i]}
    			movetofinished ${_processdir} ${TERMINATEDPROCDIRSUFFIX}
			else
			   echo "Warning: Process ${sessrunningprocesses[$i]} (PID: ${i}) was not killed."
			fi
		done
	# ----- PARAM MODULES (CONFIGS) ---------------------------------------------
  	else
  		for i in ${parmodules[@]}; do
  		   local _found=0
       	for j in ${!sessrpconfigurations[@]}; do
       	   if [[ ${i} == ${sessrpconfigurations[$j]} ]]; then
					_found=1
           		local _processdir=${SESSIONDIRECTORY}/${sessrpconfigurations[$j]}
					echo "Killing process ${sessrunningprocesses[$j]}, configuration ${sessrpconfigurations[$j]} (PID: ${j})."
					kill $j
					sleep 5
					local _running
			   	proc_is_running _running $j ${sessrunningprocesses[$j]}
					if (( $_running == 0 )); then
					   echo "`date` Process ${sessrunningprocesses[$j]}, configuration ${sessrpconfigurations[$j]}(PID: ${j}) killed." >> ${NEMEASESSIONLOG}
      				_terminatedprocesses[$j]=${sessrpconfigurations[$j]}
      				movetofinished ${_processdir} ${TERMINATEDPROCDIRSUFFIX}
					else
					   echo "Warning: Process ${sessrunningprocesses[$j]} (PID: ${j}) was not killed."
					fi
				fi
			done
			if (( _found == 0 )); then
			   echo "Warning: No process ${i} is running in session ${sessionid}."
			fi
		done
  	fi
  	if [ ! -z "$_terminatedprocesses[@]" ]; then
		updatesession ${_terminatedprocesses[@]}
	fi
  	if [ -z "$_justrestarting" ]; then
  	   echo
  	   local _f=${SESSIONDIRECTORY}/${SESSIONPROCESSES}
		if [ ! -f $_f ]; then
			echo "Session ${sessionid} was terminated."
			echo "Moving session ${sessionid} to inactive."
	      movetofinished ${SESSIONDIRECTORY} ${SESSFINEXTENSION}
			sleep 2
  		else
			defaultflag=1 #hack
	  		status
		fi
	fi
}

# *** Running mode function prototype ***
proto () {
	echo "FunctionPrototype"
	if (( $defaultflag == 1 )); then
		echo
  	else
		echo
  	fi
}

####################################<<
## END OF Running modes implementation
################################################################################

#check runing-mode (parameter 1)
runningmode=$1
shift
declare -i sess
declare -i nonsess
is_in_array sess "${runningmode}" "${sessionmodes[@]}"
is_in_array nonsess "${runningmode}" "${nonsessionmodes[@]}"
if (( $sess == 1 )); then
   sessionid=$1 #session ID is parameter 2 (allready shifted)
   if [ -z ${sessionid} ]; then
      echo "Error: Missing session ID (2nd parameter)."
      exit 3
	fi
   shift
	conf=${sessionid}".conf"
elif (( $nonsess == 1 )); then
	conf=${NEMEADEFAULTCONFIG}
else
	###TODO USAGE
 	#echo "Usage: $0 {list|start|stop|restart|status|check} [module]..."
   echo "Usage: $0 {start|stop|restart} [sessionid] [module]..."
	echo "If no module is specified all available modules are utilized as an argument."
	#echo "list    - list available modules"
	echo "start   - start modules"
	echo "stop    - stop running modules"
	echo "restart - stop and start modules"
	#echo "status  - print running modules"
	#echo "check   - check installed packages"
	exit 2
fi

SESSIONDIRECTORY=${NEMEASESSIONS}/${sessionid}

#load configuration file
if [ -f ${NEMEASESSIONS}/${conf} ]; then
	. ${NEMEASESSIONS}/${conf}
else
	echo "Error: Missing configuration file ${NEMEASESSIONS}/${conf} for session ${sessionid}."
	exit 15
fi

#set global session log name
NEMEASESSIONLOG=${SESSIONDIRECTORY}/${sessionid}${NEMEALOGSUFFIX}

if [ -z "$1" ]; then
	defaultflag=1
else
   defaultflag=0
   getparmodules "$@"
fi

case "${runningmode}" in
	check)
		check
		;;

	list)
		list
		;;
		
   status)
		status
		;;
		
   "status-all")
		status_all
		;;
		
   start)
  		start
		;;
		
	stop)
		stop
		;;
		
	restart)
  		stop "restart"
  		start
		;;
		
   refresh)
  		getsessrunningprocesses
		;;

   restore)
      mv ${SESSIONDIRECTORY}/${SESSIONPROCESSES}${BACKUPSUFFIX} ${SESSIONDIRECTORY}/${SESSIONPROCESSES}
  		#getsessrunningprocesses
		;;

	*)
		echo "Error: Unexpected error during running-mode execution."
		exit 10
		;;
esac
#Everything OK - exit with success
exit 0
