#!/bin/bash

# Copyright 2014-2020 Security Onion Solutions, LLC

#    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/>.

cd "$(dirname "$0")" || exit 255
source ./so-functions
source ./so-common-functions
source ./so-whiptail
source ./so-variables

# Parse command line arguments
setup_type=$1
automation=$2

while [[ $# -gt 0 ]]; do
	arg="$1"
	shift
	case "$arg" in
		"--turbo="* ) 
			export TURBO="http://${arg#*=}";;
		"--proxy="* ) 
			export {http,https,ftp,rsync,all}_proxy="${arg#*=}";;
		"--allow-role="* )
			export ALLOW_ROLE="${arg#*=}";;
		"--allow-cidr="* )
			export ALLOW_CIDR="${arg#*=}";;
		"--skip-reboot" )
			export SKIP_REBOOT=1;;
		* )
			if [[ "$arg" == "--"* ]]; then
				echo "Invalid option"
			fi
	esac
done

# Begin Installation pre-processing
echo "---- Starting setup at $(date -u) ----" >> $setup_log 2>&1

automated=no
function progress() {
	local title='Security Onion Install'
	if grep -q "ERROR" $setup_log || [[ -s /var/spool/mail/root ]]; then
		if [[ -s /var/spool/mail/root ]]; then 
			echo '[ ERROR ]  /var/spool/mail/root grew unexpectedly' >> $setup_log 2>&1
		fi

		export SO_ERROR=1
		export SKIP_REBOOT=1
		title="Error found, please check $setup_log"
	fi

	if [ $automated == no ]; then
		whiptail --title "$title" --gauge 'Please wait while installing...' 6 60 0 # append to text
	else
		cat >> $setup_log 2>&1
	fi		
}

if [[ -f automation/$automation && $(basename $automation) == $automation ]]; then
	echo "Preselecting variable values based on automated setup: $automation" >> $setup_log 2>&1
	source automation/$automation
	automated=yes

	echo "Checking network configuration" >> $setup_log 2>&1
	ip a >> $setup_log 2>&1

	attempt=1
	attempts=60
	ip a | grep "$MNIC:" | grep "state UP" >> $setup_log 2>&1
	while [ $? -ne 0 ]; do
		ip a >> $setup_log 2>&1
		if [ $attempt -gt $attempts ]; then
			echo "Network unavailable - setup cannot continue" >> $setup_log 2>&1
			exit 1
		fi
		echo "Waiting for network to come up (attempt $attempt of $attempts)" >> $setup_log 2>&1
		attempt=$((attempt + 1))
		sleep 10;
		ip a | grep "$MNIC:" | grep "state UP" >> $setup_log 2>&1
	done
	echo "Network is up on $MNIC" >> $setup_log 2>&1
fi

case "$setup_type" in
	iso | network) # Accepted values
		echo "Beginning Security Onion $setup_type install" >> $setup_log 2>&1
		;;
	*)
		echo "Invalid install type, must be 'iso' or 'network'" | tee $setup_log
		exit 1
		;;
esac

# Allow execution of SO tools during setup
export PATH=$PATH:../salt/common/tools/sbin

got_root

detect_os && detect_cloud
set_network_dev_status_list

if [ "$OS" == ubuntu ]; then
	update-alternatives --set newt-palette /etc/newt/palette.original >> $setup_log 2>&1
fi

setterm -blank 0

if [ "$setup_type" == 'iso' ] || (whiptail_you_sure); then
	true
else
	echo "User cancelled setup." | tee $setup_log
	whiptail_cancel
fi

whiptail_install_type

if [ "$install_type" = 'EVAL' ]; then
	is_node=true
	is_manager=true
	is_sensor=true
	is_eval=true
elif [ "$install_type" = 'STANDALONE' ]; then
	is_manager=true
	is_distmanager=true
	is_node=true
	is_sensor=true
elif [ "$install_type" = 'MANAGERSEARCH' ]; then
	is_manager=true
	is_distmanager=true
	is_node=true
elif [ "$install_type" = 'MANAGER' ]; then
	is_manager=true
	is_distmanager=true
elif [ "$install_type" = 'SENSOR' ]; then
	is_sensor=true
	is_minion=true
	STRELKA=1
