Handle account ID principals in cross-account trust checks

Trust policies can list 12-digit account IDs in addition to arns ARNs. Normalizing principals through a shared helper (and exercising exclusions/current-account filtering in tests) keeps --flag-trust-policies from missing those roles that are assumable from other accounts when only the account ID is supplied.

Author: bhendo

cloudsplaining/scan/assume_role_policy_document.py

@@ -7,12 +7,9 @@
 # or https://opensource.org/licenses/BSD-3-Clause
 from __future__ import annotations
 
-import contextlib
 import logging
 from typing import Any
 
-from policy_sentry.util.arns import get_account_from_arn
-
 from cloudsplaining.scan.resource_policy_document import (
     ResourcePolicyDocument,
     ResourceStatement,
@@ -22,6 +19,7 @@
     DEFAULT_EXCLUSIONS,
     Exclusions,
 )
+from cloudsplaining.shared.utils import get_account_id_from_principal
 
 logger = logging.getLogger(__name__)
 
@@ -158,19 +156,17 @@ def role_assumable_by_cross_account_principals(self) -> list[str]:
         if self.effect.lower() != "allow":
             return []
 
-        other_account_principals = []
-        for principal in self.principals:
-            # Check if this is an AWS IAM principal from another account
-            if principal.startswith("arn:aws:iam::"):
-                with contextlib.suppress(Exception):
-                    principal_account_id = get_account_from_arn(principal)
-                    # Only include if it's from a different account (or if we don't know our current account)
-                    # and the account is not in the known-accounts exclusion list
-                    if (
-                        self.current_account_id is None or principal_account_id != self.current_account_id
-                    ) and principal_account_id not in self.exclusions.known_accounts:
-                        other_account_principals.append(principal)
-        return other_account_principals
+        return [
+            principal
+            for principal in self.principals
+            if (
+                principal_account_id := get_account_id_from_principal(principal)
+            )
+            and (
+                self.current_account_id is None or principal_account_id != self.current_account_id
+            )
+            and principal_account_id not in self.exclusions.known_accounts
+        ]
 
     @property
     def role_assumable_by_any_principal(self) -> list[str]:

cloudsplaining/shared/exclusions.py

@@ -66,10 +66,10 @@ def _exclude_actions(self) -> list[str]:
         always_exclude_actions = [x.lower() for x in exclude_actions]
         return always_exclude_actions
 
-    def _known_accounts(self) -> list[str]:
+    def _known_accounts(self) -> set[str]:
         provided_known_accounts = self.config.get("known-accounts", [])
         # Normalize for comparisons - remove empty strings
-        known_accounts = [account for account in provided_known_accounts if account.strip()]
+        known_accounts = {account for account in provided_known_accounts if account.strip()}
         return known_accounts
 
     def is_action_always_included(self, action_in_question: str) -> bool | str:

cloudsplaining/shared/utils.py

@@ -7,6 +7,7 @@
 # or https://opensource.org/licenses/BSD-3-Clause
 from __future__ import annotations
 
+import contextlib
 import json
 import logging
 import os
@@ -20,6 +21,7 @@
     remove_actions_not_matching_access_level,
 )
 from policy_sentry.querying.all import get_all_service_prefixes
+from policy_sentry.util.arns import get_account_from_arn
 
 all_service_prefixes = get_all_service_prefixes()
 logger = logging.getLogger(__name__)
@@ -169,3 +171,15 @@ def write_json_to_file(file: str, content: str) -> None:
         os.remove(file)
 
     Path(file).write_text(json.dumps(content, indent=4, default=str), encoding="utf-8")
+
+
+def get_account_id_from_principal(principal: str) -> str | None:
+    """Return the AWS account ID for a principal ARN or bare 12-digit account identifier."""
+    principal_stripped = principal.strip()
+    if principal_stripped.startswith("arn:aws:iam::"):
+        with contextlib.suppress(Exception):
+            return get_account_from_arn(principal_stripped)
+        return None
+    if principal_stripped.isdigit() and len(principal_stripped) == 12:
+        return principal_stripped
+    return None

test/scanning/test_trust_policies.py

@@ -201,6 +201,38 @@ def test_assume_role_assumable_by_cross_account_principals(self):
             ["arn:aws:iam::098765432109:role/testrole"],
         )
 
+        # Test bare account ID principals
+        statement_account_ids = dict(
+            Effect="Allow",
+            Principal={"AWS": ["123456789012", "210987654321"]},
+            Action=["sts:AssumeRole"],
+        )
+        assume_role_account_ids = AssumeRoleStatement(statement_account_ids)
+        self.assertListEqual(
+            assume_role_account_ids.role_assumable_by_cross_account_principals,
+            ["123456789012", "210987654321"],
+        )
+
+        # Test bare account ID filtering by current account
+        assume_role_account_ids_filtered = AssumeRoleStatement(
+            statement_account_ids, current_account_id="123456789012"
+        )
+        self.assertListEqual(
+            assume_role_account_ids_filtered.role_assumable_by_cross_account_principals,
+            ["210987654321"],
+        )
+
+        # Test bare account ID filtering with exclusions
+        exclusions_account_ids_config = {"known-accounts": ["210987654321"]}
+        exclusions_account_ids = Exclusions(exclusions_account_ids_config)
+        assume_role_account_ids_exclusions = AssumeRoleStatement(
+            statement_account_ids, exclusions=exclusions_account_ids
+        )
+        self.assertListEqual(
+            assume_role_account_ids_exclusions.role_assumable_by_cross_account_principals,
+            ["123456789012"],
+        )
+
         # Test conditions that should return empty results
         # No sts:AssumeRole action
         statement_no_assume_role = dict(
@@ -232,14 +264,23 @@ def test_assume_role_assumable_by_cross_account_principals(self):
                     "Principal": {"AWS": "arn:aws:iam::098765432109:user/testuser"},
                     "Action": "sts:AssumeRole",
                 },
+                {
+                    "Effect": "Allow",
+                    "Principal": {"AWS": "210987654321"},
+                    "Action": "sts:AssumeRole",
+                },
             ],
         }
         policy_doc = AssumeRolePolicyDocument(policy)
-        expected_policy = ["arn:aws:iam::123456789012:root", "arn:aws:iam::098765432109:user/testuser"]
+        expected_policy = [
+            "arn:aws:iam::123456789012:root",
+            "arn:aws:iam::098765432109:user/testuser",
+            "210987654321",
+        ]
         self.assertListEqual(policy_doc.role_assumable_by_cross_account_principals, expected_policy)
 
         # Test AssumeRolePolicyDocument with exclusions
-        partial_exclusions_config_doc = {"known-accounts": ["123456789012"]}
+        partial_exclusions_config_doc = {"known-accounts": ["123456789012", "210987654321"]}
         partial_exclusions_doc = Exclusions(partial_exclusions_config_doc)
         policy_doc_with_exclusions = AssumeRolePolicyDocument(policy, exclusions=partial_exclusions_doc)
         expected_policy_with_exclusions = ["arn:aws:iam::098765432109:user/testuser"]

test/shared/test_utils.py

@@ -1,5 +1,9 @@
 import unittest
-from cloudsplaining.shared.utils import remove_wildcard_only_actions, remove_read_level_actions
+from cloudsplaining.shared.utils import (
+    get_account_id_from_principal,
+    remove_read_level_actions,
+    remove_wildcard_only_actions,
+)
 
 
 class TestUtils(unittest.TestCase):
@@ -20,3 +24,11 @@ def test_remove_read_level_actions(self):
         result = remove_read_level_actions(actions)
         expected_result = ["ecr:PutImage"]
         self.assertListEqual(result, expected_result)
+
+    def test_get_account_id_from_principal(self):
+        self.assertEqual(
+            get_account_id_from_principal("arn:aws:iam::123456789012:role/test"),
+            "123456789012",
+        )
+        self.assertEqual(get_account_id_from_principal(" 210987654321"), "210987654321")
+        self.assertIsNone(get_account_id_from_principal("invalid"))