[DSEC-963] add retry loop

zap/src/drive_upload.py

@@ -120,7 +120,7 @@ def find_subfolder(folder_structure, target_name, target_folder=None):
     Finds a specific subfolder by name, returns the dict for that folder.
     """
     for child in folder_structure['children']:
-        if child['name'] == target_name:
+        if target_name in child['name']:
             target_folder = child
             return target_folder
         target_folder = find_subfolder(child, target_name)

zap/src/scan.py

@@ -22,7 +22,7 @@
 from google.cloud import storage
 from slack_sdk.web import WebClient as SlackClient
 
-from zap import ScanType, zap_compliance_scan, zap_connect
+from zap import ScanType, zap_compliance_scan, zap_shutdown
 
 
 def fetch_dojo_product_name(defect_dojo, defect_dojo_user, defect_dojo_key, product_id):
@@ -34,7 +34,7 @@ def fetch_dojo_product_name(defect_dojo, defect_dojo_user, defect_dojo_key, prod
     defect_dojo, defect_dojo_key, defect_dojo_user, debug=False, timeout=120)
     max_retries = int(getenv("MAX_RETRIES", '6'))
     retry_delay = 30
-    for attempt in range(max_retries):
+    for _ in range(max_retries):
         try:
             product = dojo.get_product(product_id=product_id)
             return product.data["name"]
@@ -106,11 +106,21 @@ def defectdojo_upload(product_id: int, zap_filename: str, defect_dojo_key: str,
     date = datetime.today().strftime("%Y%m%d%H:%M")
     lead_id = fetch_dojo_lead_id(dojo, defect_dojo_user)
 
-    engagement=dojo.create_engagement( name=date, product_id=product_id, lead_id=lead_id,
-        target_start=datetime.today().strftime("%Y-%m-%d"),
-        target_end=datetime.today().strftime("%Y-%m-%d"), status="In Progress",
-        active='True',deduplication_on_engagement='False')
-    engagement_id=engagement.data["id"]
+# The call to create_engagement sometimes fails.
+    retry_delay = 20
+    max_retries = int(getenv("MAX_RETRIES", '5'))
+    for attempt in range(max_retries):
+        try:
+            engagement=dojo.create_engagement( name=date, product_id=product_id, lead_id=lead_id,
+                target_start=datetime.today().strftime("%Y-%m-%d"),
+                target_end=datetime.today().strftime("%Y-%m-%d"), status="In Progress",
+                active='True',deduplication_on_engagement='False')
+            engagement_id=engagement.data["id"]
+            break
+        except Exception: # pylint: disable=broad-except
+            sleep(retry_delay)
+            if attempt == max_retries-1:
+                raise RuntimeError("Maximum retry attempts reached for closing engagement")
 
     dojo_upload = dojo.upload_scan(engagement_id=engagement_id,
                      scan_type="ZAP Scan",
@@ -242,7 +252,6 @@ def get_codedx_initial_report(
         file_name=report_file,
         filters=filters,
     )
-
     return report_file
 
 
@@ -372,7 +381,7 @@ def upload_googledrive(scan_type, zap_filename, codedx_project, report_file, sla
     """
     root_id = os.getenv('DRIVE_ROOT_ID', None)
     drive_id = os.getenv('DRIVE_ID', None)
-    if scan_type in (ScanType.BASELINE): 
+    if scan_type in (ScanType.BASELINE):
         return
     try:
         logging.info('Setting up the google drive API service for uploading reports.')
@@ -389,7 +398,6 @@ def upload_googledrive(scan_type, zap_filename, codedx_project, report_file, sla
         date = datetime.today()
         date = drivehelper.adjust_date(date)
         _, xml_folder_dict, zap_raw_folder = drivehelper.get_upload_folders(folder_structure, date)
-
         file = drivehelper.upload_file_to_drive(zap_filename,
                                                     xml_folder_dict.get('id'),
                                                     drive_id,
@@ -508,10 +516,9 @@ def main(): # pylint: disable=too-many-locals
                     dd,
                     target_url
                 )
-                zap = zap_connect()
-                zap.core.shutdown()
+                zap_shutdown()
                 return
-            
+
             # upload its results to Code Dx
             cdx = CodeDx(codedx_url, codedx_api_key)
 
@@ -533,8 +540,7 @@ def main(): # pylint: disable=too-many-locals
             logging.info("ready to upload to google drive")
             upload_googledrive(scan_type, zap_filename, codedx_project, cdx_filename, slack_token, slack_channel)
 
-            zap = zap_connect()
-            zap.core.shutdown()
+            zap_shutdown()
             return
         except Exception as error: # pylint: disable=broad-except
             error_message = f"[RETRY-{ attempt }] Exception running Zap Scans: { error }"
@@ -546,8 +552,7 @@ def main(): # pylint: disable=too-many-locals
                 except:
                     logging.error(f"Slack could not post to {slack_channel}")
                 try:
-                    zap = zap_connect()
-                    zap.core.shutdown()
+                    zap_shutdown()
                 except Exception as zap_e: # pylint: disable=broad-except
                     error_message = f"Error shutting down zap: { zap_e }"
                     error_slack_alert(

zap/src/zap.py

@@ -21,7 +21,7 @@
 from zap_scan_type import ScanType
 from zapv2 import ZAPv2
 
-TIMEOUT_MINS = 5
+TIMEOUT_MINS = 6
 zap_port = int(os.getenv("ZAP_PORT", ""))
 PROXY = f"http://localhost:{zap_port}"
 
@@ -47,7 +47,7 @@ def zap_init(target_url: str):
     enabled = True
     matchregex = True
     matchstring = ".*"
-    requestspersecond = 5
+    requestspersecond = 4
     groupby = "host"
     zap.network.add_rate_limit_rule(description,
                                     enabled,
@@ -60,6 +60,31 @@ def zap_init(target_url: str):
 
     return zap
 
+def zap_shutdown():
+    """
+    Connect to the Zap instance and shut it down.
+    """
+    proxies = {
+        'http': PROXY,
+        'https': PROXY,
+        }
+    logging.info("Attempting to shutdown ZAP")
+    apikey = os.getenv("ZAP_API_KEY", "")
+    zap = ZAPv2(proxies={"http": PROXY, "https": PROXY}, apikey=apikey)
+    sleep_time = 20
+    attempts = int((TIMEOUT_MINS*60)/20)
+    for attempt in range(attempts):
+        try:
+            shutdown_endpoint = zap.base + "core/action/shutdown/"
+            headers = { "X-ZAP-API-Key" : apikey }
+            resp = requests.get(shutdown_endpoint, proxies=proxies, headers=headers,
+                                timeout=int(TIMEOUT_MINS*60)/2)
+            logging.info(f"Response code from requesting ZAP shutdown: {resp.status_code}")
+            if int(resp.status_code) == 200:
+                return
+        except Exception as e:
+            logging.error(f"An error occured while shutting down zap. Attempt {attempt+1}")
+        time.sleep(sleep_time)
 
 def parse_url(url):
     """
@@ -320,7 +345,7 @@ def get_hail_token():
         "aud": "https://www.googleapis.com/oauth2/v4/token",
         "iat": now,
         "scope": scope,
-        "exp": now + 300,  # 5m
+        "exp": now + 3000,  # 50m
         "iss": creds['client_email'],
     }
     encoded_assertion = jwt.encode(assertion, creds['private_key'], algorithm='RS256')