elif [[ "$install_type" =~ ^('SEARCHNODE'|'HOTNODE'|'WARMNODE')$ ]]; then
	is_node=true
	is_minion=true
elif [ "$install_type" = 'HEAVYNODE' ]; then
	is_node=true
	is_minion=true
	is_sensor=true
elif [ "$install_type" = 'FLEET' ]; then
	is_minion=true
	is_fleet_standalone=true
	OSQUERY=1
elif [ "$install_type" = 'HELIXSENSOR' ]; then
	is_helix=true
fi

if [[ $is_master && $is_sensor ]]; then
	check_requirements "standalone"
elif [[ $is_fleet_standalone ]]; then
	check_requirements "dist" "fleet"
elif [[ $is_sensor && ! $is_eval ]]; then
	check_requirements "dist" "sensor"
elif [[ $is_distmanager || $is_minion ]]; then
	check_requirements "dist"
fi

whiptail_patch_schedule

case "$setup_type" in
	'iso')
		whiptail_set_hostname
		whiptail_management_nic
		whiptail_dhcp_or_static

		if [ "$address_type" != 'DHCP' ]; then
			whiptail_management_interface_ip
			whiptail_management_interface_mask
			whiptail_management_interface_gateway
			whiptail_management_interface_dns
			whiptail_management_interface_dns_search
		fi

		collect_adminuser_inputs
		;;
	'network')
		whiptail_network_notice
		whiptail_set_hostname
		whiptail_management_nic
		;;
esac

short_name=$(echo "$HOSTNAME" | awk -F. '{print $1}')

MINION_ID=$(echo "${short_name}_${install_type}" | tr '[:upper:]' '[:lower:]')
export MINION_ID

echo "MINION_ID = $MINION_ID" >> $setup_log 2>&1

minion_type=$(get_minion_type)

# Set any variables needed
set_default_log_size >> $setup_log 2>&1

if [[ $is_helix ]]; then
	RULESETUP=ETOPEN
	NSMSETUP=BASIC
	HNSENSOR=inherit
	MANAGERUPDATES=0
fi

if [[ $is_helix || ( $is_manager && $is_node ) ]]; then
	RULESETUP=ETOPEN
	NSMSETUP=BASIC
fi

if [[ $is_manager && $is_node ]]; then
	LSPIPELINEWORKERS=1
	LSPIPELINEBATCH=125
	LSINPUTTHREADS=1
	LSINPUTBATCHCOUNT=125
	NIDS=Suricata
	BROVERSION=ZEEK
fi

if [[ $is_node ]]; then
	CURCLOSEDAYS=30
fi

# Start user prompts
if [[ $is_helix || $is_sensor ]]; then
	whiptail_sensor_nics
	calculate_useable_cores
fi

if [[ $is_helix ||  $is_manager ]]; then
	whiptail_homenet_manager
fi

if [[ $is_helix || $is_manager || $is_node ]]; then
	set_base_heapsizes
fi

if [[ $is_manager && ! $is_eval ]]; then
	whiptail_manager_adv
	whiptail_bro_version
	whiptail_nids
	whiptail_rule_setup

	if [ "$RULESETUP" != 'ETOPEN' ]; then
		whiptail_oinkcode
	fi

	if [ "$MANAGERADV" = 'ADVANCED' ] && [ "$BROVERSION" != 'SURICATA' ]; then
		whiptail_manager_adv_service_brologs
	fi
fi

if [[ $is_manager ]]; then
	whiptail_components_adv_warning
	whiptail_enable_components
	if [[ $STRELKA == 1 ]]; then
	    whiptail_strelka_rules
	fi
	collect_webuser_inputs
	get_redirect
fi

if [[ $is_distmanager || ( $is_sensor || $is_node || $is_fleet_standalone ) && ! $is_eval ]]; then
	whiptail_manager_updates
	if [[ $setup_type == 'network' && $MANAGERUPDATES == 1 ]]; then
		whiptail_manager_updates_warning
	fi
fi

if [[ $is_minion ]]; then
	whiptail_management_server
fi

if [[ $is_distmanager ]]; then
	collect_soremote_inputs
fi

