From 3e923f671b9daa6570907eba844a135eb13b6de4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81lex=20Ruiz?= <alejandro.ruiz.becerra@wazuh.com>
Date: Wed, 8 Nov 2023 18:29:03 +0100
Subject: [PATCH] Fix bug with -i option (#51)

* Fix bug with -i option

* Improve error handling
---
 scripts/indexer-ism-init.sh | 72 +++++++++++++++++++++++++++----------
 1 file changed, 54 insertions(+), 18 deletions(-)

diff --git a/scripts/indexer-ism-init.sh b/scripts/indexer-ism-init.sh
index eabdb81ca041a..15998ed84fdaf 100644
--- a/scripts/indexer-ism-init.sh
+++ b/scripts/indexer-ism-init.sh
@@ -85,10 +85,15 @@ function generate_rollover_template() {
 function load_templates() {
     # Note: the wazuh-template.json could also be loaded here.
     for alias in "${aliases[@]}"; do
-        echo "TEMPLATES AND POLICIES - Uploading ${alias} template"
-        generate_rollover_template "${alias}" | curl -s -k ${C_AUTH} \
-            -X PUT "${INDEXER_URL}/_template/${alias}-rollover" -o /dev/null \
-            -H 'Content-Type: application/json' -d @-
+        generate_rollover_template "${alias}" |
+            if ! curl -s -k ${C_AUTH} \
+                -X PUT "${INDEXER_URL}/_template/${alias}-rollover" -o /dev/null \
+                -H 'Content-Type: application/json' -d @-; then
+                echo "Error uploading ${alias} template"
+                return 1
+            else
+                echo "${alias} template uploaded"
+            fi
     done
 }
 
@@ -110,15 +115,21 @@ function upload_rollover_policy() {
 
     # Check if the ${POLICY_NAME} ISM policy was loaded (404 error if not found)
     if [[ "${policy_exists}" == "404" ]]; then
-        echo "TEMPLATES AND POLICIES - Uploading ${POLICY_NAME} ISM policy"
-        curl -s -k ${C_AUTH} -o /dev/null \
+        if ! curl -s -k ${C_AUTH} -o /dev/null \
             -X PUT "${INDEXER_URL}/_plugins/_ism/policies/${POLICY_NAME}" \
-            -H 'Content-Type: application/json' -d "$(generate_rollover_policy)"
+            -H 'Content-Type: application/json' \
+            -d "$(generate_rollover_policy)"; then
+            echo "Error uploading ${POLICY_NAME} policy"
+            return 1
+        else
+            echo "${POLICY_NAME} policy uploaded"
+        fi
     else
         if [[ "${policy_exists}" == "200" ]]; then
-            echo "TEMPLATES AND POLICIES - ${POLICY_NAME} policy already exists"
+            echo "${POLICY_NAME} policy already exists"
         else
-            echo "TEMPLATES AND POLICIES - Error uploading ${POLICY_NAME} policy"
+            echo "Error checking if ${POLICY_NAME} exists"
+            return 1
         fi
     fi
 }
@@ -158,9 +169,15 @@ function generate_write_index_alias() {
 #   1. The alias. String.
 #########################################################################
 function create_write_index() {
-    curl -s -k ${C_AUTH} -o /dev/null \
+    if ! curl -s -k ${C_AUTH} -o /dev/null \
         -X PUT "$INDEXER_URL/%3C${1}-4.x-%7Bnow%2Fd%7D-000001%3E?pretty" \
-        -H 'Content-Type: application/json' -d "$(generate_write_index_alias "${1}")"
+        -H 'Content-Type: application/json' \
+        -d "$(generate_write_index_alias "${1}")"; then
+        echo "Error creating ${1} write index"
+        exit 1
+    else
+        echo "${1} write index created"
+    fi
 }
 
 #########################################################################
@@ -169,7 +186,6 @@ function create_write_index() {
 #   1. List of aliases to initialize.
 #########################################################################
 function create_indices() {
-    echo "TEMPLATES AND POLICIES - Creating write indices"
     for alias in "${aliases[@]}"; do
         # Check if there are any write indices for the current alias
         write_index_exists=$(check_for_write_index "${alias}")
@@ -181,7 +197,6 @@ function create_indices() {
     done
 }
 
-
 #########################################################################
 # Shows usage help.
 #########################################################################
@@ -209,9 +224,15 @@ function show_help() {
     echo -e "        -p, --indexer-password <password>"
     echo -e "                Specifies the Wazuh indexer admin user password."
     echo -e ""
+    echo -e "        -P, --priority <priority>"
+    echo -e "                Specifies the policy's priority."
+    echo -e ""
     echo -e "        -s, --min-shard-size <shard-size>"
     echo -e "                Set the minimum shard size in GB. By default 25."
     echo -e ""
+    echo -e "        -v, --verbose"
+    echo -e "                Set verbose mode. Prints more information."
+    echo -e ""
     exit 1
 }
 
@@ -252,6 +273,7 @@ function main() {
                 show_help
             else
                 INDEXER_HOSTNAME="${2}"
+                INDEXER_URL="https://${INDEXER_HOSTNAME}:9200"
                 shift 2
             fi
             ;;
@@ -274,6 +296,19 @@ function main() {
                 shift 2
             fi
             ;;
+        "-P" | "--priority")
+            if [ -z "${2}" ]; then
+                echo "Error on arguments. Probably missing <priority> after -P|--priority"
+                show_help
+            else
+                ISM_PRIORITY="${2}"
+                shift 2
+            fi
+            ;;
+        "-v" | "--verbose")
+            set -x
+            shift
+            ;;
         *)
             echo "Unknow option: ${1}"
             show_help
@@ -282,13 +317,14 @@ function main() {
     done
 
     # Load the Wazuh Indexer templates
-    load_templates
-
     # Upload the rollover policy
-    upload_rollover_policy
-
     # Create the initial write indices
-    create_indices "${aliases[@]}"
+    if load_templates && upload_rollover_policy && create_indices "${aliases[@]}"; then
+        echo "Indexer ISM initialization finished successfully"
+    else
+        echo "Indexer ISM initialization failed"
+        exit 1
+    fi
 }
 
 main "$@"