#!/bin/bash

# ossec-control        This shell script takes care of starting
#                      or stopping ossec-hids
# Author: Daniel B. Cid <daniel.cid@gmail.com>
# Version: 5.0.0

LOCAL=`dirname $0`;
cd ${LOCAL}
PWD=`pwd`
DIR=`dirname $PWD`;
# AIX loader path for bundled shared libraries (libsysinfo, libsyscollector, etc.)
if [ "$(uname)" = "AIX" ]; then
    export OSSEC_HOME="${DIR}"
    export LIBPATH="${DIR}/lib:${LIBPATH}"
    export PATH="/usr/sbin:/usr/bin:/sbin:/bin:/opt/freeware/bin:${PATH}"
fi
# Detect systemd

if [[ -d /run/systemd/system ]]; then
        if [ -f /usr/lib/systemd/system/ossec-agent.service ]; then
                SYSTEMCTL=$(which systemctl)
        fi
fi

###  Do not modify bellow here ###
AUTHOR="Atomicorp Inc."
DAEMONS="ossec-modulesd ossec-logcollector ossec-syscheckd ossec-execd ossec-agentd"
INITCONF="/etc/ossec-init.conf"

[ -f ${INITCONF} ] && . ${INITCONF}  || echo "ERROR: No such file ${INITCONF}"

## Locking for the start/stop
LOCK="${DIR}/var/start-script-lock"
LOCK_PID="${LOCK}/pid"

# This number should be more than enough (even if it is
# started multiple times together). It will try for up
# to 10 attempts (or 10 seconds) to execute.
MAX_ITERATION="10"

MAX_KILL_TRIES=600
STOP_KILL_TRIES=50

function checkpid() {
    for i in ${DAEMONS}; do
        for j in `cat ${DIR}/var/run/${i}*.pid 2>/dev/null`; do
            ps -p $j > /dev/null 2>&1
            if [ ! $? = 0 ]; then
                echo "Deleting PID file '${DIR}/var/run/${i}-${j}.pid' not used..."
                rm ${DIR}/var/run/${i}-${j}.pid
            fi
        done
    done
}

function lock() {
    i=0;

    # Providing a lock.
    while [ 1 ]; do
        mkdir ${LOCK} > /dev/null 2>&1
        MSL=$?
        if [ "${MSL}" = "0" ]; then
            # Lock acquired (setting the pid)
            echo "$$" > ${LOCK_PID}
            return;
        fi

        # Waiting 1 second before trying again
        sleep 1;
        i=`expr $i + 1`;
        pid=`cat ${LOCK_PID}` 2>/dev/null

        if [ $? = 0 ]
        then
            kill -0 ${pid} >/dev/null 2>&1
            if [ ! $? = 0 ]; then
                # Pid is not present.
                # Unlocking and executing
                unlock;
                mkdir ${LOCK} > /dev/null 2>&1
                echo "$$" > ${LOCK_PID}
                return;
            fi
        fi

        # We tried 10 times to acquire the lock.
        if [ "$i" = "${MAX_ITERATION}" ]; then
            echo "ERROR: Another instance is locking this process."
            echo "If you are sure that no other instance is running, please remove ${LOCK}"
            exit 1
        fi
    done
}

function unlock() {
    rm -rf ${LOCK}
}

function help() {
    # Help message
    echo "Usage: $0 {start|stop|restart|status}";
    exit 1;
}

function status() {
    RETVAL=0
    for i in ${DAEMONS}; do
        pstatus ${i};
        if [ $? = 0 ]; then
            RETVAL=1
            echo "${i} not running..."
        else
            echo "${i} is running..."
        fi
    done
}

function testconfig() {
    # We first loop to check the config.
    for i in ${DAEMONS}; do
        ${DIR}/bin/${i} -t;
        if [ $? != 0 ]; then
            echo "${i}: Configuration error. Exiting"
            unlock;
            exit 1;
        fi
    done
}

