#!/bin/bash

# Copyright Security Onion Solutions LLC and/or licensed to Security Onion Solutions LLC under one
# or more contributor license agreements. Licensed under the Elastic License 2.0 as shown at 
# https://securityonion.net/license; you may not use this file except in compliance with the
# Elastic License 2.0.

if [ -f /usr/sbin/so-common ]; then
	. /usr/sbin/so-common
fi

if [ "$(id -u)" -ne 0 ]; then
	echo "This script must be run using sudo!"
	exit 1
fi

if [[ $# -lt 1 ]]; then
  echo "Usage: $0 -o=<operation> -m=[id]"
  echo ""
  echo " where <operation> is one of the following:"
  echo ""
  echo "     list: Lists all keys with hashes"
  echo "   accept: Accepts a new key and adds the minion files"
  echo "   delete: Removes the key and deletes the minion files"
  echo "   reject: Rejects a key"
  echo "     test: Perform minion test"
  echo ""
  exit 1
fi

for i in "$@"; do
	case $i in
		-o=*|--operation=*)
			OPERATION="${i#*=}"
			shift 
			;;
		-m=*|--minionid=*)
			MINION_ID="${i#*=}"
			shift 
			;;
		-e=*|--esheap=*)
			ES_HEAP_SIZE="${i#*=}"
			shift
			;;
		-n=*|--mgmtnic=*)
			MNIC="${i#*=}"
			shift
			;;
		-d=*|--description=*)
			NODE_DESCRIPTION="${i#*=}"
			shift
			;;
		-a=*|--monitor=*)
			INTERFACE="${i#*=}"
			shift
			;;
		-i=*|--ip=*)
			MAINIP="${i#*=}"
			shift
			;;		
		-*|--*)
			echo "Unknown option $i"
			exit 1
			;;
		*)
		;;
	esac
done

PILLARFILE=/opt/so/saltstack/local/pillar/minions/$MINION_ID.sls
ADVPILLARFILE=/opt/so/saltstack/local/pillar/minions/adv_$MINION_ID.sls

function getinstallinfo() {
	# Pull from file
	INSTALLVARS=$(sudo salt "$MINION_ID" cp.get_file_str /opt/so/install.txt --out=newline_values_only)
	source <(echo $INSTALLVARS)
}

function testminion() {
	# Always run on the host, since this is going to be the manager of a distributed grid, or an eval/standalone.
	# Distributed managers must run this in order for the sensor nodes to have access to the so-tcpreplay image.
	so-test
	result=$?

	# If this so-minion script is not running on the given minion ID, run so-test remotely on the sensor as well
	local_id=$(lookup_grain id)
	if [[ ! "$local_id" =~ "${MINION_ID}_" ]]; then
		salt "$MINION_ID" cmd.run 'so-test'
		result=$?
	fi

	exit $result
}

function listminions() {
	salt-key list -F --out=json
	exit $?
}

function rejectminion() {
	salt-key -y -r $MINION_ID
	exit $?
}

function acceptminion() {
	salt-key -y -a $MINION_ID
}

function deleteminion() {
	salt-key -y -d $MINION_ID
}

function deleteminionfiles () {
	rm -f $PILLARFILE
	rm -f $ADVPILLARFILE
}

# Create the minion file
function create_minion_files() {
	mkdir -p /opt/so/saltstack/local/pillar/minions
    touch $ADVPILLARFILE
	if [ -f "$PILLARFILE" ]; then
		rm $PILLARFILE
	fi
}

# Add Elastic settings to the minion file
function add_elasticsearch_to_minion() {
    printf '%s\n'\
		"elasticsearch:"\
		"  enabled: True"\
		"  esheap: '$ES_HEAP_SIZE'"\
		"  " >> $PILLARFILE
}


# Add Elastic Agent settings to the minion file
function add_elastic_agent_to_minion() {
    printf '%s\n'\
		"elasticagent:"\
		"  enabled: True"\
		"  " >> $PILLARFILE
}

