From 32e92d10ad0718ef87e25ecf8b58ca6d6e319cd1 Mon Sep 17 00:00:00 2001
From: Wes <wlambertts@gmail.com>
Date: Tue, 28 Mar 2023 00:55:56 +0000
Subject: [PATCH] Add new cluster space management scripts

---
 .../so-elasticsearch-cluster-space-configure  | 41 +++++++++++++++
 .../sbin/so-elasticsearch-cluster-space-total | 51 +++++++++++++++++++
 .../sbin/so-elasticsearch-cluster-space-used  | 23 +++++++++
 3 files changed, 115 insertions(+)
 create mode 100755 salt/common/tools/sbin/so-elasticsearch-cluster-space-configure
 create mode 100755 salt/common/tools/sbin/so-elasticsearch-cluster-space-total
 create mode 100755 salt/common/tools/sbin/so-elasticsearch-cluster-space-used

diff --git a/salt/common/tools/sbin/so-elasticsearch-cluster-space-configure b/salt/common/tools/sbin/so-elasticsearch-cluster-space-configure
new file mode 100755
index 000000000..70fb37e3e
--- /dev/null
+++ b/salt/common/tools/sbin/so-elasticsearch-cluster-space-configure
@@ -0,0 +1,41 @@
+#!/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
diff --git a/salt/common/tools/sbin/so-elasticsearch-cluster-space-total b/salt/common/tools/sbin/so-elasticsearch-cluster-space-total
new file mode 100755
index 000000000..962d515e2
--- /dev/null
+++ b/salt/common/tools/sbin/so-elasticsearch-cluster-space-total
@@ -0,0 +1,51 @@
+#!/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
+
+TOTAL_AVAILABLE_SPACE=0
+
+# 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
+
+# Set percentage of space to desired value, otherwise use a default value of 80 percent
+if [[ "$1" != "" ]]; then
+  PERCENTAGE=$1
+else
+  PERCENTAGE=80
+fi
+
+# Iterate through the output of _cat/allocation for each node in the cluster to determine the total available space
+for i in $(so-elasticsearch-query _cat/allocation | awk '{print $5}'); do
+  size=$(echo $i | grep -oE '[0-9]+')
+  unit=$(echo $i | grep -oE '[A-Za-z]+')
+  if [ $unit = "tb" ]; then
+    size=$(( size * 1024 ))
+  fi
+  TOTAL_AVAILABLE_SPACE=$(( TOTAL_AVAILABLE_SPACE + size ))
+done
+
+# Calculate the percentage of available space based on our previously defined value
+PERCENTAGE_AVAILABLE_SPACE=$(( TOTAL_AVAILABLE_SPACE*PERCENTAGE/100 ))
+echo "$PERCENTAGE_AVAILABLE_SPACE"
diff --git a/salt/common/tools/sbin/so-elasticsearch-cluster-space-used b/salt/common/tools/sbin/so-elasticsearch-cluster-space-used
new file mode 100755
index 000000000..3e8832ba0
--- /dev/null
+++ b/salt/common/tools/sbin/so-elasticsearch-cluster-space-used
@@ -0,0 +1,23 @@
+#!/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
+
+TOTAL_AVAILABLE_SPACE=0
+
+# Iterate through the output of _cat/allocation for each node in the cluster to determine the total available space
+for i in $(so-elasticsearch-query _cat/allocation | awk '{print $3}'); do
+  size=$(echo $i | grep -oE '[0-9].*' | awk '{print int($1+0.5)}')
+  unit=$(echo $i | grep -oE '[A-Za-z]+')
+  if [ $unit = "tb" ]; then
+    size=$(( size * 1024 ))
+  fi
+  TOTAL_AVAILABLE_SPACE=$(( TOTAL_AVAILABLE_SPACE + size ))
+done
+
+# Calculate the percentage of available space based on our previously defined value
+echo "$TOTAL_AVAILABLE_SPACE"