From ea9bbfd1aa040910fdd11bd46823c76a59d77c81 Mon Sep 17 00:00:00 2001
From: Jason Ertel <jason.ertel@securityonionsolutions.com>
Date: Tue, 24 Nov 2020 13:53:16 -0500
Subject: [PATCH] Improve wazuh agent registration with retry logic to wait for
 manager to become ready

---
 salt/wazuh/files/agent/wazuh-register-agent | 63 +++++++++++++++------
 salt/wazuh/init.sls                         | 10 ++--
 2 files changed, 52 insertions(+), 21 deletions(-)

diff --git a/salt/wazuh/files/agent/wazuh-register-agent b/salt/wazuh/files/agent/wazuh-register-agent
index da4870e47..895fbd5d1 100755
--- a/salt/wazuh/files/agent/wazuh-register-agent
+++ b/salt/wazuh/files/agent/wazuh-register-agent
@@ -57,32 +57,56 @@ register_agent() {
   echo "Adding agent:"
   echo "Executing: curl -s -u $USER:**** -k -X POST -d 'name=$AGENT_NAME&ip=$AGENT_IP' $PROTOCOL://$API_IP:$API_PORT/agents"
   API_RESULT=$(curl -s -u $USER:"$PASSWORD" -k -X POST -d 'name='$AGENT_NAME'&ip='$AGENT_IP -L $PROTOCOL://$API_IP:$API_PORT/agents)
-  echo "Result: $API_RESULT"
-  echo -e $API_RESULT | grep -q "\"error\":0" 2>&1
+  # Get agent id and key
+  AGENT_ID=$(echo "$API_RESULT" | jq -er ".data.id")
+  GOT_ID=$?
+  AGENT_KEY=$(echo "$API_RESULT" | jq -er ".data.key")
+  GOT_KEY=$?
 
-  if [ "$?" != "0" ]; then
-    echo -e $API_RESULT | sed -rn 's/.*"message":"(.+)".*/\1/p'
+  if [[ -z "$AGENT_ID" || -z "$AGENT_KEY" || $GOT_ID -ne 0 || $GOT_KEY -ne 0 ]]; then
+    echo "Failed Result: $API_RESULT"
+    return 1
   else
-    # Get agent id and agent key
-    AGENT_ID=$(echo $API_RESULT | cut -d':' -f 4 | cut -d ',' -f 1)
-    AGENT_KEY=$(echo $API_RESULT | cut -d':' -f 5 | cut -d '}' -f 1)
-
     echo "Agent '$AGENT_NAME' with ID '$AGENT_ID' added."
     echo "Key for agent '$AGENT_ID' received."
 
     # Importing key
     echo ""
     echo "Importing authentication key:"
-    echo "y" | /var/ossec/bin/manage_agents -i $AGENT_KEY
+    echo "y" | /var/ossec/bin/manage_agents -i '$AGENT_KEY'
 
     # Restarting agent
     echo ""
     echo "Restarting:"
     echo ""
     /var/ossec/bin/ossec-control restart
+    return 0
   fi
 }
 
+wait_for_manager() {
+  echo "Waiting for Wazuh manager to become ready..."
+
+  maxAttempts=$1
+  attempts=0
+  while [[ $attempts -lt $maxAttempts ]]; do
+    attempts=$((attempts+1))
+    AGENTS_OUTPUT=$(curl -s -u $USER:"$PASSWORD" -k -X GET -L $PROTOCOL://$API_IP:$API_PORT/agents)
+    MANAGER_STATUS=$(echo "$AGENTS_OUTPUT" | jq -r ".data.items[0].status")
+    if [ "$MANAGER_STATUS" == "Active" ]; then
+      echo "Wazuh manager is active, ready to proceed."
+      return 0
+    else
+      echo "Received non-Active status response: "
+      echo "$AGENTS_OUTPUT"
+      echo
+      echo "Manager is not ready after attempt $attempts of $maxAttempts, sleeping for 30 seconds."
+      sleep 30
+    fi
+  done
+  return 1
+}
+
 remove_agent() {
   echo "Found: $AGENT_ID"
   echo "Removing previous registration for '$AGENT_NAME' using ID: $AGENT_ID ..."
@@ -141,11 +165,18 @@ if [ -f /opt/so/conf/wazuh/initial_agent_registration.log ]; then
   echo "Agent $AGENT_ID already registered!"
   exit 0
 else
-  echo "Waiting before registering agent..."
-  sleep 30s
-  register_agent
-  cleanup_creds
-  echo "Initial agent $AGENT_ID with IP $AGENT_IP registered on $DATE." > /opt/so/conf/wazuh/initial_agent_registration.log
-  exit 0
+  retries=30
+  if wait_for_manager $retries; then
+    if register_agent; then
+      cleanup_creds
+      echo "Initial agent $AGENT_ID with IP $AGENT_IP registered on $DATE." > /opt/so/conf/wazuh/initial_agent_registration.log
+      exit 0
+    else
+      echo "ERROR: Failed to register agent"
+    fi
+  else
+    echo "ERROR: Wazuh manager did not become ready after $retries attempts; unable to proceed with registration"
+  fi
 fi
-#remove_agent
+
+exit 1
diff --git a/salt/wazuh/init.sls b/salt/wazuh/init.sls
index 19afa48d7..d78d7908a 100644
--- a/salt/wazuh/init.sls
+++ b/salt/wazuh/init.sls
@@ -71,7 +71,7 @@ wazuhagentconf:
 
 wazuhdir:
  file.directory:
-   - name: /nsm/wazuh
+   - name: /nsm/wazuh/etc
    - user: 945
    - group: 945
    - makedirs: True
@@ -115,6 +115,10 @@ append_so-wazuh_so-status.conf:
     - name: /opt/so/conf/so-status/so-status.conf
     - text: so-wazuh
 
+/opt/so/conf/wazuh:
+  file.symlink:
+    - target: /nsm/wazuh/etc
+
 # Register the agent
 registertheagent:
   cmd.run:
@@ -128,10 +132,6 @@ whitelistmanager:
     - name: /usr/sbin/wazuh-manager-whitelist
     - cwd: /
 
-/opt/so/conf/wazuh:
-  file.symlink:
-    - target: /nsm/wazuh/etc
-
 wazuhagentservice:
   service.running:
     - name: wazuh-agent