#!/bin/bash

WEBSHIELD="imunify360-webshield"
WSSHDICT="wsshdict"
SSL_CACHE="imunify360-webshield-ssl-cache"
JOBFILE="/etc/cron.d/imunify360-webshield"
STATEFILE="/usr/share/imunify360-webshield/.webshieldctl.status"
WEBSHIELD_CONF="/etc/imunify360-webshield/webshield.conf"
WEBSHIELD_SHDICT_CONF="/etc/imunify360-webshield/webshield-shdict.conf"
WEBSHIELD_ANTIBOT_CONF="/etc/imunify360-webshield/splashscreen-antibot.conf"

has_hosting_panel(){
    local checks=(
        /usr/local/cpanel/cpanel
        /usr/sbin/plesk
        /usr/local/directadmin/custombuild/build
    )
    for i in ${checks[@]};do
        [ -e $i ] && return 0
    done
    return 1
}

count_processes(){
    local count=$(ps aux | grep -c '[i]m360:\|[w]sshdict:\|webshield-[s]sl-cache')
    echo $count
}

check_running(){
    # for hosts with hosting panels we expect 5 processes to be running.
    # otherwise 4 ones (ssl-cache is not expected to be run on no-panel hosts)
    local num=$(count_processes)
    local expected
    if has_hosting_panel; then expected=5; else expected=4; fi
    [ $num -ge $expected ] && return 0
    return 1
}

check_stopped(){
    local num=$(count_processes)
    [ $num -eq 0 ] && return 0
    return 1
}

enable_for_systemd(){
    local RV
    systemctl -q enable $WEBSHIELD
    RV=$?
    if [ $RV -ne 0 ];then
        echo "Enabling $WEBSHIELD returned non-zero status" 1>&2
    fi
    return $RV
}

start_for_systemd(){
    local RV
    for ITEM in $WEBSHIELD $WSSHDICT;do
        systemctl -q start $ITEM;
        RV=$?
        if [ $RV -ne 0 ];then
            echo "Starting $ITEM returned non-zero status" 1>&2
            return $RV
        fi
    done
    has_hosting_panel || return $RV
    systemctl -q start $SSL_CACHE;
    RV=$?
    if [ $RV -ne 0 ];then
        echo "Starting $SSL_CACHE returned non-zero status" 1>&2
    fi
    [ $RV -ne 0 ] && return $RV
    check_running
}

activate_for_systemd(){
    local RV
    enable_for_systemd
    RV=$?
    [ $RV -ne 0 ] && return $RV
    start_for_systemd
}

disable_for_systemd(){
    local RV
    for ITEM in $WEBSHIELD $WSSHDICT $SSL_CACHE;do
        systemctl -q disable $ITEM
        RV=$?
        if [ $RV -ne 0 ];then
            echo "Disabling $ITEM returned non-zero status" 1>&2
            return $RV
        fi
    done
    return $RV
}

stop_for_systemd(){
    local RV
    for ITEM in $WEBSHIELD $WSSHDICT $SSL_CACHE;do
        systemctl -q stop $ITEM
        RV=$?
        if [ $RV -ne 0 ];then
            echo "Stopping $ITEM returned non-zero status" 1>&2
            return $RV
        fi
    done
    [ $RV -ne 0 ] && return $RV
    check_stopped
}

deactivate_for_systemd(){
    local RV
    stop_for_systemd
    RV=$?
    [ $RV -ne 0 ] && return $RV
    disable_for_systemd
}

enable_for_sysvinit(){
    local RV
    /sbin/chkconfig $WEBSHIELD on > /dev/null 2>&1
    RV=$?
    if [ $RV -ne 0 ];then
        echo "Enabling $WEBSHIELD returned non-zero status" 1>&2
        return $RV
    fi
    if [ -e $JOBFILE ];then
        sed -i -e 's/^#//g' $JOBFILE
        RV=$?
        if [ $RV -ne 0 ];then
            echo "Enabling cron job returned non-zero status" 1>&2
        fi
    fi
    return $RV
}

start_for_sysvinit(){
    local RV
    /sbin/service $WEBSHIELD start >/dev/null 2>&1
    RV=$?
    if [ $RV -ne 0 ];then
        echo "Starting $WEBSHIELD returned non-zero status" 1>&2
        return $RV
    fi
    if [ -e $JOBFILE ];then
        sed -i -e 's/^#//g' $JOBFILE
        RV=$?
        if [ $RV -ne 0 ];then
            echo "Enabling cron job returned non-zero status" 1>&2
        fi
    fi
    [ $RV -ne 0 ] && return $RV
    check_running
}

activate_for_sysvinit(){
    local RV
    enable_for_sysvinit
    RV=$?
    [ $RV -ne 0 ] && return $RV
    start_for_sysvinit
}

disable_for_sysvinit(){
    local RV
    for ITEM in $WEBSHIELD $WSSHDICT $SSL_CACHE;do
        /sbin/chkconfig --del $ITEM > /dev/null 2>&1
        RV=$?
        if [ $RV -ne 0 ];then
            echo "Disabling $ITEM returned non-zero status" 1>&2
            return $RV
        fi
    done
    if [ -e $JOBFILE ];then
        sed -i -e 's/^\([^#]\)/#\1/g' $JOBFILE
        RV=$?
        if [ $RV -ne 0 ];then
            echo "Disabling cron job returned non-zero status" 1>&2
        fi
    fi
    return $RV
}

stop_for_sysvinit(){
    local RV
    for ITEM in $WEBSHIELD $WSSHDICT $SSL_CACHE;do
        /sbin/service $ITEM stop >/dev/null 2>&1
        RV=$?
        if [ $RV -ne 0 ];then
            echo "Stopping $ITEM returned non-zero status" 1>&2
            return $RV
        fi
    done
    if [ -e $JOBFILE ];then
        sed -i -e 's/^\([^#]\)/#\1/g' $JOBFILE
        RV=$?
        if [ $RV -ne 0 ];then
            echo "Disabling cron job returned non-zero status" 1>&2
        fi
    fi
    [ $RV -ne 0 ] && return $RV
    check_stopped
}

deactivate_for_sysvinit(){
    local RV
    stop_for_sysvinit
    RV=$?
    [ $RV -ne 0 ] && return $RV
    disable_for_sysvinit
}

is_systemd(){
    if [ -e "/etc/redhat-release" ];then
        local version=$(cat /etc/redhat-release | sed -e 's/^[[:alpha:][:space:]]\+//' | head -c 1)
        if [ "$version" = "6" ];then
            echo "No"
            return 0
        fi
    fi
    echo "Yes"
}

do_enable(){
    if [ "x$(is_systemd)" = "xYes" ];then
        enable_for_systemd
    else
        enable_for_sysvinit
    fi

    local rv=$?
    if [ $rv == 0 ]; then
        echo "enabled" > $STATEFILE
    fi
    return $rv
}

do_disable(){
    if [ "x$(is_systemd)" = "xYes" ];then
        disable_for_systemd
    else
        disable_for_sysvinit
    fi

    local rv=$?
    if [ $rv == 0 ]; then
        echo "disabled" > $STATEFILE
    fi
    return $rv
}

do_start(){
    if [ "x$(is_systemd)" = "xYes" ];then
        start_for_systemd
    else
        start_for_sysvinit
    fi

    local rv=$?
    if [ $rv == 0 ]; then
        echo "started" > $STATEFILE
    fi
    return $rv
}

do_stop(){
    if [ "x$(is_systemd)" = "xYes" ];then
        stop_for_systemd
    else
        stop_for_sysvinit
    fi

    local rv=$?
    if [ $rv == 0 ]; then
        echo "stopped" > $STATEFILE
    fi
    return $rv
}

do_activate(){
    if [ "x$(is_systemd)" = "xYes" ];then
        activate_for_systemd
    else
        activate_for_sysvinit
    fi

    local rv=$?
    if [ $rv == 0 ]; then
        echo "activated" > $STATEFILE
    fi
    return $rv
}

do_deactivate(){
    if [ "x$(is_systemd)" = "xYes" ];then
        deactivate_for_systemd
    else
        deactivate_for_sysvinit
    fi

    local rv=$?
    if [ $rv == 0 ]; then
        echo "deactivated" > $STATEFILE
    fi
    return $rv
}

is_enabled_for_systemd(){
    if systemctl is-enabled $WEBSHIELD > /dev/null 2>&1; then
        echo "$WEBSHIELD is enabled"
        return 0
    else
        echo "$WEBSHIELD is disabled"
        return 1
    fi
}

is_enabled_for_sysvinit(){
    local _runlevel=$(runlevel|cut -d' ' -f2)
    local STATE=$(LANG=C /sbin/chkconfig --list $WEBSHIELD 2>/dev/null | grep "$_runlevel:on")
    if [ -n "$STATE" ];then
        echo "$WEBSHIELD is enabled"
        return 0
    else
        echo "$WEBSHIELD is disabled"
        return 1
    fi
}

is_enabled(){
    if [ "x$(is_systemd)" = "xYes" ];then
        is_enabled_for_systemd
    else
        is_enabled_for_sysvinit
    fi
}

is_active(){
    if check_running;then
        echo "$WEBSHIELD is running"
        return 0
    fi
    echo "$WEBSHIELD is not running"
    return 1
}

do_restart(){
    do_stop
    do_start
}

