From 7db6adcf91e252e3dba37f3bba65bec0a90f0700 Mon Sep 17 00:00:00 2001
From: anayeaye <alexandra@developmentseed.org>
Date: Thu, 28 Sep 2023 13:29:01 -0600
Subject: [PATCH] fix: if provided, apply permissions boundary to all roles
 created within a stack

---
 app.py                                        | 10 ++--
 .../infrastructure/construct.py               | 56 +++++++++++++++++++
 2 files changed, 62 insertions(+), 4 deletions(-)
 create mode 100644 permissions_boundary/infrastructure/construct.py

diff --git a/app.py b/app.py
index f47a2335..710204c6 100644
--- a/app.py
+++ b/app.py
@@ -1,13 +1,14 @@
 #!/usr/bin/env python3
 """ CDK Configuration for the veda-backend stack."""
 
-from aws_cdk import App, Stack, Tags, aws_iam
+from aws_cdk import App, Aspects, Stack, Tags, aws_iam
 from constructs import Construct
 
 from config import veda_app_settings
 from database.infrastructure.construct import RdsConstruct
 from domain.infrastructure.construct import DomainConstruct
 from network.infrastructure.construct import VpcConstruct
+from permissions_boundary.infrastructure.construct import PermissionsBoundaryAspect
 from raster_api.infrastructure.construct import RasterApiLambdaConstruct
 from stac_api.infrastructure.construct import StacApiLambdaConstruct
 
@@ -22,12 +23,13 @@ def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
         super().__init__(scope, construct_id, **kwargs)
 
         if veda_app_settings.permissions_boundary_policy_name:
-            permission_boundary_policy = aws_iam.Policy.from_policy_name(
+            permissions_boundary_policy = aws_iam.Policy.from_policy_name(
                 self,
-                "permission-boundary",
+                "permissions-boundary",
                 veda_app_settings.permissions_boundary_policy_name,
             )
-            aws_iam.PermissionsBoundary.of(self).apply(permission_boundary_policy)
+            aws_iam.PermissionsBoundary.of(self).apply(permissions_boundary_policy)
+            Aspects.of(self).add(PermissionsBoundaryAspect(permissions_boundary_policy))
 
 
 veda_stack = VedaStack(
diff --git a/permissions_boundary/infrastructure/construct.py b/permissions_boundary/infrastructure/construct.py
new file mode 100644
index 00000000..1fd9cbb6
--- /dev/null
+++ b/permissions_boundary/infrastructure/construct.py
@@ -0,0 +1,56 @@
+"""Class that applies permissions boundary to all the roles created within a Stack"""
+from typing import Union
+
+import jsii
+from aws_cdk import IAspect, aws_iam
+from constructs import IConstruct
+from jsii._reference_map import _refs
+from jsii._utils import Singleton
+
+
+@jsii.implements(IAspect)
+class PermissionsBoundaryAspect:
+    """
+    This aspect finds all aws_iam.Role objects in a node (ie. CDK stack) and sets permissions boundary to the given ARN.
+    """
+
+    def __init__(self, permissions_boundary: Union[aws_iam.ManagedPolicy, str]) -> None:
+        """
+        :param permissions_boundary: Either aws_iam.ManagedPolicy object or managed policy's ARN string
+        """
+        self.permissions_boundary = permissions_boundary
+
+    def visit(self, construct_ref: IConstruct) -> None:
+        """
+        construct_ref only contains a string reference to an object. To get the actual object, we need to resolve it using JSII mapping.
+        :param construct_ref: ObjRef object with string reference to the actual object.
+        :return: None
+        """
+        if isinstance(construct_ref, jsii._kernel.ObjRef) and hasattr(
+            construct_ref, "ref"
+        ):
+            kernel = Singleton._instances[
+                jsii._kernel.Kernel
+            ]  # The same object is available as: jsii.kernel
+            resolve = _refs.resolve(kernel, construct_ref)
+        else:
+            resolve = construct_ref
+
+        def _walk(obj):
+            if isinstance(obj, aws_iam.Role):
+                cfn_role = obj.node.find_child("Resource")
+                policy_arn = (
+                    self.permissions_boundary
+                    if isinstance(self.permissions_boundary, str)
+                    else self.permissions_boundary.managed_policy_arn
+                )
+                cfn_role.add_property_override("PermissionsBoundary", policy_arn)
+            else:
+                if hasattr(obj, "permissions_node"):
+                    for c in obj.permissions_node.children:
+                        _walk(c)
+                if hasattr(obj, "node") and obj.node.children:
+                    for c in obj.node.children:
+                        _walk(c)
+
+        _walk(resolve)