#!/bin/sh
VIRTUAL_ENV_PYTHON="/opt/venv3/bin/python3"
# RUNASUSER="sudo -u root"
RUNASUSER=""

lprint(){
   echo "\e[96m"$1"\e[0m";
}

activate_venv(){
. /opt/venv3/bin/activate
}

receiver()
{
    ${RUNASUSER} /opt/control.setloki/service/receiver.bin $1
}

transmitter()
{
    ${RUNASUSER} /opt/control.setloki/service/transmitter.bin $1
}

loki_fenrir()
{
    ${RUNASUSER} /opt/control.setloki/service/loki_fenrir.bin $1
}

loki_redirector()
{
    case "$1" in
        start)
            lprint "INITIALIZING"
            ${RUNASUSER} /opt/control.setloki/service/av_redirector.bin start
            lprint "LOKI REDIRECTOR INITIALIZED"
            ;;
        stop)
            lprint "STOPPING REDIRECTOR"
            ${RUNASUSER} /opt/control.setloki/service/av_redirector.bin stop
            lprint "REDIRECTOR STOPPED"
            ;;
        restart)
            lprint "RESTARTING REDIECTOR"
            ${RUNASUSER} /opt/control.setloki/service/av_redirector.bin restart
            lprint "REDIRECTOR RESTARTED"
            ;;
        *)
            lprint "loki redir [start|stop|restart]"
            ;;
    esac
}

loki_network()
{
    case "$1" in
    start)
        lprint "Starting LOKI Network"
        ${RUNASUSER} /opt/control.setloki/service/loki_network.bin start
        lprint "LOKI Network started"
        ;;
    stop)
        lprint "Stopping LOKI Network"
        ${RUNASUSER} /opt/control.setloki/service/loki_network.bin stop
        lprint "LOKI NETWORK stopped"
        ;;
    restart)
        lprint "Restarting LOKI Network"
        ${RUNASUSER} /opt/control.setloki/service/loki_network.bin restart
        lprint "LOKI NETWORK Restarted"
        ;;
    *)
        lprint "loki network [start|stop|restart]"
        ;;
    esac
}

loki_flow_log()
{
    case "$1" in
        start)
            lprint "Starting LOKI FLOW LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/flow/dpi_log_sender.py start
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/firewall/drop_log_sender.py start
            lprint "LOKI FLOW LOGGER STARTED"
            ;;
        stop)
            lprint "Stopping LOKI FLOW LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/flow/dpi_log_sender.py stop
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/firewall/drop_log_sender.py stop
            lprint "LOKI FLOW LOGGER Stopped"
            ;;
        restart)
            lprint "Restarting LOKI FLOW LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/flow/dpi_log_sender.py restart
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/firewall/drop_log_sender.py restart
            lprint "LOKI FLOW LOGGER RESTARTED"
            ;;
        *)
            lprint "loki log flow [start|stop|restart]"
            ;;
    esac
}

loki_http_log()
{
    case "$1" in
        start)
            lprint "Starting LOKI HTTP LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/http/http_log_sender.py start
            lprint "LOKI HTTP LOGGER Started"
            ;;
        stop)
            lprint "Stopping LOKI HTTP LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/http/http_log_sender.py stop
            lprint "LOKI HTTP LOGGER Stopped"
            ;;
        restart)
            lprint "Restarting LOKI HTTP LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/http/http_log_sender.py restart
            lprint "LOKI HTTP LOGGER RESTARTED"
            ;;
        *)
            lprint "loki log http [start|stop|restart]"
            ;;
    esac
}

loki_system_log()
{
    case "$1" in
        start)
            lprint "Starting LOKI SYSTEM LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/system/interface_info.py start
            lprint "LOKI SYSTEM LOGGER Started"
            ;;
        stop)
            lprint "Stopping LOKI SYSTEM LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/system/interface_info.py stop
            lprint "LOKI SYSTEM LOGGER Stopped"
            ;;
        restart)
            lprint "Restarting LOKI SYSTEM LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/system/interface_info.py restart
            lprint "LOKI SYSTEM LOGGER RESTARTED"
            ;;
        *)
            lprint "loki log system [start|stop|restart]"
            ;;
    esac
}


