Pagination refactor + Space awareness for Kibana

.gitignore

@@ -105,4 +105,9 @@ buildNumber.properties
 
 
 
-./diagnostic-output
+./diagnostic-output
+### VSCode
+.vscode/launch.json
+
+### Temporary diags
+*.zip

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

@@ -13,7 +13,6 @@
 import co.elastic.support.rest.RestClient;
 import co.elastic.support.rest.RestResult;
 import co.elastic.support.util.JsonYamlUtils;
-import co.elastic.support.util.ResourceCache;
 import co.elastic.support.util.LocalSystem;
 import co.elastic.support.util.RemoteSystem;
 import co.elastic.support.util.SystemProperties;

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

@@ -12,7 +12,6 @@
 import co.elastic.support.diagnostics.chain.DiagnosticContext;
 import co.elastic.support.rest.RestClient;
 import co.elastic.support.rest.RestEntry;
-import co.elastic.support.util.ResourceCache;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 

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

@@ -21,21 +21,21 @@
 import co.elastic.support.util.SystemUtils;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+
 import org.apache.commons.lang3.StringUtils;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.Iterator;
 
 import java.io.FileWriter;
 import java.io.IOException;
 import java.net.URI;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
+import java.nio.file.Paths;
 
 /**
  * RunKibanaQueries executes the version-dependent REST API calls against Kibana.
@@ -44,27 +44,6 @@ public class RunKibanaQueries extends BaseQuery {
 
     private static final Logger logger = LogManager.getLogger(RunKibanaQueries.class);
 
-    /**
-     * Paged actions are explicitally called out because they behave differently
-     * than normal diagnostic behaviors because they need to be called
-     * repeatedly in order to fetch all of the data.
-     */
-    private static final List<String> pagedActions = Arrays.asList(
-        new String[] {
-            "kibana_alerts",
-            "kibana_detection_engine_rules_installed",
-            "kibana_fleet_agent_policies",
-            "kibana_fleet_agents",
-            "kibana_security_endpoint_event_filters",
-            "kibana_security_endpoint_exception_items",
-            "kibana_security_endpoint_host_isolation",
-            "kibana_security_endpoint_metadata",
-            "kibana_security_endpoint_trusted_apps",
-            "kibana_security_exception_list",
-            "kibana_synthetics_monitors",
-        }
-    );
-
     /**
      * Create a new ProcessProfile object and extract the information from fileName to get PID and OS.
      * @param  tempDir where is the temporary data stored
@@ -82,6 +61,32 @@ private ProcessProfile getProfile(String tempDir, String fileName, DiagnosticCon
         return profile;
     }
 
+    private List<String> getKibanaSpacesIds(DiagnosticContext context) throws DiagnosticException {
+        RestClient restClient = context.resourceCache.getRestClient(Constants.restInputHost);
+        String url = context.fullElasticRestCalls.get("kibana_spaces").getUrl();
+        RestResult result = restClient.execQuery(url);
+
+        if (result.getStatus() != 200) {
+            throw new DiagnosticException(String.format(
+                "Kibana responded with [%d] for [%s]. Unable to proceed.",
+                result.getStatus(), url));
+        }
+
+        JsonNode spacesResponse = JsonYamlUtils.createJsonNodeFromString(result.toString());
+        if (!spacesResponse.isArray()) {
+            throw new DiagnosticException("Kibana Spaces API returned an invalid response. A list of Spaces was expected.");
+        }
+        ArrayNode arrayNode = (ArrayNode) spacesResponse;
+        List<String> spacesIds = new ArrayList<>();
+        for (JsonNode node : arrayNode) {
+            JsonNode idNode = node.path("id");
+            if(!idNode.isMissingNode()) {
+                spacesIds.add(idNode.asText());
+            }
+        }
+        return spacesIds;
+    }
+
 
     /**
      * CheckKibanaVersion (executed before) defined/set the context.elasticRestCalls.
@@ -93,39 +98,42 @@ private ProcessProfile getProfile(String tempDir, String fileName, DiagnosticCon
      * @param  context  The current diagnostic context as set in the DiagnosticService class
      * @return Number of HTTP request that will be executed.
      */
