From 55706a5ca1a111b03e78bffb9cc4a839821aabc8 Mon Sep 17 00:00:00 2001
From: Josh Carlson <jcarlson@brightcove.com>
Date: Thu, 5 Sep 2024 12:14:50 -0400
Subject: [PATCH] Revert "Seceng 1648 fix persistent scan lambda timeouts
 (#11)"

This reverts commit 8484fe930d9663f6f70ba1f8788d10e65169de0d.
---
 lambda_code/scan/scan.py         | 13 ++-------
 lambda_code/scan_ips/scan_ips.py | 11 ++------
 utils/utils_aws.py               | 45 ++++++++++++++++++--------------
 3 files changed, 29 insertions(+), 40 deletions(-)

diff --git a/lambda_code/scan/scan.py b/lambda_code/scan/scan.py
index f6b0c23..f2ad303 100644
--- a/lambda_code/scan/scan.py
+++ b/lambda_code/scan/scan.py
@@ -3,7 +3,6 @@
 import logging
 import os
 
-from utils.utils_aws import assume_role
 from utils.utils_aws import eb_susceptible
 from utils.utils_aws import get_cloudfront_s3_origin_takeover
 from utils.utils_aws import list_domains
@@ -295,19 +294,11 @@ def lambda_handler(event, context):  # pylint:disable=unused-argument
     account_id = event["Id"]
     account_name = event["Name"]
 
-    if account_id in BC_ACCT_ID_BLACKLIST:
-        logging.info("account ID found on BC account blacklist, skipping...")
-
-        return
-
-    aws_session = assume_role(account_id)
-    r53client = aws_session.client("route53")
-
-    hosted_zones = list_hosted_zones(r53client, event)
+    hosted_zones = list_hosted_zones(event)
     for hosted_zone in hosted_zones:
         print(f"Searching for vulnerable domain records in hosted zone {hosted_zone['Name']}")
 
-        record_sets = list_resource_record_sets(r53client, account_name, hosted_zone["Id"])
+        record_sets = list_resource_record_sets(account_id, account_name, hosted_zone["Id"])
         record_sets = sanitise_wildcards(record_sets)
 
         alias_cloudfront_s3(account_name, record_sets, account_id)
diff --git a/lambda_code/scan_ips/scan_ips.py b/lambda_code/scan_ips/scan_ips.py
index 9aa879f..e3048f4 100644
--- a/lambda_code/scan_ips/scan_ips.py
+++ b/lambda_code/scan_ips/scan_ips.py
@@ -2,7 +2,6 @@
 import json
 import os
 
-from utils.utils_aws import assume_role
 from utils.utils_aws import list_hosted_zones
 from utils.utils_aws import list_resource_record_sets
 from utils.utils_aws import publish_to_sns
@@ -167,19 +166,13 @@ def lambda_handler(event, context):  # pylint:disable=unused-argument
 
     get_ips(account_id, account_name)
 
-    aws_session = assume_role(account_id)
-    try:
-        r53client = aws_session.client("route53")
-    except Exception:
-        print(f"ERROR: unable to assume role in {account_name} account {account_id}")
-
-    hosted_zones = list_hosted_zones(r53client, event)
+    hosted_zones = list_hosted_zones(event)
 
     if item_count > 0:  # don't test for vulnerabilities until DynamoDB table is populated across organisation
         for hosted_zone in hosted_zones:
             print(f"Searching for vulnerable A records in hosted zone {hosted_zone['Name']}")
 
-            record_sets = list_resource_record_sets(r53client, account_name, hosted_zone["Id"])
+            record_sets = list_resource_record_sets(account_id, account_name, hosted_zone["Id"])
             record_sets = sanitise_wildcards(record_sets)
 
             a_record(account_name, record_sets, ip_prefixes)
diff --git a/utils/utils_aws.py b/utils/utils_aws.py
index 6507b2e..80ade6d 100644
--- a/utils/utils_aws.py
+++ b/utils/utils_aws.py
@@ -89,38 +89,43 @@ def list_accounts():
     return []
 
 
-def list_hosted_zones(route53, account):
+def list_hosted_zones(account):
 
+    account_id = account["Id"]
     account_name = account["Name"]
 
-    hosted_zones_list = []
-
     try:
-        paginator_zones = route53.get_paginator("list_hosted_zones")
-        pages_zones = paginator_zones.paginate()
-        for page_zones in pages_zones:
-            hosted_zones = [h for h in page_zones["HostedZones"] if not h["Config"]["PrivateZone"]]
+        boto3_session = assume_role(account_id)
+        route53 = boto3_session.client("route53")
 
-            hosted_zones_list = hosted_zones_list + hosted_zones
+        hosted_zones_list = []
 
-        return hosted_zones_list
+        try:
+            paginator_zones = route53.get_paginator("list_hosted_zones")
+            pages_zones = paginator_zones.paginate()
+            for page_zones in pages_zones:
+                hosted_zones = [h for h in page_zones["HostedZones"] if not h["Config"]["PrivateZone"]]
 
-    except exceptions.ClientError as e:
-        logging.error(
-            f"ERROR: issue when listing hosted zones in {account_name} account :: [ {e} ]"
-        )
-        # logging.error(
-        #     "ERROR: Lambda execution role requires route53:ListHostedZones permission in %a account",
-        #     account_name,
-        # )
+                hosted_zones_list = hosted_zones_list + hosted_zones
 
-    return []
+            return hosted_zones_list
+
+        except Exception:
+            logging.error(
+                "ERROR: Lambda execution role requires route53:ListHostedZones permission in %a account",
+                account_name,
+            )
 
+    except Exception:
+        logging.error("ERROR: unable to assume role in %a account %s", account_name, account_id)
+
+    return []
 
-def list_resource_record_sets(route53, account_name, hosted_zone_id):
 
-    record_set_list = []
+def list_resource_record_sets(account_id, account_name, hosted_zone_id):
 
+    boto3_session = assume_role(account_id)
+    route53 = boto3_session.client("route53")
     try:
         paginator_records = route53.get_paginator("list_resource_record_sets")
         pages_records = paginator_records.paginate(