From 49eeeac3869cf3f2315ec877bbdba47d163013cc Mon Sep 17 00:00:00 2001
From: thc202 <thc202@gmail.com>
Date: Fri, 22 Jan 2016 08:34:17 +0000
Subject: [PATCH] Use the Alerts from the Alert tree for generation of reports

---
 .../zap/extension/alert/AlertTreeModel.java   | 22 +++++++++++++++++++
 .../zap/extension/alert/ExtensionAlert.java   |  2 +-
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/org/zaproxy/zap/extension/alert/AlertTreeModel.java b/src/org/zaproxy/zap/extension/alert/AlertTreeModel.java
index 8c6fbed9652..47467c8c94a 100644
--- 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
index ddb1f696e5d..c81d03838ee 100644
--- 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++) {