-    public int runBasicQueries(RestClient client, DiagnosticContext context) throws DiagnosticException {
+    public int runBasicQueries(RestClient client, DiagnosticContext context, List<String> spacesIds) throws DiagnosticException {
         int totalRetries = 0;
         List<RestEntry> queries = new ArrayList<>();
 
         for (Map.Entry<String, RestEntry> entry : context.elasticRestCalls.entrySet()) {
-            String actionName = entry.getValue().getName().toString();
-
-            if (pagedActions.contains(actionName)) {
-                getAllPages(client, queries, context.perPage, entry.getValue());
+            RestEntry current = entry.getValue();
+            if(current.isSpaceAware()) {
+                for(String spaceId : spacesIds) {
+                    RestEntry tmp = new RestEntry(current);
+                    // The calls made for the default Space will be written without a subpath
+                    if(!spaceId.equals("default")) {
+                        tmp.url = String.format("/s/%s%s", spaceId, tmp.getUrl());
+                        // Sanitizing the spaceId to make it "safe" for a filepath
+                        tmp.subdir = Paths.get(tmp.subdir, "space_" + spaceId.replaceAll("[^a-zA-Z0-9-_]", "_")).normalize().toString();
+                    }
+                    if(current.isPageable()) {
+                        getAllPages(client, queries, context.perPage, tmp, current.getPageableFieldName());
+                    } else {
+                        queries.add(tmp);
+                    }
+                }
             } else {
-                queries.add(entry.getValue());
+                if(current.isPageable()) {
+                    getAllPages(client, queries, context.perPage, entry.getValue(), current.getPageableFieldName());
+                } else {
+                    queries.add(entry.getValue());
+                }
             }
         }
         totalRetries = runQueries(client, queries, context.tempDir, 0, 0);
 
         return totalRetries;
     }
 
-    private String getPageUrl(RestEntry action, int page, int perPage) {
+    private String getPageUrl(RestEntry action, int page, int perPage, String perPageField) {
         String actionUrl = action.getUrl();
-        String perPageField = "per_page";
-
-        if (
-            action.getName().equals("kibana_fleet_agents") ||
-            action.getName().equals("kibana_fleet_agent_policies") ||
-            action.getName().equals("kibana_synthetics_monitors")
-        ) {
-            perPageField = "perPage";
-        } else if (
-            action.getName().equals("kibana_security_endpoint_metadata")
-        ) {
-            perPageField = "pageSize";
-        }
 
         final String querystringPrefix = URI.create(actionUrl).getQuery() != null ? "&" : "?";
         final String params = "page=" + page + "&" + perPageField + "=" + perPage;
@@ -143,10 +151,11 @@ private String getPageUrl(RestEntry action, int page, int perPage) {
      * @param queries we will store the list of queries that need to be executed
      * @param perPage  Number of docusment we reques to the API
      * @param action Kibana API name we are running
+     * @param perPageField 
      */
-    public void getAllPages(RestClient client, List<RestEntry> queries, int perPage, RestEntry action) throws DiagnosticException {
+    public void getAllPages(RestClient client, List<RestEntry> queries, int perPage, RestEntry action, String perPageField) throws DiagnosticException {
         // get the values needed to the pagination (only need the total)
-        String url = getPageUrl(action, 1, 1);
+        String url = getPageUrl(action, 1, 1, perPageField);
 
         // get the values needed to the pagination.
         RestResult res = client.execQuery(url);
@@ -159,11 +168,11 @@ public void getAllPages(RestClient client, List<RestEntry> queries, int perPage,
         }
 
         // guarantee at least one page is returned regardless of total
-        queries.add(getNewEntryPage(perPage, 1, action));
+        queries.add(getNewEntryPage(perPage, 1, action, perPageField));
 
         if (totalPages > 1) {
             for (int currentPage = 2; currentPage <= totalPages; ++currentPage) {
-                queries.add(getNewEntryPage(perPage, currentPage, action));
+                queries.add(getNewEntryPage(perPage, currentPage, action, perPageField));
             }
         }
     }
@@ -174,10 +183,18 @@ public void getAllPages(RestClient client, List<RestEntry> queries, int perPage,
     * @param  perPage how many events need to be retreived in the response
     * @param  page the apge we are requesting
     * @param  action Kibana API name we are running
+    * @param perPageField 
     * @return new object with the API and params that need to be executed.
     */
-    private RestEntry getNewEntryPage(int perPage, int page, RestEntry action) {
-        return new RestEntry(String.format("%s_%s", action.getName(), page), "", ".json", false, getPageUrl(action, page, perPage), false);
+    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
+        );
     }
 
 
@@ -299,7 +316,8 @@ public void execute(DiagnosticContext context) throws DiagnosticException {
         try {
             context.perPage         = 100;
             RestClient client       = context.resourceCache.getRestClient(Constants.restInputHost);
-            int totalRetries        = runBasicQueries(client, context);
+            List<String> spacesIds  = getKibanaSpacesIds(context);
+            int totalRetries        = runBasicQueries(client, context, spacesIds);
             filterActionsHeaders(context);
             execSystemCommands(context);
 

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

@@ -9,7 +9,6 @@
 import co.elastic.support.util.JsonYamlUtils;
 import co.elastic.support.util.LocalSystem;
 import co.elastic.support.util.RemoteSystem;
-import co.elastic.support.util.ResourceCache;
 import co.elastic.support.util.SystemCommand;
 import co.elastic.support.util.SystemProperties;
 import co.elastic.support.Constants;

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

@@ -16,7 +16,7 @@ public class RestEntry {
 
     public static final String MISSING = "missing";
 
-    public RestEntry(String name, String subdir, String extension, boolean retry, String url, boolean showErrors){
+    public RestEntry(String name, String subdir, String extension, boolean retry, String url, boolean showErrors) {
         this.name = name;
         this.subdir = subdir;
         this.extension = extension;
@@ -25,6 +25,19 @@ public RestEntry(String name, String subdir, String extension, boolean retry, St
         this.showErrors = showErrors;
     }
 
+    // Copy constructor
+    public RestEntry(RestEntry other) {
+        this.name = other.name;
+        this.subdir = other.subdir;
+        this.extension = other.extension;
+        this.retry = other.retry;
+        this.url = other.url;
+        this.showErrors = other.showErrors;
+        this.isPageable = other.isPageable;
+        this.pageableFieldName = other.pageableFieldName;
+        this.isSpaceAware = other.isSpaceAware;
+    }
+
     public String name;
 
     public String getName() {
@@ -57,8 +70,32 @@ public boolean isRetry() {
 
     public boolean showErrors = true;
 
+    private String pageableFieldName = null;
+
+    private boolean isPageable = false;
+
+    private boolean isSpaceAware = false;
+
+    public boolean isSpaceAware() {
+        return isSpaceAware;
+    }
 
+    public boolean isPageable() {
+        return isPageable;
+    }
 
+    public String getPageableFieldName() {
+        return pageableFieldName;
+    }
 
+    public void setSpaceAware(boolean isSpaceAware) {
+        this.isSpaceAware = isSpaceAware;
+    }
 
+    public void setPageableFieldName(String pageableFieldName) {
+        if(pageableFieldName != null) {
+            this.pageableFieldName = pageableFieldName;
+            this.isPageable = true;
+        }
+    }
 }

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

@@ -57,27 +57,37 @@ private RestEntry build(Map.Entry<String, Object> entry) {
         if ("full".equals(mode) == false && mode.equals(values.get("tags")) == false) {
             return null;
         }
-
-        String subdir = (String) ObjectUtils.defaultIfNull(values.get("subdir"), "");
-        String extension = (String) ObjectUtils.defaultIfNull(values.get("extension"), ".json");
-        Boolean retry = (Boolean) ObjectUtils.defaultIfNull(values.get("retry"), false);
-        Boolean showErrors = (Boolean) ObjectUtils.defaultIfNull(values.get("showErrors"), true);
-        Map<String, String> versions = (Map) values.get("versions");
-
-        String url = getVersionSpecificUrl(versions);
-
-        return new RestEntry(name, subdir, extension, retry, url, showErrors);
+
+        RestEntry tmp = new RestEntry(
+            name,
+            (String) ObjectUtils.defaultIfNull(values.get("subdir"), ""),
+            (String) ObjectUtils.defaultIfNull(values.get("extension"), ".json"),
+            (Boolean) ObjectUtils.defaultIfNull(values.get("retry"), false),
+            RestEntry.MISSING,
+            (Boolean) ObjectUtils.defaultIfNull(values.get("showErrors"), true)
+        );
+        Map<String, Object> versions = (Map) values.get("versions");
+        populateVersionSpecificUrlAndModifiers(versions, tmp);
+
+        return tmp;
     }
 
-    private String getVersionSpecificUrl(Map<String, String> versions) {
-        for (Map.Entry<String, String> urlVersion : versions.entrySet()) {
+    private void populateVersionSpecificUrlAndModifiers(Map<String, Object> versions, RestEntry entry) {
+        for (Map.Entry<String, Object> urlVersion : versions.entrySet()) {
             if (semver.satisfies(urlVersion.getKey())) {
-                return urlVersion.getValue();
+                // We allow it to be String,String or String,Map(url,paginate,spaceaware)
+                if(urlVersion.getValue() instanceof Map) {
+                    Map<String, Object> info = (Map) urlVersion.getValue();
+                    entry.url = (String) ObjectUtils.defaultIfNull(info.get("url"), RestEntry.MISSING);
+                    entry.setPageableFieldName((String)ObjectUtils.defaultIfNull(info.get("paginate"), null));
+                    entry.setSpaceAware((boolean)ObjectUtils.defaultIfNull(info.get("spaceaware"), false));
+                    return;
+                } else if (urlVersion.getValue() instanceof String){
+                    entry.url = (String) urlVersion.getValue();
+                    return;
+                }
             }
         }
-
-        // This can happen if it's older
-        return RestEntry.MISSING;
     }
 
 }