#!/bin/bash
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Originally written by Bryant Treacle
# https://raw.githubusercontent.com/bryant-treacle/so-elastalert-test-rule/master/so-elastalert-test
# Modified by Doug Burks and Wes Lambert
#
# Purpose:  This script will allow you to test your elastalert rule without entering the Docker container.

. /usr/sbin/so-elastic-common

OPTIONS=""
SKIP=0
RESULTS_TO_LOG="n"
RULE_NAME=""
FILE_SAVE_LOCATION=""

usage()
{
cat <<EOF

Test Elastalert Rule
  Options:
  -h         		This message
  -a         		Trigger real alerts instead of the debug alert
  -l <path_to_file>     Write results to specified log file
  -o '<options>'	Specify Elastalert options ( Ex. --schema-only , --count-only, --days N ) 
  -r <rule_name>        Specify path/name of rule to test

EOF
}

while getopts "hal:o:r:" OPTION
do
        case $OPTION in
                h)
                        usage
                        exit 0
                        ;;
                a)
                        OPTIONS="--alert"
                        ;;
                l)
                        RESULTS_TO_LOG="y"
                        FILE_SAVE_LOCATION=$OPTARG
                        ;;
                
                o)
			OPTIONS=$OPTARG
                        ;;  
                
                r)
                        RULE_NAME=$OPTARG
                        SKIP=1
                        ;;
                *)
                        usage
                        exit 0
                        ;;
        esac
done

docker_exec(){
	if [ ${RESULTS_TO_LOG,,} = "y" ] ; then
        	docker exec -it so-elastalert bash -c "elastalert-test-rule $RULE_NAME $OPTIONS" > $FILE_SAVE_LOCATION
	else
        	docker exec -it so-elastalert bash -c "elastalert-test-rule $RULE_NAME $OPTIONS"
	fi	
}

rule_prompt(){
	CURRENT_RULES=$(find /opt/so/rules/elastalert -name "*.yaml")
        echo
        echo "This script will allow you to test an Elastalert rule."
        echo
        echo "Below is a list of active Elastalert rules:"
        echo
	echo "-----------------------------------"
        echo
        echo "$CURRENT_RULES"
        echo
	echo "-----------------------------------"
	echo 
        echo "Note: To test a rule it must be accessible by the Elastalert Docker container."
        echo
        echo "Make sure to swap the local path (/opt/so/rules/elastalert/) for the docker path (/etc/elastalert/rules/)"
        echo "Example: /opt/so/rules/elastalert/nids2hive.yaml would be /etc/elastalert/rules/nids2hive.yaml"
        echo
        while [ -z $RULE_NAME ]; do
		echo "Please enter the file path and rule name you want to test."
		read -e RULE_NAME
	done
}

log_save_prompt(){
	RESULTS_TO_LOG=""
        while [ -z $RESULTS_TO_LOG ]; do
		echo "The results can be rather long.  Would you like to write the results to a file? (Y/N)"
		read RESULTS_TO_LOG
	done
}

log_path_prompt(){
        while [ -z $FILE_SAVE_LOCATION ]; do
		echo "Please enter the file path and file name."
		read -e FILE_SAVE_LOCATION
        done
		echo "Depending on the rule this may take a while."
}

if [ $SKIP -eq 0 ]; then
	rule_prompt
	log_save_prompt
        if [ ${RESULTS_TO_LOG,,} = "y" ] ; then
        	log_path_prompt
        fi
fi

docker_exec

if [ $? -eq 0 ]; then
	echo "Test completed successfully!"
else
	echo "Something went wrong..."
fi

echo