# Add Elastic Fleet Server settings to the minion file
function add_fleet_to_minion() {

    # Create ES Token for Fleet server (Curl to Kibana API)
	# TODO: Add error handling
    ESTOKEN=$(curl -K /opt/so/conf/elasticsearch/curl.config -L -X POST "localhost:5601/api/fleet/service_tokens" -H 'kbn-xsrf: true' -H 'Content-Type: application/json' |  jq -r .value)

    # Write out settings to minion file
    printf '%s\n'\
		"elasticfleet:"\
		"  enabled: True"\
		"  config:"\
		"    server:"\
		"      es_token: '$ESTOKEN'"\
		"  " >> $PILLARFILE
}

# Add IDH Services info to the minion file
function add_idh_to_minion() {
	printf '%s\n'\
	"idh:"\
	"  enabled: True"\
	"  restrict_management_ip: $IDH_MGTRESTRICT"\
	"  services:" >> "$PILLARFILE"
	IFS=',' read -ra IDH_SERVICES_ARRAY <<< "$IDH_SERVICES"
	for service in ${IDH_SERVICES_ARRAY[@]}; do
	  echo "    - $service" | tr '[:upper:]' '[:lower:]' | tr -d '"' >> "$PILLARFILE"
	done
}

function add_logstash_to_minion() {
	# Create the logstash advanced pillar
	printf '%s\n'\
		"logstash:"\
		"  enabled: True"\
		"  config:"\
		"    pipeline_x_workers: $CPUCORES"\
		"  settings:"\
		"    lsheap: $LSHEAP"\
		"  " >> $PILLARFILE
}

# Security Onion Desktop
function add_desktop_to_minion() {
    printf '%s\n'\
		"desktop:"\
		"  gui:"\
		"    enabled: true"\ >> $PILLARFILE
}

# Add basic host info to the minion file
function add_host_to_minion() {
    printf '%s\n'\
    "host:"\
	"  mainip: '$MAINIP'"\
    "  mainint: '$MNIC'" >> $PILLARFILE
}

# Add sensoroni specific information - Can we pull node_adrees from the host pillar?
function add_sensoroni_to_minion() {
	printf '%s\n'\
    "sensoroni:"\
	"  enabled: True"\
	"  config:"\
	"    node_description: '${NODE_DESCRIPTION//\'/''}'"\
    "  " >> $PILLARFILE
}

# Add sensoroni specific information - Can we pull node_adrees from the host pillar?
function add_sensoroni_with_analyze_to_minion() {
	printf '%s\n'\
    "sensoroni:"\
	"  enabled: True"\
	"  config:"\
    "    analyze:"\
    "      enabled: True"\
	"    node_description: '${NODE_DESCRIPTION//\'/''}'"\
    "  " >> $PILLARFILE
}

# Sensor settings for the minion pillar
function add_sensor_to_minion() {
	echo "sensor:" >> $PILLARFILE
	echo "  interface: '$INTERFACE'" >> $PILLARFILE
	echo "  mtu: 9000" >> $PILLARFILE
	echo "zeek:" >> $PILLARFILE
	echo "  enabled: True" >> $PILLARFILE
	echo "  config:" >> $PILLARFILE
	echo "    node:" >> $PILLARFILE
	echo "      lb_procs: '$CORECOUNT'" >> $PILLARFILE
	echo "suricata:" >> $PILLARFILE
	echo "  enabled: True " >> $PILLARFILE
	echo "  config:" >> $PILLARFILE
	echo "    af-packet:" >> $PILLARFILE
	echo "      threads: '$CORECOUNT'" >> $PILLARFILE
	echo "pcap:" >> $PILLARFILE
	echo "  enabled: True" >> $PILLARFILE
	if [[ $is_pcaplimit ]]; then
		echo "  config:" >> $PILLARFILE
		echo "    diskfreepercentage: 60" >> $PILLARFILE
	fi
	echo "  " >> $PILLARFILE
}