function check_queue() {
    QUEUE_DIR="../queue/alerts"

    if [ ! -d $QUEUE_DIR ]
    then
        if rm -rf $QUEUE_DIR && mkdir -p $QUEUE_DIR && chown ossec:ossec $QUEUE_DIR && chmod 770 $QUEUE_DIR
        then
            echo "WARNING: missing directory 'queue/alerts'. Recreated."
        else
            echo "ERROR: missing directory 'queue/alerts', unable to recreate."
            exit 1
        fi
    fi
}

# Modules function
function modules() {
    # Modules
    if [ -f /var/ossec/modules/clamav/clam-module.sh ]; then
	/var/ossec/modules/clamav/clam-module.sh
    fi

    if [ -f /var/ossec/modules/firewall/firewall-module.sh ]; then
	/var/ossec/modules/firewall/firewall-module.sh
    fi

    if [ -f /var/ossec/modules/fapolicyd/fapolicyd-module.sh ]; then
	/var/ossec/modules/fapolicyd/fapolicyd-module.sh
    fi
}

# Start function
function start() {


    echo "Starting $NAME $VERSION..."
    checkpid;
    
    # Delete all files in temporary folder
    TO_DELETE="$DIR/tmp/*"
    rm -rf $TO_DELETE

    # We actually start them now.
    for i in ${DAEMONS}; do
        pstatus ${i};
        if [ $? = 0 ]; then
            ${DIR}/bin/${i};
            if [ $? != 0 ]; then
                echo "${i} did not start";
                unlock;
                exit 1;
            fi

            echo "Started ${i}..."
	    sleep 5
        else
            echo "${i} already running..."
        fi
    done

    # After we start we give 2 seconds for the daemons
    # to internally create their PID files.
    sleep 2;
    echo "Completed."




}

function pstatus() {
    pfile=$1;

    # pfile must be set
    if [ "X${pfile}" = "X" ]; then
        return 0;
    fi

    ls ${DIR}/var/run/${pfile}*.pid > /dev/null 2>&1
    if [ $? = 0 ]; then
        for j in `cat ${DIR}/var/run/${pfile}*.pid 2>/dev/null`; do
            kill -0 $j > /dev/null 2>&1
            if [ $? = 0 ]; then
                return 1;
            fi

            # AIX ps -p can fail for valid daemons; treat kill -0 as authoritative.
            if [ "$(uname)" != "AIX" ]; then
                ps -p $j > /dev/null 2>&1
                if [ $? = 0 ]; then
                    return 1;
                fi
            fi

            echo "${pfile}: Process $j not used by ossec, removing .."
            rm -f ${DIR}/var/run/${pfile}-$j.pid
            continue;
        done
    fi

    return 0;
}

function wait_pid() {
    local wpid=$1
    local max_tries=${2:-$MAX_KILL_TRIES}
    wp_counter=1

    while kill -0 ${wpid} 2> /dev/null
    do
        if [ "$wp_counter" = "$max_tries" ]
        then
            return 1
        else
            # sleep doesn't work in AIX
            # read doesn't work in FreeBSD
            sleep 0.1 > /dev/null 2>&1 || read -t 0.1 > /dev/null 2>&1
            wp_counter=`expr $wp_counter + 1`
        fi
    done

    return 0
}

