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

. /usr/sbin/so-common

# Determine available disk space
{% import_yaml 'elasticsearch/defaults.yaml' as ELASTICDEFAULTS %}
{% set ELASTICMERGED = salt['pillar.get']('elasticsearch:retention', ELASTICDEFAULTS.elasticsearch.retention, merge=true) %}

# Wait for ElasticSearch to initialize
#COUNT=0
ELASTICSEARCH_CONNECTED="no"
while [[ "$COUNT" -le 240 ]]; do
      so-elasticsearch-query / -k --output /dev/null --silent --head --fail
        if [ $? -eq 0 ]; then
                ELASTICSEARCH_CONNECTED="yes"
                break
        else
                ((COUNT+=1))
                sleep 1
        fi
done
if [ "$ELASTICSEARCH_CONNECTED" == "no" ]; then
        echo
        echo -e "Connection attempt timed out.  Unable to connect to ElasticSearch.  \nPlease try: \n  -checking log(s) in /var/log/elasticsearch/\n  -running 'sudo docker ps' \n  -running 'sudo so-elastic-restart'"
        echo
        exit 1
fi

AVAILABLE_SPACE=$(/usr/sbin/so-elasticsearch-cluster-space-total {{ ELASTICMERGED.retention_pct }})
ELASTICSEARCH_PILLAR="/opt/so/saltstack/local/pillar/elasticsearch/soc_elasticsearch.sls"
if grep -q log_size_limit $ELASTICSEARCH_PILLAR ; then
  sed -i s"/log_size_limit:.*/log_size_limit: $AVAILABLE_SPACE/" $ELASTICSEARCH_PILLAR
else
  echo "  retention:" >> $ELASTICSEARCH_PILLAR
  echo "    log_size_limit: $AVAILABLE_SPACE" >> $ELASTICSEARCH_PILLAR
fi