function add_playbook_to_minion() {
    printf '%s\n'\
    "playbook:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_elastalert_to_minion() {
    printf '%s\n'\
    "elastalert:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_kibana_to_minion() {
    printf '%s\n'\
    "kibana:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_redis_to_minion() {
    printf '%s\n'\
    "redis:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_strelka_to_minion() {
    printf '%s\n'\
    "strelka:"\
	"  backend:"\
    "    enabled: True"\
	"  filestream:"\
    "    enabled: True"\
	"  frontend:"\
    "    enabled: True"\
	"  manager:"\
    "    enabled: True"\
	"  coordinator:"\
    "    enabled: True"\
	"  gatekeeper:"\
    "    enabled: True"\
	"  " >> $PILLARFILE
}

function add_curator_to_minion() {
    printf '%s\n'\
    "curator:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_telegraf_to_minion() {
    printf '%s\n'\
    "telegraf:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_influxdb_to_minion() {
    printf '%s\n'\
    "influxdb:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_nginx_to_minion() {
    printf '%s\n'\
    "nginx:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_soctopus_to_minion() {
    printf '%s\n'\
    "soctopus:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_soc_to_minion() {
    printf '%s\n'\
    "soc:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_registry_to_minion() {
    printf '%s\n'\
    "registry:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_mysql_to_minion() {
    printf '%s\n'\
    "mysql:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_kratos_to_minion() {
    printf '%s\n'\
    "kratos:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_idstools_to_minion() {
    printf '%s\n'\
    "idstools:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function add_elastic_fleet_package_registry_to_minion() {
    printf '%s\n'\
    "elastic_fleet_package_registry:"\
    "  enabled: True"\
	"  " >> $PILLARFILE
}

function create_fleet_policy() {

	JSON_STRING=$( jq -n \
					--arg NAME "FleetServer_$LSHOSTNAME" \
					--arg DESC "Fleet Server - $LSHOSTNAME" \
					'{"name": $NAME,"id":$NAME,"description":$DESC,"namespace":"default","monitoring_enabled":["logs"],"inactivity_timeout":1209600,"has_fleet_server":true}'
					)

	# Create Fleet Sever Policy
	curl -K /opt/so/conf/elasticsearch/curl.config -L -X POST "localhost:5601/api/fleet/agent_policies" -H 'kbn-xsrf: true' -H 'Content-Type: application/json' -d "$JSON_STRING"

	JSON_STRING_UPDATE=$( jq -n \
					--arg NAME "FleetServer_$LSHOSTNAME" \
					--arg DESC "Fleet Server - $LSHOSTNAME" \
					'{"name":$NAME,"description":$DESC,"namespace":"default","monitoring_enabled":["logs"],"inactivity_timeout":120,"data_output_id":"so-manager_elasticsearch"}'
					)

	# Update Fleet Policy - ES Output
	curl -K /opt/so/conf/elasticsearch/curl.config -L -X PUT "localhost:5601/api/fleet/agent_policies/FleetServer_$LSHOSTNAME" -H 'kbn-xsrf: true' -H 'Content-Type: application/json' -d "$JSON_STRING_UPDATE"
}

function update_fleet_host_urls() {
	# Query for current Fleet Host URLs & append New Fleet Node Hostname & IP
    JSON_STRING=$(curl -K /opt/so/conf/elasticsearch/curl.config 'http://localhost:5601/api/fleet/fleet_server_hosts/grid-default' | jq --arg HOSTNAME "https://$LSHOSTNAME:8220" --arg IP "https://$MAINIP:8220" '.item.host_urls += [ $HOSTNAME, $IP ]  | {"name":"grid-default","is_default":true,"host_urls": .item.host_urls}')

	# Update Fleet Host URLs
	curl -K /opt/so/conf/elasticsearch/curl.config -L -X PUT "localhost:5601/api/fleet/fleet_server_hosts/grid-default" -H 'kbn-xsrf: true' -H 'Content-Type: application/json' -d "$JSON_STRING"
}

