Exim - Site export / prune

Signed-off-by: Simon Bennetts <[email protected]>

addOns/exim/CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ## Unreleased
 ### Added
 - Add Automation Framework job to export data (e.g. HAR, URLs).
+- Support for Sites Tree export and prune.
 
 ### Changed
 - Update dependency.

addOns/exim/src/main/java/org/zaproxy/addon/exim/Exporter.java

@@ -19,6 +19,7 @@
  */
 package org.zaproxy.addon.exim;
 
+import java.io.BufferedWriter;
 import java.io.IOException;
 import java.io.Writer;
 import java.nio.charset.StandardCharsets;
@@ -30,7 +31,10 @@
 import org.parosproxy.paros.db.DatabaseException;
 import org.parosproxy.paros.model.HistoryReference;
 import org.parosproxy.paros.model.Model;
+import org.zaproxy.addon.exim.ExporterOptions.Source;
+import org.zaproxy.addon.exim.ExporterOptions.Type;
 import org.zaproxy.addon.exim.har.HarExporter;
+import org.zaproxy.addon.exim.sites.SitesTreeHandler;
 import org.zaproxy.zap.model.Context;
 import org.zaproxy.zap.utils.Stats;
 
@@ -82,25 +86,24 @@ private ExporterResult exportImpl(ExporterOptions options) {
                         StandardOpenOption.CREATE,
                         StandardOpenOption.TRUNCATE_EXISTING)) {
 
-            List<Integer> historyIds =
-                    model.getDb()
-                            .getTableHistory()
-                            .getHistoryIdsOfHistType(
-                                    model.getSession().getSessionId(), getHistoryTypes(options));
-
-            ExporterType type = createExporterType(options);
-            type.begin(writer);
-            Context context = options.getContext();
-            for (Integer id : historyIds) {
-                HistoryReference ref = new HistoryReference(id, true);
-                if (context != null && !context.isInContext(ref)) {
-                    continue;
+            if (Source.SITESTREE.equals(options.getSource())) {
+                if (!Type.YAML.equals(options.getType())) {
+                    result.addError(
+                            Constant.messages.getString(
+                                    "exim.exporter.error.type.sitestree", options.getType()));
+                } else {
+                    SitesTreeHandler.exportSitesTree(writer, result);
+                }
+            } else {
+                if (Type.YAML.equals(options.getType())) {
+                    result.addError(
+                            Constant.messages.getString(
+                                    "exim.exporter.error.type.messages", options.getSource()));
+                } else {
+                    exportMessagesImpl(writer, result, options);
                 }
-
-                result.incrementCount();
-                type.write(writer, ref);
             }
-            type.end(writer);
+
         } catch (IOException e) {
             result.addError(
                     Constant.messages.getString("exim.exporter.error.io", e.getLocalizedMessage()),
@@ -114,6 +117,30 @@ private ExporterResult exportImpl(ExporterOptions options) {
         return result;
     }
 
+    private void exportMessagesImpl(
+            BufferedWriter writer, ExporterResult result, ExporterOptions options)
+            throws DatabaseException, IOException {
+        List<Integer> historyIds =
+                model.getDb()
+                        .getTableHistory()
+                        .getHistoryIdsOfHistType(
+                                model.getSession().getSessionId(), getHistoryTypes(options));
+
+        ExporterType type = createExporterType(options);
+        type.begin(writer);
+        Context context = options.getContext();
+        for (Integer id : historyIds) {
+            HistoryReference ref = new HistoryReference(id, true);
+            if (context != null && !context.isInContext(ref)) {
+                continue;
+            }
+
+            result.incrementCount();
+            type.write(writer, ref);
+        }
+        type.end(writer);
+    }
+
     private static boolean isValid(Path file, ExporterResult result) {
         if (Files.exists(file)) {
             if (!Files.isRegularFile(file)) {

addOns/exim/src/main/java/org/zaproxy/addon/exim/ExporterOptions.java

@@ -168,7 +168,9 @@ public enum Type {
         /** The messages are exported as an HAR. */
         HAR,
         /** The messages are exported as URLs. */
-        URL;
+        URL,
+        /** The SiteTree will be exported as YAML. */
+        YAML;
 
         private String id;
         private String name;
@@ -200,6 +202,9 @@ public static Type fromString(String value) {
             if (URL.id.equalsIgnoreCase(value)) {
                 return URL;
             }
+            if (YAML.id.equalsIgnoreCase(value)) {
+                return YAML;
+            }
             return HAR;
         }
     }
@@ -209,7 +214,9 @@ public enum Source {
         /** Exports the messages proxied and manually accessed by the user. */
         HISTORY,
         /** Exports all messages accessed, includes temporary messages. */
-        ALL;
+        ALL,
+        /** Exports the Sites tree, only in yaml format */
+        SITESTREE;
 
         private String id;
         private String name;
@@ -241,6 +248,9 @@ public static Source fromString(String value) {
             if (ALL.id.equalsIgnoreCase(value)) {
                 return ALL;
             }
+            if (SITESTREE.id.equalsIgnoreCase(value)) {
+                return SITESTREE;
+            }
             return HISTORY;
         }
     }

addOns/exim/src/main/java/org/zaproxy/addon/exim/ExporterResult.java

@@ -44,7 +44,7 @@ public int getCount() {
         return count;
     }
 
-    void incrementCount() {
+    public void incrementCount() {
         this.count++;
     }
 

addOns/exim/src/main/java/org/zaproxy/addon/exim/ExtensionExim.java

@@ -37,6 +37,8 @@
 import org.zaproxy.addon.exim.har.PopupMenuItemSaveHarMessage;
 import org.zaproxy.addon.exim.log.MenuItemImportLogs;
 import org.zaproxy.addon.exim.pcap.MenuItemImportPcap;
+import org.zaproxy.addon.exim.sites.MenuPruneSites;
+import org.zaproxy.addon.exim.sites.MenuSaveSites;
 import org.zaproxy.addon.exim.urls.MenuItemImportUrls;
 
 public class ExtensionExim extends ExtensionAdaptor {
@@ -110,6 +112,9 @@ public void hook(ExtensionHook extensionHook) {
             MainMenuBar menuBar = getView().getMainFrame().getMainMenuBar();
             menuBar.add(getMenuExport(), menuBar.getMenuCount() - 2); // Before Online and Help
 
+            getMenuExport().add(new MenuSaveSites());
+            extensionHook.getHookMenu().addToolsMenuItem(new MenuPruneSites());
+
             extensionHook.getHookMenu().addImportMenuItem(new MenuImportHar());
             extensionHook.getHookMenu().addImportMenuItem(new MenuItemImportUrls());
             extensionHook.getHookMenu().addImportMenuItem(new MenuItemImportLogs());

addOns/exim/src/main/java/org/zaproxy/addon/exim/ImportExportApi.java

@@ -26,7 +26,9 @@
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import net.sf.json.JSONObject;
 import org.apache.commons.httpclient.URI;
 import org.apache.logging.log4j.LogManager;
@@ -45,6 +47,8 @@
 import org.zaproxy.addon.exim.har.HarImporter;
 import org.zaproxy.addon.exim.har.HarUtils;
 import org.zaproxy.addon.exim.log.LogsImporter;
+import org.zaproxy.addon.exim.sites.PruneSiteResult;
+import org.zaproxy.addon.exim.sites.SitesTreeHandler;
 import org.zaproxy.addon.exim.urls.UrlsImporter;
 import org.zaproxy.zap.extension.api.API;
 import org.zaproxy.zap.extension.api.ApiAction;
@@ -54,6 +58,7 @@
 import org.zaproxy.zap.extension.api.ApiOther;
 import org.zaproxy.zap.extension.api.ApiResponse;
 import org.zaproxy.zap.extension.api.ApiResponseElement;
+import org.zaproxy.zap.extension.api.ApiResponseSet;
 import org.zaproxy.zap.network.HttpRedirectionValidator;
 import org.zaproxy.zap.network.HttpRequestConfig;
 import org.zaproxy.zap.utils.ApiUtils;
@@ -76,6 +81,8 @@ public class ImportExportApi extends ApiImplementor {
     private static final String ACTION_IMPORT_URLS = "importUrls";
     private static final String ACTION_IMPORT_ZAP_LOGS = "importZapLogs";
     private static final String ACTION_IMPORT_MODSEC2_LOGS = "importModsec2Logs";
+    private static final String ACTION_EXPORT_SITES = "exportSitesTree";
+    private static final String ACTION_PRUNE_SITES = "pruneSitesTree";
 
     private static final String OTHER_EXPORT_HAR = "exportHar";
     private static final String OTHER_EXPORT_HAR_BY_ID = "exportHarById";
@@ -90,6 +97,8 @@ public ImportExportApi() {
         this.addApiAction(new ApiAction(ACTION_IMPORT_ZAP_LOGS, new String[] {PARAM_FILE_PATH}));
         this.addApiAction(
                 new ApiAction(ACTION_IMPORT_MODSEC2_LOGS, new String[] {PARAM_FILE_PATH}));
+        this.addApiAction(new ApiAction(ACTION_EXPORT_SITES, new String[] {PARAM_FILE_PATH}));
+        this.addApiAction(new ApiAction(ACTION_PRUNE_SITES, new String[] {PARAM_FILE_PATH}));
 
         this.addApiOthers(
                 new ApiOther(
@@ -133,6 +142,24 @@ public ApiResponse handleApiAction(String name, JSONObject params) throws ApiExc
                 LogsImporter logsImporter =
                         new LogsImporter(file, LogsImporter.LogType.MOD_SECURITY_2);
                 return handleFileImportResponse(logsImporter.isSuccess(), file);
+            case ACTION_EXPORT_SITES:
+                file = new File(ApiUtils.getNonEmptyStringParam(params, PARAM_FILE_PATH));
+                try {
+                    SitesTreeHandler.exportSitesTree(file, new ExporterResult());
+                    return handleFileImportResponse(true, file);
+                } catch (IOException e) {
+                    LOGGER.error(e.getMessage(), e);
+                    return handleFileImportResponse(false, file);
+                }
+            case ACTION_PRUNE_SITES:
+                file = new File(ApiUtils.getNonEmptyStringParam(params, PARAM_FILE_PATH));
+                PruneSiteResult res = SitesTreeHandler.pruneSiteNodes(file);
+
+                Map<String, String> map = new HashMap<>();
+                map.put("readNodes", Integer.toString(res.getReadNodes()));
+                map.put("deletedNodes", Integer.toString(res.getDeletedNodes()));
+                map.put("errors", res.getError() == null ? "" : res.getError());
+                return new ApiResponseSet<>(name, map);
             default:
                 throw new ApiException(Type.BAD_ACTION);
         }

addOns/exim/src/main/java/org/zaproxy/addon/exim/automation/ExportJob.java

@@ -20,6 +20,7 @@
 package org.zaproxy.addon.exim.automation;
 
 import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import lombok.Getter;
@@ -79,6 +80,23 @@ public void verifyParameters(AutomationProgress progress) {
                 this.getName(),
                 null,
                 progress);
+
+        // Check for invalid combinations
+        if (Source.SITESTREE.equals(this.parameters.getSource())
+                && !Type.YAML.equals(this.parameters.getType())) {
+            progress.error(
+                    Constant.messages.getString(
+                            "exim.automation.export.error.sitestree.type",
+                            this.getName(),
+                            this.parameters.getType()));
+        } else if (!Source.SITESTREE.equals(this.parameters.getSource())
+                && Type.YAML.equals(this.parameters.getType())) {
+            progress.error(
+                    Constant.messages.getString(
+                            "exim.automation.export.error.messages.type",
+                            this.getName(),
+                            this.parameters.getSource()));
+        }
     }
 
     @Override
@@ -112,18 +130,23 @@ public void runJob(AutomationEnvironment env, AutomationProgress progress) {
             return;
         }
 
+        Path path = JobUtils.getFile(fileName, getPlan()).toPath();
+
         ExporterOptions options =
                 ExporterOptions.builder()
                         .setContext(contextWrapper.getContext())
-                        .setOutputFile(JobUtils.getFile(fileName, getPlan()).toPath())
+                        .setOutputFile(path)
                         .setType(getParameters().getType())
                         .setSource(getParameters().getSource())
                         .build();
 
         ExporterResult result = extension.getExporter().export(options);
         progress.info(
                 Constant.messages.getString(
-                        "exim.automation.export.exportcount", getName(), result.getCount()));
+                        "exim.automation.export.exportcount",
+                        getName(),
+                        result.getCount(),
+                        path.toAbsolutePath()));
         result.getErrors()
                 .forEach(
                         error ->

addOns/exim/src/main/java/org/zaproxy/addon/exim/automation/ExportJobDialog.java

@@ -24,6 +24,7 @@
 import java.util.stream.Stream;
 import javax.swing.DefaultComboBoxModel;
 import javax.swing.JFileChooser;
+import org.parosproxy.paros.Constant;
 import org.parosproxy.paros.view.View;
 import org.zaproxy.addon.exim.ExporterOptions;
 import org.zaproxy.zap.utils.DisplayUtils;
@@ -89,7 +90,16 @@ public void save() {
 
     @Override
     public String validateFields() {
-        // Nothing to do
+        if (ExporterOptions.Source.SITESTREE.equals(sourceOptionModel.getSelectedItem())
+                && !ExporterOptions.Type.YAML.equals(typeOptionModel.getSelectedItem())) {
+            return Constant.messages.getString(
+                    "exim.automation.export.dialog.error.sitestree.type");
+        } else if (!ExporterOptions.Source.SITESTREE.equals(sourceOptionModel.getSelectedItem())
+                && ExporterOptions.Type.YAML.equals(typeOptionModel.getSelectedItem())) {
+            return Constant.messages.getString(
+                    "exim.automation.export.dialog.error.messages.type",
+                    sourceOptionModel.getSelectedItem());
+        }
         return null;
     }
 }

addOns/exim/src/main/java/org/zaproxy/addon/exim/automation/ExtensionEximAutomation.java

@@ -37,6 +37,7 @@ public class ExtensionEximAutomation extends ExtensionAdaptor {
 
     private ImportJob importJob;
     private ExportJob exportJob;
+    private PruneJob pruneJob;
 
     public ExtensionEximAutomation() {
         super(NAME);
@@ -55,6 +56,8 @@ public void hook(ExtensionHook extensionHook) {
         extAuto.registerAutomationJob(importJob);
         exportJob = new ExportJob(getExtension(ExtensionExim.class));
         extAuto.registerAutomationJob(exportJob);
+        pruneJob = new PruneJob();
+        extAuto.registerAutomationJob(pruneJob);
     }
 
     private static <T extends Extension> T getExtension(Class<T> clazz) {
@@ -72,6 +75,7 @@ public void unload() {
 
         extAuto.unregisterAutomationJob(importJob);
         extAuto.unregisterAutomationJob(exportJob);
+        extAuto.unregisterAutomationJob(pruneJob);
     }
 
     @Override