loki_alert_log()
{
    case "$1" in
        start)
            lprint "Starting LOKI ALERT LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/ids/alert_generator.py start
            lprint "LOKI ALERT LOGGER Started"
            ;;
        stop)
            lprint "Stopping LOKI ALERT LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/ids/alert_generator.py stop
            lprint "LOKI ALERT LOGGER Stopped"
            ;;
        restart)
            lprint "Restarting LOKI ALERT LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/ids/alert_generator.py restart
            lprint "LOKI ALERT LOGGER RESTARTED"
            ;;
        *)
            lprint "loki log alert [start|stop|restart]"
            ;;
    esac
}

loki_dhcrelay_log()
{
    case "$1" in
        start)
            lprint "Starting LOKI DHCRELAY LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/dhcrelay/dhcrelay_log_daemon.py start
            lprint "LOKI DHCRELAY LOGGER Started"
            ;;
        stop)
            lprint "Stopping LOKI DHCRELAY LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/dhcrelay/dhcrelay_log_daemon.py stop
            lprint "LOKI DHCRELAY LOGGER Stopped"
            ;;
        restart)
            lprint "Restarting LOKI DHCRELAY LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/dhcrelay/dhcrelay_log_daemon.py restart
            lprint "LOKI DHCRELAY LOGGER RESTARTED"
            ;;
        *)
            lprint "loki log dhcrelay [start|stop|restart]"
            ;;
    esac
}

loki_log()
{
    case "$1" in
        flow)
            loki_flow_log $2
            ;;
        http)
            loki_http_log $2
            ;;
        system)
            loki_system_log $2
            ;;
        alert)
            loki_alert_log $2
            ;;
        dhcrelay)
            loki_dhcrelay_log $2
            ;;
        start)
            loki_flow_log start
            loki_http_log start
            loki_system_log start
            loki_alert_log start
            loki_dhcrelay_log start
            ;;
        stop)
            loki_flow_log stop
            loki_http_log stop
            loki_system_log stop
            loki_alert_log stop
            loki_dhcrelay_log stop
            ;;
        restart)
            loki_flow_log restart
            loki_http_log restart
            loki_system_log restart
            loki_alert_log restart
            loki_dhcrelay_log restart
            ;;
        *)
            lprint "loki log [start|stop|restart|flow|http|system|alert|device|dhcrelay]"
            ;;
    esac
}

loki_defender(){
     case "$1" in
          start)
              lprint "Starting LOKI ALERTER"
              ${RUNASUSER} /opt/control.setloki/service/loki_defender.bin start
              lprint "LOKI LOKI ALERTER Started"
              ;;
          stop)
              lprint "Stopping LOKI ALERTER LOGGER"
              ${RUNASUSER} /opt/control.setloki/service/loki_defender.bin stop
              lprint "LOKI LOKI ALERTER Stopped"
              ;;
          restart)
              lprint "Restarting LOKI ALERTER LOGGER"
              ${RUNASUSER} /opt/control.setloki/service/loki_defender.bin restart
              lprint "LOKI LOKI ALERTER RESTARTED"
              ;;
          *)
              lprint "loki defender [start|stop|restart]"
              ;;
      esac
}


loki_app()
{
    case "$1" in
        start)
            lprint "Starting LOKI CONNECT APP"
            ${RUNASUSER} /opt/control.setloki/service/setup_app.bin start
            lprint "LOKI CONNECT APP Started"
            ;;
        stop)
            lprint "Stopping LOKI CONNECT APP"
            ${RUNASUSER} /opt/control.setloki/service/setup_app.bin stop
            lprint "LOKI CONNECT APP Stopped"
            ;;
        restart)
            lprint "Restarting LOKI CONNECT APP"
            # ${RUNASUSER} /opt/control.setloki/service/setup_app.bin restart
            # XXX: After 2 restart command, pid file disappears. So process does not stop, continue working.
            ${RUNASUSER} /opt/control.setloki/service/setup_app.bin stop
            ${RUNASUSER} /opt/control.setloki/service/setup_app.bin start
            lprint "LOKI CONNECT APP Restarted"
            ;;
        *)
            lprint "loki app [start|stop|restart]"
            ;;
    esac
}

loki_ha_heartbeat()
{
    case "$1" in
        start)
            lprint "Starting LOKI HA HEARTBEAT APP"
            /opt/control.setloki/service/ha_heartbeat.bin start
            lprint "LOKI HA HEARTBEAT APP Started"
            ;;
        stop)
            lprint "Stopping LOKI HA HEARTBEAT APP"
            /opt/control.setloki/service/ha_heartbeat.bin stop
            lprint "LOKI HA HEARTBEAT APP Stopped"
            ;;
        restart)
            lprint "Restarting LOKI HA HEARTBEAT APP"
            /opt/control.setloki/service/ha_heartbeat.bin restart
            lprint "LOKI HA HEARTBEAT APP Restarted"
            ;;
        *)
            lprint "loki app [start|stop|restart]"
            ;;
    esac
}

loki_vrrp_listener()
{
    case "$2" in
        start)
            lprint "Starting LOKI HA HEARTBEAT APP"
            /opt/control.setloki/service/vrrp_listener.bin $1 start
            lprint "LOKI HA HEARTBEAT APP Started"
            ;;
        stop)
            lprint "Stopping LOKI HA HEARTBEAT APP"
            /opt/control.setloki/service/vrrp_listener.bin $1 stop
            lprint "LOKI HA HEARTBEAT APP Stopped"
            ;;
        restart)
            lprint "Restarting LOKI HA HEARTBEAT APP"
            /opt/control.setloki/service/vrrp_listener.bin $1 restart
            lprint "LOKI HA HEARTBEAT APP Restarted"
            ;;
        *)
            lprint "loki app [start|stop|restart]"
            ;;
    esac
}

