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

… access (#408)

Logic chanced slightly so that enabling cross account access is performed while reducing the number of API calls in the deployment account.

Before it would update the IAM policies per region (3 x number of regions), increasing the chance of being throttled.
This changes it to 3 API calls per account (or how many role/policy mappings are passed in).

Changes include backwards compatibility, just in case customers have downstream modifications to this lambda function.

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"] 
+