From db60c27da23ecb1c4f2d632fc0e578cb932e55c3 Mon Sep 17 00:00:00 2001
From: reyesj2 <94730068+reyesj2@users.noreply.github.com>
Date: Tue, 22 Sep 2026 12:24:23 -0500
Subject: [PATCH] fail cleanly when endpoints-initial is missing or has
 multiple enrollment tokens

---
 .../tools/sbin/so-elastic-fleet-common        | 61 ++++++++++++++++---
 ...ic-fleet-integration-policy-elastic-defend |  5 +-
 .../so-elastic-fleet-integration-policy-load  |  4 ++
 .../so-elastic-agent-gen-installers           | 18 ++++--
 .../tools/sbin_jinja/so-elastic-fleet-setup   | 23 +------
 5 files changed, 74 insertions(+), 37 deletions(-)

diff --git a/salt/elasticfleet/tools/sbin/so-elastic-fleet-common b/salt/elasticfleet/tools/sbin/so-elastic-fleet-common
index 855a28510..2e5ad084c 100644
--- a/salt/elasticfleet/tools/sbin/so-elastic-fleet-common
+++ b/salt/elasticfleet/tools/sbin/so-elastic-fleet-common
@@ -30,6 +30,49 @@ fleet_api() {
     curl -sK /opt/so/conf/elasticsearch/curl.config -L "localhost:5601/api/fleet/${QUERYPATH}" "$@" --retry 3 --retry-delay 10 --fail 2>/dev/null
 }
 
+elastic_fleet_require_agent_policy() {
+    local AGENT_POLICY=$1
+    local POLICY_JSON
+
+    POLICY_JSON=$(fleet_api "agent_policies/$AGENT_POLICY")
+    if ! jq -e '.item.package_policies' <<<"$POLICY_JSON" >/dev/null 2>&1; then
+        echo "Error: Agent policy '$AGENT_POLICY' was not found or is not visible in the current Kibana space." >&2
+        return 1
+    fi
+
+    echo "$POLICY_JSON"
+}
+
+elastic_fleet_active_enrollment_token() {
+    local POLICY_ID=$1
+    local RESP TOKEN_COUNT API_KEY
+
+    if ! RESP=$(fleet_api "enrollment_api_keys?perPage=100" -H 'kbn-xsrf: true' -H 'Content-Type: application/json'); then
+        echo "Error: Failed to retrieve enrollment tokens for agent policy '$POLICY_ID'." >&2
+        return 1
+    fi
+
+    if ! jq -e '.list' <<<"$RESP" >/dev/null 2>&1; then
+        echo "Error: Invalid enrollment token response for agent policy '$POLICY_ID'." >&2
+        return 1
+    fi
+
+    TOKEN_COUNT=$(jq --arg pid "$POLICY_ID" '[.list[] | select(.policy_id == $pid and .active == true)] | length' <<<"$RESP")
+
+    if [ "$TOKEN_COUNT" -eq 0 ]; then
+        echo "Error: No active enrollment token found for agent policy '$POLICY_ID'." >&2
+        return 1
+    fi
+
+    if [ "$TOKEN_COUNT" -gt 1 ]; then
+        echo "Error: Found $TOKEN_COUNT active enrollment tokens for agent policy '$POLICY_ID'; expected exactly one." >&2
+        return 1
+    fi
+
+    API_KEY=$(jq -r --arg pid "$POLICY_ID" '.list[] | select(.policy_id == $pid and .active == true) | .api_key' <<<"$RESP")
+    echo "$API_KEY"
+}
+
 # Max number of concurrent Fleet write jobs (create/update). Override via env if needed.
 MAX_FLEET_JOBS=${MAX_FLEET_JOBS:-10}
 