loki_celery()
{
    CELERY_WORKING_DIR="/opt/control.setloki"
    CELERY_APP="tasks.task_manager"
    CELERYD_STATE_DIR="/var/run/celery"
    CELERYD_PID_FILE="${CELERYD_STATE_DIR}/celery.pid"
    CELERYD_LOG_LEVEL="INFO"

    case "$1" in
        start)
            lprint "Starting LOKI CELERY WORKER"

            MODEL_NAME=$(redis-cli -n 1 GET "model" 2>/dev/null)
            if [ -z "${MODEL_NAME}" ]; then
                MODEL_NAME="virtual"
                lprint "The model name will default to 'virtual' since it cannot be accessed from Redis."
            fi

            CELERY_CONCURRENCY=$(jq -r .system.celery_concurrency "/etc/loki/models/${MODEL_NAME}.json" 2>/dev/null)
            if [ -z "${CELERY_CONCURRENCY}" ]; then
                CELERY_CONCURRENCY=32
                lprint "The celery concurrency will default to 32 since it cannot be accessed from /etc/loki/models/${MODEL_NAME}.json"
            fi

            mkdir -p "${CELERYD_STATE_DIR}"

            is_celery_started=false
            retry_count_celery=0
            MAX_RETRIES_CELERY=3
            while [ "${retry_count_celery}" -lt "${MAX_RETRIES_CELERY}" ]; do

                ${RUNASUSER} celery --workdir "${CELERY_WORKING_DIR}" -A "${CELERY_APP}" worker \
                -Ofair --without-heartbeat --without-gossip --without-mingle \
                --pidfile="${CELERYD_PID_FILE}" --loglevel="${CELERYD_LOG_LEVEL}" \
                --concurrency="${CELERY_CONCURRENCY}" --pool=gevent 2>&1 | logger -t "celery_daemon" &

                MAX_RETRIES=15
                RETRY_INTERVAL=1
                retry_count=0

                while [ "${retry_count}" -lt "${MAX_RETRIES}" ]; do

                    if ${RUNASUSER} celery --workdir "${CELERY_WORKING_DIR}" -A "${CELERY_APP}" status; then
                        lprint "LOKI CELERY WORKER Started"
                        is_celery_started=true
                        return 0
                    else
                        echo "Waiting ${RETRY_INTERVAL} seconds to start celery..."
                        sleep "${RETRY_INTERVAL}"
                        retry_count=$((retry_count+1))
                    fi
                done

                retry_count_celery=$((retry_count_celery+1))

            done

            if ! $is_celery_started ; then
                lprint "LOKI CELERY WORKER Not Started"
                return 1
            fi

            ;;
        stop)
            lprint "Stopping LOKI CELERY WORKER"

            if [ -f "${CELERYD_PID_FILE}" ]; then
              CELERYD_PID="$(cat ${CELERYD_PID_FILE})"
              kill -TERM "${CELERYD_PID}"
            fi

            MAX_RETRIES=15
            RETRY_INTERVAL=1
            retry_count=0
            while [ "${retry_count}" -lt "${MAX_RETRIES}" ]; do

                if ${RUNASUSER} celery --workdir "${CELERY_WORKING_DIR}" -A "${CELERY_APP}" status; then

                    lprint "LOKI CELERY WORKER Already Running"
                    echo "Waiting ${RETRY_INTERVAL} seconds to stop celery..."
                    sleep "${RETRY_INTERVAL}"
                    retry_count=$((retry_count+1))
                else
                    lprint "Already Running LOKI CELERY WORKER Stopped"
                    break
                fi
            done

            if [ -n "${CELERYD_PID}" ] && [ -n "$(ps -p ${CELERYD_PID} -o pid=)" ]; then
              kill -9 "${CELERYD_PID}" \
              && echo "Killing Celery Worker pid:${CELERYD_PID}" || { echo "Cannot Killing Celery Worker pid:${CELERYD_PID}" ; return 1; }

              MAX_RETRIES=5
              RETRY_INTERVAL=1
              retry_count=0
              is_celery_pid_killed=false
              while [ "${retry_count}" -lt "${MAX_RETRIES}" ]; do
                if [ -n "$(ps -p ${CELERYD_PID} -o pid=)" ]; then
                  echo "Waiting ${RETRY_INTERVAL} seconds to kill celery pid:${CELERYD_PID}..."
                  sleep "${RETRY_INTERVAL}"
                  retry_count=$((retry_count+1))
                else
                  is_celery_pid_killed=true
                  break
                fi
              done

              if ! $is_celery_pid_killed ; then
                  lprint "Celery Worker pid:${CELERYD_PID} Not Killed"
                  return 1
              fi

              lprint "Killed Celery Worker pid:${CELERYD_PID}"
            fi

            if pgrep -f celeryd; then
              pkill -9 -f celeryd \
              && echo "Killing All Stuck Celery Workers" || { echo "Cannot Killing Stuck Celery Workers" ; return 1; }

              MAX_RETRIES=5
              RETRY_INTERVAL=1
              retry_count=0
              is_celeryd_killed=false
              while [ "${retry_count}" -lt "${MAX_RETRIES}" ]; do
                if pgrep -f celeryd; then
                  echo "Waiting ${RETRY_INTERVAL} seconds to kill celeryd..."
                  sleep "${RETRY_INTERVAL}"
                  retry_count=$((retry_count+1))
                else
                  is_celeryd_killed=true
                  break
                fi
              done

              if ! $is_celeryd_killed ; then
                  lprint "Stuck Celery Workers Not Killed"
                  return 1
              fi

              lprint "Killed All Stuck Celery Workers"

            fi

            rm -rf "${CELERYD_PID_FILE}"
            lprint "LOKI CELERY WORKER Stopped"
            return 0

            ;;
        restart)
            lprint "Restarting LOKI CELERY WORKER"

            if loki_celery stop && loki_celery start; then
                lprint "LOKI CELERY WORKER Restarted"
                return 0
            else
                lprint "LOKI CELERY WORKER Not Restarted"
                return 1
            fi

            ;;
        *)
            lprint "loki celery [start|stop|restart]"
            ;;
    esac
}

