Reducing the number of IAM API calls made during enable cross-account…

src/lambda_codebase/initial_commit/bootstrap_repository/adf-bootstrap/deployment/lambda_codebase/enable_cross_account_access.py

@@ -26,11 +26,11 @@
 LOGGER = configure_logger(__name__)
 
 
-def update_iam(role, s3_bucket, kms_key_arn, role_policies):
-    iam = IAM(role)
+def update_iam(role, s3_buckets, kms_key_arns, role_policies):
+    iam = IAM(role.client("iam"))
     iam.update_iam_roles(
-        s3_bucket,
-        kms_key_arn,
+        s3_buckets,
+        kms_key_arns,
         role_policies
     )
 
@@ -55,14 +55,17 @@ def lambda_handler(event, _):
         role=boto3
     )
     account_id = event.get("account_id")
+    kms_key_arns = []
+    s3_buckets = []
     for region in list(set([event.get('deployment_account_region')] + event.get("regions", []))):
         kms_key_arn = parameter_store.fetch_parameter(
             "/cross_region/kms_arn/{0}".format(region)
         )
+        kms_key_arns.append(kms_key_arn)
         s3_bucket = parameter_store.fetch_parameter(
             "/cross_region/s3_regional_bucket/{0}".format(region)
         )
-        update_iam(boto3, s3_bucket, kms_key_arn, role_policies)
+        s3_buckets.append(s3_bucket)
         try:
             role = sts.assume_cross_account_role(
                 f'arn:{partition}:iam::{account_id}:role/adf-cloudformation-deployment-role',
@@ -74,4 +77,6 @@ def lambda_handler(event, _):
             LOGGER.debug("%s could not be assumed (%s), continuing", account_id, err, exc_info=True)
             continue
 
+    update_iam(boto3, s3_buckets, kms_key_arns, role_policies)
+
     return event

src/lambda_codebase/initial_commit/bootstrap_repository/adf-build/shared/python/iam.py

@@ -18,22 +18,29 @@
 class IAM:
     """Class used for modeling IAM."""
 
-    def __init__(self, role):
-        self.client = role.client('iam')
+    def __init__(self, client):
+        self.client = client
         self.role_name = None
         self.policy_name = None
         self.policy = None
 
     def update_iam_roles(
             self,
-            s3_bucket,
-            kms_key_arn,
+            s3_buckets,
+            kms_key_arns,
             role_policies
         ):
+        if not isinstance(s3_buckets, list):
+            s3_buckets = [s3_buckets]
+        if not isinstance(kms_key_arns, list):
+            kms_key_arns = [kms_key_arns]
+
         for role_name, policy_name in role_policies.items():
             self._fetch_policy_document(role_name, policy_name)
-            self._update_iam_policy_bucket(s3_bucket)
-            self._update_iam_cfn(kms_key_arn)
+            for s3_bucket in s3_buckets:
+                self._update_iam_policy_bucket(s3_bucket)
+            for kms_key_arn in kms_key_arns:
+                self._update_iam_cfn(kms_key_arn)
             self._put_role_policy()
 
     def _get_policy(self):

src/lambda_codebase/initial_commit/bootstrap_repository/adf-build/shared/python/tests/test_iam.py

@@ -44,3 +44,27 @@ def test_update_iam_cfn(cls):
     for policy in cls.policy.get('Statement'):
         if policy["Sid"] == "KMS":
             assert 'kms::12345678910::some_arn' in policy["Resource"]
+
+
+def test_update_iam_roles_with_lists(cls):
+    cls.update_iam_roles(["bucket1", "bucket2"], ["kms::12345678910::kms_key_1", "kms::12345678910::kms_key_2"], {'some_role_name':'some_policy_name'})
+    for policy in cls.policy.get('Statement'):
+        if policy["Sid"] == "KMS":
+            assert 'kms::12345678910::kms_key_1' in policy["Resource"]
+            assert 'kms::12345678910::kms_key_2' in policy["Resource"]
+        if policy["Sid"] == "S3":
+            assert 'arn:aws:s3:::bucket2' in policy["Resource"]
+            assert 'arn:aws:s3:::bucket2/*' in policy["Resource"]
+            assert 'arn:aws:s3:::bucket1' in policy["Resource"]
+            assert 'arn:aws:s3:::bucket1/*' in policy["Resource"]
+
+
+def test_update_iam_roles_is_backwards_compatible(cls):
+    cls.update_iam_roles("bucket1", "kms::12345678910::kms_key_1", {'some_role_name':'some_policy_name'})
+    for policy in cls.policy.get('Statement'):
+        if policy["Sid"] == "KMS":
+            assert 'kms::12345678910::kms_key_1' in policy["Resource"]
+        if policy["Sid"] == "S3":  
+            assert 'arn:aws:s3:::bucket1' in policy["Resource"]
+            assert 'arn:aws:s3:::bucket1/*' in policy["Resource"] 
+