Add per request functionality to add HTTP request headers

src/main/java/co/elastic/support/diagnostics/commands/BaseQuery.java

@@ -78,6 +78,8 @@ List<RestEntry> execQueryList(RestClient restClient, List<RestEntry> calls, Stri
                     }
                 }
                 String fileName = subdir + SystemProperties.fileSeparator + entry.getName() + entry.getExtension();
+                restClient.setCurrentRestEntry(entry);
+                logger.debug("Execute RestClient: {} with extra headers: {}", entry.getName(), entry.getExtraHeaders());
                 RestResult restResult = restClient.execQuery(entry.getUrl(), fileName);
                 if (restResult.isValid()) {
                     logger.info(Constants.CONSOLE, "Results written to: {}", fileName);

src/main/java/co/elastic/support/diagnostics/commands/RunClusterQueries.java

@@ -35,6 +35,10 @@ public void execute(DiagnosticContext context) throws DiagnosticException {
             entries.addAll(context.elasticRestCalls.values());
             RestClient client;
             client = context.resourceCache.getRestClient(Constants.restInputHost);
+
+            // Set headers for each RestEntry
+            entries.forEach(client::setCurrentRestEntry);
+
 /*            if(ResourceCache.resourceExists(Constants.restTargetHost)){
                 client = ResourceCache.getRestClient(Constants.restTargetHost);
             }

src/main/java/co/elastic/support/diagnostics/commands/RunKibanaQueries.java

@@ -119,13 +119,13 @@ public int runBasicQueries(RestClient client, DiagnosticContext context, List<St
                     }
 
                     if (restEntry.isPageable()) {
-                        getAllPages(client, queries, context.perPage, restEntry, current.getPageableFieldName());
+                        getAllPages(client, queries, context.perPage, restEntry, current.getPageableFieldName(), context);
                     } else {
                         queries.add(restEntry);
                     }
                 }
             } else if (current.isPageable()) {
-                getAllPages(client, queries, context.perPage, entry.getValue(), current.getPageableFieldName());
+                getAllPages(client, queries, context.perPage, entry.getValue(), current.getPageableFieldName(), context);
             } else {
                 queries.add(entry.getValue());
             }
@@ -155,11 +155,12 @@ private String getPageUrl(RestEntry action, int page, int perPage, String perPag
      * @param action Kibana API name we are running
      * @param perPageField 
      */
-    public void getAllPages(RestClient client, List<RestEntry> queries, int perPage, RestEntry action, String perPageField) throws DiagnosticException {
+    public void getAllPages(RestClient client, List<RestEntry> queries, int perPage, RestEntry action, String perPageField, DiagnosticContext context) throws DiagnosticException {
         // get the values needed to the pagination (only need the total)
         String url = getPageUrl(action, 1, 1, perPageField);
 
         // get the values needed to the pagination.
+        client.setCurrentRestEntry(action);
         RestResult res = client.execQuery(url);
 
         int totalPages = 0;
@@ -177,6 +178,7 @@ public void getAllPages(RestClient client, List<RestEntry> queries, int perPage,
                 queries.add(getNewEntryPage(perPage, currentPage, action, perPageField));
             }
         }
+        client.setCurrentRestEntry(null);
     }
 
    /**
@@ -189,14 +191,16 @@ public void getAllPages(RestClient client, List<RestEntry> queries, int perPage,
     * @return new object with the API and params that need to be executed.
     */
     private RestEntry getNewEntryPage(int perPage, int page, RestEntry action, String perPageField) {
-        return new RestEntry(
-            String.format("%s_%s", action.getName(), page),
-            action.getSubdir(),
-            action.getExtension(),
-            false,
-            getPageUrl(action, page, perPage, perPageField),
-            false
+        RestEntry newEntry = new RestEntry(
+           String.format("%s_%s", action.getName(), page),
+           action.getSubdir(),
+           action.getExtension(),
+           false,
+           getPageUrl(action, page, perPage, perPageField),
+           false
         );
+        newEntry.setExtraHeaders(action.getExtraHeaders());
+        return newEntry;
     }
 
 

src/main/java/co/elastic/support/diagnostics/commands/RunLogstashQueries.java

@@ -47,6 +47,8 @@ public void execute(DiagnosticContext context) throws DiagnosticException {
             Map<String, RestEntry> entries = builder.buildEntryMap(restCalls);
 
             List<RestEntry> queries = new ArrayList<>();
+            // Set headers for each RestEntry
+            queries.forEach(client::setCurrentRestEntry);
             queries.addAll(entries.values());
             runQueries(client, queries, context.tempDir, 0, 0);
 

src/main/java/co/elastic/support/rest/RestClient.java

@@ -41,6 +41,7 @@
 import java.io.FileInputStream;
 import java.io.UnsupportedEncodingException;
 import java.security.KeyStore;
+import java.util.Arrays;
 import java.util.Map;
 
 public class RestClient implements Closeable {
@@ -54,6 +55,9 @@ public class RestClient implements Closeable {
 
     private Map<String, String> extraHeaders;
 
+    private Map<String, RestEntry> restEntries;
+    private RestEntry currentRestEntry;
+
     public RestClient(CloseableHttpClient client, HttpHost httpHost, HttpClientContext context,
             Map<String, String> extraHeaders) {
         this.client = client;
@@ -63,17 +67,46 @@ public RestClient(CloseableHttpClient client, HttpHost httpHost, HttpClientConte
     }
 
     public RestResult execQuery(String url) {
-        return new RestResult(execGet(url), url);
+        HttpRequestBase request = new HttpGet(url);
+
+        if (currentRestEntry != null && currentRestEntry.getExtraHeaders() != null) {
+            for (Map.Entry<String, String> header : currentRestEntry.getExtraHeaders().entrySet()) {
+                request.addHeader(header.getKey(), header.getValue());
+            }
+        }
+        logger.debug("Executing query: {} with currentRestEntry headers: {}", url, currentRestEntry != null ? currentRestEntry.getExtraHeaders() : "null");
+
+        return new RestResult(execGet(url, currentRestEntry), url);
     }
 
     public RestResult execQuery(String url, String fileName) {
-        return new RestResult(execGet(url), fileName, url);
+        HttpRequestBase request = new HttpGet(url);
+
+        if (currentRestEntry != null && currentRestEntry.getExtraHeaders() != null) {
+            for (Map.Entry<String, String> header : currentRestEntry.getExtraHeaders().entrySet()) {
+                request.addHeader(header.getKey(), header.getValue());
+            }
+        }
+        logger.debug("Executing query: {} with fileName: {} and currentRestEntry headers: {}", url, fileName, currentRestEntry != null ? currentRestEntry.getExtraHeaders() : "null");
+        return new RestResult(execGet(url, currentRestEntry), fileName, url);
     }
 
-    public HttpResponse execGet(String query) {
+    public HttpResponse execGet(String query, RestEntry restEntry) {
+        setCurrentRestEntry(restEntry);
         HttpGet httpGet = new HttpGet(query);
         logger.debug(query);
-        return execRequest(httpGet);
+        HttpResponse response = execRequest(httpGet);
+        setCurrentRestEntry(null);
+        return response;
+    }
+
+    public void setRestEntries(Map<String, RestEntry> restEntries) {
+        this.restEntries = restEntries;
+    }
+
+    public void setCurrentRestEntry(RestEntry restEntry) {
+        this.currentRestEntry = restEntry;
+        logger.debug("Setting currentRestEntry with headers: {}", restEntry != null ? restEntry.getExtraHeaders() : "null");
     }
 
     private HttpResponse execRequest(HttpRequestBase httpRequest) {
@@ -82,6 +115,17 @@ private HttpResponse execRequest(HttpRequestBase httpRequest) {
                 httpRequest.addHeader(entry.getKey(), entry.getValue());
             }
         }
+
+        // Then, add (or overwrite) with the per-call extra headers if RestEntry is provided
+        if (currentRestEntry != null && currentRestEntry.getExtraHeaders() != null) {
+            logger.debug("Adding extra headers from currentRestEntry: {}", currentRestEntry.getExtraHeaders());
+            for (Map.Entry<String, String> entry : currentRestEntry.getExtraHeaders().entrySet()) {
+                httpRequest.setHeader(entry.getKey(), entry.getValue());
+            }
+        }
+
+        logger.debug("Executing request with headers: " + Arrays.toString(httpRequest.getAllHeaders()) + " url " + httpRequest);
+
         try {
             return client.execute(httpHost, httpRequest, httpContext);
         } catch (HttpHostConnectException e) {
@@ -93,26 +137,44 @@ private HttpResponse execRequest(HttpRequestBase httpRequest) {
         }
     }
 
-    public HttpResponse execPost(String uri, String payload) {
+    public HttpResponse execPost(String uri, String payload, RestEntry restEntry) {
         try {
+            setCurrentRestEntry(restEntry);
             HttpPost httpPost = new HttpPost(uri);
             StringEntity entity = new StringEntity(payload);
             httpPost.setEntity(entity);
             httpPost.setHeader("Accept", "application/json");
             httpPost.setHeader("Content-type", "application/json");
             logger.debug(uri + SystemProperties.fileSeparator + payload);
-            return execRequest(httpPost);
+            HttpResponse response = execRequest(httpPost);
+            setCurrentRestEntry(null);
+            return response;
         } catch (UnsupportedEncodingException e) {
             logger.error(Constants.CONSOLE, "Error with json body.", e);
             throw new RuntimeException("Could not complete post request.");
         }
     }
 
-    public HttpResponse execDelete(String uri) {
+    public HttpResponse execDelete(String uri, RestEntry restEntry) {
+        setCurrentRestEntry(restEntry);
         HttpDelete httpDelete = new HttpDelete(uri);
         logger.debug(uri);
+        HttpResponse response = execRequest(httpDelete);
+        setCurrentRestEntry(null);
+        return response;
+    }
 
-        return execRequest(httpDelete);
+    // Add these methods to maintain backwards compatibility
+    public HttpResponse execGet(String query) {
+        return execGet(query, null);
+    }
+
+    public HttpResponse execPost(String uri, String payload) {
+        return execPost(uri, payload, null);
+    }
+
+    public HttpResponse execDelete(String uri) {
+        return execDelete(uri, null);
     }
 
     public void close() {

src/main/java/co/elastic/support/rest/RestEntry.java

@@ -7,6 +7,7 @@
 package co.elastic.support.rest;
 
 import lombok.Getter;
+import java.util.Map;
 
 @Getter
 public class RestEntry {
@@ -21,6 +22,7 @@ public class RestEntry {
     private final String pageableFieldName;
     private final boolean pageable;
     private final boolean spaceAware;
+    private Map<String, String> extraHeaders;
 
     public RestEntry(String name, String subdir, String extension, boolean retry, String url, boolean showErrors) {
         this(name, subdir, extension, retry, url, showErrors, null, false);
@@ -47,6 +49,14 @@ public RestEntry(
         this.spaceAware = spaceAware;
     }
 
+    public Map<String, String> getExtraHeaders() {
+        return extraHeaders;
+    }
+
+    public void setExtraHeaders(Map<String, String> extraHeaders) {
+        this.extraHeaders = extraHeaders;
+    }
+
     public RestEntry copyWithNewUrl(String url, String subdir) {
         return new RestEntry(name, subdir, extension, retry, url, showErrors, pageableFieldName, spaceAware);
     }

src/main/java/co/elastic/support/rest/RestEntryConfig.java

@@ -74,24 +74,35 @@ private RestEntry buildRestEntryForVersion(String name, Map<String, Object> entr
             if (semver.satisfies(urlVersion.getKey())) {
                 if (urlVersion.getValue() instanceof String) {
                     return new RestEntry(name, subdir, extension, retry, (String) urlVersion.getValue(), showErrors);
-                    // We allow it to be String,String or String,Map(url,paginate,spaceaware)
+
                 } else if (urlVersion.getValue() instanceof Map) {
                     Map<String, Object> info = (Map<String, Object>) urlVersion.getValue();
-
                     String url = (String) ObjectUtils.defaultIfNull(info.get("url"), null);
-
                     if (url == null) {
                         throw new RuntimeException("Undefined URL for REST entry (route)");
                     }
-
                     String pageableFieldName = (String) ObjectUtils.defaultIfNull(info.get("paginate"), null);
                     boolean spaceAware = (boolean) ObjectUtils.defaultIfNull(info.get("spaceaware"), false);
 
-                    return new RestEntry(name, subdir, extension, retry, url, showErrors, pageableFieldName, spaceAware);
+                    // Construct the RestEntry object first
+                    RestEntry restEntry = new RestEntry(name, subdir, extension, retry, url, showErrors, pageableFieldName, spaceAware);
+
+                    // Logic for handling extra-headers, if they exist, after RestEntry is created
+                    if (info.containsKey("extra-headers")) {
+                        Object extraHeadersObj = info.get("extra-headers");
+                        if (extraHeadersObj instanceof Map) {
+                            Map<String, String> extraHeaders = (Map<String, String>) extraHeadersObj;
+                            // Add the extra headers to the restEntry (assuming this method is present on RestEntry)
+                            restEntry.setExtraHeaders(extraHeaders);
+                        }
+                    }
+
+                    // Return the constructed RestEntry with (potentially) extra headers
+                    return restEntry;
                 }
             }
         }
 
         return null;
     }
-}
+}

src/main/resources/kibana-rest.yml

@@ -44,14 +44,20 @@ kibana_data_views:
 # https://github.com/elastic/kibana/blob/main/x-pack/plugins/security_solution/common/detection_engine/rule_monitoring/api/detection_engine_health/README.md
 kibana_detection_engine_health_cluster:
   versions:
-    ">= 8.8.2": "/internal/detection_engine/health/_cluster"
+    ">= 8.8.2":
+      url: "/internal/detection_engine/health/_cluster"
+      extra-headers:
+        elastic-api-version: "1"
 
 # Calculates health of Detection Engine and returns a health snapshot.
 # Scope: all detection rules in the current/default Kibana space.
 # https://github.com/elastic/kibana/blob/main/x-pack/plugins/security_solution/common/detection_engine/rule_monitoring/api/detection_engine_health/README.md
 kibana_detection_engine_health_space:
   versions:
-    ">= 8.8.2": "/internal/detection_engine/health/_space"
+    ">= 8.8.2":
+      url: "/internal/detection_engine/health/_space"
+      extra-headers:
+        elastic-api-version: "1"
 
 kibana_detection_engine_privileges:
   versions:
@@ -71,7 +77,10 @@ kibana_detection_engine_rules_prebuilt_status:
   versions:
     ">= 7.10.0 < 7.15.0": "/api/detection_engine/prepackaged"
     ">= 7.15.0 < 8.9.0": "/api/detection_engine/rules/prepackaged/_status"
-    ">= 8.9.0": "/internal/detection_engine/prebuilt_rules/status"
+    ">= 8.9.0":
+      url: "/internal/detection_engine/prebuilt_rules/status"
+      extra-headers:
+        elastic-api-version: "1"
 
 kibana_fleet_agents:
   versions:
@@ -96,7 +105,10 @@ kibana_fleet_agent_status:
 
 kibana_fleet_agents_current_upgrades:
   versions:
-    ">= 8.3.0": "/api/fleet/agents/current_upgrades"
+    ">= 8.3.0":
+      url: "/api/fleet/agents/current_upgrades"
+      extra-headers:
+        elastic-api-version: "2023-10-31"
 
 kibana_lists_privileges:
   versions: