Use the Alerts from the Alert tree for generation of reports

thc202 · Jan 22, 2016 · 49eeeac · 49eeeac
1 parent 5f72b0e
commit 49eeeac
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/src/org/zaproxy/zap/extension/alert/AlertTreeModel.java b/src/org/zaproxy/zap/extension/alert/AlertTreeModel.java
@@ -20,7 +20,10 @@
 package org.zaproxy.zap.extension.alert;
 
 import java.awt.EventQueue;
+import java.util.ArrayList;
 import java.util.Comparator;
+import java.util.Enumeration;
+import java.util.List;
 
 import javax.swing.tree.DefaultTreeModel;
 
@@ -196,6 +199,25 @@ public synchronized void deletePath(Alert alert) {
             this.nodeChanged(parent);
         }
     }
+
+    public List<Alert> getAlerts(String cleanNodeName) {
+        ArrayList<Alert> alerts = new ArrayList<>();
+        @SuppressWarnings("unchecked")
+        Enumeration<AlertNode> groupAlertNodes = ((AlertNode) getRoot()).children();
+        while (groupAlertNodes.hasMoreElements()) {
+            AlertNode groupAlertNode = groupAlertNodes.nextElement();
+            @SuppressWarnings("unchecked")
+            Enumeration<AlertNode> alertNodes = groupAlertNode.children();
+            while (alertNodes.hasMoreElements()) {
+                AlertNode alertNode = alertNodes.nextElement();
+                if (alertNode.getUserObject().getHistoryRef().getURI().toString().startsWith(cleanNodeName)) {
+                    alerts.add(alertNode.getUserObject());
+                }
+            }
+        }
+        alerts.trimToSize();
+        return alerts;
+    }
 
     private static class GroupAlertChildNodeComparator implements Comparator<AlertNode> {
 

diff --git a/src/org/zaproxy/zap/extension/alert/ExtensionAlert.java b/src/org/zaproxy/zap/extension/alert/ExtensionAlert.java
@@ -662,7 +662,7 @@ public List<Alert> getAllAlerts() {
     public String getXml(SiteNode site) {
         StringBuilder xml = new StringBuilder();
         xml.append("<alerts>");
-        List<Alert> alerts = site.getAlerts();
+        List<Alert> alerts = treeModel.getAlerts(site.getCleanNodeName());
         SortedSet<String> handledAlerts = new TreeSet<String>(); 
 
         for (int i=0; i < alerts.size(); i++) {