Patch scan.py, scan_ips.py, and utils_aws.py to reuse the generated route53 clients

Relates To: SECENG-1648

Author: bc-jcarlson

lambda_code/scan/scan.py

@@ -292,13 +292,17 @@ def lambda_handler(event, context):  # pylint:disable=unused-argument
     account_name = event["Name"]
 
     aws_session = assume_role(account_id)
-
-    hosted_zones = list_hosted_zones(event)
+    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)
 
     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(account_id, account_name, hosted_zone["Id"])
+        record_sets = list_resource_record_sets(r53client, account_name, hosted_zone["Id"])
         record_sets = sanitise_wildcards(record_sets)
 
         alias_cloudfront_s3(account_name, record_sets, account_id)

lambda_code/scan_ips/scan_ips.py

@@ -2,6 +2,7 @@
 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
@@ -166,13 +167,19 @@ def lambda_handler(event, context):  # pylint:disable=unused-argument
 
     get_ips(account_id, account_name)
 
-    hosted_zones = list_hosted_zones(event)
+    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)
 
     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(account_id, account_name, hosted_zone["Id"])
+            record_sets = list_resource_record_sets(r53client, account_name, hosted_zone["Id"])
             record_sets = sanitise_wildcards(record_sets)
 
             a_record(account_name, record_sets, ip_prefixes)

utils/utils_aws.py

@@ -86,70 +86,53 @@ def list_accounts():
     return []
 
 
-def list_hosted_zones(account):
+def list_hosted_zones(route53, account):
 
-    account_id = account["Id"]
     account_name = account["Name"]
 
-    try:
-        boto3_session = assume_role(account_id)
-        route53 = boto3_session.client("route53")
-
-        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"]]
+    hosted_zones_list = []
 
-                hosted_zones_list = hosted_zones_list + hosted_zones
+    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"]]
 
-            return hosted_zones_list
+            hosted_zones_list = hosted_zones_list + hosted_zones
 
-        except Exception:
-            logging.error(
-                "ERROR: Lambda execution role requires route53:ListHostedZones permission in %a account",
-                account_name,
-            )
+        return hosted_zones_list
 
     except Exception:
-        logging.error("ERROR: unable to assume role in %a account %s", account_name, account_id)
+        logging.error(
+            "ERROR: Lambda execution role requires route53:ListHostedZones permission in %a account",
+            account_name,
+        )
 
     return []
 
 
-def list_resource_record_sets(account_id, account_name, hosted_zone_id):
+def list_resource_record_sets(route53, account_name, hosted_zone_id):
 
     try:
-        boto3_session = assume_role(account_id)
-        route53 = boto3_session.client("route53")
-
-        record_set_list = []
-
-        try:
-            paginator_records = route53.get_paginator("list_resource_record_sets")
-            pages_records = paginator_records.paginate(
-                HostedZoneId=hosted_zone_id,
-                StartRecordName="_",
-                StartRecordType="NS",
-            )
+        paginator_records = route53.get_paginator("list_resource_record_sets")
+        pages_records = paginator_records.paginate(
+            HostedZoneId=hosted_zone_id,
+            StartRecordName="_",
+            StartRecordType="NS",
+        )
 
-            for page_records in pages_records:
-                record_sets = page_records["ResourceRecordSets"]
+        for page_records in pages_records:
+            record_sets = page_records["ResourceRecordSets"]
 
-                record_set_list = record_set_list + record_sets
+            record_set_list = record_set_list + record_sets
 
-            return record_set_list
-
-        except Exception:
-            logging.exception(
-                "ERROR: Lambda execution role requires route53:ListResourceRecordSets permission in %a account",
-                account_name,
-            )
+        return record_set_list
 
     except Exception:
-        logging.error("ERROR: unable to assume role in %a account %s", account_name, account_id)
+        logging.exception(
+            "ERROR: Lambda execution role requires route53:ListResourceRecordSets permission in %a account",
+            account_name,
+        )
 
     return []