@@ -62,15 +105,7 @@ elastic_fleet_load_integrations_dir() {
     i=0
 
     # Fetch the agent policy a single time; we look up integration ids locally below.
-    if ! POLICY_JSON=$(fleet_api "agent_policies/$AGENT_POLICY"); then
-        echo "Error: Failed to retrieve agent policy '$AGENT_POLICY'."
-        rm -f "$FAIL_FILE"
-        rm -rf "$OUT_DIR"
-        return 1
-    fi
-
-    if ! jq -e '.item.package_policies' <<<"$POLICY_JSON" >/dev/null 2>&1; then
-        echo "Error: Invalid agent policy response for '$AGENT_POLICY'."
+    if ! POLICY_JSON=$(elastic_fleet_require_agent_policy "$AGENT_POLICY"); then
         rm -f "$FAIL_FILE"
         rm -rf "$OUT_DIR"
         return 1
@@ -125,8 +160,14 @@ elastic_fleet_integration_check() {
     JSON_STRING=$2
 
     NAME=$(jq -r .name $JSON_STRING)
+    INTEGRATION_ID=""
 
-    INTEGRATION_ID=$(/usr/sbin/so-elastic-fleet-agent-policy-view "$AGENT_POLICY" | jq -r '.item.package_policies[] | select(.name=="'"$NAME"'") | .id')
+    local POLICY_JSON
+    if ! POLICY_JSON=$(elastic_fleet_require_agent_policy "$AGENT_POLICY"); then
+        return 1
+    fi
+
+    INTEGRATION_ID=$(jq -r --arg n "$NAME" '.item.package_policies[]? | select(.name==$n) | .id' <<<"$POLICY_JSON")
 
 }
 
diff --git a/salt/elasticfleet/tools/sbin/so-elastic-fleet-integration-policy-elastic-defend b/salt/elasticfleet/tools/sbin/so-elastic-fleet-integration-policy-elastic-defend
index d036f0d94..013a8089d 100755
--- a/salt/elasticfleet/tools/sbin/so-elastic-fleet-integration-policy-elastic-defend
+++ b/salt/elasticfleet/tools/sbin/so-elastic-fleet-integration-policy-elastic-defend
@@ -13,7 +13,10 @@ ERROR=false
 for INTEGRATION in /opt/so/conf/elastic-fleet/integrations/elastic-defend/*.json
 do
   printf "\n\nInitial Endpoints Policy - Loading $INTEGRATION\n"
-  elastic_fleet_integration_check "endpoints-initial" "$INTEGRATION"
+  if ! elastic_fleet_integration_check "endpoints-initial" "$INTEGRATION"; then
+    ERROR=true
+    continue
+  fi
   if [ -n "$INTEGRATION_ID" ]; then
       printf "\n\nIntegration $NAME exists - Upgrading integration policy\n"
       if ! elastic_fleet_integration_policy_upgrade "$INTEGRATION_ID"; then
diff --git a/salt/elasticfleet/tools/sbin/so-elastic-fleet-integration-policy-load b/salt/elasticfleet/tools/sbin/so-elastic-fleet-integration-policy-load
index 7c2aeb006..57569b56b 100644
--- a/salt/elasticfleet/tools/sbin/so-elastic-fleet-integration-policy-load
+++ b/salt/elasticfleet/tools/sbin/so-elastic-fleet-integration-policy-load
@@ -10,6 +10,10 @@ RETURN_CODE=0
 
 if [ ! -f /opt/so/state/eaintegrations.txt ]; then
 
+  for AGENT_POLICY in endpoints-initial so-grid-nodes_general so-grid-nodes_heavy; do
+    elastic_fleet_require_agent_policy "$AGENT_POLICY" >/dev/null || exit 1
+  done
+
   # update Fleet Server policies
   /usr/sbin/so-elastic-fleet-integration-policy-elastic-fleet-server
 
diff --git a/salt/elasticfleet/tools/sbin_jinja/so-elastic-agent-gen-installers b/salt/elasticfleet/tools/sbin_jinja/so-elastic-agent-gen-installers
index de342657f..344f2bdc7 100755
--- a/salt/elasticfleet/tools/sbin_jinja/so-elastic-agent-gen-installers
+++ b/salt/elasticfleet/tools/sbin_jinja/so-elastic-agent-gen-installers
@@ -22,12 +22,12 @@ NUM_RUNNING=$(pgrep -cf "/bin/bash /sbin/so-elastic-agent-gen-installers")
 
 for i in {1..30}
 do
-   ENROLLMENTOKEN=$(curl -K /opt/so/conf/elasticsearch/curl.config -L "localhost:5601/api/fleet/enrollment_api_keys?perPage=100" -H 'kbn-xsrf: true' -H 'Content-Type: application/json'  | jq .list | jq -r -c '.[] | select(.policy_id | contains("endpoints-initial")) | .api_key')
+   ENROLLMENTOKEN=$(elastic_fleet_active_enrollment_token "endpoints-initial")
    FLEETHOST=$(curl -K /opt/so/conf/elasticsearch/curl.config 'http://localhost:5601/api/fleet/fleet_server_hosts/grid-default' | jq -r  '.item.host_urls[]' | paste -sd ',')
 if [[ $FLEETHOST ]] && [[  $ENROLLMENTOKEN ]]; then break; else sleep 10; fi
 done
 
-if [[ -z $FLEETHOST ]] || [[ -z $ENROLLMENTOKEN ]]; then 
+if [[ -z "$FLEETHOST" ]] || [[ -z "$ENROLLMENTOKEN" ]]; then
   printf "\nFleet Host URL, Enrollment Token or Elastic Version empty - exiting..."
   printf "\nFleet Host: $FLEETHOST, Enrollment Token: $ENROLLMENTOKEN\n"
   exit 1
@@ -67,19 +67,25 @@ for GOOS in "${GOTARGETOS[@]}"; do
     GOARCH="amd64"
     if [[ $GOOS == 'darwin/arm64' ]]; then GOOS="darwin" && GOARCH="arm64"; fi
     printf "\n\n### Generating $GOOS/$GOARCH Installer...\n"
-    docker run -e CGO_ENABLED=0 -e GOOS=$GOOS -e GOARCH=$GOARCH \
+    if ! docker run -e CGO_ENABLED=0 -e GOOS=$GOOS -e GOARCH=$GOARCH \
     --mount type=bind,source=/etc/pki/tls/certs/,target=/workspace/files/cert/ \
     --mount type=bind,source=/nsm/elastic-agent-workspace/,target=/workspace/files/elastic-agent/ \
     --mount type=bind,source=/opt/so/saltstack/local/salt/elasticfleet/files/,target=/output/ \
-    {{ GLOBALS.registry_host }}:5000/{{ GLOBALS.image_repo }}/so-elastic-agent-builder:{{ GLOBALS.so_version }} go build -ldflags "-X main.fleetHostURLsList=$FLEETHOST -X main.enrollmentToken=$ENROLLMENTOKEN"  -o /output/so-elastic-agent_${GOOS}_${GOARCH}
+    {{ GLOBALS.registry_host }}:5000/{{ GLOBALS.image_repo }}/so-elastic-agent-builder:{{ GLOBALS.so_version }} go build -ldflags "-X main.fleetHostURLsList=$FLEETHOST -X main.enrollmentToken=$ENROLLMENTOKEN"  -o /output/so-elastic-agent_${GOOS}_${GOARCH}; then
+        printf "\n### ERROR: Failed to generate $GOOS/$GOARCH installer. Exiting...\n"
+        exit 1
+    fi
     printf "\n### $GOOS/$GOARCH Installer Generated...\n"
 done
 
 printf "\n\n### Generating MSI...\n"
 cp /opt/so/saltstack/local/salt/elasticfleet/files/so-elastic-agent_windows_amd64 /opt/so/saltstack/local/salt/elasticfleet/files/so-elastic-agent_windows_amd64.exe
-docker run \
+if ! docker run \
 --mount type=bind,source=/opt/so/saltstack/local/salt/elasticfleet/files/,target=/output/ -w /output \
-{{ GLOBALS.registry_host }}:5000/{{ GLOBALS.image_repo }}/so-elastic-agent-builder:{{ GLOBALS.so_version }} wixl -o so-elastic-agent_windows_amd64_msi --arch x64 /workspace/so-elastic-agent.wxs
+{{ GLOBALS.registry_host }}:5000/{{ GLOBALS.image_repo }}/so-elastic-agent-builder:{{ GLOBALS.so_version }} wixl -o so-elastic-agent_windows_amd64_msi --arch x64 /workspace/so-elastic-agent.wxs; then
+    printf "\n### ERROR: Failed to generate MSI. Exiting...\n"
+    exit 1
+fi
 printf "\n### MSI Generated...\n"
 
 # Verify installers were created
diff --git a/salt/elasticfleet/tools/sbin_jinja/so-elastic-fleet-setup b/salt/elasticfleet/tools/sbin_jinja/so-elastic-fleet-setup
index 22e0c7554..77c45c16e 100755
--- a/salt/elasticfleet/tools/sbin_jinja/so-elastic-fleet-setup
+++ b/salt/elasticfleet/tools/sbin_jinja/so-elastic-fleet-setup
@@ -202,26 +202,9 @@ fi
 ### Finalization ###
 
 # Query for Enrollment Tokens for default policies
-if ENDPOINTSENROLLMENTOKEN_RAW=$(fleet_api "enrollment_api_keys" -H 'kbn-xsrf: true' -H 'Content-Type: application/json'); then
-    ENDPOINTSENROLLMENTOKEN=$(echo "$ENDPOINTSENROLLMENTOKEN_RAW" | jq .list | jq -r -c '.[] | select(.policy_id | contains("endpoints-initial")) | .api_key')
-else
-    echo -e "\nFailed to query for Endpoints enrollment token"
-    exit 1
-fi
-
-if GRIDNODESENROLLMENTOKENGENERAL_RAW=$(fleet_api "enrollment_api_keys" -H 'kbn-xsrf: true' -H 'Content-Type: application/json'); then
-    GRIDNODESENROLLMENTOKENGENERAL=$(echo "$GRIDNODESENROLLMENTOKENGENERAL_RAW" | jq .list | jq -r -c '.[] | select(.policy_id | contains("so-grid-nodes_general")) | .api_key')
-else
-    echo -e "\nFailed to query for Grid nodes - General enrollment token"
-    exit 1
-fi
-
-if GRIDNODESENROLLMENTOKENHEAVY_RAW=$(fleet_api "enrollment_api_keys" -H 'kbn-xsrf: true' -H 'Content-Type: application/json'); then
-    GRIDNODESENROLLMENTOKENHEAVY=$(echo "$GRIDNODESENROLLMENTOKENHEAVY_RAW" | jq .list | jq -r -c '.[] | select(.policy_id | contains("so-grid-nodes_heavy")) | .api_key')
-else
-    echo -e "\nFailed to query for Grid nodes - Heavy enrollment token"
-    exit 1
-fi
+ENDPOINTSENROLLMENTOKEN=$(elastic_fleet_active_enrollment_token "endpoints-initial") || exit 1
+GRIDNODESENROLLMENTOKENGENERAL=$(elastic_fleet_active_enrollment_token "so-grid-nodes_general") || exit 1
+GRIDNODESENROLLMENTOKENHEAVY=$(elastic_fleet_active_enrollment_token "so-grid-nodes_heavy") || exit 1
 
 # Store needed data in minion pillar
 pillar_file=/opt/so/saltstack/local/pillar/minions/{{ GLOBALS.minion_id }}.sls