From d6a9da91d73e6441339f7dc56a8e8aa374b150c8 Mon Sep 17 00:00:00 2001
From: Marc Greisen <marc@greisen.org>
Date: Fri, 23 Jul 2021 13:30:25 -0700
Subject: [PATCH] Track zap complete state (#227)

---
 Scripts/Tests/bvt-petstore3.py   |  7 +++++--
 cli/raft-tools/tools/ZAP/scan.py | 29 ++++++++++++++++++++++++++---
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/Scripts/Tests/bvt-petstore3.py b/Scripts/Tests/bvt-petstore3.py
index 604669d..5140eb8 100644
--- a/Scripts/Tests/bvt-petstore3.py
+++ b/Scripts/Tests/bvt-petstore3.py
@@ -151,8 +151,11 @@ def bvt(cli, definitions, subs):
         print('Validating that bugs posted events matches total bugs found in job status')
         total_bugs_found = 0
         for r in job_status_events:
-            if r['Data']['State'] == 'Completed' and r['Data']['AgentName'] != r['Data']['JobId'] and r['Data']['Tool'] == 'RESTler':
-                total_bugs_found += r['Data']['Metrics']['TotalBugBucketsCount']
+            if r['Data']['State'] == 'Completed' and r['Data']['AgentName'] != r['Data']['JobId']:
+                if r['Data']['Tool'] == 'RESTler':
+                    total_bugs_found += r['Data']['Metrics']['TotalBugBucketsCount']
+                elif  r['Data']['Tool'] == 'ZAP':
+                    total_bugs_found += int(r['Data']['Details']['totalBugCount'])
 
         print(f'Total bugs found: {total_bugs_found}')
         print(f'Number of Bug found events: {len(bug_found_events)}')
diff --git a/cli/raft-tools/tools/ZAP/scan.py b/cli/raft-tools/tools/ZAP/scan.py
index 91d1b46..fc676bb 100644
--- a/cli/raft-tools/tools/ZAP/scan.py
+++ b/cli/raft-tools/tools/ZAP/scan.py
@@ -30,6 +30,12 @@ class StatusReporter(StreamHandler):
             if i != -1:
                 self.details["Scan progress"] = txt[i :]
                 raftUtils.report_status_running(self.details)
+            else:
+                progress='Passive scanning complete'
+                i = txt.find(progress)
+                if i != -1:
+                    self.details["Scan progress"] = "Active and Passive Scan progress %100"
+                    raftUtils.report_status_running(self.details)
 
 zap = __import__("zap-api-scan")
 
@@ -57,6 +63,18 @@ def post_bugs(target_index):
     else:
         print(f'File {target_index}-report.json does NOT exist.')
 
+def count_bugs(target_index):
+    bugCount = 0
+    if os.path.exists(f'/zap/wrk/{target_index}-report.json'):
+        with open(f'/zap/wrk/{target_index}-report.json') as f:
+            reportData = json.load(f)
+
+        # Every alert is a bug
+        for site in reportData['site']:
+            bugCount = len(site['alerts'])
+
+    return bugCount
+
 def run_zap(target_index, targets_total, host, target, token):
     if token:
         raftUtils.log_trace('Authentication token is set')
@@ -85,7 +103,7 @@ def run_zap(target_index, targets_total, host, target, token):
         pass
 
     try:
-        details = {"targetIndex": target_index, "numberOfTargets" : targets_total, "target": target}
+        details = {"targetIndex": target_index, "numberOfTargets" : targets_total, "target": target, "totalBugCount": 0}
         print(f"Starting ZAP target: {target} host_config: {host_config}")
 
         if os.path.exists(target):
@@ -94,9 +112,11 @@ def run_zap(target_index, targets_total, host, target, token):
 
         raftUtils.log_trace(f"Starting ZAP")
         raftUtils.report_status_running(details)
+
         status_reporter = StatusReporter(details)
         logger = logging.getLogger()
         logger.addHandler(status_reporter)
+
         zap.main([ '-t', target,
                    '-f', 'openapi',
                    '-J', f'{target_index}-report.json',
@@ -104,14 +124,17 @@ def run_zap(target_index, targets_total, host, target, token):
                    '-w', f'{target_index}-report.md',
                    '-x', f'{target_index}-report.xml',
                    '-d'] + zap_auth_config + host_config)
-        details["Scan progress"] = "Active scan progress %: 100"
-        raftUtils.report_status_running(details)
 
     except SystemExit as e:
         r = e.code
 
     raftUtils.log_trace(f"ZAP exited with exit code: {r}")
     shutil.copy('/zap/zap.out', f'/zap/wrk/{target_index}-zap.out')
+
+    # Update the status with the total bug count.
+    details["totalBugCount"] = count_bugs(target_index)
+    raftUtils.report_status_running(details)
+
     post_bugs(target_index)
 
     if r <= 2: