extend target_ids to support AWS Organization Units (#30)

* extend target_ids to support AWS Organization Units * refactor validation of accounts & ou's Co-authored-by: Dave Connell <[email protected]> * use build_full_result instead of paginator Co-authored-by: Dave Connell <[email protected]> * minor code quality adjustments * extend ignore_ids to support ou's * fix variable naming * refactor _resolve_target_accounts() * fix variable after merging from master * adjust _gather_ignored_accounts() * add tests for wrong target/ignore id type * add test cases for targeting/ignoring ou's * make target/ignore id wrong type msg more specific Co-authored-by: [email protected] <[email protected]> Co-authored-by: Dave Connell <[email protected]>
connelldave · Mar 4, 2022 · 8b2cff3 · 8b2cff3
1 parent 283513e
commit 8b2cff3
Show file tree

Hide file tree

Showing 5 changed files with 249 additions and 64 deletions.
diff --git a/README.md b/README.md
@@ -125,8 +125,8 @@ Equivalent to:
 
 `target_ids`: List[str]
 
-A list of AWS accounts as strings to attempt to assume role in to. When unset,
-default attempts to use every available account ID in an AWS organization.
+A list of AWS accounts and/or AWS Organization Units as strings to attempt to assume role in to. When unset,
+default attempts to use every available account ID in an AWS organization. When specifing ou's it will recursivly fetch all child ou's as well.
 
 `ignore_ids`: List[str]
 

diff --git a/botocove/cove_host_account.py b/botocove/cove_host_account.py
@@ -1,5 +1,6 @@
 import logging
-from typing import Any, Iterable, List, Literal, Optional, Sequence, Set, Union
+import re
+from typing import Any, Iterable, List, Literal, Optional, Sequence, Set, Tuple, Union
 
 import boto3
 from boto3.session import Session
@@ -114,31 +115,111 @@ def _get_boto3_sts_client(self, assuming_session: Optional[Session]) -> STSClien
         return client
 
     def _resolve_target_accounts(self, target_ids: Optional[List[str]]) -> Set[str]:
-        # Ensure we never run botocove on the account it's being run from
-        validated_ignore_ids = {self.host_account_id}
+        accounts_to_ignore = self._gather_ignored_accounts()
+        logger.info(f"Ignoring account IDs: {accounts_to_ignore=}")
+        accounts_to_target = self._gather_target_accounts(target_ids)
+        final_accounts: Set = accounts_to_target - accounts_to_ignore
+        if len(final_accounts) < 1:
+            raise ValueError(
+                "There are no eligible account ids to run decorated func against"
+            )
+        return final_accounts
 
-        if self.provided_ignore_ids:
-            validated_ignore_ids.update(self._parse_ignore_ids())
-        logger.info(f"Ignoring account IDs: {validated_ignore_ids=}")
+    def _gather_ignored_accounts(self) -> Set[str]:
+        ignored_accounts = {self.host_account_id}
 
-        if target_ids is None:
-            # No target_ids passed, get all accounts in org
-            target_accounts = self._get_active_org_accounts()
+        if self.provided_ignore_ids:
+            accs, ous = self._get_validated_ids(self.provided_ignore_ids)
+            ignored_accounts.update(accs)
+            if ous:
+                accs_from_ous = self._get_all_accounts_by_organization_units(ous)
+                ignored_accounts.update(accs_from_ous)
+
+        return ignored_accounts
+
+    def _gather_target_accounts(self, targets: Optional[List[str]]) -> Set[str]:
+        if targets:
+            accs, ous = self._get_validated_ids(targets)
+            if ous:
+                accs_from_ous = self._get_all_accounts_by_organization_units(ous)
+                accs.extend(accs_from_ous)
+            return set(accs)
         else:
-            # Specific list of IDs passed
-            target_accounts = set(target_ids)
-
-        return target_accounts - validated_ignore_ids
-
-    def _parse_ignore_ids(self) -> Set[str]:
-        if not isinstance(self.provided_ignore_ids, list):
-            raise TypeError("ignore_ids must be a list of account IDs")
-        for account_id in self.provided_ignore_ids:
-            if len(account_id) != 12:
-                raise TypeError("All ignore_id in list must be 12 character strings")
-            if not isinstance(account_id, str):
-                raise TypeError("All ignore_id list entries must be strings")
-        return set(self.provided_ignore_ids)
+            # No target_ids passed, getting all accounts in org
+            return self._get_active_org_accounts()
+
+    def _get_validated_ids(self, ids: List[str]) -> Tuple[List[str], List[str]]:
+
+        accounts: List[str] = []
+        ous: List[str] = []
+
+        for current_id in ids:
+            if not isinstance(current_id, str):
+                raise TypeError(
+                    f"{current_id} is an incorrect type: all account and ou id's must be strings not {type(current_id)}"  # noqa E501
+                )
+            if re.match(r"^\d{12}$", current_id):
+                accounts.append(current_id)
+                continue
+            if re.match(r"^ou-[0-9a-z]{4,32}-[a-z0-9]{8,32}$", current_id):
+                ous.append(current_id)
+                continue
+            raise ValueError(
+                f"provided id is neither an aws account nor an ou: {current_id}"
+            )
+
+        return accounts, ous
+
+    def _get_all_accounts_by_organization_units(
+        self, target_ous: List[str]
+    ) -> List[str]:
+
+        account_list: List[str] = []
+
+        for parent_ou in target_ous:
+
+            current_ou_list: List[str] = []
+
+            # current_ou_list is mutated and recursivly populated with all childs
+            self._get_all_child_ous(parent_ou, current_ou_list)
+
+            # for complete list add parent ou as well to list of child ous
+            current_ou_list.append(parent_ou)
+
+            account_list.extend(
+                self._get_accounts_by_organization_units(current_ou_list)
+            )
+
+        return account_list
+
+    def _get_all_child_ous(self, parent_ou: str, ou_list: List[str]) -> None:
+
+        child_ous = (
+            self.org_client.get_paginator("list_children")
+            .paginate(ChildType="ORGANIZATIONAL_UNIT", ParentId=parent_ou)
+            .build_full_result()
+        )
+        child_ous_list = [ou["Id"] for ou in child_ous["Children"]]
+        ou_list.extend(child_ous_list)
+
+        for ou in child_ous_list:
+            self._get_all_child_ous(ou, ou_list)
+
+    def _get_accounts_by_organization_units(
+        self, organization_units: List[str]
+    ) -> List[str]:
+
+        account_list: List[str] = []
+
+        for ou in organization_units:
+            ou_children = (
+                self.org_client.get_paginator("list_children")
+                .paginate(ChildType="ACCOUNT", ParentId=ou)
+                .build_full_result()
+            )
+            account_list.extend(acc["Id"] for acc in ou_children["Children"])
+
+        return account_list
 
     def _get_active_org_accounts(self) -> Set[str]:
         all_org_accounts = (

diff --git a/tests/moto_mock_org/test_host_account/test_resolve_targets.py b/tests/moto_mock_org/test_host_account/test_resolve_targets.py
@@ -1,5 +1,5 @@
 from botocove.cove_host_account import CoveHostAccount
-from tests.moto_mock_org.moto_models import SmallOrg
+from tests.moto_mock_org.moto_models import LargeOrg, SmallOrg
 
 
 def test_target_all_in_org(mock_small_org: SmallOrg) -> None:
@@ -45,6 +45,88 @@ def test_target_all_in_org_ignore_one(mock_small_org: SmallOrg) -> None:
     assert ignore_acc not in account_ids
 
 
+def test_target_all_in_org_ignore_one_ou(mock_small_org: SmallOrg) -> None:
+
+    ignore_ou = mock_small_org.new_org3
+    ignore_ou_accounts = mock_small_org.account_group_one
+
+    host_account = CoveHostAccount(
+        target_ids=None,
+        ignore_ids=[ignore_ou],
+        rolename=None,
+        role_session_name=None,
+        policy=None,
+        policy_arns=None,
+        org_master=True,
+        assuming_session=None,
+        regions=None,
+        thread_workers=20,
+    )
+    sessions = host_account.get_cove_sessions()
+    account_ids = [acc_id["Id"] for acc_id in sessions]
+
+    assert not any(acc in account_ids for acc in ignore_ou_accounts)
+    assert mock_small_org.master_acc_id not in account_ids
+    assert set(account_ids) == (
+        set(mock_small_org.all_accounts[1:]) - set(ignore_ou_accounts)
+    )
+
+
+def test_large_org_target_all_in_org_ignore_one_ou(mock_large_org: LargeOrg) -> None:
+
+    ignore_ou = mock_large_org.ou_B
+    ignore_ou_accounts = mock_large_org.account_group_one
+
+    host_account = CoveHostAccount(
+        target_ids=None,
+        ignore_ids=[ignore_ou],
+        rolename=None,
+        role_session_name=None,
+        policy=None,
+        policy_arns=None,
+        org_master=True,
+        assuming_session=None,
+        regions=None,
+        thread_workers=20,
+    )
+    sessions = host_account.get_cove_sessions()
+    account_ids = [acc_id["Id"] for acc_id in sessions]
+
+    assert not any(acc in account_ids for acc in ignore_ou_accounts)
+    assert mock_large_org.master_acc_id not in account_ids
+    assert set(account_ids) == (
+        set(mock_large_org.all_accounts[1:]) - set(ignore_ou_accounts)
+    )
+
+
+def test_large_org_target_all_in_org_ignore_two_ous(mock_large_org: LargeOrg) -> None:
+
+    ignore_ou = [mock_large_org.ou_B, mock_large_org.ou_C]
+    ignore_ou_accounts = mock_large_org.account_group_one
+    ignore_ou_accounts.extend(mock_large_org.account_group_two)
+
+    host_account = CoveHostAccount(
+        target_ids=None,
+        ignore_ids=ignore_ou,
+        rolename=None,
+        role_session_name=None,
+        policy=None,
+        policy_arns=None,
+        org_master=True,
+        assuming_session=None,
+        regions=None,
+        thread_workers=20,
+    )
+    sessions = host_account.get_cove_sessions()
+    account_ids = [acc_id["Id"] for acc_id in sessions]
+
+    assert not any(acc in account_ids for acc in ignore_ou_accounts)
+    assert mock_large_org.master_acc_id not in account_ids
+    assert set(account_ids) == (
+        set(mock_large_org.all_accounts[1:]) - set(ignore_ou_accounts)
+    )
+
+
 def test_target_targets_and_ignores(mock_small_org: SmallOrg) -> None:
 
     target_accs = mock_small_org.all_accounts[1:]

diff --git a/tests/test_decorator.py b/tests/test_decorator.py
@@ -59,40 +59,6 @@ def simple_func(session: CoveSession) -> str:
     assert len(cove_output["Results"]) == 2
 
 
-def test_target_ids(mock_boto3_session: MagicMock) -> None:
-    @cove(assuming_session=mock_boto3_session, target_ids=["1"])
-    def simple_func(session: CoveSession) -> str:
-        return "hello"
-
-    cove_output = simple_func()
-    # One account in target_ids, two in mock response.
-    # simple_func calls == one mock AWS accounts
-    assert len(cove_output["Results"]) == 1
-
-
-def test_empty_target_ids(mock_boto3_session: MagicMock) -> None:
-    @cove(assuming_session=mock_boto3_session, target_ids=[])
-    def simple_func(session: CoveSession) -> str:
-        return "hello"
-
-    with pytest.raises(
-        ValueError,
-        match="There are no eligible account ids to run decorated func against",  # noqa: E501
-    ):
-        simple_func()
-
-
-def test_ignore_ids(mock_boto3_session: MagicMock) -> None:
-    @cove(assuming_session=mock_boto3_session, ignore_ids=["123123123123"])
-    def simple_func(session: CoveSession) -> str:
-        return "hello"
-
-    cove_output = simple_func()
-    # Two in mock response, one ignored.
-    # simple_func calls == one mock AWS accounts
-    assert len(cove_output["Results"]) == 1
-
-
 def test_target_and_ignore_ids(mock_boto3_session: MagicMock) -> None:
     @cove(
         assuming_session=mock_boto3_session,

diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py
@@ -67,13 +67,13 @@ def simple_func(session: CoveSession) -> str:
 
 
 def test_handled_exception_in_wrapped_func(mock_boto3_session: MagicMock) -> None:
-    @cove(assuming_session=mock_boto3_session, target_ids=["123"])
+    @cove(assuming_session=mock_boto3_session, target_ids=["456456456456"])
     def simple_func(session: CoveSession) -> None:
         raise Exception("oh no")
 
     results = simple_func()
     expected = {
-        "Id": "123",
+        "Id": "456456456456",
         "RoleName": "OrganizationAccountAccessRole",
         "AssumeRoleSuccess": True,
         "Arn": "hello-arn",
@@ -93,7 +93,11 @@ def simple_func(session: CoveSession) -> None:
 
 
 def test_raised_exception_in_wrapped_func(mock_boto3_session: MagicMock) -> None:
-    @cove(assuming_session=mock_boto3_session, target_ids=["123"], raise_exception=True)
+    @cove(
+        assuming_session=mock_boto3_session,
+        target_ids=["456456456456"],
+        raise_exception=True,
+    )
     def simple_func(session: CoveSession) -> None:
         raise Exception("oh no")
 
@@ -110,8 +114,60 @@ def test_malformed_ignore_ids(mock_boto3_session: MagicMock) -> None:
     def simple_func(session: CoveSession) -> str:
         return "hello"
 
+    with pytest.raises(
+        ValueError,
+        match=("provided id is neither an aws account nor an ou: cat"),
+    ):
+        simple_func()
+
+
+def test_malformed_ignore_ids_type(mock_boto3_session: MagicMock) -> None:
+    @cove(
+        assuming_session=mock_boto3_session,
+        target_ids=None,
+        ignore_ids=[456456456456],  # type: ignore
+    )
+    def simple_func(session: CoveSession) -> str:
+        return "hello"
+
+    with pytest.raises(
+        TypeError,
+        match=(
+            "456456456456 is an incorrect type: all account and ou id's must be strings not <class 'int'>"  # noqa E501
+        ),
+    ):
+        simple_func()
+
+
+def test_malformed_target_id(mock_boto3_session: MagicMock) -> None:
+    @cove(
+        assuming_session=mock_boto3_session,
+        target_ids=["xu-gzxu-393a2l5b"],
+        ignore_ids=["456456456456"],
+    )
+    def simple_func(session: CoveSession) -> str:
+        return "hello"
+
+    with pytest.raises(
+        ValueError,
+        match=("provided id is neither an aws account nor an ou: xu-gzxu-393a2l5b"),
+    ):
+        simple_func()
+
+
+def test_malformed_target_id_type(mock_boto3_session: MagicMock) -> None:
+    @cove(
+        assuming_session=mock_boto3_session,
+        target_ids=[456456456456],  # type: ignore
+        ignore_ids=[],
+    )
+    def simple_func(session: CoveSession) -> str:
+        return "hello"
+
     with pytest.raises(
         TypeError,
-        match=("All ignore_id in list must be 12 character strings"),
+        match=(
+            "456456456456 is an incorrect type: all account and ou id's must be strings not <class 'int'>"  # noqa E501
+        ),
     ):
         simple_func()