#!/bin/bash
### Copyright 1999-2023. Plesk International GmbH. All rights reserved.
# vim:ft=sh:

set -e

usage()
{
	if [ -n "$*" ]; then
		echo "drwebmng: $*" >&2
		echo "Use '--help' option to get help on the utility usage" >&2
		exit 2
	fi

	cat <<-EOT
		Usage: drwebmng COMMAND [OPTIONS]
		
		  --add-check --mailname=<mail1,mail2,...> --type={"from"|"to"|"any"|"incoming"|"outgoing"} [--no-restart-service]
		                            Add rule check into antivirus configuration file
		  --remove-check --mailname=<mail1,mail2,...> [--no-restart-service]
		                            Delete rule check from antivirus configuration file
		  -p, --start               Start drweb service
		  -o, --stop                Stop drweb service
		  -s, --status              Show drweb service status
		  -r, --restart             Restart drweb service
		  -l, --reload              Reload drweb service
		  -e, --on                  Turn on drweb
		  -d, --off                 Turn off drweb
		  -u, --update              Reload, turn on or turn off drweb depending on the handlers state
		  -t, --setup-trial-vhost   Set up support for drweb trial license in sw-cp-server
		  -h, --help                Display this help and exit
		EOT
	exit 2
}

mail_handler_add()
{
	local direction="$1"
	local mailname="$2"

	if [ "$direction" = "in" ]; then
		"$MAIL_HANDLERS_BIN" --add --name "$DRWEB_HOOK_NAME" --mailname "$mailname" \
			--priority "$DRWEB_HOOK_PRIORITY" --executable "$HOOK_DIR/$DRWEB_HOOK_NAME" \
			--enabled --dont-preserve-on-restore \
			--queue "before-queue" --type "recipient"
	elif [ "$direction" = "out" ]; then
		"$MAIL_HANDLERS_BIN" --add --name "$DRWEB_HOOK_NAME" --mailname "$mailname" \
			--priority "$DRWEB_HOOK_PRIORITY" --executable "$HOOK_DIR/$DRWEB_HOOK_NAME" \
			--enabled --dont-preserve-on-restore \
			--queue "before-remote" --type "sender"
	fi
}

mail_handler_remove()
{
	local direction="$1"
	local mailname="$2"

	if [ "$direction" = "in" ]; then
		"$MAIL_HANDLERS_BIN" --remove --name "$DRWEB_HOOK_NAME" --mailname "$mailname" \
			--queue "before-queue" --type "recipient"
	elif [ "$direction" = "out" ]; then
		"$MAIL_HANDLERS_BIN" --remove --name "$DRWEB_HOOK_NAME" --mailname "$mailname" \
			--queue "before-remote" --type "sender"
	fi
}

mail_handler_set()
{
	local direction="$1"
	local mailname="$2"

	if [ "$direction" = "none" ]; then
		mail_handler_remove "in"  "$mailname"
		mail_handler_remove "out" "$mailname"
	elif [ "$direction" = "any" ]; then
		mail_handler_add "in"  "$mailname"
		mail_handler_add "out" "$mailname"
	elif [ "$direction" = "incoming" -o "$direction" = "to" ]; then
		mail_handler_add "in" "$mailname"
		mail_handler_remove "out" "$mailname"
	elif [ "$direction" = "outgoing" -o "$direction" = "from" ]; then
		mail_handler_remove "in" "$mailname"
		mail_handler_add "out" "$mailname"
	fi
}

drweb_service()
{
	local action="$1"
	/bin/systemctl "$action" "$DRWEBD_SERVICE.service"
}

do_start()
{
	drweb_service start
}

do_stop()
{
	drweb_service stop
}

do_status()
{
	if drweb_service is-active > /dev/null; then
		echo "is running"
	else
		echo "is stopped"
	fi
}

do_restart()
{
	drweb_service restart
}

do_reload()
{
	if drweb_service is-active > /dev/null; then
		drweb_service reload
	else
		echo "Skip drweb antivirus daemon reloading - service is stopped" >&2
	fi
}

