Merge pull request #5916 from thc202/sequence/af-ascan-alerts

sequence: provide alert result data
zaproxy · Nov 14, 2024 · 9cfe4f8 · 9cfe4f8
2 parents f051bdf + 15ea645
commit 9cfe4f8
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 8 deletions.
diff --git a/addOns/sequence/CHANGELOG.md b/addOns/sequence/CHANGELOG.md
@@ -5,8 +5,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## Unreleased
 ### Added
-- Add `sequence-import` Automation Framework job.
-- Initial sequence-activeScan implementation.
+- Add Automation Framework jobs:
+  - `sequence-import` to import HARs as sequences.
+  - `sequence-activeScan` to active scan sequences.
 - Data for reporting.
 - Sequence active scan policy.
 

diff --git a/...c/main/java/org/zaproxy/zap/extension/sequence/automation/SequenceAScanJobResultData.java b/...c/main/java/org/zaproxy/zap/extension/sequence/automation/SequenceAScanJobResultData.java
@@ -20,17 +20,30 @@
 package org.zaproxy.zap.extension.sequence.automation;
 
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import lombok.Getter;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.parosproxy.paros.core.scanner.Alert;
+import org.parosproxy.paros.db.DatabaseException;
+import org.parosproxy.paros.db.RecordAlert;
+import org.parosproxy.paros.model.Model;
 import org.zaproxy.addon.automation.JobResultData;
 import org.zaproxy.zap.extension.sequence.StdActiveScanRunner.SequenceStepData;
 
 @Getter
 public class SequenceAScanJobResultData extends JobResultData {
 
+    private static final Logger LOGGER = LogManager.getLogger(SequenceAScanJobResultData.class);
+
     public static final String KEY = "seqAScanData";
 
     private List<SequenceData> seqData = new ArrayList<>();
+    private Map<Integer, Alert> alertDataMap;
 
     public SequenceAScanJobResultData(String jobName) {
         super(jobName);
@@ -45,6 +58,42 @@ public String getKey() {
         return KEY;
     }
 
+    @Override
+    public Alert getAlertData(int alertId) {
+        return getAlertDataMap().get(alertId);
+    }
+
+    @Override
+    public Collection<Alert> getAllAlertData() {
+        return getAlertDataMap().values();
+    }
+
+    private Map<Integer, Alert> getAlertDataMap() {
+        if (alertDataMap == null) {
+            Map<Integer, Alert> data = new HashMap<Integer, Alert>();
+            seqData.forEach(
+                    sequenceData ->
+                            sequenceData.getSteps().forEach(step -> addStepAlerts(data, step)));
+
+            alertDataMap = Collections.unmodifiableMap(data);
+        }
+        return alertDataMap;
+    }
+
+    private static void addStepAlerts(Map<Integer, Alert> data, SequenceStepData step) {
+        List<Integer> alertIds = step.getAlertIds();
+        for (int id : alertIds) {
+            try {
+                RecordAlert recordAlert = Model.getSingleton().getDb().getTableAlert().read(id);
+                if (recordAlert != null) {
+                    data.put(id, new Alert(recordAlert));
+                }
+            } catch (DatabaseException e) {
+                LOGGER.error("Could not read alert with ID {} from the database:", id, e);
+            }
+        }
+    }
+
     @Getter
     public static class SequenceData {
         private String sequenceName;

diff --git a/...ce/src/main/java/org/zaproxy/zap/extension/sequence/automation/SequenceActiveScanJob.java b/...ce/src/main/java/org/zaproxy/zap/extension/sequence/automation/SequenceActiveScanJob.java
@@ -44,7 +44,6 @@
 import org.zaproxy.addon.automation.AutomationProgress;
 import org.zaproxy.addon.automation.ContextWrapper;
 import org.zaproxy.addon.automation.JobResultData;
-import org.zaproxy.addon.automation.jobs.ActiveScanJobResultData;
 import org.zaproxy.addon.automation.jobs.JobData;
 import org.zaproxy.addon.automation.jobs.JobUtils;
 import org.zaproxy.addon.automation.jobs.PolicyDefinition;
@@ -97,7 +96,7 @@ public boolean supportsAlertTests() {
 
     @Override
     public String getKeyAlertTestsResultData() {
-        return ActiveScanJobResultData.KEY;
+        return SequenceAScanJobResultData.KEY;
     }
 
     @Override
@@ -200,10 +199,20 @@ public void runJob(AutomationEnvironment env, AutomationProgress progress) {
                 contextSpecificObjects.add(scanPolicy);
             }
 
+            Map<String, List<SequenceStepData>> result = new HashMap<>();
             Stream<ZestScriptWrapper> sequenceZestScripts = getSequenceScripts();
             if (StringUtils.isEmpty(parameters.getSequence())) {
                 sequenceZestScripts.forEach(
-                        e -> scanSequence(e, context, user, contextSpecificObjects, progress));
+                        e ->
+                                result.put(
+                                        e.getName(),
+                                        scanSequence(
+                                                e,
+                                                context,
+                                                user,
+                                                contextSpecificObjects,
+                                                progress)));
+
             } else {
                 Optional<ZestScriptWrapper> scriptWrapper =
                         sequenceZestScripts
@@ -218,8 +227,13 @@ public void runJob(AutomationEnvironment env, AutomationProgress progress) {
                     return;
                 }
 
-                scanSequence(scriptWrapper.get(), context, user, contextSpecificObjects, progress);
+                ZestScriptWrapper sw = scriptWrapper.get();
+                result.put(
+                        sw.getName(),
+                        scanSequence(sw, context, user, contextSpecificObjects, progress));
             }
+
+            progress.addJobResultData(createJobResultData(result));
         } finally {
             extAScan.setPanelSwitch(true);
         }
@@ -231,7 +245,7 @@ private Stream<ZestScriptWrapper> getSequenceScripts() {
                 .map(ZestScriptWrapper.class::cast);
     }
 
-    private static void scanSequence(
+    private static List<SequenceStepData> scanSequence(
             ZestScriptWrapper script,
             ContextWrapper contextWrapper,
             User user,
@@ -244,19 +258,25 @@ private static void scanSequence(
         try {
             zzr.run(null, null);
             ascans.put(script.getName(), zzr.getSteps());
+            return zzr.getSteps();
         } catch (Exception e) {
             progress.error(
                     Constant.messages.getString(
                             "automation.error.unexpected.internal", e.getMessage()));
             LOGGER.error(e.getMessage(), e);
         }
+        return List.of();
     }
 
     @Override
     public List<JobResultData> getJobResultData() {
+        return createJobResultData(ascans);
+    }
+
+    private List<JobResultData> createJobResultData(Map<String, List<SequenceStepData>> result) {
         List<JobResultData> list = new ArrayList<>();
         SequenceAScanJobResultData data = new SequenceAScanJobResultData(this.getName());
-        ascans.entrySet().stream()
+        result.entrySet().stream()
                 .forEach(entry -> data.addSequenceData(entry.getKey(), entry.getValue()));
         list.add(data);
         return list;