if [[ $is_sensor && ! $is_eval ]]; then
	whiptail_homenet_sensor
	whiptail_sensor_config
	if [ $NSMSETUP == 'ADVANCED' ]; then
		whiptail_bro_pins
		whiptail_suricata_pins
		whiptail_bond_nics_mtu
	else
		whiptail_basic_bro
		whiptail_basic_suri
	fi
fi

if [[ $is_node && ! $is_eval ]]; then
	whiptail_node_advanced
	if [ "$NODESETUP" == 'NODEADVANCED' ]; then
		whiptail_node_es_heap
		whiptail_node_ls_heap
		whiptail_node_ls_pipeline_worker
		whiptail_node_ls_pipline_batchsize
		whiptail_node_ls_input_threads
		whiptail_node_ls_input_batch_count
		whiptail_cur_close_days
		whiptail_log_size_limit
	else
		NODE_ES_HEAP_SIZE=$ES_HEAP_SIZE
		NODE_LS_HEAP_SIZE=$LS_HEAP_SIZE
		LSPIPELINEWORKERS=$num_cpu_cores
		LSPIPELINEBATCH=125
		LSINPUTTHREADS=1
		LSINPUTBATCHCOUNT=125
	fi
fi

if [ "$install_type" == 'FLEET' ]; then
	collect_fleetuser_inputs
	collect_fleet_custom_hostname_inputs
else
	FLEETNODEUSER=$WEBUSER
	FLEETNODEPASSWD1=$WEBPASSWD1
fi

if [[ $is_master ]]; then whiptail_so_allow; fi

whiptail_make_changes

if [[ -n "$TURBO" ]]; then
	use_turbo_proxy
fi

if [[ "$setup_type" == 'iso' ]]; then
	# Init networking so rest of install works
	set_hostname_iso
	set_management_interface

	add_admin_user
	disable_onion_user
fi

{
	set_hostname;
	set_version;
	clear_manager;
} >> $setup_log 2>&1


if [[ $is_manager ]]; then
	{
		generate_passwords;
		secrets_pillar;
		add_socore_user_manager;
	} >> $setup_log 2>&1	
fi

if [[ $is_manager && ! $is_eval ]]; then
	add_soremote_user_manager >> $setup_log 2>&1
fi

set_main_ip >> $setup_log 2>&1

host_pillar >> $setup_log 2>&1

if [[ $is_minion ]]; then
	set_updates >> $setup_log 2>&1
	copy_ssh_key >> $setup_log 2>&1
fi