function update_logstash_outputs() {
	# Query for current Logstash outputs & append New Fleet Node Hostname & IP
    JSON_STRING=$(curl -K /opt/so/conf/elasticsearch/curl.config 'http://localhost:5601/api/fleet/outputs/so-manager_logstash' | jq --arg HOSTNAME "$LSHOSTNAME:5055" --arg IP "$MAINIP:5055" -r '.item.hosts += [ $HOSTNAME, $IP ] | {"name":"grid-logstash","type":"logstash","hosts": .item.hosts,"is_default":true,"is_default_monitoring":true,"config_yaml":""}')

	# Update Logstash Outputs
    curl -K /opt/so/conf/elasticsearch/curl.config -L -X PUT "localhost:5601/api/fleet/outputs/so-manager_logstash" -H 'kbn-xsrf: true' -H 'Content-Type: application/json' -d "$JSON_STRING"
}

function checkMine() {
	local func=$1
	# make sure the minion sees itself in the mine since it needs to see itself for states as opposed to using salt-run
	retry 20 1 "salt '$MINION_ID' mine.get '\*' '$func'" "$MINION_ID"

}

function updateMine() {
	retry 20 1 "salt '$MINION_ID' mine.update" True
}

function createEVAL() {
	is_pcaplimit=true
	add_elasticsearch_to_minion
	add_sensor_to_minion
	add_strelka_to_minion
	add_playbook_to_minion
	add_elastalert_to_minion
	add_kibana_to_minion
	add_curator_to_minion
	add_telegraf_to_minion
	add_influxdb_to_minion
	add_nginx_to_minion
	add_soctopus_to_minion
	add_soc_to_minion
	add_registry_to_minion
	add_mysql_to_minion
	add_kratos_to_minion
	add_idstools_to_minion
	add_elastic_fleet_package_registry_to_minion
}

function createSTANDALONE() {
	is_pcaplimit=true
	add_elasticsearch_to_minion
	add_logstash_to_minion
	add_sensor_to_minion
	add_strelka_to_minion
	add_playbook_to_minion
	add_elastalert_to_minion
	add_kibana_to_minion
	add_redis_to_minion
	add_curator_to_minion
	add_telegraf_to_minion
	add_influxdb_to_minion
	add_nginx_to_minion
	add_soctopus_to_minion
	add_soc_to_minion
	add_registry_to_minion
	add_mysql_to_minion
	add_kratos_to_minion
	add_idstools_to_minion
	add_elastic_fleet_package_registry_to_minion
}

function createMANAGER() {
	add_elasticsearch_to_minion
	add_logstash_to_minion
	add_playbook_to_minion
	add_elastalert_to_minion
	add_kibana_to_minion
	add_redis_to_minion
	add_curator_to_minion
	add_telegraf_to_minion
	add_influxdb_to_minion
	add_nginx_to_minion
	add_soctopus_to_minion
	add_soc_to_minion
	add_registry_to_minion
	add_mysql_to_minion
	add_kratos_to_minion
	add_idstools_to_minion
	add_elastic_fleet_package_registry_to_minion
}

function createMANAGERSEARCH() {
	add_elasticsearch_to_minion
	add_logstash_to_minion
	add_playbook_to_minion
	add_elastalert_to_minion
	add_kibana_to_minion
	add_redis_to_minion
	add_curator_to_minion
	add_telegraf_to_minion
	add_influxdb_to_minion
	add_nginx_to_minion
	add_soctopus_to_minion
	add_soc_to_minion
	add_registry_to_minion
	add_mysql_to_minion
	add_kratos_to_minion
	add_idstools_to_minion
	add_elastic_fleet_package_registry_to_minion
}