setup_swcpserver()
{
	local action="$1"
	[ "$action" = "enable" -o "$action" = "disable" ]

	if [ "$action" = "enable" ]; then
		# allow TLSv1 connections in sw-cp-server if there is no key in the drwebd directory
		[ ! -r "$DRWEB_SWCPSERVER_CONF" ] || return 0

		cp -fT "$DRWEB_SWCPSERVER_CONF_TPL" "$DRWEB_SWCPSERVER_CONF"
		# listen to ipv6 only if it is available
		if [ -n "`cat /proc/net/if_inet6`" ]; then
			cp -fT "$DRWEB_SWCPSERVER_IPV6_TPL" "$DRWEB_SWCPSERVER_IPV6"
		fi

		/usr/sbin/sw-cp-serverd -q -t && "$PLESKRC_BIN" sw-cp-server try-reload || {
			echo "Unable to setup sw-cp-server to receive connections from drweb." >&2
			rm -f "$DRWEB_SWCPSERVER_IPV6"
			rm -f "$DRWEB_SWCPSERVER_CONF"
			return 1
		}
	elif [ -r "$DRWEB_SWCPSERVER_CONF" ]; then
		rm -f "$DRWEB_SWCPSERVER_IPV6"
		rm -f "$DRWEB_SWCPSERVER_CONF"

		"$PLESKRC_BIN" sw-cp-server try-reload
	fi
}

register_service()
{
	local action="$1"
	[ "$action" = "enable" -o "$action" = "disable" ]

	local REGISTER_SERVICE_BIN="$PRODUCT_ROOT_D/admin/sbin/register_service"
	"$REGISTER_SERVICE_BIN" "--$action" "$DRWEBD_SERVICE" || {
		echo "Unable to $action drweb antivirus daemon" >&2
		return 1
	}
}

watchdog_monit_service()
{
	local action="$1"
	[ "$action" = "monit" -o "$action" = "unmonit" ]

	local WATCHDOG_BIN="$PRODUCT_ROOT_D/admin/bin/modules/watchdog/wd"

	[ -x "$WATCHDOG_BIN" ] || {
		echo "Skipped setting up monitoring in WatchDog as it is not accessible" >&2
		return 0
	}

	"$WATCHDOG_BIN" "--$action-service=$DRWEB_WD_SERVICE" || {
		echo "Unable to $action drweb in WatchDog" >&2
		return 1
	}
}

do_on()
{
	do_setup_trial_vhost
	register_service enable
	do_start
	watchdog_monit_service monit
}

do_off()
{
	setup_swcpserver disable
	watchdog_monit_service unmonit
	do_stop
	register_service disable
}