do_enable_splashscreen(){
    local ss_state=$(awk '$1 == "splashscreen_antibot" {gsub(";","",$2);print $2}' $WEBSHIELD_ANTIBOT_CONF)
    if [ "$ss_state" = on ];then
        echo "splashscreen is already enabled"
        return 0
    fi
    local old_shm_size=$(awk -F= '$1 ~ /data_size/ && $1 !~ /#/ {gsub(" ","",$2);print $2}' $WEBSHIELD_SHDICT_CONF)
    local old_ws_val=$(awk '$1 == "shared_storage" {gsub(";","",$3);print $3}' $WEBSHIELD_CONF)
    local ws_suff=$(echo $old_ws_val | awk '{idx=match($0,/[^0-9]/);print substr($0,idx)}')
    local old_ws_size=$(echo $old_ws_val | awk '{sub(/[^0-9]/,"",$0);print $0}')
    if [ "$ws_suff" = "$old_ws_size" ];then
        ws_suff=""
    fi
    local new_shm_size=$((old_shm_size*2))
    local new_ws_size=$((old_ws_size*2))
    local new_ws_val="$new_ws_size$ws_suff"
    sed -i -e "/splashscreen_antibot/ {s/off/on/}" $WEBSHIELD_ANTIBOT_CONF
    sed -i -e "/^data_size/ {s/$old_shm_size/$new_shm_size/}" $WEBSHIELD_SHDICT_CONF
    sed -i -e "/shared_storage/ {s/$old_ws_val/$new_ws_val/}" $WEBSHIELD_CONF
    do_restart
}

do_disable_splashscreen(){
    local ss_state=$(awk '$1 == "splashscreen_antibot" {gsub(";","",$2);print $2}' $WEBSHIELD_ANTIBOT_CONF)
    if [ "$ss_state" = off ];then
        echo "splashscreen is already disabled"
        return 0
    fi
    local old_shm_size=$(awk -F= '$1 ~ /data_size/ && $1 !~ /#/ {gsub(" ","",$2);print $2}' $WEBSHIELD_SHDICT_CONF)
    local old_ws_val=$(awk '$1 == "shared_storage" {gsub(";","",$3);print $3}' $WEBSHIELD_CONF)
    local ws_suff=$(echo $old_ws_val | awk '{idx=match($0,/[^0-9]/);print substr($0,idx)}')
    local old_ws_size=$(echo $old_ws_val | awk '{sub(/[^0-9]/,"",$0);print $0}')
    if [ "$ws_suff" = "$old_ws_size" ];then
        ws_suff=""
    fi
    local new_shm_size=$((old_shm_size/2))
    local new_ws_size=$((old_ws_size/2))
    if [ $new_shm_size -eq 0 ] || [ $new_ws_size -eq 0 ];then
        echo "Got improper shared memory size while disabling splashscreen. All left unchanged. Please contact support."
        return 1
    fi
    local new_ws_val="$new_ws_size$ws_suff"
    sed -i -e "/splashscreen_antibot/ {s/on/off/}" $WEBSHIELD_ANTIBOT_CONF
    sed -i -e "/^data_size/ {s/$old_shm_size/$new_shm_size/}" $WEBSHIELD_SHDICT_CONF
    sed -i -e "/shared_storage/ {s/$old_ws_val/$new_ws_val/}" $WEBSHIELD_CONF
    do_restart
}

print_help(){
    echo "enable                : enables webshield starting on boot (without actully starting it)"
    echo "is-enabled            : shows if the webshield is enabled to start on boot"
    echo "is-active             : shows if the webshield is running now"
    echo "disable               : disables webshield starting on boot (without actully stopping it)"
    echo "start                 : starts webshield (without enabling its starting on boot)"
    echo "stop                  : stops webshield (without disabling its starting on boot)"
    echo "activate              : enables webshield starting on boot and starts it right away"
    echo "deactivate            : stops webshield right away and disables its starting on boot"
    echo "enable-splashscreen   : enables splashscreen functionality for webshield"
    echo "disable-splashscreen  : disables splashscreen functionality for webshield"
}

case "$1" in
    enable)
        do_enable
        ;;
    disable)
        do_disable
        ;;
    is-enabled)
        is_enabled
        ;;
    is-active)
        is_active
        ;;
    start)
        do_start
        ;;
    stop)
        do_stop
        ;;
    activate)
        do_activate
        ;;
    deactivate)
        do_deactivate
        ;;
    enable-splashscreen)
        do_enable_splashscreen
        ;;
    disable-splashscreen)
        do_disable_splashscreen
        ;;
    help)
        print_help
        ;;
    *)
        echo "Usage: $0 {enable|disable|start|stop|activate|deactivate|is-enabled|is-active|enable-splashscreen|disable-splashscreen|help}"
        exit 2
esac

