Release: 1.6.5

VERSION

@@ -1 +1 @@
-1.6.4
+1.6.5

sources/aft-lambda-layer/aft_common/account_provisioning_framework.py

@@ -298,12 +298,12 @@ def persist_metadata(
 def get_ssm_parameters_names_by_path(session: Session, path: str) -> List[str]:
 
     client = session.client("ssm")
-    response = client.get_parameters_by_path(Path=path, Recursive=True)
-    logger.debug(response)
+    paginator = client.get_paginator("get_parameters_by_path")
+    pages = paginator.paginate(Path=path, Recursive=True)
 
     parameter_names = []
-    for p in response["Parameters"]:
-        parameter_names.append(p["Name"])
+    for page in pages:
+        parameter_names.extend([param["Name"] for param in page["Parameters"]])
 
     return parameter_names
 
@@ -313,18 +313,16 @@ def delete_ssm_parameters(session: Session, parameters: Sequence[str]) -> None:
     if len(parameters) > 0:
         client = session.client("ssm")
         response = client.delete_parameters(Names=parameters)
-        logger.info(response)
 
 
-def create_ssm_parameters(session: Session, parameters: Dict[str, str]) -> None:
+def put_ssm_parameters(session: Session, parameters: Dict[str, str]) -> None:
 
     client = session.client("ssm")
 
     for key, value in parameters.items():
         response = client.put_parameter(
             Name=SSM_PARAMETER_PATH + key, Value=value, Type="String", Overwrite=True
         )
-        logger.info(response)
 
 
 def tag_account(

sources/aft-lambda-layer/aft_common/aft_utils.py

@@ -481,72 +481,6 @@ def get_all_aft_account_ids(session: Session) -> List[str]:
     return aft_account_ids
 
 
-def get_account_ids_in_ous(
-    session: Session, ou_names: List[str]
-) -> Optional[List[str]]:
-    client: OrganizationsClient = session.client("organizations")
-    logger.info("Getting Account IDs in the following OUs: " + str(ou_names))
-    ou_ids = []
-    account_ids = []
-    for n in ou_names:
-        ou_ids.append(get_org_ou_id(session, n))
-    logger.info("OU IDs: " + str(ou_ids))
-    for ou_id in ou_ids:
-        if ou_id is not None:
-            logger.info("Listing accounts in the OU ID " + ou_id)
-
-            response = client.list_children(ParentId=ou_id, ChildType="ACCOUNT")
-            children = response["Children"]
-            while "NextToken" in response:
-                response = client.list_children(
-                    ParentId=ou_id,
-                    ChildType="ACCOUNT",
-                    NextToken=response["NextToken"],
-                )
-                children.extend(response["Children"])
-
-            logger.info(str(children))
-
-            for a in children:
-                account_ids.append(a["Id"])
-        else:
-            logger.info("OUs in " + str(ou_names) + " was not found")
-    logger.info("Account IDs: " + str(account_ids))
-    if len(account_ids) > 0:
-        return account_ids
-    else:
-        return None
-
-
-def get_org_ou_id(session: Session, ou_name: str) -> Optional[str]:
-    client: OrganizationsClient = session.client("organizations")
-    logger.info("Listing Org Roots")
-    list_roots_response = client.list_roots(MaxResults=1)
-    logger.info(list_roots_response)
-    root_id = list_roots_response["Roots"][0]["Id"]
-    logger.info("Root ID is " + root_id)
-
-    logger.info("Listing OUs in the Organization")
-
-    list_ou_response = client.list_organizational_units_for_parent(ParentId=root_id)
-    ous = list_ou_response["OrganizationalUnits"]
-    while "NextToken" in list_ou_response:
-        list_ou_response = client.list_organizational_units_for_parent(
-            ParentId=root_id, NextToken=list_ou_response["NextToken"]
-        )
-        ous.extend(list_ou_response["OrganizationalUnits"])
-
-    logger.info(ous)
-
-    for ou in ous:
-        if ou["Name"] == ou_name:
-            ou_id: str = ou["Id"]
-            logger.info("OU ID for " + ou_name + " is " + ou_id)
-            return ou_id
-
-    return None
-
-
 def get_accounts_by_tags(
     aft_mgmt_session: Session, ct_mgmt_session: Session, tags: List[Dict[str, str]]
 ) -> Optional[List[str]]:

sources/aft-lambda-layer/aft_common/customizations.py

@@ -8,6 +8,7 @@
 
 import aft_common.aft_utils as utils
 import jsonschema
+from aft_common.organizations import OrganizationsAgent
 from boto3.session import Session
 
 CUSTOMIZATIONS_PIPELINE_PATTERN = "^\d\d\d\d\d\d\d\d\d\d\d\d-.*$"
@@ -201,11 +202,10 @@ def get_included_accounts(
             core_accounts = get_core_accounts(session)
             included_accounts.extend(core_accounts)
         if d["type"] == "ous":
-            ou_accounts = utils.get_account_ids_in_ous(
-                ct_mgmt_session, d["target_value"]
+            orgs_agent = OrganizationsAgent(ct_mgmt_session)
+            included_accounts.extend(
+                orgs_agent.get_account_ids_in_ous(ou_names=d["target_value"])
             )
-            if ou_accounts is not None:
-                included_accounts.extend(ou_accounts)
         if d["type"] == "tags":
             tag_accounts = utils.get_accounts_by_tags(
                 session, ct_mgmt_session, d["target_value"]
@@ -234,11 +234,10 @@ def get_excluded_accounts(
             core_accounts = get_core_accounts(session)
             excluded_accounts.extend(core_accounts)
         if d["type"] == "ous":
-            ou_accounts = utils.get_account_ids_in_ous(
-                ct_mgmt_session, d["target_value"]
+            orgs_agent = OrganizationsAgent(ct_mgmt_session)
+            excluded_accounts.extend(
+                orgs_agent.get_account_ids_in_ous(ou_names=d["target_value"])
             )
-            if ou_accounts is not None:
-                excluded_accounts.extend(ou_accounts)
         if d["type"] == "tags":
             tag_accounts = utils.get_accounts_by_tags(
                 session, ct_mgmt_session, d["target_value"]

sources/aft-lambda-layer/aft_common/organizations.py

@@ -1,8 +1,11 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
-from typing import TYPE_CHECKING, List, Optional
+import re
+from copy import deepcopy
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
+from aft_common.aft_utils import get_logger
 from boto3.session import Session
 
 if TYPE_CHECKING:
@@ -21,45 +24,151 @@
     OrganizationalUnitTypeDef = object
     AccountTypeDef = object
 
+logger = get_logger()
+
 
 class OrganizationsAgent:
     ROOT_OU = "Root"
+    # https://docs.aws.amazon.com/organizations/latest/APIReference/API_OrganizationalUnit.html
+    # Ex: Sandbox (ou-1234-zxcv)
+    OU_ID_PATTERN = r"\(ou-.*\)"
+    OU_NAME_PATTERN = r".{1,128}"
+    NESTED_OU_NAME_PATTERN = (
+        rf"{OU_NAME_PATTERN}\s{OU_ID_PATTERN}"  # <Name> space (<Id>)
+    )
 
     def __init__(self, ct_management_session: Session):
         self.orgs_client: OrganizationsClient = ct_management_session.client(
             "organizations"
         )
 
+        # Memoize expensive all-org traversal
+        self.org_ous: Optional[List[OrganizationalUnitTypeDef]] = None
+
+    @staticmethod
+    def ou_name_is_nested_format(ou_name: str) -> bool:
+        pattern = re.compile(OrganizationsAgent.NESTED_OU_NAME_PATTERN)
+        if pattern.match(ou_name) is not None:
+            return True
+        return False
+
+    @staticmethod
+    def get_name_and_id_from_nested_ou(
+        nested_ou_name: str,
+    ) -> Optional[Tuple[str, str]]:
+        if not OrganizationsAgent.ou_name_is_nested_format(ou_name=nested_ou_name):
+            return None
+
+        pattern = re.compile(OrganizationsAgent.OU_ID_PATTERN)
+        match = pattern.search(nested_ou_name)
+        if match is None:
+            return None
+        first_id_idx, last_id_idx = match.span()
+
+        # Grab the matched ID from the nested-ou-string using the span,
+        id = nested_ou_name[first_id_idx:last_id_idx]
+        id = id.strip("()")
+
+        # The name is what remains of the nested OU without the ID, minus
+        # the whitespace between the name and ID
+        name = nested_ou_name[: first_id_idx - 1]
+        return (name, id)
+
     def get_root_ou_id(self) -> str:
         return self.orgs_client.list_roots()["Roots"][0]["Id"]
 
     def get_ous_for_root(self) -> List[OrganizationalUnitTypeDef]:
+        return self.get_children_ous_from_parent_id(parent_id=self.get_root_ou_id())
+
+    def get_all_org_ous(self) -> List[OrganizationalUnitTypeDef]:
+        # Memoize calls / cache previous results
+        # Cache is not shared between invocations so staleness due to org updates is unlikely
+        if self.org_ous is not None:
+            return self.org_ous
+
+        # Including the root OU
+        list_root_response = self.orgs_client.list_roots()
+        root_ou: OrganizationalUnitTypeDef = {
+            "Id": list_root_response["Roots"][0]["Id"],
+            "Arn": list_root_response["Roots"][0]["Arn"],
+            "Name": list_root_response["Roots"][0]["Name"],
+        }
+
+        org_ous = [root_ou]
+
+        # Get the children OUs of the root as the first pass
+        root_children = self.get_ous_for_root()
+        org_ous.extend(root_children)
+
+        # Exclude root to avoid double counting children
+        ous_to_query = deepcopy(root_children)
+
+        # Recursively search all children OUs for further children
+        while len(ous_to_query) > 0:
+            parent_id: str = ous_to_query.pop()["Id"]
+            children_ous = self.get_children_ous_from_parent_id(parent_id=parent_id)
+            org_ous.extend(children_ous)
+            ous_to_query.extend(children_ous)
+
+        self.org_ous = org_ous
+
+        return self.org_ous
+
+    def get_children_ous_from_parent_id(
+        self, parent_id: str
+    ) -> List[OrganizationalUnitTypeDef]:
+
         paginator = self.orgs_client.get_paginator(
             "list_organizational_units_for_parent"
         )
-        pages = paginator.paginate(ParentId=self.get_root_ou_id())
-        ous_under_root = []
+        pages = paginator.paginate(ParentId=parent_id)
+        children_ous = []
         for page in pages:
-            ous_under_root.extend(page["OrganizationalUnits"])
-        return ous_under_root
+            children_ous.extend(page["OrganizationalUnits"])
+        return children_ous
 
-    def list_tags_for_resource(self, resource: str) -> List[TagTypeDef]:
-        return self.orgs_client.list_tags_for_resource(ResourceId=resource)["Tags"]
+    def get_ou_ids_from_ou_names(self, target_ou_names: List[str]) -> List[str]:
+        ous = self.get_all_org_ous()
+        ou_map = {}
+
+        # Convert list of OUs to name->id map for constant time lookups
+        for ou in ous:
+            ou_map[ou["Name"]] = ou["Id"]
 
-    def get_ou_id_for_account_id(
+        # Search the map for every target exactly once
+        matched_ou_ids = []
+        for target_name in target_ou_names:
+            # Only match nested OU targets if both name and ID are the same
+            nested_parsed = OrganizationsAgent.get_name_and_id_from_nested_ou(
+                nested_ou_name=target_name
+            )
+            if nested_parsed is not None:  # Nested OU pattern matched!
+                target_name, target_id = nested_parsed
+                if ou_map[target_name] == target_id:
+                    matched_ou_ids.append(ou_map[target_name])
+            else:
+                if target_name in ou_map:
+                    matched_ou_ids.append(ou_map[target_name])
+
+        return matched_ou_ids
+
+    def get_ou_from_account_id(
         self, account_id: str
-    ) -> Optional[DescribeOrganizationalUnitResponseTypeDef]:
+    ) -> Optional[OrganizationalUnitTypeDef]:
         if self.account_id_is_member_of_root(account_id=account_id):
-            return self.orgs_client.describe_organizational_unit(
-                OrganizationalUnitId=self.get_root_ou_id()
-            )
-        ou_ids = [ou["Id"] for ou in self.get_ous_for_root()]
-        for ou_id in ou_ids:
-            ou_accounts = self.get_accounts_for_ou(ou_id=ou_id)
-            if account_id in [account_object["Id"] for account_object in ou_accounts]:
-                return self.orgs_client.describe_organizational_unit(
-                    OrganizationalUnitId=ou_id
-                )
+            list_root_response = self.orgs_client.list_roots()
+            root_ou: OrganizationalUnitTypeDef = {
+                "Id": list_root_response["Roots"][0]["Id"],
+                "Arn": list_root_response["Roots"][0]["Arn"],
+                "Name": list_root_response["Roots"][0]["Name"],
+            }
+            return root_ou
+
+        ous = self.get_all_org_ous()
+        for ou in ous:
+            account_ids = [acct["Id"] for acct in self.get_accounts_for_ou(ou["Id"])]
+            if account_id in account_ids:
+                return ou
         return None
 
     def get_accounts_for_ou(self, ou_id: str) -> List[AccountTypeDef]:
@@ -70,6 +179,15 @@ def get_accounts_for_ou(self, ou_id: str) -> List[AccountTypeDef]:
             accounts.extend(page["Accounts"])
         return accounts
 
+    def get_account_ids_in_ous(self, ou_names: List[str]) -> List[str]:
+        ou_ids = self.get_ou_ids_from_ou_names(target_ou_names=ou_names)
+        account_ids = []
+        for ou_id in ou_ids:
+            account_ids.extend(
+                [acct["Id"] for acct in self.get_accounts_for_ou(ou_id=ou_id)]
+            )
+        return account_ids
+
     def account_id_is_member_of_root(self, account_id: str) -> bool:
         root_id = self.get_root_ou_id()
         accounts_under_root = self.get_accounts_for_ou(ou_id=root_id)
@@ -78,9 +196,11 @@ def account_id_is_member_of_root(self, account_id: str) -> bool:
     def ou_contains_account(self, ou_name: str, account_id: str) -> bool:
         if ou_name == OrganizationsAgent.ROOT_OU:
             return self.account_id_is_member_of_root(account_id=account_id)
-        current_ou = self.get_ou_id_for_account_id(account_id=account_id)
+        current_ou = self.get_ou_from_account_id(account_id=account_id)
         if current_ou:
-            current_ou_name = current_ou["OrganizationalUnit"]["Name"]
-            if current_ou_name == ou_name:
+            if ou_name == current_ou["Name"]:
                 return True
         return False
+
+    def list_tags_for_resource(self, resource: str) -> List[TagTypeDef]:
+        return self.orgs_client.list_tags_for_resource(ResourceId=resource)["Tags"]

src/aft_lambda/aft_account_provisioning_framework/aft_account_provisioning_framework_account_metadata_ssm.py

@@ -11,9 +11,9 @@
 from aft_common.account_provisioning_framework import (
     SSM_PARAMETER_PATH,
     ProvisionRoles,
-    create_ssm_parameters,
     delete_ssm_parameters,
     get_ssm_parameters_names_by_path,
+    put_ssm_parameters,
 )
 from aft_common.auth import AuthClient
 
@@ -71,7 +71,7 @@ def lambda_handler(event: Dict[str, Any], context: LambdaContext) -> None:
 
         # Update / Add SSM parameters for custom fields provided
         logger.info(message=f"Adding/Updating SSM params: {custom_fields}")
-        create_ssm_parameters(target_account_session, custom_fields)
+        put_ssm_parameters(target_account_session, custom_fields)
 
     except Exception as error:
         notifications.send_lambda_failure_sns_message(