do_update()
{
	local allow_reload="$1"
	local ignore_current_state="$2"

	local handlers_count="`
		"$MAIL_HANDLERS_BIN" --list --json | python3 -c "
import sys, json
handlers = json.load(sys.stdin)['handlers']
print(sum(1 for h in handlers if h['name'] == '$DRWEB_HOOK_NAME' and h['enabled']))
		"
	`"
	local was_running=
	! drweb_service is-active > /dev/null || was_running="yes"

	local needs_to_be_running=
	! [ "$handlers_count" -gt 0 ] || needs_to_be_running="yes"

	if [ "$was_running" != "$needs_to_be_running" -o -n "$ignore_current_state" ]; then
		if [ -n "$needs_to_be_running" ]; then
			do_on
		else
			do_off
		fi
	fi

	if [ -n "$allow_reload" -a -n "$was_running" -a -n "$needs_to_be_running" ]; then
		drweb_service reload
	fi
}

do_setup_trial_vhost()
{
	setup_swcpserver "`[ -r "$DRWEB_KEY" ] && echo disable || echo enable`"
}

# --- environment and constants setup ---

unset GREP_OPTIONS
umask 022

PRODUCT_ROOT_D='/opt/psa'
MAIL_HANDLERS_BIN="$PRODUCT_ROOT_D/admin/sbin/mail_handlers_control"
PLESKRC_BIN="$PRODUCT_ROOT_D/admin/sbin/pleskrc"

DRWEBD_SERVICE=drwebd
DRWEB_HOOK_NAME='drweb'
DRWEB_HOOK_PRIORITY='20'
DRWEB_KEY='/opt/drweb/drweb32.key'
DRWEB_SWCPSERVER_CONF='/etc/sw-cp-server/conf.d/plesk-drweb-local.conf'
DRWEB_SWCPSERVER_CONF_TPL='/etc/drweb/swcpserver.plesk-drweb-local.conf.tpl'
DRWEB_SWCPSERVER_IPV6='/etc/sw-cp-server/conf.d/plesk-drweb-local.ipv6.inc'
DRWEB_SWCPSERVER_IPV6_TPL='/etc/drweb/swcpserver.plesk-drweb-local.ipv6.inc.tpl'
DRWEB_WD_SERVICE='drweb'
HOOK_DIR='/opt/psa/handlers/hooks'

# --- parse args ---

opt_mailnames=()
opt_type=
opt_allow_restart="yes"
opt_cmd=

TEMP=`getopt -o posrleduth \
	--long add-check,remove-check,mailname:,type:,no-restart-service,start,stop,status,restart,reload,on,off,update,setup-trial-vhost,help \
	-n drwebmng -- "$@"` || exit 2
eval set -- "$TEMP"

declare -A CMD_SHORT_TO_LONG=(
	[-p]=--start
	[-o]=--stop
	[-s]=--status
	[-r]=--restart
	[-l]=--reload
	[-e]=--on
	[-d]=--off
	[-u]=--update
	[-t]=--setup-trial-vhost
)

while [ "$#" -gt 0 ]; do
	case "$1" in
		-h|--help)
			usage
			;;
		--mailname)
			IFS=, opt_mailnames+=($2)
			shift
			;;
		--type)
			[ -z "$opt_type" ] || usage "Use '--type' only once"
			opt_type="$2"
			shift
			;;
		--no-restart-service)
			opt_allow_restart=
			;;
		--add-check|--remove-check)
			[ -z "$opt_cmd" ] || usage "Specify only one command, got '$opt_cmd' and '$1'"
			opt_cmd="$1"
			;;
		--[a-z]*)
			[[ "$1" =~ ^(`echo "${CMD_SHORT_TO_LONG[@]}" | tr " " "|"`)$ ]] ||
				usage "Unrecognized option '$1'"
			[ -z "$opt_cmd" ] || usage "Specify only one command, got '$opt_cmd' and '$1'"
			opt_cmd="$1"
			;;
		-[a-z])
			[[ "$1" =~ ^(`echo "${!CMD_SHORT_TO_LONG[@]}" | tr " " "|"`)$ ]] ||
				usage "Unrecognized option '$1'"
			[ -z "$opt_cmd" ] || usage "Specify only one command, got '$opt_cmd' and '$1'"
			opt_cmd="${CMD_SHORT_TO_LONG[$1]}"
			;;
		--)
			shift
			break
			;;
		*)
			usage "Unrecognized option '$1'"
			;;
	esac
	shift
done

[ -z "$*" ] ||
	usage "Positional arguments are not accepted"
[ -n "$opt_cmd" ] ||
	usage "Command is empty"
[ "$opt_cmd" != "--add-check" ] || [ -n "${opt_mailnames[*]}" -a -n "$opt_type" ] ||
	usage "Command '$opt_cmd' requires '--mailname' and '--type' options"
[ "$opt_cmd" != "--remove-check" ] || [ -n "${opt_mailnames[*]}" ] ||
	usage "Command '$opt_cmd' requires '--mailname' option"
[ -n "$opt_allow_restart" ] || [ "$opt_cmd" = "--add-check" -o "$opt_cmd" = "--remove-check" ] ||
	usage "Option '--no-restart-service' is accepted only by '--add-check' and '--remove-check' commands"
[ -z "$opt_type" ] || [ "$opt_type" = "from" -o "$opt_type" = "to" -o "$opt_type" = "any" \
		-o "$opt_type" = "incoming" -o "$opt_type" = "outgoing" ] ||
	usage "Unsupported '--type' option value: '$opt_type'"

# --- execute ---

case "$opt_cmd" in
	--add-check)
		for mailname in "${opt_mailnames[@]}"; do
			mail_handler_set "$opt_type" "$mailname"
		done
		do_update "$opt_allow_restart"
		;;
	--remove-check)
		for mailname in "${opt_mailnames[@]}"; do
			mail_handler_set "none" "$mailname"
		done
		do_update "$opt_allow_restart"
		;;
	--start)
		do_start
		;;
	--stop)
		do_stop
		;;
	--status)
		do_status
		;;
	--restart)
		do_restart
		;;
	--reload)
		do_reload
		;;
	--on)
		do_on
		;;
	--off)
		do_off
		;;
	--update)
		do_update --allow-reload --ignore-current-state
		;;
	--setup-trial-vhost)
		do_setup_trial_vhost
		;;
esac