Update list_hosted_zones() and list_resource_record_sets() to use…

… preconfigured route53 client (#15)

This will prevent them from having to configure it for every run

lambda_code/scan/scan.py

@@ -2,6 +2,7 @@
 import json
 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
@@ -290,12 +291,15 @@ def lambda_handler(event, context):  # pylint:disable=unused-argument
     account_id = event["Id"]
     account_name = event["Name"]
 
-    hosted_zones = list_hosted_zones(event)
+    boto3_session = assume_role(account_id)
+    route53 = boto3_session.client("route53")
+
+    hosted_zones = list_hosted_zones(event, route53)
 
     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(account_id, account_name, hosted_zone["Id"], route53)
         record_sets = sanitise_wildcards(record_sets)
 
         alias_cloudfront_s3(account_name, record_sets, account_id)

utils/utils_aws.py

@@ -86,14 +86,15 @@ def list_accounts():
     return []
 
 
-def list_hosted_zones(account):
+def list_hosted_zones(account, route53):
 
     account_id = account["Id"]
     account_name = account["Name"]
 
     try:
-        boto3_session = assume_role(account_id)
-        route53 = boto3_session.client("route53")
+        if not route53:
+            boto3_session = assume_role(account_id)
+            route53 = boto3_session.client("route53")
 
         hosted_zones_list = []
 
@@ -119,11 +120,12 @@ def list_hosted_zones(account):
     return []
 
 
-def list_resource_record_sets(account_id, account_name, hosted_zone_id):
+def list_resource_record_sets(account_id, account_name, hosted_zone_id, route53):
 
     try:
-        boto3_session = assume_role(account_id)
-        route53 = boto3_session.client("route53")
+        if not route53:
+            boto3_session = assume_role(account_id)
+            route53 = boto3_session.client("route53")
 
         record_set_list = []