diff --git a/src/main/java/co/elastic/support/diagnostics/commands/BaseQuery.java b/src/main/java/co/elastic/support/diagnostics/commands/BaseQuery.java index 3b50cad5..bfd090ff 100644 --- a/src/main/java/co/elastic/support/diagnostics/commands/BaseQuery.java +++ b/src/main/java/co/elastic/support/diagnostics/commands/BaseQuery.java @@ -78,6 +78,8 @@ List execQueryList(RestClient restClient, List 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); diff --git a/src/main/java/co/elastic/support/diagnostics/commands/RunClusterQueries.java b/src/main/java/co/elastic/support/diagnostics/commands/RunClusterQueries.java index ed6eaec8..6b6b59bf 100644 --- a/src/main/java/co/elastic/support/diagnostics/commands/RunClusterQueries.java +++ b/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); } diff --git a/src/main/java/co/elastic/support/diagnostics/commands/RunKibanaQueries.java b/src/main/java/co/elastic/support/diagnostics/commands/RunKibanaQueries.java index 136cd779..d1af9e77 100644 --- a/src/main/java/co/elastic/support/diagnostics/commands/RunKibanaQueries.java +++ b/src/main/java/co/elastic/support/diagnostics/commands/RunKibanaQueries.java @@ -119,13 +119,13 @@ public int runBasicQueries(RestClient client, DiagnosticContext context, List queries, int perPage, RestEntry action, String perPageField) throws DiagnosticException { + public void getAllPages(RestClient client, List 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 queries, int perPage, queries.add(getNewEntryPage(perPage, currentPage, action, perPageField)); } } + client.setCurrentRestEntry(null); } /** @@ -189,14 +191,16 @@ public void getAllPages(RestClient client, List 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; } diff --git a/src/main/java/co/elastic/support/diagnostics/commands/RunLogstashQueries.java b/src/main/java/co/elastic/support/diagnostics/commands/RunLogstashQueries.java index c186d2d8..51c667c8 100644 --- a/src/main/java/co/elastic/support/diagnostics/commands/RunLogstashQueries.java +++ b/src/main/java/co/elastic/support/diagnostics/commands/RunLogstashQueries.java @@ -47,6 +47,8 @@ public void execute(DiagnosticContext context) throws DiagnosticException { Map entries = builder.buildEntryMap(restCalls); List queries = new ArrayList<>(); + // Set headers for each RestEntry + queries.forEach(client::setCurrentRestEntry); queries.addAll(entries.values()); runQueries(client, queries, context.tempDir, 0, 0); diff --git a/src/main/java/co/elastic/support/rest/RestClient.java b/src/main/java/co/elastic/support/rest/RestClient.java index 85feb33b..a3f46086 100644 --- a/src/main/java/co/elastic/support/rest/RestClient.java +++ b/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 extraHeaders; + private Map restEntries; + private RestEntry currentRestEntry; + public RestClient(CloseableHttpClient client, HttpHost httpHost, HttpClientContext context, Map 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 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 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 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 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() { diff --git a/src/main/java/co/elastic/support/rest/RestEntry.java b/src/main/java/co/elastic/support/rest/RestEntry.java index f820e23e..39b2062b 100644 --- a/src/main/java/co/elastic/support/rest/RestEntry.java +++ b/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 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 getExtraHeaders() { + return extraHeaders; + } + + public void setExtraHeaders(Map extraHeaders) { + this.extraHeaders = extraHeaders; + } + public RestEntry copyWithNewUrl(String url, String subdir) { return new RestEntry(name, subdir, extension, retry, url, showErrors, pageableFieldName, spaceAware); } diff --git a/src/main/java/co/elastic/support/rest/RestEntryConfig.java b/src/main/java/co/elastic/support/rest/RestEntryConfig.java index 83124ef9..f9d2b89a 100644 --- a/src/main/java/co/elastic/support/rest/RestEntryConfig.java +++ b/src/main/java/co/elastic/support/rest/RestEntryConfig.java @@ -74,24 +74,35 @@ private RestEntry buildRestEntryForVersion(String name, Map 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 info = (Map) 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 extraHeaders = (Map) 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; } -} +} \ No newline at end of file diff --git a/src/main/resources/kibana-rest.yml b/src/main/resources/kibana-rest.yml index 993b4220..327106c3 100644 --- a/src/main/resources/kibana-rest.yml +++ b/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: