#!/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-whiptail"
source "./so-variables"
source "./so-common-functions"

setup_type=$1
export setup_type

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

# Allow execution of SO tools during setup
PATH=$PATH:"$(dirname "$0")"/../salt/common/tools/sbin
export PATH

date -u > $setup_log 2>&1

got_root
detect_os

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

setterm -blank 0

if ! (whiptail_you_sure); then
	echo "User cancelled setup.">> $setup_log 2>&1
	whiptail_cancel
fi

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

		# Init networking so rest of install works
		set_hostname_iso
		set_management_interface

		collect_adminuser_inputs
		add_admin_user
		disable_onion_user

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

whiptail_install_type

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

whiptail_patch_schedule

filter_unused_nics
minion_type=$(get_minion_type)

if [ "$install_type" = 'EVAL' ]; then
	is_node=true
	is_master=true
	is_sensor=true
	is_eval=true
elif [ "$install_type" = 'MASTERSEARCH' ]; then
	is_master=true
	is_distmaster=true
	is_node=true
elif [ "$install_type" = 'MASTER' ]; then
	is_master=true
	is_distmaster=true
elif [ "$install_type" = 'SENSOR' ]; then
	is_sensor=true
	STRELKA=1
elif [[ "$install_type" =~ ^('SEARCHNODE'|'PARSINGNODE'|'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
	OSQUERY=1
elif [ "$install_type" = 'HELIXSENSOR' ]; then
	is_helix=true
fi

# Set any constants needed
if [[ $is_helix ]]; then
	RULESETUP=ETOPEN
	NSMSETUP=BASIC
	HNSENSOR=inherit
	MASTERUPDATES=0
fi
if [[ $is_helix || ( $is_master && $is_node ) ]]; then
	RULESETUP=ETOPEN
	NSMSETUP=BASIC
fi
if [[ $is_master && $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_bond_nics
	calculate_useable_cores
fi

if [[ $is_helix ||  $is_master ]]; then
	whiptail_homenet_master
fi

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

if [[ $is_master && ! $is_eval ]]; then
	whiptail_master_adv
	whiptail_bro_version
	whiptail_nids
	whiptail_rule_setup

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

	if [ "$MASTERADV" = 'ADVANCED' ] && [ "$BROVERSION" != 'SURICATA' ]; then
		whiptail_master_adv_service_brologs
	fi

	get_redirect
fi

if [[ $is_master ]]; then
	whiptail_components_adv_warning
	whiptail_enable_components
fi

if [[ $is_distmaster || ( $is_sensor || $is_node ) && ! $is_eval ]]; then
	whiptail_master_updates
fi

if [[ $is_minion ]]; then
	whiptail_management_server
fi

if [[ $is_distmaster ]]; then
	whiptail_create_soremote_user
	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=$cpu_cores
		LSPIPELINEBATCH=125
		LSINPUTTHREADS=1
		LSINPUTBATCHCOUNT=125
	fi
fi

whiptail_make_changes
set_hostname
set_version
clear_master

if [[ $is_master && ! $is_eval ]]; then
	generate_passwords
	secrets_pillar
	add_socore_user_master
	add_soremote_user_master
fi

if [[ $is_helix || $is_master ]]; then
	set_main_ip
fi

if [[ $is_minion ]]; then
	set_updates
fi

# Begin install
{
	progress_str_printer &

	if [[ $is_sensor || $is_helix ]]; then
		set_progress_str 1 'Creating bond interface'
		create_sensor_bond
		
		set_progress_str 2 'Generating the sensor pillar'
		sensor_pillar
	fi

	set_progress_str 2 'Installing Salt and dependencies'
	saltify
	
	set_progress_str 10 'Installing Docker and dependencies'
	docker_install
	
	set_progress_str 12 'Initializing Salt minion'
	configure_minion "$minion_type"

	if [[ $is_master || $is_helix ]]; then
		set_progress_str 13 'Configuring Salt master'
		copy_master_config
		setup_salt_master_dirs
		
		set_progress_str 15 'Updating sudoers file for soremote user'
		update_sudoers
		
		set_progress_str 16 'Generating master static pillar'
		master_static
		
		set_progress_str 17 'Generating master pillar'
		master_pillar

		set_progress_str 18 'Accepting Salt key'
		salt-key -ya "$MINION_ID" >> "$setup_log" 2>&1
	fi

	if [[ $is_helix ]]; then
		set_progress_str 19 'Generating the FireEye pillar'
		fireeye_pillar
	fi

	if [[ $is_minion ]]; then
		set_progress_str 20 'Accepting salt key on master'
		accept_salt_key_remote
	fi

	set_progress_str 20 'Copying minion pillars to master'
	copy_minion_tmp_files

	set_progress_str 21 'Running intial Salt highstate'
	salt-call state.highstate -l info >> "$setup_log" 2>&1

	set_progress_str 23 'Configuring firewall'
	set_initial_firewall_policy

	set_progress_str 25 'Generating CA'
	gen_ca

	if [[ $is_node ]]; then
		set_progress_str 25 'Setting node type'
		set_node_type

		set_progress_str 26 'Generating search node pillar'
		node_pillar

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

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

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

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

	if [[ $is_master || $is_helix ]]; then
		set_progress_str 30 'Downloading containers from the internet'
		salt-call state.apply -l info registry >> "$setup_log" 2>&1
		docker_seed_registry # ~ 65% when finished
	fi

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

	if [[ "$OSQUERY" = 1 ]]; then
		set_progress_str 66 "$(print_salt_state_apply 'fleet')"
		salt-call state.apply -l info fleet >> $setup_log 2>&1

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

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

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

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

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

	if [[ $is_master || $is_helix ]]; then
		set_progress_str 72 "$(print_salt_state_apply 'master')"
		salt-call state.apply -l info master >> $setup_log 2>&1

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

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

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

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

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

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

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

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

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

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

	set_progress_str 90 'Applying finishing touches'
	filter_unused_nics
	network_setup

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

} | whiptail --title "Hybrid Hunter Install" --gauge "Please wait while installing" 6 60 0

success=$(tail -10 $setup_log | grep Failed | awk '{ print $2}')
if [[ "$success" = 0 ]]; then
	whiptail_setup_complete
	if [[ $THEHIVE == '1' ]]; then
		check_hive_init_then_reboot
	else
		shutdown -r now
	fi
else
	whiptail_setup_failed
	shutdown -r now
fi