loki_celery_beat()
{
    CELERY_WORKING_DIR="/opt/control.setloki"
    CELERY_APP="tasks.task_manager"
    CELERYD_LOG_LEVEL="INFO"
    CELERYBEAT_STATE_DIR="/var/run/celery"
    CELERYBEAT_PID_FILE="${CELERYBEAT_STATE_DIR}/beat.pid"

    case "$1" in
        start)
            lprint "Starting LOKI CELERY BEAT SCHEDULER"
            mkdir -p ${CELERYBEAT_STATE_DIR}
            ${RUNASUSER} celery --workdir "${CELERY_WORKING_DIR}" -A "${CELERY_APP}" beat \
            --pidfile="${CELERYBEAT_PID_FILE}" --loglevel="${CELERYD_LOG_LEVEL}" 2>&1 | logger -t "celery_beat" &
            lprint "LOKI CELERY BEAT SCHEDULER Started"
            ;;
        stop)
            lprint "Stopping LOKI CELERY BEAT SCHEDULER"

            if [ -f "${CELERYBEAT_PID_FILE}" ]; then
                kill -TERM "$(cat ${CELERYBEAT_PID_FILE})"
                rm -rf "${CELERYBEAT_PID_FILE}"
            fi

            lprint "LOKI CELERY BEAT SCHEDULER Stopped"
            ;;
        restart)
            lprint "Restarting LOKI CELERY BEAT SCHEDULER"

            loki_celery_beat stop
            loki_celery_beat start

            lprint "LOKI CELERY BEAT SCHEDULER Restarted"
            ;;
        *)
            lprint "loki celery_beat [start|stop|restart]"
            ;;
    esac
}


case "$1" in
    fenrir)
        activate_venv
        loki_fenrir $2
        ;;

    log)
        loki_log $2 $3 $4
        ;;

    network)
        activate_venv
        loki_network $2
        ;;

    defender)
        activate_venv
        loki_defender $2
        ;;

    redir)
        loki_redirector $2
        ;;

    ha_heartbeat)
        activate_venv
        loki_ha_heartbeat $2
        ;;

    vrrp_listener)
        activate_venv
        loki_vrrp_listener $2 $3
        ;;
    app)
        activate_venv
        loki_app $2
        ;;
    start)
        loki_log start
        activate_venv
        loki_defender start
        loki_redirector start
        loki_app start
        loki_fenrir start
        receiver start
        transmitter start
        loki_celery start
        loki_celery_beat start
        ;;
    stop)
        activate_venv
        loki_fenrir stop
        loki_log stop
        loki_defender stop
        loki_redirector stop
        receiver stop
        transmitter stop
        loki_celery stop
        loki_celery_beat stop
        ;;
    restart)
        loki_log restart
        activate_venv
        loki_defender restart
        loki_redirector restart
        loki_app restart
        loki_fenrir restart
        receiver restart
        transmitter restart
        loki_celery restart
        loki_celery_beat restart
        ;;
    setup)
        activate_venv
        loki_fenrir start
        loki_log start
        loki_redirector start
        receiver start
        transmitter start
        ;;
    celery)
        activate_venv
        loki_celery $2
        ;;
    celery_beat)
        activate_venv
        loki_celery_beat $2
        ;;
    receiver)
        activate_venv
        receiver $2
        ;;
    transmitter)
        activate_venv
        transmitter $2
        ;;
    *)
        echo "--------------------------------------------------------"
        lprint "LOKI TERMINAL APP USAGE"
        echo "--------------------------------------------------------"
        lprint "loki [service] [subservice] [ start | stop | restart ]"
        lprint "Services:\n\tapp : Captive Application\n\tfenrir: Consistency Service\n\treceiver : Consumer service\n\ttransmitter : Producer service\n\tcron : CRON Daemon\n\tlog : Log Services\n\tnetwork : Newtork Daemon\n\tDefender : Defender Service\n\tcelery : Celery Service\n\tcelery_beat : Celery Beat Service\n"
        lprint "SUB-SERVICES:\n\tlog:\n\t    http : HTTP log collection and sending service\n\t    flow : FLOW log collection and sending service\n\t    system : SYSTEM informational log service\n\t    device: Client Device Information Log\n\t    alert: Alert Logger\n"
        ;;
esac