# Collect running PIDs for a daemon (pid files + exact binary path for orphans).
collect_daemon_pids() {
    local name=$1
    local bin="${DIR}/bin/${name}"
    local pid=""
    local pid_list=""

    for pidfile in ${DIR}/var/run/${name}*.pid; do
        [ -f "${pidfile}" ] || continue
        for pid in `cat "${pidfile}" 2>/dev/null`; do
            if kill -0 ${pid} 2>/dev/null; then
                pid_list="${pid_list} ${pid}"
            fi
        done
    done

    if [ -x "${bin}" ]; then
        if command -v pgrep >/dev/null 2>&1; then
            for pid in `pgrep -f "${bin}([[:space:]]|$)" 2>/dev/null`; do
                if kill -0 ${pid} 2>/dev/null; then
                    pid_list="${pid_list} ${pid}"
                fi
            done
        else
            for pid in `ps -eo pid,args 2>/dev/null | awk -v bin="${bin}" '$2 == bin { print $1 }'`; do
                if kill -0 ${pid} 2>/dev/null; then
                    pid_list="${pid_list} ${pid}"
                fi
            done
        fi
    fi

    echo ${pid_list}
}

daemon_still_running() {
    local name=$1
    local pids=`collect_daemon_pids "${name}"`
    [ -n "${pids}" ]
}

kill_daemon() {
    local name=$1
    local pid=""
    local pids=""

    pids=`collect_daemon_pids "${name}"`
    if [ -z "${pids}" ]; then
        rm -f ${DIR}/var/run/${name}*.pid
        return 0
    fi

    echo "Killing ${name}... "
    for pid in ${pids}; do
        kill -9 ${pid} 2>/dev/null
    done

    for pid in ${pids}; do
        if ! wait_pid ${pid} ${STOP_KILL_TRIES}; then
            echo "Process ${name} couldn't be killed."
            return 1
        fi
    done

    if daemon_still_running "${name}"; then
        echo "Process ${name} couldn't be killed."
        return 1
    fi

    rm -f ${DIR}/var/run/${name}*.pid
    return 0
}

function stopa() {
    local failed=0

    checkpid;
    for i in ${DAEMONS}; do
        if daemon_still_running "${i}"; then
            kill_daemon "${i}" || failed=1
        else
            echo "${i} not running..."
            rm -f ${DIR}/var/run/${i}*.pid
        fi
     done

    if [ "${failed}" = "1" ]; then
        echo "ERROR: One or more OSSEC processes are still running."
        exit 1
    fi

    echo "$NAME $VERSION Stopped"
}

function disable_scan_on_start() {
	cp /var/ossec/etc/ossec.conf /var/ossec/etc/ossec.conf.$(date +%s)
	sed -i 's/<start-on-scan>yes<\/start-on-scan>/<start-on-scan>no<\/start-on-scan>/g' /var/ossec/etc/ossec.conf
	sed -i 's/<scan_on_start>yes<\/>/<scan_on_start>no<\/>/g' /var/ossec/etc/ossec.conf
	cp /var/ossec/etc/ossec.conf.$(date +%s) /var/ossec/etc/ossec.conf
    	rm -f /var/ossec/etc/ossec.conf.$(date +%s)
}


### MAIN HERE ###

case "$1" in
start)
    testconfig
    check_queue
    modules
    if [[ ! -z "${SYSTEMCTL}" ]]; then
       	${SYSTEMCTL} start ossec-hids
       	exit 0
    fi
    lock
    start
    unlock
    ;;
stop)
    if [[ ! -z "${SYSTEMCTL}" ]]; then
        ${SYSTEMCTL} stop ossec-hids
	exit 0
    fi
    lock
    stopa
    unlock
    ;;
restart)
    testconfig
    modules
    if [[ ! -z "${SYSTEMCTL}" ]]; then
        ${SYSTEMCTL} restart ossec-hids
	exit 0
    fi 
    lock
    stopa
    sleep 1
    start
    unlock
    ;;
reload)
    DAEMONS=$(echo $DAEMONS | sed 's/ossec-execd//')
    disable_scan_on_start
    if [[ ! -z "${SYSTEMCTL}" ]]; then
        ${SYSTEMCTL} restart ossec-hids
	exit 0
    fi 
    lock
    stopa
    sleep 1
    start
    unlock
    ;;
status)
    lock
    status
    unlock
    ;;
help)
    help
    ;;
*)
    help
esac

exit $RETVAL
