diff --git a/zap/src/main/java/org/zaproxy/zap/extension/alert/AlertNode.java b/zap/src/main/java/org/zaproxy/zap/extension/alert/AlertNode.java index 8a92679a9bf..1cfd7935d3e 100644 --- a/zap/src/main/java/org/zaproxy/zap/extension/alert/AlertNode.java +++ b/zap/src/main/java/org/zaproxy/zap/extension/alert/AlertNode.java @@ -32,21 +32,31 @@ public class AlertNode extends DefaultMutableTreeNode { private final Comparator childComparator; private String nodeName = null; + private String alertRef; private int risk = -1; private Alert alert; private boolean systemic; public AlertNode(int risk, String nodeName) { - this(risk, nodeName, null); + this(risk, nodeName, null, null); } public AlertNode(int risk, String nodeName, Comparator childComparator) { + this(risk, nodeName, null, childComparator); + } + + AlertNode(int risk, String nodeName, String alertRef, Comparator childComparator) { super(); this.nodeName = nodeName; + this.alertRef = alertRef; this.setRisk(risk); this.childComparator = new AlertNodeComparatorWrapper(childComparator); } + String getAlertRef() { + return alertRef; + } + /** Sets an alert for this node. The {@link #setAlert(Alert)} method should be used instead. */ @Override public void setUserObject(Object userObject) { diff --git a/zap/src/main/java/org/zaproxy/zap/extension/alert/AlertTreeModel.java b/zap/src/main/java/org/zaproxy/zap/extension/alert/AlertTreeModel.java index ce466261169..b356eb34735 100644 --- a/zap/src/main/java/org/zaproxy/zap/extension/alert/AlertTreeModel.java +++ b/zap/src/main/java/org/zaproxy/zap/extension/alert/AlertTreeModel.java @@ -120,7 +120,9 @@ public AlertNode getAlertNode(Alert alert) { risk = -1; } - AlertNode needle = new AlertNode(risk, alert.getName(), GROUP_ALERT_CHILD_COMPARATOR); + AlertNode needle = + new AlertNode( + risk, alert.getName(), alert.getAlertRef(), GROUP_ALERT_CHILD_COMPARATOR); needle.setAlert(alert); int idx = parent.findIndex(needle); if (idx < 0) { @@ -194,7 +196,7 @@ private AlertNode findAndAddGroup(AlertNode parent, String nodeName, Alert alert risk = -1; } - AlertNode node = new AlertNode(risk, nodeName, ALERT_CHILD_COMPARATOR); + AlertNode node = new AlertNode(risk, nodeName, alert.getAlertRef(), ALERT_CHILD_COMPARATOR); int idx = parent.findIndex(node); if (idx < 0) { idx = -(idx + 1); @@ -214,7 +216,8 @@ private AlertNode addLeaf(AlertNode parent, String nodeName, Alert alert) { risk = -1; } - AlertNode needle = new AlertNode(risk, nodeName, ALERT_CHILD_COMPARATOR); + AlertNode needle = + new AlertNode(risk, nodeName, alert.getAlertRef(), ALERT_CHILD_COMPARATOR); needle.setAlert(alert); int idx = parent.findIndex(needle); if (idx < 0) { @@ -268,7 +271,11 @@ public int compare(AlertNode alertNode, AlertNode anotherAlertNode) { return -1; } - return alertNode.getNodeName().compareTo(anotherAlertNode.getNodeName()); + int res = alertNode.getNodeName().compareTo(anotherAlertNode.getNodeName()); + if (res == 0) { + return alertNode.getAlertRef().compareTo(anotherAlertNode.getAlertRef()); + } + return res; } } diff --git a/zap/src/test/java/org/zaproxy/zap/extension/alert/AlertTreeModelUnitTest.java b/zap/src/test/java/org/zaproxy/zap/extension/alert/AlertTreeModelUnitTest.java index a7298cbe434..97705ac853d 100644 --- a/zap/src/test/java/org/zaproxy/zap/extension/alert/AlertTreeModelUnitTest.java +++ b/zap/src/test/java/org/zaproxy/zap/extension/alert/AlertTreeModelUnitTest.java @@ -73,11 +73,22 @@ void shouldAddUniqueAlerts() { "https://www.example.com", Alert.RISK_MEDIUM, Alert.CONFIDENCE_MEDIUM); + Alert a4 = + newAlert( + 1, + "1-2", + 3, + "Alert A", + "https://www.example.com", + "https://www.example.com", + Alert.RISK_MEDIUM, + Alert.CONFIDENCE_MEDIUM); // When atModel.addPath(a1); atModel.addPath(a3); atModel.addPath(a2); + atModel.addPath(a4); // Then @@ -89,12 +100,15 @@ void shouldAddUniqueAlerts() { - Medium: Alert A - GET:https://www.example.com - GET:https://www.example.net + - Medium: Alert A + - GET:https://www.example.com """, TextAlertTree.toString(atModel)); assertEquals(a1, atModel.getRoot().getChildAt(0).getChildAt(0).getAlert()); assertEquals(a3, atModel.getRoot().getChildAt(1).getChildAt(0).getAlert()); assertEquals(a2, atModel.getRoot().getChildAt(1).getChildAt(1).getAlert()); + assertEquals(a4, atModel.getRoot().getChildAt(2).getChildAt(0).getAlert()); } @Test @@ -177,11 +191,22 @@ void shouldFindDuplicateAlerts() { "https://www.example.com?a=3", Alert.RISK_MEDIUM, Alert.CONFIDENCE_MEDIUM); + Alert a4 = + newAlert( + 1, + "1-2", + 3, + "Alert A", + "https://www.example.com(a)", + "https://www.example.com?a=4", + Alert.RISK_MEDIUM, + Alert.CONFIDENCE_MEDIUM); // When atModel.addPath(a1); atModel.addPath(a2); atModel.addPath(a3); + atModel.addPath(a4); AlertNode an1 = atModel.getAlertNode(a1); AlertNode an2 = atModel.getAlertNode(a2); @@ -191,6 +216,7 @@ void shouldFindDuplicateAlerts() { assertEquals("GET:https://www.example.com(a)", an1.getNodeName()); assertEquals("GET:https://www.example.com(a)", an2.getNodeName()); assertEquals("GET:https://www.example.com(a)", an3.getNodeName()); + assertEquals("GET:https://www.example.com(a)", atModel.getAlertNode(a4).getNodeName()); } @Test @@ -225,16 +251,27 @@ void shouldChangeDuplicateAlerts() { "https://www.example.com?a=3", Alert.RISK_MEDIUM, Alert.CONFIDENCE_MEDIUM); + Alert a4 = + newAlert( + 1, + "1-2", + 3, + "Alert A", + "https://www.example.com(a)", + "https://www.example.com?a=4", + Alert.RISK_MEDIUM, + Alert.CONFIDENCE_MEDIUM); // When atModel.addPath(a1); atModel.addPath(a2); atModel.addPath(a3); + atModel.addPath(a4); a1.setRisk(Alert.RISK_HIGH); atModel.updatePath(a1); // Then - assertEquals(1, atModel.getRoot().getChildCount()); + assertEquals(2, atModel.getRoot().getChildCount()); // Only child - Medium risk assertEquals("Alert A", atModel.getRoot().getChildAt(0).getNodeName()); @@ -245,6 +282,11 @@ void shouldChangeDuplicateAlerts() { "GET:https://www.example.com(a)", atModel.getRoot().getChildAt(0).getChildAt(0).getNodeName()); assertEquals(Alert.RISK_HIGH, atModel.getRoot().getChildAt(0).getChildAt(0).getRisk()); + + assertEquals( + "GET:https://www.example.com(a)", + atModel.getRoot().getChildAt(1).getChildAt(0).getNodeName()); + assertEquals(Alert.RISK_MEDIUM, atModel.getRoot().getChildAt(1).getChildAt(0).getRisk()); } @Test @@ -280,16 +322,27 @@ void shouldDeleteUniqueAlert() { "https://www.example.net", Alert.RISK_MEDIUM, Alert.CONFIDENCE_MEDIUM); + Alert a4 = + newAlert( + 1, + "1-2", + 3, + "Alert A", + "https://www.example.com/a1", + "https://www.example.com/a1", + Alert.RISK_MEDIUM, + Alert.CONFIDENCE_MEDIUM); // When atModel.addPath(a1); atModel.addPath(a2); atModel.addPath(a3); + atModel.addPath(a4); atModel.deletePath(a1); // Then - assertEquals(1, atModel.getRoot().getChildCount()); + assertEquals(2, atModel.getRoot().getChildCount()); assertEquals( """ @@ -297,11 +350,14 @@ void shouldDeleteUniqueAlert() { - Medium: Alert A - GET:https://www.example.com/a2 - GET:https://www.example.net + - Medium: Alert A + - GET:https://www.example.com/a1 """, TextAlertTree.toString(atModel)); assertEquals(a2, atModel.getRoot().getChildAt(0).getChildAt(0).getAlert()); assertEquals(a3, atModel.getRoot().getChildAt(0).getChildAt(1).getAlert()); + assertEquals(a4, atModel.getRoot().getChildAt(1).getChildAt(0).getAlert()); } @Test @@ -337,11 +393,22 @@ void shouldChangeUniqueAlert() { "https://www.example.com/a2", Alert.RISK_MEDIUM, Alert.CONFIDENCE_MEDIUM); + Alert a4 = + newAlert( + 1, + "1-2", + 3, + "Alert A", + "https://www.example.com/a1", + "https://www.example.com/a1", + Alert.RISK_MEDIUM, + Alert.CONFIDENCE_MEDIUM); // When atModel.addPath(a1); atModel.addPath(a2); atModel.addPath(a3); + atModel.addPath(a4); a1.setRisk(Alert.RISK_HIGH); atModel.updatePath(a1); @@ -355,12 +422,15 @@ void shouldChangeUniqueAlert() { - Medium: Alert A - GET:https://www.example.com/a2 - GET:https://www.example.net + - Medium: Alert A + - GET:https://www.example.com/a1 """, TextAlertTree.toString(atModel)); assertEquals(a1, atModel.getRoot().getChildAt(0).getChildAt(0).getAlert()); assertEquals(a3, atModel.getRoot().getChildAt(1).getChildAt(0).getAlert()); assertEquals(a2, atModel.getRoot().getChildAt(1).getChildAt(1).getAlert()); + assertEquals(a4, atModel.getRoot().getChildAt(2).getChildAt(0).getAlert()); } @Test @@ -393,18 +463,30 @@ void shouldDeleteNodeWhenNoAlertsLeft() { "https://www.example.com?a=3", Alert.RISK_MEDIUM, Alert.CONFIDENCE_MEDIUM); + Alert a4 = + newAlert( + 1, + "1-2", + 3, + "Alert A", + "https://www.example.com(a)", + "https://www.example.com?a=4", + Alert.RISK_MEDIUM, + Alert.CONFIDENCE_MEDIUM); // When atModel.addPath(a1); atModel.addPath(a2); atModel.addPath(a3); + atModel.addPath(a4); atModel.deletePath(a1); atModel.deletePath(a3); atModel.deletePath(a2); // Then - assertEquals(0, atModel.getRoot().getChildCount()); + assertEquals(1, atModel.getRoot().getChildCount()); + assertEquals(a4, atModel.getRoot().getChildAt(0).getChildAt(0).getAlert()); } private static Alert newAlert( @@ -415,7 +497,22 @@ private static Alert newAlert( String uri, int risk, int confidence) { + return newAlert(pluginId, null, id, name, nodeName, uri, risk, confidence); + } + + private static Alert newAlert( + int pluginId, + String alertRef, + int id, + String name, + String nodeName, + String uri, + int risk, + int confidence) { Alert alert = new Alert(pluginId, risk, confidence, name); + if (alertRef != null) { + alert.setAlertRef(alertRef); + } alert.setUri(uri); alert.setAlertId(id); alert.setNodeName(nodeName);