#!/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_ipsec_log()
{
    case "$1" in
        start)
            lprint "Starting LOKI IPSEC LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/ipsec/ipsec_log_sender.py start
            lprint "LOKI IPSEC LOGGER Started"
            ;;
        stop)
            lprint "Stopping LOKI IPSEC LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/ipsec/ipsec_log_sender.py stop
            lprint "LOKI IPSEC LOGGER Stopped"
            ;;
        restart)
            lprint "Restarting LOKI IPSEC LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/ipsec/ipsec_log_sender.py restart
            lprint "LOKI IPSEC LOGGER RESTARTED"
            ;;
        *)
            lprint "loki log ipsec [start|stop|restart]"
            ;;
    esac
}

loki_clamav_log()
{
    case "$1" in
        start)
            lprint "Starting LOKI CLAMAV LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/http/clamav_log_sender.py start
            lprint "LOKI CLAMAV LOGGER Started"
            ;;
        stop)
            lprint "Stopping LOKI CLAMAV LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/http/clamav_log_sender.py stop
            lprint "LOKI CLAMAV LOGGER Stopped"
            ;;
        restart)
            lprint "Restarting LOKI CLAMAV LOGGER"
            ${RUNASUSER} ${VIRTUAL_ENV_PYTHON} /opt/control.setloki/service/logger/http/clamav_log_sender.py restart
            lprint "LOKI CLAMAV LOGGER RESTARTED"
            ;;
        *)
            lprint "loki log clamav [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
            ;;
        ipsec)
            loki_ipsec_log $2
            ;;
        clamav)
            loki_clamav_log $2
            ;;
        start)
            loki_flow_log start
            loki_http_log start
            loki_system_log start
            loki_dhcrelay_log start
            loki_ipsec_log start
            loki_clamav_log start
            ;;
        stop)
            loki_flow_log stop
            loki_http_log stop
            loki_system_log stop
            loki_dhcrelay_log stop
            loki_ipsec_log stop
            loki_clamav_log stop
            ;;
        restart)
            loki_flow_log restart
            loki_http_log restart
            loki_system_log restart
            loki_dhcrelay_log restart
            loki_ipsec_log restart
            loki_clamav_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)
            activated=$(redis-cli -n 1 get activated)
            if [ "${activated}" = "1" ]; then
              lprint "LOKI CONNECT APP will not start, gateway is activated"
            else
              lprint "Starting LOKI CONNECT APP"
              ${RUNASUSER} /opt/control.setloki/service/setup_app.bin start
              lprint "LOKI CONNECT APP Started"
            fi
            ;;
        stop)
            lprint "Stopping LOKI CONNECT APP"
            ${RUNASUSER} /opt/control.setloki/service/setup_app.bin stop
            lprint "LOKI CONNECT APP Stopped"
            ;;
        restart)
            activated=$(redis-cli -n 1 get activated)
            if [ "${activated}" = "1" ]; then
              lprint "LOKI CONNECT APP will not restart, gateway is activated"
            else
              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"
            fi
            ;;
        *)
            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/admin/celery"
    celeryd_log_level="INFO"
    worker_name="$2"

    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
    model_file="/etc/loki/models/${model_name}.json"

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

            mkdir -p "${celeryd_state_dir}"

            for process in $(jq -r .system.celery "$model_file" | jq -c '.processes[]'); do
                current_worker_name=$(echo "$process" | jq -r .worker_name)
                
                # Skip if worker_name is provided and doesn't match current worker
                if [ -n "${worker_name}" ] && [ "${worker_name}" != "${current_worker_name}" ]; then
                    continue
                fi

                queue_name=$(echo "$process" | jq -r '.queues | map(tostring) | join(",")')
                celery_concurrency=$(echo "$process" | jq -r .concurrency)

                if timeout 5s ${RUNASUSER} celery --workdir "${celery_working_dir}" -A "${celery_app}" status | grep -w "${current_worker_name}"; then
                    lprint "Celery ${current_worker_name} already running"
                    continue
                fi

                celery_pid_file="${celeryd_state_dir}/${current_worker_name}.pid"
                celery_state_file="${celeryd_state_dir}/${current_worker_name}.state"

                is_celery_started=0
                retry_count_celery=0
                max_retries_celery=3
                while [ "${retry_count_celery}" -lt "${max_retries_celery}" ] && [ ${is_celery_started} -eq 0 ]; do
                    
                   ${RUNASUSER} celery --workdir "${celery_working_dir}" -A "${celery_app}" worker \
                    -Ofair --without-heartbeat --without-gossip --without-mingle \
                    --pidfile="${celery_pid_file}" --statedb="${celery_state_file}" \
                    --loglevel="${celeryd_log_level}" --concurrency="${celery_concurrency}" \
                    -n "${current_worker_name}" --pool=gevent -Q "${queue_name}" 2>&1 | logger -t "celery_daemon" &

                    max_retries=15
                    retry_interval=1
                    retry_count=0

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

                        if timeout 5s ${RUNASUSER} celery --workdir "${celery_working_dir}" -A "${celery_app}" status | grep -w "${current_worker_name}"; then
                            lprint "LOKI CELERY WORKER ${current_worker_name} Started"
                            is_celery_started=1
                            break
                        else
                            echo "Waiting ${retry_interval} seconds to start celery ${current_worker_name}..."
                            sleep "${retry_interval}"
                            retry_count=$((retry_count+1))
                        fi
                    done

                    retry_count_celery=$((retry_count_celery+1))
                done

                if [ ${is_celery_started} -eq 0 ]; then
                    lprint "LOKI CELERY WORKER ${current_worker_name} Not Started"
                    return 1
                fi
            done
            ;;
        stop)
            lprint "Stopping 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

            for process in $(jq -r .system.celery $model_file | jq -c '.processes[]'); do
                current_worker_name=$(echo "$process" | jq -r .worker_name)
                

                # Skip if worker_name is provided and doesn't match current worker
                if [ -n "${worker_name}" ] && [ "${worker_name}" != "${current_worker_name}" ]; then
                    continue
                fi

                celery_pid_file="${celeryd_state_dir}/${current_worker_name}.pid"

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

                max_retries=15
                retry_interval=1
                retry_count=0
                while [ "${retry_count}" -lt "${max_retries}" ]; do
                    if timeout 5s ${RUNASUSER} celery --workdir "${celery_working_dir}" -A "${celery_app}" status | grep -w "${current_worker_name}"; then
                        lprint "LOKI CELERY WORKER ${current_worker_name} 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 ${current_worker_name} Stopped"
                        break
                    fi
                done

                if [ -n "${celery_pid}" ] && [ -n "$(ps -p "${celery_pid}" -o pid=)" ]; then
                    kill -9 "${celery_pid}" \
                    && echo "Killing Celery Worker pid:${celery_pid}" || { echo "Cannot Killing Celery Worker pid:${celery_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 "${celery_pid}" -o pid=)" ]; then
                        echo "Waiting ${retry_interval} seconds to kill celery pid:${celery_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:${celery_pid} Not Killed"
                    else
                        lprint "Killed Celery Worker pid:${celery_pid}"
                    fi
                fi

                rm -rf "${celery_pid_file}"
            done

            # Only kill all stuck workers if no specific worker was specified
            if [ -z "${worker_name}" ] && 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

            lprint "LOKI CELERY WORKER Stopped"
            return 0
            ;;
        restart)
            lprint "Restarting LOKI CELERY WORKER"

            if loki_celery stop "$2" && loki_celery start "$2"; then
                lprint "LOKI CELERY WORKER Restarted"
                return 0
            else
                lprint "LOKI CELERY WORKER Not Restarted"
                return 1
            fi
            ;;
        *)
            lprint "loki celery [start|stop|restart] [worker_name]"
            ;;
    esac
}

loki_celery_beat()
{
    celery_working_dir="/opt/control.setloki"
    celery_app="tasks.task_scheduler"
    celery_log_level="INFO"
    celerybeat_state_dir="/var/run/admin/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="${celery_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
        ;;

    app)
        activate_venv
        loki_app $2
        ;;
    start)
        activate_venv
        loki_log start
        loki_redirector start
        loki_app start
        receiver start
        transmitter start
        loki_celery start $2
        loki_celery_beat start
        loki_fenrir start
        ;;
    stop)
        activate_venv
        loki_fenrir stop
        loki_log stop
        loki_redirector stop
        receiver stop
        transmitter stop
        loki_celery stop $2
        loki_celery_beat stop
        ;;
    restart)
        activate_venv
        loki_fenrir stop
        loki_log restart
        loki_redirector restart
        loki_app restart
        receiver restart
        transmitter restart
        loki_celery restart $2
        loki_celery_beat restart
        loki_fenrir start
        ;;
    setup)
        activate_venv
        loki_log start
        loki_redirector start
        receiver start
        transmitter start
        loki_fenrir start
        ;;
    celery)
        activate_venv
        loki_celery $2 $3
        ;;
    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