function createIMPORT() {
	add_elasticsearch_to_minion
	add_sensor_to_minion
	add_kibana_to_minion
	add_telegraf_to_minion
	add_influxdb_to_minion
	add_nginx_to_minion
	add_soc_to_minion
	add_registry_to_minion
	add_kratos_to_minion
	add_idstools_to_minion
	add_elastic_fleet_package_registry_to_minion
}

function createFLEET() {
	add_fleet_to_minion
	add_logstash_to_minion
	create_fleet_policy
	update_fleet_host_urls
	#update_logstash_outputs
	add_telegraf_to_minion
	add_nginx_to_minion
}

function createIDH() {
	add_idh_to_minion
	add_telegraf_to_minion
}

function createHEAVYNODE() {
	is_pcaplimit=true
	add_elasticsearch_to_minion
	add_elastic_agent_to_minion
	add_logstash_to_minion
	add_sensor_to_minion
	add_strelka_to_minion
	add_redis_to_minion
	add_curator_to_minion
	add_telegraf_to_minion
}

function createSENSOR() {
	add_sensor_to_minion
	add_strelka_to_minion
	add_telegraf_to_minion
}

function createSEARCHNODE() {
	add_elasticsearch_to_minion
	add_logstash_to_minion
	add_telegraf_to_minion
}

function createRECEIVER() {
	add_logstash_to_minion
	add_redis_to_minion
	add_telegraf_to_minion
}

function createDESKTOP() {
	add_desktop_to_minion
	add_telegraf_to_minion
}

function testConnection() {
	# the minion should be trying to auth every 10 seconds so 15 seconds should be more than enough time to see this in the log
	# this retry was put in because it is possible that a minion is attempted to be pinged before it has authenticated and connected to the Salt master
	# causing the first ping to fail and typically wouldn't be successful until the second ping
	# this check may pass without the minion being authenticated if it was previously connected and the line exists in the log
	retry 15 1 "grep 'Authentication accepted from $MINION_ID' /opt/so/log/salt/master"
	local retauth=$?
	if [[ $retauth != 0 ]]; then
		echo "The Minion did not authenticate with the Salt master in the allotted time"
		echo "Deleting the key"
		deleteminion
		exit 1
	fi

	retry 15 3 "salt '$MINION_ID' test.ping" True
	local ret=$?
	if [[ $ret != 0 ]]; then
		echo "The Minion has been accepted but is not online. Try again later"
		echo "Deleting the key"
		deleteminion
		exit 1
	fi
}

if [[ "$OPERATION" = 'list' ]]; then
	listminions
fi

if [[ "$OPERATION" = 'delete' ]]; then
	deleteminionfiles
	deleteminion
fi

if [[ "$OPERATION" == 'add' || "$OPERATION" == 'setup' ]]; then
	# Skip this if its setup
	if [[ $OPERATION == 'add' ]]; then
		# Accept the salt key
		acceptminion
		# Test to see if the minion was accepted
		testConnection
		# Pull the info from the file to build what is needed
		getinstallinfo
	fi
	# Check to see if nodetype is set
	if [ -z $NODETYPE ]; then
		echo "No node type specified"
		exit 1
	fi
	create_minion_files
	add_host_to_minion
	managers=("EVAL" "STANDALONE" "IMPORT" "MANAGER" "MANAGERSEARCH")
	if echo "${managers[@]}" | grep -qw "$NODETYPE"; then
		add_sensoroni_with_analyze_to_minion
	else
		add_sensoroni_to_minion
	fi

	create$NODETYPE
	echo "Minion file created for $MINION_ID"	

	if [[ "$OPERATION" == 'add' ]]; then
		# tell the minion to populate the mine with data from mine_functions which is populated during setup
		# this only needs to happen on non managers since they handle this during setup
		# and they need to wait for ca creation to update the mine
		updateMine
		checkMine "network.ip_addrs"
		# run this async so the cli doesn't wait for a return
		salt "$MINION_ID" state.highstate --async
	fi
fi

if [[ "$OPERATION" = 'test' ]]; then
	testminion
fi