# Begin install
{
	# Set initial percentage to 0
	export percentage=0

	if [[ $is_minion ]]; then
		set_progress_str 1 'Configuring firewall'
		set_initial_firewall_policy >> $setup_log 2>&1
	fi

	set_progress_str 2 'Updating packages'
	update_packages >> $setup_log 2>&1

	if [[ $is_sensor || $is_helix ]]; then
		set_progress_str 3 'Configuring sensor interface'
		configure_network_sensor >> $setup_log 2>&1
		set_progress_str 4 'Generating sensor pillar'
		sensor_pillar >> $setup_log 2>&1
	fi

	set_progress_str 5 'Installing Salt and dependencies'
	saltify 2>> $setup_log

	set_progress_str 7 'Installing Docker and dependencies'
	docker_install >> $setup_log 2>&1
	
	set_progress_str 8 'Generating patch pillar'
	patch_pillar >> $setup_log 2>&1

	set_progress_str 9 'Initializing Salt minion'
	configure_minion "$minion_type" >> $setup_log 2>&1

	if [[ $is_manager || $is_helix ]]; then
		set_progress_str 10 'Configuring Salt master'
		{
			create_local_directories;
			addtotab_generate_templates;
			copy_salt_master_config;
			setup_salt_master_dirs;
			firewall_generate_templates;
		} >> $setup_log 2>&1
		
		set_progress_str 11 'Updating sudoers file for soremote user'
		update_sudoers >> $setup_log 2>&1
		
		set_progress_str 12 'Generating manager static pillar'
		manager_static >> $setup_log 2>&1
		
		set_progress_str 13 'Generating manager pillar'
		manager_pillar >> $setup_log 2>&1
	fi


	set_progress_str 16 'Running first Salt checkin'
	salt_firstcheckin >> $setup_log 2>&1

	if [[ $is_helix ]]; then
		set_progress_str 17 'Generating the FireEye pillar'
		fireeye_pillar >> $setup_log 2>&1
	fi
	
	if [[ $is_node ]]; then
		set_progress_str 18 'Setting node type'
		set_node_type >> $setup_log 2>&1

		set_progress_str 19 'Generating search node pillar'
		elasticsearch_pillar >> $setup_log 2>&1
	fi

	if [[ $is_minion ]]; then
		set_progress_str 20 'Accepting Salt key on manager'
		accept_salt_key_remote >> $setup_log 2>&1
	fi

	if [[ $is_manager ]]; then
		set_progress_str 20 'Accepting Salt key'
		salt-key -ya "$MINION_ID" >> $setup_log 2>&1
	fi

	set_progress_str 21 'Copying minion pillars to manager'
	copy_minion_tmp_files >> $setup_log 2>&1

	set_progress_str 22 'Generating CA and checking in'
	salt_checkin >> $setup_log 2>&1

	if [[ $is_manager || $is_helix ]]; then
		set_progress_str 25 'Configuring firewall'
		set_initial_firewall_policy >> $setup_log 2>&1
		
		if [[ "$setup_type" == 'iso' ]]; then
			set_progress_str 26 'Copying containers from iso'
		else
			set_progress_str 26 'Downloading containers from the internet'
		fi

		salt-call state.apply -l info registry >> $setup_log 2>&1
		docker_seed_registry  2>> "$setup_log" # ~ 60% when finished
		
		set_progress_str 60 "$(print_salt_state_apply 'manager')"
		salt-call state.apply -l info manager >> $setup_log 2>&1

		set_progress_str 61 "$(print_salt_state_apply 'idstools')"
		salt-call state.apply -l info idstools >> $setup_log 2>&1

		set_progress_str 61 "$(print_salt_state_apply 'suricata.manager')"
		salt-call state.apply -l info suricata.manager >> $setup_log 2>&1

	fi

	set_progress_str 62 "$(print_salt_state_apply 'firewall')"
	salt-call state.apply -l info firewall >> $setup_log 2>&1

	if [ $OS = 'centos' ]; then
		set_progress_str 63 'Installing Yum utilities'
		salt-call state.apply -l info yum.packages >> $setup_log 2>&1
	fi

	set_progress_str 63 "$(print_salt_state_apply 'common')"
	salt-call state.apply -l info common >> $setup_log 2>&1

	set_progress_str 64 "$(print_salt_state_apply 'nginx')"
	salt-call state.apply -l info nginx >> $setup_log 2>&1

	if [[ $is_manager || $is_node ]]; then
		set_progress_str 64 "$(print_salt_state_apply 'elasticsearch')"
		salt-call state.apply -l info elasticsearch >> $setup_log 2>&1
	fi

	if [[ $is_sensor ]]; then
		set_progress_str 65 "$(print_salt_state_apply 'pcap')"
		salt-call state.apply -l info pcap >> $setup_log 2>&1

		set_progress_str 66 "$(print_salt_state_apply 'suricata')"
		salt-call state.apply -l info suricata >> $setup_log 2>&1

		set_progress_str 67 "$(print_salt_state_apply 'zeek')"
		salt-call state.apply -l info zeek >> $setup_log 2>&1
	fi

	if [[ $is_node ]]; then
		set_progress_str 68 "$(print_salt_state_apply 'curator')"
		salt-call state.apply -l info curator >> $setup_log 2>&1
	fi

	if [[ $is_manager ]]; then
		set_progress_str 69 "$(print_salt_state_apply 'soc')"
		salt-call state.apply -l info soc >> $setup_log 2>&1

		set_progress_str 70 "$(print_salt_state_apply 'kibana')"
		salt-call state.apply -l info kibana >> $setup_log 2>&1

		set_progress_str 71 "$(print_salt_state_apply 'elastalert')"
		salt-call state.apply -l info elastalert >> $setup_log 2>&1

		set_progress_str 72 "$(print_salt_state_apply 'soctopus')"
		salt-call state.apply -l info soctopus >> $setup_log 2>&1
	fi

	if [[ "$PLAYBOOK" = 1 ]]; then
		set_progress_str 73 "$(print_salt_state_apply 'mysql')"
		salt-call state.apply -l info mysql >> $setup_log 2>&1
		
		set_progress_str 73 "$(print_salt_state_apply 'playbook')"
		salt-call state.apply -l info playbook >> $setup_log 2>&1
		so-playbook-ruleupdate >> /root/setup_playbook_rule_update.log 2>&1 &
	fi
	
	if [[ "$OSQUERY" = 1 ]]; then
		if [[ "$PLAYBOOK" != 1 ]]; then
		set_progress_str 74 "$(print_salt_state_apply 'mysql')"
		salt-call state.apply -l info mysql >> $setup_log 2>&1
		fi

		set_progress_str 75 "$(print_salt_state_apply 'fleet')"
		salt-call state.apply -l info fleet >> $setup_log 2>&1

		set_progress_str 76 "$(print_salt_state_apply 'redis')"
		salt-call state.apply -l info redis >> $setup_log 2>&1

		if [[ $is_fleet_standalone && $FLEETCUSTOMHOSTNAME != '' ]]; then
			set_progress_str 77 "$(print_salt_state_apply 'fleet.event_update-custom-hostname')"
			pillar_override="{\"static\":{\"fleet_custom_hostname\": \"$FLEETCUSTOMHOSTNAME\"}}"
			salt-call state.apply -l info fleet.event_update-custom-hostname pillar="$pillar_override" >> $setup_log 2>&1
		fi
	
		set_progress_str 74 "$(print_salt_state_apply 'so-fleet-setup')"	
		so-fleet-setup "$FLEETNODEUSER" "$FLEETNODEPASSWD1" >> $setup_log 2>&1

	fi

	if [[ "$WAZUH" = 1 ]]; then
		set_progress_str 78 "$(print_salt_state_apply 'wazuh')"
		salt-call state.apply -l info wazuh >> $setup_log 2>&1
	fi

	if [[ "$THEHIVE" = 1 ]]; then
		set_progress_str 79 "$(print_salt_state_apply 'thehive')"
		salt-call state.apply -l info thehive >> $setup_log 2>&1
	fi

	if [[ "$STRELKA" = 1 ]]; then
		set_progress_str 80 "$(print_salt_state_apply 'strelka')"
		salt-call state.apply -l info strelka >> $setup_log 2>&1
		if [[ $STRELKARULES == 1 ]]; then
			/usr/sbin/so-yara-update >> $setup_log 2>&1	
		fi
	fi

	if [[ $is_manager || $is_helix ]]; then		
		set_progress_str 81 "$(print_salt_state_apply 'utility')"
		salt-call state.apply -l info utility >> $setup_log 2>&1
	fi

	if [[ ( $is_helix || $is_manager || $is_node ) && ! $is_eval ]]; then
		set_progress_str 82 "$(print_salt_state_apply 'logstash')"
		salt-call state.apply -l info logstash >> $setup_log 2>&1

		set_progress_str 83 "$(print_salt_state_apply 'filebeat')"
		salt-call state.apply -l info filebeat >> $setup_log 2>&1
	fi

	set_progress_str 85 'Applying finishing touches'
	filter_unused_nics >> $setup_log 2>&1
	network_setup >> $setup_log 2>&1

	if [[ $is_manager ]]; then
		set_progress_str 87 'Adding user to SOC'
		add_web_user >> $setup_log 2>&1
	fi

	set_progress_str 90 'Enabling checkin at boot'
	checkin_at_boot >> $setup_log 2>&1

	set_progress_str 95 'Verifying setup'
	salt-call -l info  state.highstate >> $setup_log 2>&1

} | progress

success=$(tail -10 $setup_log | grep Failed | awk '{ print $2}')

if [[ -n $ALLOW_ROLE && -n $ALLOW_CIDR ]]; then 
	IP=$ALLOW_CIDR so-allow -$ALLOW_ROLE >> $setup_log 2>&1
fi

if [[ $success != 0 || -n $SO_ERROR ]]; then whiptail_setup_failed
else 
	whiptail_setup_complete
	if [[ $THEHIVE == 1 ]]; then check_hive_init; fi
fi

install_cleanup >> $setup_log 2>&1

if [[ -z $SKIP_REBOOT ]]; then shutdown -r now; else exit; fi