From 752909bbfdb8d70086c14933c49543680539342a Mon Sep 17 00:00:00 2001
From: Fanny Jiang <fanny.jiang@datadoghq.com>
Date: Wed, 8 Jul 2026 14:12:16 -0400
Subject: [PATCH 1/6] TON-710: prefer structured policy_statements over legacy
 permission chunks

Fetch full API attributes, render ${AccountId}/${Partition} placeholders,
and attach rendered statement chunks directly, falling back to the legacy
per-chunk action lists when policy_statements isn't present.
---
 .../attach_integration_permissions.py         | 53 +++++++++++++---
 .../attach_integration_permissions_test.py    | 62 +++++++++++++++++++
 .../datadog_integration_permissions.yaml      | 53 +++++++++++++---
 3 files changed, 148 insertions(+), 20 deletions(-)

diff --git a/aws_quickstart/attach_integration_permissions.py b/aws_quickstart/attach_integration_permissions.py
index 7102f889..9e963693 100644
--- a/aws_quickstart/attach_integration_permissions.py
+++ b/aws_quickstart/attach_integration_permissions.py
@@ -39,7 +39,7 @@ class DatadogAPIError(Exception):
     pass
 
 
-def fetch_permissions_from_datadog(api_url):
+def fetch_attributes_from_datadog(api_url):
     headers = {
         "Dd-Aws-Api-Call-Source": API_CALL_SOURCE_HEADER_VALUE,
     }
@@ -53,7 +53,11 @@ def fetch_permissions_from_datadog(api_url):
         error_message = error_body.get('errors', ['Unknown error'])[0]
         raise DatadogAPIError(f"Datadog API error: {error_message}") from e
 
-    return json.loads(response.read())["data"]["attributes"]["permissions"]
+    return json.loads(response.read())["data"]["attributes"]
+
+
+def fetch_permissions_from_datadog(api_url):
+    return fetch_attributes_from_datadog(api_url)["permissions"]
 
 
 def parse_resource_types(raw):
@@ -72,6 +76,31 @@ def build_instrumentation_permissions_url(datadog_site, resource_types):
     return f"https://api.{datadog_site}{INSTRUMENTATION_PERMISSIONS_API_PATH}?{query}"
 
 
+def render_placeholders(value, account_id, partition):
+    if isinstance(value, str):
+        return value.replace("${AccountId}", account_id).replace("${Partition}", partition)
+    if isinstance(value, list):
+        return [render_placeholders(v, account_id, partition) for v in value]
+    if isinstance(value, dict):
+        return {k: render_placeholders(v, account_id, partition) for k, v in value.items()}
+    return value
+
+
+def legacy_chunk_to_statement(chunk):
+    return {"Effect": "Allow", "Action": list(chunk), "Resource": "*"}
+
+
+def resolve_instrumentation_statement_chunks(attributes, account_id, partition):
+    # Prefer structured policy_statements when the API provides them; permissions is the legacy,
+    # pre-chunked fallback. Chunk boundaries and standard-permission filtering land in later work.
+    policy_statements = attributes.get("policy_statements")
+    if policy_statements is not None:
+        rendered = [render_placeholders(s, account_id, partition) for s in policy_statements]
+        return [rendered]
+
+    return [[legacy_chunk_to_statement(chunk)] for chunk in attributes.get("permissions", [])]
+
+
 def _detach_and_delete_policy(iam_client, role_name, policy_arn, policy_name):
     # Detach + delete are both no-ops if the entity is already gone, so callers can blindly
     # iterate the policy-name space without first checking what actually exists.
@@ -139,14 +168,17 @@ def attach_standard_permissions(iam_client, role_name):
 
 
 def _create_and_attach_policy(iam_client, role_name, policy_name, actions):
+    _create_and_attach_statement_policy(
+        iam_client, role_name, policy_name, [{"Effect": "Allow", "Action": actions, "Resource": "*"}]
+    )
+
+
+def _create_and_attach_statement_policy(iam_client, role_name, policy_name, statements):
     policy_json = json.dumps(
-        {
-            "Version": "2012-10-17",
-            "Statement": [{"Effect": "Allow", "Action": actions, "Resource": "*"}],
-        },
+        {"Version": "2012-10-17", "Statement": statements},
         separators=(',', ':'),
     )
-    LOGGER.info(f"Creating policy {policy_name} with {len(actions)} permissions ({len(policy_json)} characters)")
+    LOGGER.info(f"Creating policy {policy_name} with {len(statements)} statements ({len(policy_json)} characters)")
     policy = iam_client.create_policy(PolicyName=policy_name, PolicyDocument=policy_json)
     iam_client.attach_role_policy(RoleName=role_name, PolicyArn=policy['Policy']['Arn'])
 
@@ -179,7 +211,8 @@ def attach_instrumentation_permissions(iam_client, role_name, account_id, partit
     try:
         url = build_instrumentation_permissions_url(datadog_site, resource_types)
         LOGGER.info(f"Fetching instrumentation permissions for {resource_types} from {url}")
-        permission_chunks = fetch_permissions_from_datadog(url)
+        attributes = fetch_attributes_from_datadog(url)
+        chunks = resolve_instrumentation_statement_chunks(attributes, account_id, partition)
     except Exception as e:
         if fail_on_error:
             raise
@@ -190,10 +223,10 @@ def attach_instrumentation_permissions(iam_client, role_name, account_id, partit
         return
 
     cleanup_instrumentation_policies(iam_client, role_name, account_id, partition)
-    for i, chunk in enumerate(permission_chunks):
+    for i, chunk in enumerate(chunks):
         policy_name = f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-{role_name}-{i+1}"
         try:
-            _create_and_attach_policy(iam_client, role_name, policy_name, chunk)
+            _create_and_attach_statement_policy(iam_client, role_name, policy_name, chunk)
         except Exception as e:
             if fail_on_error:
                 raise
diff --git a/aws_quickstart/attach_integration_permissions_test.py b/aws_quickstart/attach_integration_permissions_test.py
index 191ef671..2dae1e72 100644
--- a/aws_quickstart/attach_integration_permissions_test.py
+++ b/aws_quickstart/attach_integration_permissions_test.py
@@ -22,6 +22,9 @@
     cleanup_legacy_base_policies,
     handle_create_update,
     handle_delete,
+    render_placeholders,
+    legacy_chunk_to_statement,
+    resolve_instrumentation_statement_chunks,
     POLICY_NAME_STANDARD,
     BASE_POLICY_PREFIX_INSTRUMENTATION,
     BASE_POLICY_PREFIX_RESOURCE_COLLECTION,
@@ -393,5 +396,64 @@ def test_instrumentation_names_disjoint_from_legacy(self):
         )
 
 
+class TestRenderPlaceholders(unittest.TestCase):
+    def test_renders_in_string(self):
+        self.assertEqual(
+            render_placeholders("arn:${Partition}:iam::${AccountId}:role/x", "123456789012", "aws"),
+            "arn:aws:iam::123456789012:role/x",
+        )
+
+    def test_renders_recursively_in_list_and_dict(self):
+        value = {
+            "Resource": ["arn:${Partition}:iam::${AccountId}:role/x", "arn:${Partition}:s3:::bucket"],
+            "Condition": {"StringEquals": {"iam:PassedToService": ["ec2.${Partition}.example"]}},
+        }
+        rendered = render_placeholders(value, "123456789012", "aws-cn")
+        self.assertEqual(
+            rendered["Resource"],
+            ["arn:aws-cn:iam::123456789012:role/x", "arn:aws-cn:s3:::bucket"],
+        )
+        self.assertEqual(
+            rendered["Condition"]["StringEquals"]["iam:PassedToService"],
+            ["ec2.aws-cn.example"],
+        )
+
+    def test_leaves_non_placeholder_values_unchanged(self):
+        self.assertEqual(render_placeholders("Allow", "123456789012", "aws"), "Allow")
+        self.assertEqual(render_placeholders(42, "123456789012", "aws"), 42)
+
+
+class TestResolveInstrumentationStatementChunks(unittest.TestCase):
+    def test_legacy_fallback_preserves_chunk_boundaries(self):
+        attributes = {"permissions": [["ec2:DescribeInstances"], ["iam:GetRole", "iam:PassRole"]]}
+        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws")
+        self.assertEqual(len(chunks), 2)
+        self.assertEqual(chunks[0], [legacy_chunk_to_statement(["ec2:DescribeInstances"])])
+        self.assertEqual(chunks[1], [legacy_chunk_to_statement(["iam:GetRole", "iam:PassRole"])])
+
+    def test_empty_policy_statements_list_is_not_treated_as_absent(self):
+        # An explicitly present but empty policy_statements must not fall back to the legacy
+        # permissions chunks — presence, not truthiness, decides which path is authoritative.
+        attributes = {"permissions": [["ec2:DescribeInstances"]], "policy_statements": []}
+        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws")
+        self.assertEqual(chunks, [[]])
+
+    def test_prefers_structured_statements_over_legacy(self):
+        attributes = {
+            "permissions": [["ec2:DescribeInstances"]],
+            "policy_statements": [
+                {
+                    "Effect": "Allow",
+                    "Action": ["iam:PassRole"],
+                    "Resource": ["arn:${Partition}:iam::${AccountId}:role/datadog-ssm-*"],
+                }
+            ],
+        }
+        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws")
+        statements = [s for chunk in chunks for s in chunk]
+        self.assertEqual(len(statements), 1)
+        self.assertEqual(statements[0]["Resource"], ["arn:aws:iam::123456789012:role/datadog-ssm-*"])
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/aws_quickstart/datadog_integration_permissions.yaml b/aws_quickstart/datadog_integration_permissions.yaml
index 5bf11c80..f07ff51e 100644
--- a/aws_quickstart/datadog_integration_permissions.yaml
+++ b/aws_quickstart/datadog_integration_permissions.yaml
@@ -137,7 +137,7 @@ Resources:
               pass
 
 
-          def fetch_permissions_from_datadog(api_url):
+          def fetch_attributes_from_datadog(api_url):
               headers = {
                   "Dd-Aws-Api-Call-Source": API_CALL_SOURCE_HEADER_VALUE,
               }
@@ -151,7 +151,11 @@ Resources:
                   error_message = error_body.get('errors', ['Unknown error'])[0]
                   raise DatadogAPIError(f"Datadog API error: {error_message}") from e
 
-              return json.loads(response.read())["data"]["attributes"]["permissions"]
+              return json.loads(response.read())["data"]["attributes"]
+
+
+          def fetch_permissions_from_datadog(api_url):
+              return fetch_attributes_from_datadog(api_url)["permissions"]
 
 
           def parse_resource_types(raw):
@@ -170,6 +174,31 @@ Resources:
               return f"https://api.{datadog_site}{INSTRUMENTATION_PERMISSIONS_API_PATH}?{query}"
 
 
+          def render_placeholders(value, account_id, partition):
+              if isinstance(value, str):
+                  return value.replace("${AccountId}", account_id).replace("${Partition}", partition)
+              if isinstance(value, list):
+                  return [render_placeholders(v, account_id, partition) for v in value]
+              if isinstance(value, dict):
+                  return {k: render_placeholders(v, account_id, partition) for k, v in value.items()}
+              return value
+
+
+          def legacy_chunk_to_statement(chunk):
+              return {"Effect": "Allow", "Action": list(chunk), "Resource": "*"}
+
+
+          def resolve_instrumentation_statement_chunks(attributes, account_id, partition):
+              # Prefer structured policy_statements when the API provides them; permissions is the legacy,
+              # pre-chunked fallback. Chunk boundaries and standard-permission filtering land in later work.
+              policy_statements = attributes.get("policy_statements")
+              if policy_statements is not None:
+                  rendered = [render_placeholders(s, account_id, partition) for s in policy_statements]
+                  return [rendered]
+
+              return [[legacy_chunk_to_statement(chunk)] for chunk in attributes.get("permissions", [])]
+
+
           def _detach_and_delete_policy(iam_client, role_name, policy_arn, policy_name):
               # Detach + delete are both no-ops if the entity is already gone, so callers can blindly
               # iterate the policy-name space without first checking what actually exists.
@@ -237,14 +266,17 @@ Resources:
 
 
           def _create_and_attach_policy(iam_client, role_name, policy_name, actions):
+              _create_and_attach_statement_policy(
+                  iam_client, role_name, policy_name, [{"Effect": "Allow", "Action": actions, "Resource": "*"}]
+              )
+
+
+          def _create_and_attach_statement_policy(iam_client, role_name, policy_name, statements):
               policy_json = json.dumps(
-                  {
-                      "Version": "2012-10-17",
-                      "Statement": [{"Effect": "Allow", "Action": actions, "Resource": "*"}],
-                  },
+                  {"Version": "2012-10-17", "Statement": statements},
                   separators=(',', ':'),
               )
-              LOGGER.info(f"Creating policy {policy_name} with {len(actions)} permissions ({len(policy_json)} characters)")
+              LOGGER.info(f"Creating policy {policy_name} with {len(statements)} statements ({len(policy_json)} characters)")
               policy = iam_client.create_policy(PolicyName=policy_name, PolicyDocument=policy_json)
               iam_client.attach_role_policy(RoleName=role_name, PolicyArn=policy['Policy']['Arn'])
 
@@ -277,7 +309,8 @@ Resources:
               try:
                   url = build_instrumentation_permissions_url(datadog_site, resource_types)
                   LOGGER.info(f"Fetching instrumentation permissions for {resource_types} from {url}")
-                  permission_chunks = fetch_permissions_from_datadog(url)
+                  attributes = fetch_attributes_from_datadog(url)
+                  chunks = resolve_instrumentation_statement_chunks(attributes, account_id, partition)
               except Exception as e:
                   if fail_on_error:
                       raise
@@ -288,10 +321,10 @@ Resources:
                   return
 
               cleanup_instrumentation_policies(iam_client, role_name, account_id, partition)
-              for i, chunk in enumerate(permission_chunks):
+              for i, chunk in enumerate(chunks):
                   policy_name = f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-{role_name}-{i+1}"
                   try:
-                      _create_and_attach_policy(iam_client, role_name, policy_name, chunk)
+                      _create_and_attach_statement_policy(iam_client, role_name, policy_name, chunk)
                   except Exception as e:
                       if fail_on_error:
                           raise

From f2ec4035427f8c1b91274edcda744d3e2cb695c5 Mon Sep 17 00:00:00 2001
From: Fanny Jiang <fanny.jiang@datadoghq.com>
Date: Thu, 9 Jul 2026 12:42:07 -0400
Subject: [PATCH 2/6] TON-710: bump aws_quickstart version to v4.15.2

---
 aws_quickstart/version.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/aws_quickstart/version.txt b/aws_quickstart/version.txt
index 1597f02b..0bacf6ee 100644
--- a/aws_quickstart/version.txt
+++ b/aws_quickstart/version.txt
@@ -1 +1 @@
-v4.15.1
+v4.15.2

From 77ec128cbeab4ca4a68e94cc8104388e8de010fa Mon Sep 17 00:00:00 2001
From: Fanny Jiang <fanny.jiang@datadoghq.com>
Date: Wed, 8 Jul 2026 14:16:06 -0400
Subject: [PATCH 3/6] TON-711: filter instrumentation statements already
 covered by standard permissions

Only drop an instrumentation action when an existing standard-integration
statement grants an equivalent-or-broader permission across Action,
Resource, and Condition. The role-creation path filters against the
permissions it just attached; the add-on path can't verify what's on the
role, so it passes no standard statements and filtering is a no-op.
---
 .../attach_integration_permissions.py         |  86 +++++++++--
 .../attach_integration_permissions_test.py    | 133 +++++++++++++++++-
 .../datadog_integration_permissions.yaml      |  86 +++++++++--
 3 files changed, 285 insertions(+), 20 deletions(-)

diff --git a/aws_quickstart/attach_integration_permissions.py b/aws_quickstart/attach_integration_permissions.py
index 9e963693..bd9bff14 100644
--- a/aws_quickstart/attach_integration_permissions.py
+++ b/aws_quickstart/attach_integration_permissions.py
@@ -90,15 +90,70 @@ def legacy_chunk_to_statement(chunk):
     return {"Effect": "Allow", "Action": list(chunk), "Resource": "*"}
 
 
-def resolve_instrumentation_statement_chunks(attributes, account_id, partition):
-    # Prefer structured policy_statements when the API provides them; permissions is the legacy,
-    # pre-chunked fallback. Chunk boundaries and standard-permission filtering land in later work.
+def resolve_instrumentation_statement_chunks(attributes, account_id, partition, standard_statements):
+    # Prefer structured policy_statements when the API provides them; permissions is pre-chunked
+    # server-side for the legacy fallback — filter each chunk but preserve its boundaries rather
+    # than re-chunking, since dd-source already sized it to fit. Chunking the (unchunked)
+    # structured statements by size lands in later work.
     policy_statements = attributes.get("policy_statements")
     if policy_statements is not None:
         rendered = [render_placeholders(s, account_id, partition) for s in policy_statements]
-        return [rendered]
+        filtered = filter_instrumentation_statements(rendered, standard_statements)
+        return [filtered] if filtered else []
 
-    return [[legacy_chunk_to_statement(chunk)] for chunk in attributes.get("permissions", [])]
+    legacy_statements = [legacy_chunk_to_statement(chunk) for chunk in attributes.get("permissions", [])]
+    filtered = filter_instrumentation_statements(legacy_statements, standard_statements)
+    return [[statement] for statement in filtered]
+
+
+def _resource_covers(standard_resource, instrumentation_resource):
+    standard_set = set(standard_resource) if isinstance(standard_resource, list) else {standard_resource}
+    if "*" in standard_set:
+        return True
+    instrumentation_set = (
+        set(instrumentation_resource) if isinstance(instrumentation_resource, list) else {instrumentation_resource}
+    )
+    return instrumentation_set.issubset(standard_set)
+
+
+def _condition_covers(standard_condition, instrumentation_condition):
+    # An absent condition on the standard statement is broader than any condition on the
+    # instrumentation side; a present condition must match exactly, otherwise it's not
+    # equivalent-or-broader and the instrumentation action must be kept.
+    if not standard_condition:
+        return True
+    return standard_condition == instrumentation_condition
+
+
+def _action_covered_by_standard(action, resource, condition, standard_statements):
+    for standard_statement in standard_statements:
+        if action not in standard_statement.get("Action", []):
+            continue
+        if not _resource_covers(standard_statement.get("Resource", "*"), resource):
+            continue
+        if not _condition_covers(standard_statement.get("Condition"), condition):
+            continue
+        return True
+    return False
+
+
+def filter_instrumentation_statements(statements, standard_statements):
+    # Only drop an instrumentation action when an existing standard-integration statement
+    # grants an equivalent-or-broader permission across Action, Resource, and Condition.
+    # Action-name-only filtering would silently drop scoped statements just because the same
+    # action name also appears, unconditioned, in the standard policy.
+    filtered = []
+    for statement in statements:
+        resource = statement.get("Resource", "*")
+        condition = statement.get("Condition")
+        kept_actions = [
+            action
+            for action in statement.get("Action", [])
+            if not _action_covered_by_standard(action, resource, condition, standard_statements)
+        ]
+        if kept_actions:
+            filtered.append({**statement, "Action": kept_actions})
+    return filtered
 
 
 def _detach_and_delete_policy(iam_client, role_name, policy_arn, policy_name):
@@ -165,6 +220,7 @@ def attach_standard_permissions(iam_client, role_name):
         PolicyName=POLICY_NAME_STANDARD,
         PolicyDocument=json.dumps(policy_document, separators=(',', ':')),
     )
+    return permissions
 
 
 def _create_and_attach_policy(iam_client, role_name, policy_name, actions):
@@ -194,13 +250,24 @@ def attach_resource_collection_permissions(iam_client, role_name):
         )
 
 
-def attach_instrumentation_permissions(iam_client, role_name, account_id, partition, datadog_site, resource_types, previous_resource_types, fail_on_error=False):
+def attach_instrumentation_permissions(
+    iam_client, role_name, account_id, partition, datadog_site, resource_types, previous_resource_types,
+    *, standard_statements=(), fail_on_error=False,
+):
     # Best-effort by default: instrumentation permissions are additive convenience on top of the
     # integration, so any failure is logged and swallowed rather than blocking install. The
     # post-setup add-on passes fail_on_error=True because attaching these policies is the stack's
     # whole purpose, so a silent SUCCESS that attached nothing would be worse than a visible failure.
     # Fetch before cleanup so that a transient API failure on an Update leaves the
     # previously-attached policies in place instead of silently revoking them.
+    #
+    # standard_statements defaults to empty (no filtering): filtering is only safe when the caller
+    # knows what's actually attached to this role. The role-creation path passes the exact
+    # permissions it just attached via attach_standard_permissions in this same invocation; the
+    # add-on path (ManageBasePermissions=false) doesn't manage or verify the role's standard
+    # permissions, so it can't assume the *current* Datadog standard permission list reflects
+    # what's really on the role — filtering against that would risk dropping instrumentation
+    # actions the role doesn't actually have covered.
     if not resource_types:
         # Only clean up if the previous Update had instrumentation enabled — avoids running
         # delete calls on stacks that never opted in to instrumentation in the first place.
@@ -212,7 +279,7 @@ def attach_instrumentation_permissions(iam_client, role_name, account_id, partit
         url = build_instrumentation_permissions_url(datadog_site, resource_types)
         LOGGER.info(f"Fetching instrumentation permissions for {resource_types} from {url}")
         attributes = fetch_attributes_from_datadog(url)
-        chunks = resolve_instrumentation_statement_chunks(attributes, account_id, partition)
+        chunks = resolve_instrumentation_statement_chunks(attributes, account_id, partition, standard_statements)
     except Exception as e:
         if fail_on_error:
             raise
@@ -266,15 +333,18 @@ def handle_create_update(event, context):
 
     try:
         iam_client = boto3.client('iam')
+        standard_statements = []
         if manage_base_permissions:
             cleanup_legacy_base_policies(iam_client, role_name, account_id, partition)
             cleanup_existing_policies(iam_client, role_name, account_id, partition)
-            attach_standard_permissions(iam_client, role_name)
+            standard_permissions = attach_standard_permissions(iam_client, role_name)
+            standard_statements = [{"Action": standard_permissions, "Resource": "*"}]
             if should_install_security_audit_policy:
                 attach_resource_collection_permissions(iam_client, role_name)
         attach_instrumentation_permissions(
             iam_client, role_name, account_id, partition,
             datadog_site, instrumentation_resource_types, previous_instrumentation_resource_types,
+            standard_statements=standard_statements,
             fail_on_error=fail_on_instrumentation_error,
         )
         cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData={})
diff --git a/aws_quickstart/attach_integration_permissions_test.py b/aws_quickstart/attach_integration_permissions_test.py
index 2dae1e72..40fe7c15 100644
--- a/aws_quickstart/attach_integration_permissions_test.py
+++ b/aws_quickstart/attach_integration_permissions_test.py
@@ -25,6 +25,7 @@
     render_placeholders,
     legacy_chunk_to_statement,
     resolve_instrumentation_statement_chunks,
+    filter_instrumentation_statements,
     POLICY_NAME_STANDARD,
     BASE_POLICY_PREFIX_INSTRUMENTATION,
     BASE_POLICY_PREFIX_RESOURCE_COLLECTION,
@@ -214,6 +215,45 @@ def test_fail_on_error_raises_on_attach_failure(self, mock_urlopen):
                 ["aws:ec2:instance"], (), fail_on_error=True,
             )
 
+    def _mock_statements_response(self, statements):
+        body = json.dumps({"data": {"attributes": {"policy_statements": statements}}}).encode()
+        resp = Mock()
+        resp.read.return_value = body
+        return resp
+
+    @patch("attach_integration_permissions.urllib.request.urlopen")
+    def test_default_standard_statements_means_no_filtering(self, mock_urlopen):
+        # attach_instrumentation_permissions must not silently fetch/apply live standard
+        # permissions for filtering unless the caller explicitly supplies standard_statements —
+        # the add-on path (ManageBasePermissions=false) can't verify what's actually on the role.
+        mock_urlopen.return_value = self._mock_statements_response(
+            [{"Effect": "Allow", "Action": ["iam:PassRole"], "Resource": "*"}]
+        )
+
+        attach_instrumentation_permissions(
+            self.iam, self.role_name, self.account_id, self.partition, self.site,
+            ["aws:ec2:instance"], (), fail_on_error=True,
+        )
+
+        self.iam.create_policy.assert_called_once()
+        policy_document = json.loads(self.iam.create_policy.call_args.kwargs["PolicyDocument"])
+        self.assertEqual(policy_document["Statement"][0]["Action"], ["iam:PassRole"])
+
+    @patch("attach_integration_permissions.urllib.request.urlopen")
+    def test_caller_supplied_standard_statements_are_used_for_filtering(self, mock_urlopen):
+        mock_urlopen.return_value = self._mock_statements_response(
+            [{"Effect": "Allow", "Action": ["iam:PassRole"], "Resource": "*"}]
+        )
+        standard_statements = [{"Action": ["iam:PassRole"], "Resource": "*"}]
+
+        attach_instrumentation_permissions(
+            self.iam, self.role_name, self.account_id, self.partition, self.site,
+            ["aws:ec2:instance"], (), standard_statements=standard_statements, fail_on_error=True,
+        )
+
+        # The only action was fully covered by standard_statements, so nothing is left to attach.
+        self.iam.create_policy.assert_not_called()
+
 
 class TestCleanup(unittest.TestCase):
     def setUp(self):
@@ -297,6 +337,11 @@ def test_create_manage_base_true_attaches_base(
         mock_rc.assert_called_once()
         mock_instr.assert_called_once()
         mock_legacy.assert_called_once()
+        # The role-creation path must filter against the exact permissions it just attached.
+        self.assertEqual(
+            mock_instr.call_args.kwargs["standard_statements"],
+            [{"Action": mock_standard.return_value, "Resource": "*"}],
+        )
 
     @patch("attach_integration_permissions.cleanup_legacy_base_policies")
     @patch("attach_integration_permissions.boto3.client")
@@ -315,6 +360,8 @@ def test_create_manage_base_false_only_instrumentation(
         mock_instr.assert_called_once()
         # Add-on mode must not touch the role stack's standard/resource-collection policies.
         mock_legacy.assert_not_called()
+        # It also can't verify the role's actual standard permissions, so it must not filter.
+        self.assertEqual(mock_instr.call_args.kwargs["standard_statements"], [])
 
     @patch("attach_integration_permissions.boto3.client")
     @patch("attach_integration_permissions.cleanup_instrumentation_policies")
@@ -426,7 +473,7 @@ def test_leaves_non_placeholder_values_unchanged(self):
 class TestResolveInstrumentationStatementChunks(unittest.TestCase):
     def test_legacy_fallback_preserves_chunk_boundaries(self):
         attributes = {"permissions": [["ec2:DescribeInstances"], ["iam:GetRole", "iam:PassRole"]]}
-        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws")
+        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
         self.assertEqual(len(chunks), 2)
         self.assertEqual(chunks[0], [legacy_chunk_to_statement(["ec2:DescribeInstances"])])
         self.assertEqual(chunks[1], [legacy_chunk_to_statement(["iam:GetRole", "iam:PassRole"])])
@@ -435,8 +482,8 @@ def test_empty_policy_statements_list_is_not_treated_as_absent(self):
         # An explicitly present but empty policy_statements must not fall back to the legacy
         # permissions chunks — presence, not truthiness, decides which path is authoritative.
         attributes = {"permissions": [["ec2:DescribeInstances"]], "policy_statements": []}
-        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws")
-        self.assertEqual(chunks, [[]])
+        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+        self.assertEqual(chunks, [])
 
     def test_prefers_structured_statements_over_legacy(self):
         attributes = {
@@ -449,11 +496,89 @@ def test_prefers_structured_statements_over_legacy(self):
                 }
             ],
         }
-        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws")
+        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
         statements = [s for chunk in chunks for s in chunk]
         self.assertEqual(len(statements), 1)
         self.assertEqual(statements[0]["Resource"], ["arn:aws:iam::123456789012:role/datadog-ssm-*"])
 
+    def test_filters_against_standard_statements(self):
+        attributes = {"policy_statements": [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]}
+        standard = [{"Action": ["ec2:DescribeInstances"], "Resource": "*"}]
+        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", standard)
+        self.assertEqual(chunks, [])
+
+
+class TestFilterInstrumentationStatements(unittest.TestCase):
+    def test_removes_action_when_standard_is_equivalent_or_broader(self):
+        statements = [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]
+        standard = [{"Action": ["ec2:DescribeInstances"], "Resource": "*"}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), [])
+
+    def test_list_form_wildcard_resource_is_treated_as_global(self):
+        # A standard statement can express "*" either as a bare string or as a single-element
+        # list; both must be treated as covering any scoped instrumentation resource.
+        statements = [
+            {"Effect": "Allow", "Action": ["iam:PassRole"], "Resource": ["arn:aws:iam::123:role/datadog-ssm-*"]}
+        ]
+        standard = [{"Action": ["iam:PassRole"], "Resource": ["*"]}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), [])
+
+    def test_keeps_action_when_resource_differs(self):
+        statements = [
+            {"Effect": "Allow", "Action": ["iam:PassRole"], "Resource": ["arn:aws:iam::123:role/datadog-ssm-*"]}
+        ]
+        standard = [{"Action": ["iam:PassRole"], "Resource": ["arn:aws:iam::123:role/other-role"]}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), statements)
+
+    def test_keeps_action_when_condition_differs(self):
+        statements = [
+            {
+                "Effect": "Allow",
+                "Action": ["iam:PassRole"],
+                "Resource": "*",
+                "Condition": {"StringEquals": {"iam:PassedToService": ["ec2.amazonaws.com"]}},
+            }
+        ]
+        standard = [
+            {
+                "Action": ["iam:PassRole"],
+                "Resource": "*",
+                "Condition": {"StringEquals": {"iam:PassedToService": ["ecs.amazonaws.com"]}},
+            }
+        ]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), statements)
+
+    def test_keeps_scoped_statement_whose_action_is_not_in_standard(self):
+        statements = [
+            {
+                "Effect": "Allow",
+                "Action": ["iam:PassRole"],
+                "Resource": ["arn:aws:iam::123:role/datadog-ssm-*"],
+                "Condition": {"StringEquals": {"iam:PassedToService": ["ec2.amazonaws.com"]}},
+            }
+        ]
+        standard = [{"Action": ["ec2:DescribeInstances"], "Resource": "*"}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), statements)
+
+    def test_drops_statement_when_every_action_is_covered(self):
+        statements = [
+            {"Effect": "Allow", "Action": ["ec2:DescribeInstances", "ec2:DescribeTags"], "Resource": "*"}
+        ]
+        standard = [{"Action": ["ec2:DescribeInstances", "ec2:DescribeTags"], "Resource": "*"}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), [])
+
+    def test_absent_standard_condition_covers_conditioned_instrumentation(self):
+        statements = [
+            {
+                "Effect": "Allow",
+                "Action": ["iam:PassRole"],
+                "Resource": "*",
+                "Condition": {"StringEquals": {"iam:PassedToService": ["ec2.amazonaws.com"]}},
+            }
+        ]
+        standard = [{"Action": ["iam:PassRole"], "Resource": "*"}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), [])
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/aws_quickstart/datadog_integration_permissions.yaml b/aws_quickstart/datadog_integration_permissions.yaml
index f07ff51e..39655126 100644
--- a/aws_quickstart/datadog_integration_permissions.yaml
+++ b/aws_quickstart/datadog_integration_permissions.yaml
@@ -188,15 +188,70 @@ Resources:
               return {"Effect": "Allow", "Action": list(chunk), "Resource": "*"}
 
 
-          def resolve_instrumentation_statement_chunks(attributes, account_id, partition):
-              # Prefer structured policy_statements when the API provides them; permissions is the legacy,
-              # pre-chunked fallback. Chunk boundaries and standard-permission filtering land in later work.
+          def resolve_instrumentation_statement_chunks(attributes, account_id, partition, standard_statements):
+              # Prefer structured policy_statements when the API provides them; permissions is pre-chunked
+              # server-side for the legacy fallback — filter each chunk but preserve its boundaries rather
+              # than re-chunking, since dd-source already sized it to fit. Chunking the (unchunked)
+              # structured statements by size lands in later work.
               policy_statements = attributes.get("policy_statements")
               if policy_statements is not None:
                   rendered = [render_placeholders(s, account_id, partition) for s in policy_statements]
-                  return [rendered]
+                  filtered = filter_instrumentation_statements(rendered, standard_statements)
+                  return [filtered] if filtered else []
 
-              return [[legacy_chunk_to_statement(chunk)] for chunk in attributes.get("permissions", [])]
+              legacy_statements = [legacy_chunk_to_statement(chunk) for chunk in attributes.get("permissions", [])]
+              filtered = filter_instrumentation_statements(legacy_statements, standard_statements)
+              return [[statement] for statement in filtered]
+
+
+          def _resource_covers(standard_resource, instrumentation_resource):
+              standard_set = set(standard_resource) if isinstance(standard_resource, list) else {standard_resource}
+              if "*" in standard_set:
+                  return True
+              instrumentation_set = (
+                  set(instrumentation_resource) if isinstance(instrumentation_resource, list) else {instrumentation_resource}
+              )
+              return instrumentation_set.issubset(standard_set)
+
+
+          def _condition_covers(standard_condition, instrumentation_condition):
+              # An absent condition on the standard statement is broader than any condition on the
+              # instrumentation side; a present condition must match exactly, otherwise it's not
+              # equivalent-or-broader and the instrumentation action must be kept.
+              if not standard_condition:
+                  return True
+              return standard_condition == instrumentation_condition
+
+
+          def _action_covered_by_standard(action, resource, condition, standard_statements):
+              for standard_statement in standard_statements:
+                  if action not in standard_statement.get("Action", []):
+                      continue
+                  if not _resource_covers(standard_statement.get("Resource", "*"), resource):
+                      continue
+                  if not _condition_covers(standard_statement.get("Condition"), condition):
+                      continue
+                  return True
+              return False
+
+
+          def filter_instrumentation_statements(statements, standard_statements):
+              # Only drop an instrumentation action when an existing standard-integration statement
+              # grants an equivalent-or-broader permission across Action, Resource, and Condition.
+              # Action-name-only filtering would silently drop scoped statements just because the same
+              # action name also appears, unconditioned, in the standard policy.
+              filtered = []
+              for statement in statements:
+                  resource = statement.get("Resource", "*")
+                  condition = statement.get("Condition")
+                  kept_actions = [
+                      action
+                      for action in statement.get("Action", [])
+                      if not _action_covered_by_standard(action, resource, condition, standard_statements)
+                  ]
+                  if kept_actions:
+                      filtered.append({**statement, "Action": kept_actions})
+              return filtered
 
 
           def _detach_and_delete_policy(iam_client, role_name, policy_arn, policy_name):
@@ -263,6 +318,7 @@ Resources:
                   PolicyName=POLICY_NAME_STANDARD,
                   PolicyDocument=json.dumps(policy_document, separators=(',', ':')),
               )
+              return permissions
 
 
           def _create_and_attach_policy(iam_client, role_name, policy_name, actions):
@@ -292,13 +348,24 @@ Resources:
                   )
 
 
-          def attach_instrumentation_permissions(iam_client, role_name, account_id, partition, datadog_site, resource_types, previous_resource_types, fail_on_error=False):
+          def attach_instrumentation_permissions(
+              iam_client, role_name, account_id, partition, datadog_site, resource_types, previous_resource_types,
+              *, standard_statements=(), fail_on_error=False,
+          ):
               # Best-effort by default: instrumentation permissions are additive convenience on top of the
               # integration, so any failure is logged and swallowed rather than blocking install. The
               # post-setup add-on passes fail_on_error=True because attaching these policies is the stack's
               # whole purpose, so a silent SUCCESS that attached nothing would be worse than a visible failure.
               # Fetch before cleanup so that a transient API failure on an Update leaves the
               # previously-attached policies in place instead of silently revoking them.
+              #
+              # standard_statements defaults to empty (no filtering): filtering is only safe when the caller
+              # knows what's actually attached to this role. The role-creation path passes the exact
+              # permissions it just attached via attach_standard_permissions in this same invocation; the
+              # add-on path (ManageBasePermissions=false) doesn't manage or verify the role's standard
+              # permissions, so it can't assume the *current* Datadog standard permission list reflects
+              # what's really on the role — filtering against that would risk dropping instrumentation
+              # actions the role doesn't actually have covered.
               if not resource_types:
                   # Only clean up if the previous Update had instrumentation enabled — avoids running
                   # delete calls on stacks that never opted in to instrumentation in the first place.
@@ -310,7 +377,7 @@ Resources:
                   url = build_instrumentation_permissions_url(datadog_site, resource_types)
                   LOGGER.info(f"Fetching instrumentation permissions for {resource_types} from {url}")
                   attributes = fetch_attributes_from_datadog(url)
-                  chunks = resolve_instrumentation_statement_chunks(attributes, account_id, partition)
+                  chunks = resolve_instrumentation_statement_chunks(attributes, account_id, partition, standard_statements)
               except Exception as e:
                   if fail_on_error:
                       raise
@@ -364,15 +431,18 @@ Resources:
 
               try:
                   iam_client = boto3.client('iam')
+                  standard_statements = []
                   if manage_base_permissions:
                       cleanup_legacy_base_policies(iam_client, role_name, account_id, partition)
                       cleanup_existing_policies(iam_client, role_name, account_id, partition)
-                      attach_standard_permissions(iam_client, role_name)
+                      standard_permissions = attach_standard_permissions(iam_client, role_name)
+                      standard_statements = [{"Action": standard_permissions, "Resource": "*"}]
                       if should_install_security_audit_policy:
                           attach_resource_collection_permissions(iam_client, role_name)
                   attach_instrumentation_permissions(
                       iam_client, role_name, account_id, partition,
                       datadog_site, instrumentation_resource_types, previous_instrumentation_resource_types,
+                      standard_statements=standard_statements,
                       fail_on_error=fail_on_instrumentation_error,
                   )
                   cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData={})

From 8434897f8b331aa794bc90dd1d3046398476938b Mon Sep 17 00:00:00 2001
From: Fanny Jiang <fanny.jiang@datadoghq.com>
Date: Wed, 8 Jul 2026 14:22:02 -0400
Subject: [PATCH 4/6] TON-712: Add size-based chunking and preflight quota
 validation for instrumentation policies

Chunk structured policy_statements by managed-policy size limit, validate the
full candidate chunk set (size + attachment quota) against the role's existing
non-instrumentation policies before cleaning up and re-attaching.
---
 .../attach_integration_permissions.py         | 105 +++++++++-
 .../attach_integration_permissions_test.py    | 194 ++++++++++++++++++
 .../datadog_integration_permissions.yaml      | 106 +++++++++-
 3 files changed, 387 insertions(+), 18 deletions(-)

diff --git a/aws_quickstart/attach_integration_permissions.py b/aws_quickstart/attach_integration_permissions.py
index bd9bff14..6e6985fa 100644
--- a/aws_quickstart/attach_integration_permissions.py
+++ b/aws_quickstart/attach_integration_permissions.py
@@ -33,12 +33,20 @@
 STANDARD_PERMISSIONS_API_URL = "https://api.datadoghq.com/api/v2/integration/aws/iam_permissions/standard"
 RESOURCE_COLLECTION_PERMISSIONS_API_URL = "https://api.datadoghq.com/api/v2/integration/aws/iam_permissions/resource_collection?chunked=true"
 INSTRUMENTATION_PERMISSIONS_API_PATH = "/api/unstable/instrumenter/aws/iam_permissions"
+# AWS hard limit on a customer managed policy document (non-whitespace characters).
+MAX_MANAGED_POLICY_SIZE = 6144
+# AWS default managed-policies-per-role quota; matches the cleanup loops' max_policies bound.
+MAX_ATTACHED_MANAGED_POLICIES = 10
 
 
 class DatadogAPIError(Exception):
     pass
 
 
+class InstrumentationPolicyLimitError(Exception):
+    pass
+
+
 def fetch_attributes_from_datadog(api_url):
     headers = {
         "Dd-Aws-Api-Call-Source": API_CALL_SOURCE_HEADER_VALUE,
@@ -90,16 +98,18 @@ def legacy_chunk_to_statement(chunk):
     return {"Effect": "Allow", "Action": list(chunk), "Resource": "*"}
 
 
-def resolve_instrumentation_statement_chunks(attributes, account_id, partition, standard_statements):
-    # Prefer structured policy_statements when the API provides them; permissions is pre-chunked
-    # server-side for the legacy fallback — filter each chunk but preserve its boundaries rather
-    # than re-chunking, since dd-source already sized it to fit. Chunking the (unchunked)
-    # structured statements by size lands in later work.
+def resolve_instrumentation_statement_chunks(
+    attributes, account_id, partition, standard_statements, max_size=MAX_MANAGED_POLICY_SIZE
+):
+    # Prefer structured policy_statements when the API provides them; policy_statements is
+    # always unchunked, so QuickStart renders, filters, and chunks it by size (Task 5).
+    # permissions is pre-chunked server-side for the legacy fallback — filter each chunk but
+    # preserve its boundaries rather than re-chunking, since dd-source already sized it to fit.
     policy_statements = attributes.get("policy_statements")
     if policy_statements is not None:
         rendered = [render_placeholders(s, account_id, partition) for s in policy_statements]
         filtered = filter_instrumentation_statements(rendered, standard_statements)
-        return [filtered] if filtered else []
+        return chunk_statements(filtered, max_size=max_size)
 
     legacy_statements = [legacy_chunk_to_statement(chunk) for chunk in attributes.get("permissions", [])]
     filtered = filter_instrumentation_statements(legacy_statements, standard_statements)
@@ -156,6 +166,79 @@ def filter_instrumentation_statements(statements, standard_statements):
     return filtered
 
 
+def _policy_document_size(statements):
+    return len(json.dumps({"Version": "2012-10-17", "Statement": statements}, separators=(',', ':')))
+
+
+def chunk_statements(statements, max_size=MAX_MANAGED_POLICY_SIZE):
+    # Chunk by measured minified-document size, never splitting a statement's Action away from
+    # its Resource/Condition. If a single statement alone exceeds max_size it still gets its own
+    # (oversize) chunk; validate_statement_chunks is what catches that before anything is deleted.
+    chunks = []
+    current = []
+    for statement in statements:
+        candidate = current + [statement]
+        if current and _policy_document_size(candidate) > max_size:
+            chunks.append(current)
+            current = [statement]
+        else:
+            current = candidate
+    if current:
+        chunks.append(current)
+    return chunks
+
+
+def validate_statement_chunks(
+    chunks,
+    resource_types,
+    existing_attached_count=0,
+    max_size=MAX_MANAGED_POLICY_SIZE,
+    max_policies=MAX_ATTACHED_MANAGED_POLICIES,
+):
+    # Local size checks are preflight estimates, not a substitute for AWS's own validation on
+    # create_policy; this only exists to fail loudly, before touching any existing policy, rather
+    # than partially attaching a replacement set that AWS would reject partway through.
+    # existing_attached_count is the role's other (non-instrumentation) managed-policy attachments
+    # — those aren't being replaced, so they still count against the per-role attachment quota
+    # alongside the new instrumentation chunks.
+    total = existing_attached_count + len(chunks)
+    if total > max_policies:
+        raise InstrumentationPolicyLimitError(
+            f"Instrumentation permissions for resource types {resource_types} would attach "
+            f"{len(chunks)} policies on top of {existing_attached_count} already attached to the "
+            f"role, totaling {total} and exceeding the {max_policies}-policy attachment limit."
+        )
+    for i, chunk in enumerate(chunks):
+        size = _policy_document_size(chunk)
+        if size > max_size:
+            raise InstrumentationPolicyLimitError(
+                f"Instrumentation permissions for resource types {resource_types} produced "
+                f"policy chunk {i + 1}/{len(chunks)} of {size} characters, exceeding the "
+                f"{max_size}-character managed-policy size limit."
+            )
+
+
+def _chunked_policy_names(role_name, prefix, max_policies=MAX_ATTACHED_MANAGED_POLICIES):
+    return [f"{prefix}-{role_name}-{i+1}" for i in range(max_policies)]
+
+
+def _count_non_instrumentation_attached_policies(iam_client, role_name):
+    # Exact-name match against what cleanup_instrumentation_policies actually deletes, not a
+    # substring check — a stale/unrelated policy whose name merely contains the instrumentation
+    # prefix isn't touched by cleanup and must still count against the attachment quota.
+    instrumentation_names = set(_chunked_policy_names(role_name, BASE_POLICY_PREFIX_INSTRUMENTATION))
+    count = 0
+    kwargs = {"RoleName": role_name}
+    while True:
+        response = iam_client.list_attached_role_policies(**kwargs)
+        for policy in response.get("AttachedPolicies", []):
+            if policy.get("PolicyName", "") not in instrumentation_names:
+                count += 1
+        if not response.get("IsTruncated"):
+            return count
+        kwargs["Marker"] = response["Marker"]
+
+
 def _detach_and_delete_policy(iam_client, role_name, policy_arn, policy_name):
     # Detach + delete are both no-ops if the entity is already gone, so callers can blindly
     # iterate the policy-name space without first checking what actually exists.
@@ -177,8 +260,7 @@ def _detach_and_delete_policy(iam_client, role_name, policy_arn, policy_name):
 
 
 def _cleanup_chunked_policies(iam_client, role_name, account_id, partition, prefix, max_policies=10):
-    for i in range(max_policies):
-        policy_name = f"{prefix}-{role_name}-{i+1}"
+    for policy_name in _chunked_policy_names(role_name, prefix, max_policies):
         policy_arn = f"arn:{partition}:iam::{account_id}:policy/{policy_name}"
         _detach_and_delete_policy(iam_client, role_name, policy_arn, policy_name)
 
@@ -280,11 +362,16 @@ def attach_instrumentation_permissions(
         LOGGER.info(f"Fetching instrumentation permissions for {resource_types} from {url}")
         attributes = fetch_attributes_from_datadog(url)
         chunks = resolve_instrumentation_statement_chunks(attributes, account_id, partition, standard_statements)
+        existing_attached_count = _count_non_instrumentation_attached_policies(iam_client, role_name)
+        # Validate the full candidate replacement set before cleanup_instrumentation_policies
+        # below deletes anything currently attached — a set that can't fit must fail here,
+        # leaving the previously-attached policies untouched.
+        validate_statement_chunks(chunks, resource_types, existing_attached_count)
     except Exception as e:
         if fail_on_error:
             raise
         LOGGER.warning(
-            f"Failed to fetch instrumentation permissions for {resource_types}: {e}. "
+            f"Failed to prepare instrumentation permissions for {resource_types}: {e}. "
             "Leaving any previously-attached instrumentation policies in place."
         )
         return
diff --git a/aws_quickstart/attach_integration_permissions_test.py b/aws_quickstart/attach_integration_permissions_test.py
index 40fe7c15..108dc892 100644
--- a/aws_quickstart/attach_integration_permissions_test.py
+++ b/aws_quickstart/attach_integration_permissions_test.py
@@ -26,6 +26,12 @@
     legacy_chunk_to_statement,
     resolve_instrumentation_statement_chunks,
     filter_instrumentation_statements,
+    chunk_statements,
+    validate_statement_chunks,
+    _count_non_instrumentation_attached_policies,
+    InstrumentationPolicyLimitError,
+    MAX_MANAGED_POLICY_SIZE,
+    MAX_ATTACHED_MANAGED_POLICIES,
     POLICY_NAME_STANDARD,
     BASE_POLICY_PREFIX_INSTRUMENTATION,
     BASE_POLICY_PREFIX_RESOURCE_COLLECTION,
@@ -41,6 +47,7 @@ def make_iam_mock(cleanup_side_effects=True):
     if cleanup_side_effects:
         iam.detach_role_policy.side_effect = iam.exceptions.NoSuchEntityException
         iam.delete_policy.side_effect = iam.exceptions.NoSuchEntityException
+    iam.list_attached_role_policies.return_value = {"AttachedPolicies": [], "IsTruncated": False}
     return iam
 
 
@@ -255,6 +262,94 @@ def test_caller_supplied_standard_statements_are_used_for_filtering(self, mock_u
         self.iam.create_policy.assert_not_called()
 
 
+class TestAttachInstrumentationPermissionsPreflight(unittest.TestCase):
+    def setUp(self):
+        self.iam = make_iam_mock()
+        self.iam.create_policy.return_value = {"Policy": {"Arn": "arn:aws:iam::123:policy/X"}}
+        self.role_name = "DatadogIntegrationRole"
+        self.account_id = "123456789012"
+        self.partition = "aws"
+        self.site = "datadoghq.com"
+
+    def _mock_statements_response(self, statements):
+        body = json.dumps({"data": {"attributes": {"policy_statements": statements}}}).encode()
+        resp = Mock()
+        resp.read.return_value = body
+        return resp
+
+    @patch("attach_integration_permissions.urllib.request.urlopen")
+    def test_oversize_candidate_set_preserves_existing_policies(self, mock_urlopen):
+        # A single statement whose Action list alone exceeds the managed-policy size limit can
+        # never be validly chunked; the preflight check must fail before any existing
+        # instrumentation policy is detached, so the previous (working) set stays attached.
+        huge_action_list = [f"service{i}:Action{i}" for i in range(2000)]
+        mock_urlopen.return_value = self._mock_statements_response(
+            [{"Effect": "Allow", "Action": huge_action_list, "Resource": "*"}]
+        )
+
+        attach_instrumentation_permissions(
+            self.iam, self.role_name, self.account_id, self.partition, self.site,
+            ["aws:ec2:instance"], (),
+        )
+
+        self.iam.create_policy.assert_not_called()
+        self.iam.detach_role_policy.assert_not_called()
+        self.iam.delete_policy.assert_not_called()
+
+    @patch("attach_integration_permissions.urllib.request.urlopen")
+    def test_fail_on_error_raises_on_preflight_failure(self, mock_urlopen):
+        huge_action_list = [f"service{i}:Action{i}" for i in range(2000)]
+        mock_urlopen.return_value = self._mock_statements_response(
+            [{"Effect": "Allow", "Action": huge_action_list, "Resource": "*"}]
+        )
+
+        with self.assertRaises(InstrumentationPolicyLimitError):
+            attach_instrumentation_permissions(
+                self.iam, self.role_name, self.account_id, self.partition, self.site,
+                ["aws:ec2:instance"], (), fail_on_error=True,
+            )
+
+    @patch("attach_integration_permissions.urllib.request.urlopen")
+    def test_existing_non_instrumentation_policies_count_against_the_quota(self, mock_urlopen):
+        mock_urlopen.return_value = self._mock_statements_response(
+            [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]
+        )
+        self.iam.list_attached_role_policies.return_value = {
+            "AttachedPolicies": [
+                {"PolicyName": f"SomeOtherPolicy{i}"} for i in range(MAX_ATTACHED_MANAGED_POLICIES)
+            ],
+            "IsTruncated": False,
+        }
+
+        with self.assertRaises(InstrumentationPolicyLimitError):
+            attach_instrumentation_permissions(
+                self.iam, self.role_name, self.account_id, self.partition, self.site,
+                ["aws:ec2:instance"], (), fail_on_error=True,
+            )
+
+    @patch("attach_integration_permissions.urllib.request.urlopen")
+    def test_existing_instrumentation_policies_do_not_count_against_the_quota(self, mock_urlopen):
+        # The instrumentation policies about to be replaced by cleanup_instrumentation_policies
+        # must not count against the quota for the replacement set.
+        mock_urlopen.return_value = self._mock_statements_response(
+            [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]
+        )
+        self.iam.list_attached_role_policies.return_value = {
+            "AttachedPolicies": [
+                {"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-{self.role_name}-{i+1}"}
+                for i in range(MAX_ATTACHED_MANAGED_POLICIES)
+            ],
+            "IsTruncated": False,
+        }
+
+        attach_instrumentation_permissions(
+            self.iam, self.role_name, self.account_id, self.partition, self.site,
+            ["aws:ec2:instance"], (), fail_on_error=True,
+        )
+
+        self.iam.create_policy.assert_called_once()
+
+
 class TestCleanup(unittest.TestCase):
     def setUp(self):
         self.iam = make_iam_mock()
@@ -507,6 +602,21 @@ def test_filters_against_standard_statements(self):
         chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", standard)
         self.assertEqual(chunks, [])
 
+    def test_structured_statements_are_size_chunked(self):
+        # Unlike the legacy permissions path (pre-chunked server-side), policy_statements is
+        # always unchunked and must be split locally to respect the managed-policy size limit.
+        statements = [
+            {"Effect": "Allow", "Action": [f"service{i}:Action{i}"], "Resource": "*"} for i in range(50)
+        ]
+        attributes = {"policy_statements": statements}
+        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [], max_size=200)
+        self.assertGreater(len(chunks), 1)
+        flattened = [s for chunk in chunks for s in chunk]
+        self.assertEqual(len(flattened), 50)
+        for chunk in chunks:
+            size = len(json.dumps({"Version": "2012-10-17", "Statement": chunk}, separators=(',', ':')))
+            self.assertLessEqual(size, 200)
+
 
 class TestFilterInstrumentationStatements(unittest.TestCase):
     def test_removes_action_when_standard_is_equivalent_or_broader(self):
@@ -580,5 +690,89 @@ def test_absent_standard_condition_covers_conditioned_instrumentation(self):
         self.assertEqual(filter_instrumentation_statements(statements, standard), [])
 
 
+class TestChunkAndValidateStatements(unittest.TestCase):
+    def test_chunk_statements_never_exceeds_max_size(self):
+        statements = [
+            {"Effect": "Allow", "Action": [f"service{i}:Action{i}"], "Resource": "*"} for i in range(30)
+        ]
+        chunks = chunk_statements(statements, max_size=200)
+        for chunk in chunks:
+            size = len(json.dumps({"Version": "2012-10-17", "Statement": chunk}, separators=(',', ':')))
+            self.assertLessEqual(size, 200)
+        self.assertEqual(sum(len(c) for c in chunks), len(statements))
+
+    def test_chunk_statements_empty_input_returns_no_chunks(self):
+        self.assertEqual(chunk_statements([]), [])
+
+    def test_chunk_statements_oversize_single_statement_gets_its_own_chunk(self):
+        huge = {"Effect": "Allow", "Action": [f"service{i}:Action{i}" for i in range(500)], "Resource": "*"}
+        chunks = chunk_statements([huge], max_size=200)
+        self.assertEqual(chunks, [[huge]])
+
+    def test_validate_statement_chunks_passes_within_limits(self):
+        chunks = [[{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]]
+        validate_statement_chunks(chunks, ["aws:ec2:instance"], existing_attached_count=0)
+
+    def test_validate_statement_chunks_raises_on_oversize_chunk(self):
+        huge = [{"Effect": "Allow", "Action": [f"service{i}:Action{i}" for i in range(500)], "Resource": "*"}]
+        with self.assertRaises(InstrumentationPolicyLimitError):
+            validate_statement_chunks([huge], ["aws:ec2:instance"], existing_attached_count=0, max_size=200)
+
+    def test_validate_statement_chunks_raises_when_over_policy_attachment_limit(self):
+        chunks = [[{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}] for _ in range(5)]
+        with self.assertRaises(InstrumentationPolicyLimitError):
+            validate_statement_chunks(chunks, ["aws:ec2:instance"], existing_attached_count=8, max_policies=10)
+
+    def test_validate_statement_chunks_accounts_for_existing_attached_count(self):
+        chunks = [[{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]]
+        validate_statement_chunks(chunks, ["aws:ec2:instance"], existing_attached_count=9, max_policies=10)
+
+
+class TestCountNonInstrumentationAttachedPolicies(unittest.TestCase):
+    def setUp(self):
+        self.iam = make_iam_mock()
+        self.role_name = "DatadogIntegrationRole"
+
+    def test_counts_only_non_instrumentation_policies(self):
+        self.iam.list_attached_role_policies.return_value = {
+            "AttachedPolicies": [
+                {"PolicyName": POLICY_NAME_STANDARD},
+                {"PolicyName": f"{BASE_POLICY_PREFIX_RESOURCE_COLLECTION}-{self.role_name}-1"},
+                {"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-{self.role_name}-1"},
+                {"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-{self.role_name}-2"},
+            ],
+            "IsTruncated": False,
+        }
+        self.assertEqual(_count_non_instrumentation_attached_policies(self.iam, self.role_name), 2)
+
+    def test_paginates_through_truncated_results(self):
+        self.iam.list_attached_role_policies.side_effect = [
+            {
+                "AttachedPolicies": [{"PolicyName": "SomePolicy1"}],
+                "IsTruncated": True,
+                "Marker": "page2",
+            },
+            {
+                "AttachedPolicies": [
+                    {"PolicyName": "SomePolicy2"},
+                    {"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-{self.role_name}-1"},
+                ],
+                "IsTruncated": False,
+            },
+        ]
+        self.assertEqual(_count_non_instrumentation_attached_policies(self.iam, self.role_name), 2)
+        second_call_kwargs = self.iam.list_attached_role_policies.call_args_list[1].kwargs
+        self.assertEqual(second_call_kwargs.get("Marker"), "page2")
+
+    def test_unrelated_policy_containing_instrumentation_prefix_still_counts(self):
+        # Exact-name match, not substring — a stale policy that merely contains the
+        # instrumentation prefix isn't touched by cleanup, so it must still count.
+        self.iam.list_attached_role_policies.return_value = {
+            "AttachedPolicies": [{"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-extra-unrelated"}],
+            "IsTruncated": False,
+        }
+        self.assertEqual(_count_non_instrumentation_attached_policies(self.iam, self.role_name), 1)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/aws_quickstart/datadog_integration_permissions.yaml b/aws_quickstart/datadog_integration_permissions.yaml
index 39655126..31496ab2 100644
--- a/aws_quickstart/datadog_integration_permissions.yaml
+++ b/aws_quickstart/datadog_integration_permissions.yaml
@@ -76,6 +76,7 @@ Resources:
                   - iam:AttachRolePolicy
                   - iam:DetachRolePolicy
                   - iam:PutRolePolicy
+                  - iam:ListAttachedRolePolicies
                 Resource:
                   # Wildcards cover both the v2 names this template creates and the un-suffixed legacy
                   # names it cleans up on an in-place upgrade.
@@ -131,12 +132,20 @@ Resources:
           STANDARD_PERMISSIONS_API_URL = "https://api.datadoghq.com/api/v2/integration/aws/iam_permissions/standard"
           RESOURCE_COLLECTION_PERMISSIONS_API_URL = "https://api.datadoghq.com/api/v2/integration/aws/iam_permissions/resource_collection?chunked=true"
           INSTRUMENTATION_PERMISSIONS_API_PATH = "/api/unstable/instrumenter/aws/iam_permissions"
+          # AWS hard limit on a customer managed policy document (non-whitespace characters).
+          MAX_MANAGED_POLICY_SIZE = 6144
+          # AWS default managed-policies-per-role quota; matches the cleanup loops' max_policies bound.
+          MAX_ATTACHED_MANAGED_POLICIES = 10
 
 
           class DatadogAPIError(Exception):
               pass
 
 
+          class InstrumentationPolicyLimitError(Exception):
+              pass
+
+
           def fetch_attributes_from_datadog(api_url):
               headers = {
                   "Dd-Aws-Api-Call-Source": API_CALL_SOURCE_HEADER_VALUE,
@@ -188,16 +197,18 @@ Resources:
               return {"Effect": "Allow", "Action": list(chunk), "Resource": "*"}
 
 
-          def resolve_instrumentation_statement_chunks(attributes, account_id, partition, standard_statements):
-              # Prefer structured policy_statements when the API provides them; permissions is pre-chunked
-              # server-side for the legacy fallback — filter each chunk but preserve its boundaries rather
-              # than re-chunking, since dd-source already sized it to fit. Chunking the (unchunked)
-              # structured statements by size lands in later work.
+          def resolve_instrumentation_statement_chunks(
+              attributes, account_id, partition, standard_statements, max_size=MAX_MANAGED_POLICY_SIZE
+          ):
+              # Prefer structured policy_statements when the API provides them; policy_statements is
+              # always unchunked, so QuickStart renders, filters, and chunks it by size (Task 5).
+              # permissions is pre-chunked server-side for the legacy fallback — filter each chunk but
+              # preserve its boundaries rather than re-chunking, since dd-source already sized it to fit.
               policy_statements = attributes.get("policy_statements")
               if policy_statements is not None:
                   rendered = [render_placeholders(s, account_id, partition) for s in policy_statements]
                   filtered = filter_instrumentation_statements(rendered, standard_statements)
-                  return [filtered] if filtered else []
+                  return chunk_statements(filtered, max_size=max_size)
 
               legacy_statements = [legacy_chunk_to_statement(chunk) for chunk in attributes.get("permissions", [])]
               filtered = filter_instrumentation_statements(legacy_statements, standard_statements)
@@ -254,6 +265,79 @@ Resources:
               return filtered
 
 
+          def _policy_document_size(statements):
+              return len(json.dumps({"Version": "2012-10-17", "Statement": statements}, separators=(',', ':')))
+
+
+          def chunk_statements(statements, max_size=MAX_MANAGED_POLICY_SIZE):
+              # Chunk by measured minified-document size, never splitting a statement's Action away from
+              # its Resource/Condition. If a single statement alone exceeds max_size it still gets its own
+              # (oversize) chunk; validate_statement_chunks is what catches that before anything is deleted.
+              chunks = []
+              current = []
+              for statement in statements:
+                  candidate = current + [statement]
+                  if current and _policy_document_size(candidate) > max_size:
+                      chunks.append(current)
+                      current = [statement]
+                  else:
+                      current = candidate
+              if current:
+                  chunks.append(current)
+              return chunks
+
+
+          def validate_statement_chunks(
+              chunks,
+              resource_types,
+              existing_attached_count=0,
+              max_size=MAX_MANAGED_POLICY_SIZE,
+              max_policies=MAX_ATTACHED_MANAGED_POLICIES,
+          ):
+              # Local size checks are preflight estimates, not a substitute for AWS's own validation on
+              # create_policy; this only exists to fail loudly, before touching any existing policy, rather
+              # than partially attaching a replacement set that AWS would reject partway through.
+              # existing_attached_count is the role's other (non-instrumentation) managed-policy attachments
+              # — those aren't being replaced, so they still count against the per-role attachment quota
+              # alongside the new instrumentation chunks.
+              total = existing_attached_count + len(chunks)
+              if total > max_policies:
+                  raise InstrumentationPolicyLimitError(
+                      f"Instrumentation permissions for resource types {resource_types} would attach "
+                      f"{len(chunks)} policies on top of {existing_attached_count} already attached to the "
+                      f"role, totaling {total} and exceeding the {max_policies}-policy attachment limit."
+                  )
+              for i, chunk in enumerate(chunks):
+                  size = _policy_document_size(chunk)
+                  if size > max_size:
+                      raise InstrumentationPolicyLimitError(
+                          f"Instrumentation permissions for resource types {resource_types} produced "
+                          f"policy chunk {i + 1}/{len(chunks)} of {size} characters, exceeding the "
+                          f"{max_size}-character managed-policy size limit."
+                      )
+
+
+          def _chunked_policy_names(role_name, prefix, max_policies=MAX_ATTACHED_MANAGED_POLICIES):
+              return [f"{prefix}-{role_name}-{i+1}" for i in range(max_policies)]
+
+
+          def _count_non_instrumentation_attached_policies(iam_client, role_name):
+              # Exact-name match against what cleanup_instrumentation_policies actually deletes, not a
+              # substring check — a stale/unrelated policy whose name merely contains the instrumentation
+              # prefix isn't touched by cleanup and must still count against the attachment quota.
+              instrumentation_names = set(_chunked_policy_names(role_name, BASE_POLICY_PREFIX_INSTRUMENTATION))
+              count = 0
+              kwargs = {"RoleName": role_name}
+              while True:
+                  response = iam_client.list_attached_role_policies(**kwargs)
+                  for policy in response.get("AttachedPolicies", []):
+                      if policy.get("PolicyName", "") not in instrumentation_names:
+                          count += 1
+                  if not response.get("IsTruncated"):
+                      return count
+                  kwargs["Marker"] = response["Marker"]
+
+
           def _detach_and_delete_policy(iam_client, role_name, policy_arn, policy_name):
               # Detach + delete are both no-ops if the entity is already gone, so callers can blindly
               # iterate the policy-name space without first checking what actually exists.
@@ -275,8 +359,7 @@ Resources:
 
 
           def _cleanup_chunked_policies(iam_client, role_name, account_id, partition, prefix, max_policies=10):
-              for i in range(max_policies):
-                  policy_name = f"{prefix}-{role_name}-{i+1}"
+              for policy_name in _chunked_policy_names(role_name, prefix, max_policies):
                   policy_arn = f"arn:{partition}:iam::{account_id}:policy/{policy_name}"
                   _detach_and_delete_policy(iam_client, role_name, policy_arn, policy_name)
 
@@ -378,11 +461,16 @@ Resources:
                   LOGGER.info(f"Fetching instrumentation permissions for {resource_types} from {url}")
                   attributes = fetch_attributes_from_datadog(url)
                   chunks = resolve_instrumentation_statement_chunks(attributes, account_id, partition, standard_statements)
+                  existing_attached_count = _count_non_instrumentation_attached_policies(iam_client, role_name)
+                  # Validate the full candidate replacement set before cleanup_instrumentation_policies
+                  # below deletes anything currently attached — a set that can't fit must fail here,
+                  # leaving the previously-attached policies untouched.
+                  validate_statement_chunks(chunks, resource_types, existing_attached_count)
               except Exception as e:
                   if fail_on_error:
                       raise
                   LOGGER.warning(
-                      f"Failed to fetch instrumentation permissions for {resource_types}: {e}. "
+                      f"Failed to prepare instrumentation permissions for {resource_types}: {e}. "
                       "Leaving any previously-attached instrumentation policies in place."
                   )
                   return

From 54f88d40aa7f22c1a85f8f65e774293d75f107b8 Mon Sep 17 00:00:00 2001
From: Fanny Jiang <fanny.jiang@datadoghq.com>
Date: Thu, 9 Jul 2026 13:30:51 -0400
Subject: [PATCH 5/6] TON-710: simplify and speed up instrumentation policy
 helpers

- use boto3 paginator for list_attached_role_policies, matching the
  existing pattern in datadog_agentless_api_call.py
- inline the single-caller _create_and_attach_policy wrapper
- chunk_statements tracks running document size incrementally instead
  of re-serializing the whole accumulated chunk each iteration
- filter_instrumentation_statements precomputes standard Action sets
  once instead of scanning a list per action per standard statement
---
 .../attach_integration_permissions.py         | 52 +++++++-------
 .../attach_integration_permissions_test.py    | 67 +++++++++----------
 .../datadog_integration_permissions.yaml      | 52 +++++++-------
 3 files changed, 88 insertions(+), 83 deletions(-)

diff --git a/aws_quickstart/attach_integration_permissions.py b/aws_quickstart/attach_integration_permissions.py
index 6e6985fa..c82c9179 100644
--- a/aws_quickstart/attach_integration_permissions.py
+++ b/aws_quickstart/attach_integration_permissions.py
@@ -135,13 +135,13 @@ def _condition_covers(standard_condition, instrumentation_condition):
     return standard_condition == instrumentation_condition
 
 
-def _action_covered_by_standard(action, resource, condition, standard_statements):
-    for standard_statement in standard_statements:
-        if action not in standard_statement.get("Action", []):
+def _action_covered_by_standard(action, resource, condition, standard_index):
+    for action_set, standard_resource, standard_condition in standard_index:
+        if action not in action_set:
             continue
-        if not _resource_covers(standard_statement.get("Resource", "*"), resource):
+        if not _resource_covers(standard_resource, resource):
             continue
-        if not _condition_covers(standard_statement.get("Condition"), condition):
+        if not _condition_covers(standard_condition, condition):
             continue
         return True
     return False
@@ -152,6 +152,11 @@ def filter_instrumentation_statements(statements, standard_statements):
     # grants an equivalent-or-broader permission across Action, Resource, and Condition.
     # Action-name-only filtering would silently drop scoped statements just because the same
     # action name also appears, unconditioned, in the standard policy.
+    # Action lists are converted to sets once here so each instrumentation action's membership
+    # check is O(1) instead of a fresh linear scan per action, per standard statement.
+    standard_index = [
+        (set(s.get("Action", [])), s.get("Resource", "*"), s.get("Condition")) for s in standard_statements
+    ]
     filtered = []
     for statement in statements:
         resource = statement.get("Resource", "*")
@@ -159,7 +164,7 @@ def filter_instrumentation_statements(statements, standard_statements):
         kept_actions = [
             action
             for action in statement.get("Action", [])
-            if not _action_covered_by_standard(action, resource, condition, standard_statements)
+            if not _action_covered_by_standard(action, resource, condition, standard_index)
         ]
         if kept_actions:
             filtered.append({**statement, "Action": kept_actions})
@@ -174,15 +179,23 @@ def chunk_statements(statements, max_size=MAX_MANAGED_POLICY_SIZE):
     # Chunk by measured minified-document size, never splitting a statement's Action away from
     # its Resource/Condition. If a single statement alone exceeds max_size it still gets its own
     # (oversize) chunk; validate_statement_chunks is what catches that before anything is deleted.
+    # Tracks the running document size incrementally (base wrapper + each statement's own
+    # serialized size + one comma per additional array element) instead of re-serializing the
+    # whole accumulated chunk on every statement, which would be O(n^2) in the chunk size.
+    base_size = _policy_document_size([])
     chunks = []
     current = []
+    current_size = base_size
     for statement in statements:
-        candidate = current + [statement]
-        if current and _policy_document_size(candidate) > max_size:
+        statement_size = len(json.dumps(statement, separators=(',', ':')))
+        candidate_size = current_size + statement_size + (1 if current else 0)
+        if current and candidate_size > max_size:
             chunks.append(current)
             current = [statement]
+            current_size = base_size + statement_size
         else:
-            current = candidate
+            current.append(statement)
+            current_size = candidate_size
     if current:
         chunks.append(current)
     return chunks
@@ -228,15 +241,12 @@ def _count_non_instrumentation_attached_policies(iam_client, role_name):
     # prefix isn't touched by cleanup and must still count against the attachment quota.
     instrumentation_names = set(_chunked_policy_names(role_name, BASE_POLICY_PREFIX_INSTRUMENTATION))
     count = 0
-    kwargs = {"RoleName": role_name}
-    while True:
-        response = iam_client.list_attached_role_policies(**kwargs)
-        for policy in response.get("AttachedPolicies", []):
+    paginator = iam_client.get_paginator("list_attached_role_policies")
+    for page in paginator.paginate(RoleName=role_name):
+        for policy in page.get("AttachedPolicies", []):
             if policy.get("PolicyName", "") not in instrumentation_names:
                 count += 1
-        if not response.get("IsTruncated"):
-            return count
-        kwargs["Marker"] = response["Marker"]
+    return count
 
 
 def _detach_and_delete_policy(iam_client, role_name, policy_arn, policy_name):
@@ -305,12 +315,6 @@ def attach_standard_permissions(iam_client, role_name):
     return permissions
 
 
-def _create_and_attach_policy(iam_client, role_name, policy_name, actions):
-    _create_and_attach_statement_policy(
-        iam_client, role_name, policy_name, [{"Effect": "Allow", "Action": actions, "Resource": "*"}]
-    )
-
-
 def _create_and_attach_statement_policy(iam_client, role_name, policy_name, statements):
     policy_json = json.dumps(
         {"Version": "2012-10-17", "Statement": statements},
@@ -324,11 +328,11 @@ def _create_and_attach_statement_policy(iam_client, role_name, policy_name, stat
 def attach_resource_collection_permissions(iam_client, role_name):
     permission_chunks = fetch_permissions_from_datadog(RESOURCE_COLLECTION_PERMISSIONS_API_URL)
     for i, chunk in enumerate(permission_chunks):
-        _create_and_attach_policy(
+        _create_and_attach_statement_policy(
             iam_client,
             role_name,
             f"{BASE_POLICY_PREFIX_RESOURCE_COLLECTION}-{role_name}-{i+1}",
-            chunk,
+            [{"Effect": "Allow", "Action": chunk, "Resource": "*"}],
         )
 
 
diff --git a/aws_quickstart/attach_integration_permissions_test.py b/aws_quickstart/attach_integration_permissions_test.py
index 108dc892..366207b2 100644
--- a/aws_quickstart/attach_integration_permissions_test.py
+++ b/aws_quickstart/attach_integration_permissions_test.py
@@ -47,10 +47,15 @@ def make_iam_mock(cleanup_side_effects=True):
     if cleanup_side_effects:
         iam.detach_role_policy.side_effect = iam.exceptions.NoSuchEntityException
         iam.delete_policy.side_effect = iam.exceptions.NoSuchEntityException
-    iam.list_attached_role_policies.return_value = {"AttachedPolicies": [], "IsTruncated": False}
+    set_attached_policies(iam, [])
     return iam
 
 
+def set_attached_policies(iam, attached_policies, extra_pages=()):
+    pages = [{"AttachedPolicies": attached_policies}, *extra_pages]
+    iam.get_paginator.return_value.paginate.return_value = pages
+
+
 def detached_arns(iam):
     return [c.kwargs["PolicyArn"] for c in iam.detach_role_policy.call_args_list]
 
@@ -314,12 +319,9 @@ def test_existing_non_instrumentation_policies_count_against_the_quota(self, moc
         mock_urlopen.return_value = self._mock_statements_response(
             [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]
         )
-        self.iam.list_attached_role_policies.return_value = {
-            "AttachedPolicies": [
-                {"PolicyName": f"SomeOtherPolicy{i}"} for i in range(MAX_ATTACHED_MANAGED_POLICIES)
-            ],
-            "IsTruncated": False,
-        }
+        set_attached_policies(
+            self.iam, [{"PolicyName": f"SomeOtherPolicy{i}"} for i in range(MAX_ATTACHED_MANAGED_POLICIES)]
+        )
 
         with self.assertRaises(InstrumentationPolicyLimitError):
             attach_instrumentation_permissions(
@@ -334,13 +336,13 @@ def test_existing_instrumentation_policies_do_not_count_against_the_quota(self,
         mock_urlopen.return_value = self._mock_statements_response(
             [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]
         )
-        self.iam.list_attached_role_policies.return_value = {
-            "AttachedPolicies": [
+        set_attached_policies(
+            self.iam,
+            [
                 {"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-{self.role_name}-{i+1}"}
                 for i in range(MAX_ATTACHED_MANAGED_POLICIES)
             ],
-            "IsTruncated": False,
-        }
+        )
 
         attach_instrumentation_permissions(
             self.iam, self.role_name, self.account_id, self.partition, self.site,
@@ -734,43 +736,38 @@ def setUp(self):
         self.role_name = "DatadogIntegrationRole"
 
     def test_counts_only_non_instrumentation_policies(self):
-        self.iam.list_attached_role_policies.return_value = {
-            "AttachedPolicies": [
+        set_attached_policies(
+            self.iam,
+            [
                 {"PolicyName": POLICY_NAME_STANDARD},
                 {"PolicyName": f"{BASE_POLICY_PREFIX_RESOURCE_COLLECTION}-{self.role_name}-1"},
                 {"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-{self.role_name}-1"},
                 {"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-{self.role_name}-2"},
             ],
-            "IsTruncated": False,
-        }
+        )
         self.assertEqual(_count_non_instrumentation_attached_policies(self.iam, self.role_name), 2)
 
     def test_paginates_through_truncated_results(self):
-        self.iam.list_attached_role_policies.side_effect = [
-            {
-                "AttachedPolicies": [{"PolicyName": "SomePolicy1"}],
-                "IsTruncated": True,
-                "Marker": "page2",
-            },
-            {
-                "AttachedPolicies": [
-                    {"PolicyName": "SomePolicy2"},
-                    {"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-{self.role_name}-1"},
-                ],
-                "IsTruncated": False,
-            },
-        ]
+        set_attached_policies(
+            self.iam,
+            [{"PolicyName": "SomePolicy1"}],
+            extra_pages=[
+                {
+                    "AttachedPolicies": [
+                        {"PolicyName": "SomePolicy2"},
+                        {"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-{self.role_name}-1"},
+                    ],
+                },
+            ],
+        )
         self.assertEqual(_count_non_instrumentation_attached_policies(self.iam, self.role_name), 2)
-        second_call_kwargs = self.iam.list_attached_role_policies.call_args_list[1].kwargs
-        self.assertEqual(second_call_kwargs.get("Marker"), "page2")
+        self.iam.get_paginator.assert_called_once_with("list_attached_role_policies")
+        self.iam.get_paginator.return_value.paginate.assert_called_once_with(RoleName=self.role_name)
 
     def test_unrelated_policy_containing_instrumentation_prefix_still_counts(self):
         # Exact-name match, not substring — a stale policy that merely contains the
         # instrumentation prefix isn't touched by cleanup, so it must still count.
-        self.iam.list_attached_role_policies.return_value = {
-            "AttachedPolicies": [{"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-extra-unrelated"}],
-            "IsTruncated": False,
-        }
+        set_attached_policies(self.iam, [{"PolicyName": f"{BASE_POLICY_PREFIX_INSTRUMENTATION}-extra-unrelated"}])
         self.assertEqual(_count_non_instrumentation_attached_policies(self.iam, self.role_name), 1)
 
 
diff --git a/aws_quickstart/datadog_integration_permissions.yaml b/aws_quickstart/datadog_integration_permissions.yaml
index 31496ab2..17e0b096 100644
--- a/aws_quickstart/datadog_integration_permissions.yaml
+++ b/aws_quickstart/datadog_integration_permissions.yaml
@@ -234,13 +234,13 @@ Resources:
               return standard_condition == instrumentation_condition
 
 
-          def _action_covered_by_standard(action, resource, condition, standard_statements):
-              for standard_statement in standard_statements:
-                  if action not in standard_statement.get("Action", []):
+          def _action_covered_by_standard(action, resource, condition, standard_index):
+              for action_set, standard_resource, standard_condition in standard_index:
+                  if action not in action_set:
                       continue
-                  if not _resource_covers(standard_statement.get("Resource", "*"), resource):
+                  if not _resource_covers(standard_resource, resource):
                       continue
-                  if not _condition_covers(standard_statement.get("Condition"), condition):
+                  if not _condition_covers(standard_condition, condition):
                       continue
                   return True
               return False
@@ -251,6 +251,11 @@ Resources:
               # grants an equivalent-or-broader permission across Action, Resource, and Condition.
               # Action-name-only filtering would silently drop scoped statements just because the same
               # action name also appears, unconditioned, in the standard policy.
+              # Action lists are converted to sets once here so each instrumentation action's membership
+              # check is O(1) instead of a fresh linear scan per action, per standard statement.
+              standard_index = [
+                  (set(s.get("Action", [])), s.get("Resource", "*"), s.get("Condition")) for s in standard_statements
+              ]
               filtered = []
               for statement in statements:
                   resource = statement.get("Resource", "*")
@@ -258,7 +263,7 @@ Resources:
                   kept_actions = [
                       action
                       for action in statement.get("Action", [])
-                      if not _action_covered_by_standard(action, resource, condition, standard_statements)
+                      if not _action_covered_by_standard(action, resource, condition, standard_index)
                   ]
                   if kept_actions:
                       filtered.append({**statement, "Action": kept_actions})
@@ -273,15 +278,23 @@ Resources:
               # Chunk by measured minified-document size, never splitting a statement's Action away from
               # its Resource/Condition. If a single statement alone exceeds max_size it still gets its own
               # (oversize) chunk; validate_statement_chunks is what catches that before anything is deleted.
+              # Tracks the running document size incrementally (base wrapper + each statement's own
+              # serialized size + one comma per additional array element) instead of re-serializing the
+              # whole accumulated chunk on every statement, which would be O(n^2) in the chunk size.
+              base_size = _policy_document_size([])
               chunks = []
               current = []
+              current_size = base_size
               for statement in statements:
-                  candidate = current + [statement]
-                  if current and _policy_document_size(candidate) > max_size:
+                  statement_size = len(json.dumps(statement, separators=(',', ':')))
+                  candidate_size = current_size + statement_size + (1 if current else 0)
+                  if current and candidate_size > max_size:
                       chunks.append(current)
                       current = [statement]
+                      current_size = base_size + statement_size
                   else:
-                      current = candidate
+                      current.append(statement)
+                      current_size = candidate_size
               if current:
                   chunks.append(current)
               return chunks
@@ -327,15 +340,12 @@ Resources:
               # prefix isn't touched by cleanup and must still count against the attachment quota.
               instrumentation_names = set(_chunked_policy_names(role_name, BASE_POLICY_PREFIX_INSTRUMENTATION))
               count = 0
-              kwargs = {"RoleName": role_name}
-              while True:
-                  response = iam_client.list_attached_role_policies(**kwargs)
-                  for policy in response.get("AttachedPolicies", []):
+              paginator = iam_client.get_paginator("list_attached_role_policies")
+              for page in paginator.paginate(RoleName=role_name):
+                  for policy in page.get("AttachedPolicies", []):
                       if policy.get("PolicyName", "") not in instrumentation_names:
                           count += 1
-                  if not response.get("IsTruncated"):
-                      return count
-                  kwargs["Marker"] = response["Marker"]
+              return count
 
 
           def _detach_and_delete_policy(iam_client, role_name, policy_arn, policy_name):
@@ -404,12 +414,6 @@ Resources:
               return permissions
 
 
-          def _create_and_attach_policy(iam_client, role_name, policy_name, actions):
-              _create_and_attach_statement_policy(
-                  iam_client, role_name, policy_name, [{"Effect": "Allow", "Action": actions, "Resource": "*"}]
-              )
-
-
           def _create_and_attach_statement_policy(iam_client, role_name, policy_name, statements):
               policy_json = json.dumps(
                   {"Version": "2012-10-17", "Statement": statements},
@@ -423,11 +427,11 @@ Resources:
           def attach_resource_collection_permissions(iam_client, role_name):
               permission_chunks = fetch_permissions_from_datadog(RESOURCE_COLLECTION_PERMISSIONS_API_URL)
               for i, chunk in enumerate(permission_chunks):
-                  _create_and_attach_policy(
+                  _create_and_attach_statement_policy(
                       iam_client,
                       role_name,
                       f"{BASE_POLICY_PREFIX_RESOURCE_COLLECTION}-{role_name}-{i+1}",
-                      chunk,
+                      [{"Effect": "Allow", "Action": chunk, "Resource": "*"}],
                   )
 
 

From 553aa438cf468bf70864213f6770fce5af973751 Mon Sep 17 00:00:00 2001
From: Fanny Jiang <fanny.jiang@datadoghq.com>
Date: Thu, 9 Jul 2026 16:46:30 -0400
Subject: [PATCH 6/6] TON-710: harden instrumentation policy schema validation
 and IAM wildcard matching

Validate structured policy_statements (Effect/Action/Resource/Condition, including
blank-string checks) before they can trigger cleanup of previously-attached
instrumentation policies, and replace fnmatch with an IAM-correct action-glob
matcher to avoid over-matching on bracket-class syntax fnmatch grants but IAM
does not. Also trims duplicate validation/policy-JSON logic and a couple of
overly-verbose comments.
---
 .../attach_integration_permissions.py         | 123 +++++++++---
 .../attach_integration_permissions_test.py    | 180 ++++++++++++++++++
 .../datadog_integration_permissions.yaml      | 123 +++++++++---
 3 files changed, 366 insertions(+), 60 deletions(-)

diff --git a/aws_quickstart/attach_integration_permissions.py b/aws_quickstart/attach_integration_permissions.py
index c82c9179..630c11b7 100644
--- a/aws_quickstart/attach_integration_permissions.py
+++ b/aws_quickstart/attach_integration_permissions.py
@@ -1,5 +1,6 @@
 import json
 import logging
+import re
 from urllib.request import Request
 import urllib.error
 import urllib.parse
@@ -98,20 +99,77 @@ def legacy_chunk_to_statement(chunk):
     return {"Effect": "Allow", "Action": list(chunk), "Resource": "*"}
 
 
+def _is_nonblank_string_or_list(value):
+    # .strip() only tests for blankness — it never rewrites the value, so a legitimately-formatted
+    # field passes through untouched.
+    if isinstance(value, str):
+        return bool(value.strip())
+    if isinstance(value, list):
+        return bool(value) and all(isinstance(v, str) and v.strip() for v in value)
+    return False
+
+
+def _validate_action_field(action, context):
+    # A statement/chunk with a missing, null, or empty Action silently disappears in
+    # filter_instrumentation_statements (kept_actions ends up empty), which can collapse the
+    # entire candidate set to [] and trigger the same wipe-previously-attached-policies failure
+    # mode as a missing top-level field — just one level deeper.
+    if not _is_nonblank_string_or_list(action):
+        raise DatadogAPIError(f"Datadog API returned a statement with an invalid Action field: {action!r} ({context})")
+
+
+def _validate_policy_statement(statement, context):
+    # A statement missing Effect/Resource, or with a non-"Allow" Effect, passes prep and
+    # reaches cleanup_instrumentation_policies before create_policy rejects it as malformed —
+    # the same wipe-previously-attached-policies failure mode as an invalid Action.
+    if not isinstance(statement, dict):
+        raise DatadogAPIError(f"Datadog API returned a non-dict statement: {statement!r} ({context})")
+    if statement.get("Effect") != "Allow":
+        raise DatadogAPIError(
+            f"Datadog API returned a statement with an invalid Effect field: {statement.get('Effect')!r} ({context})"
+        )
+    _validate_action_field(statement.get("Action"), context)
+    resource = statement.get("Resource")
+    if not _is_nonblank_string_or_list(resource):
+        raise DatadogAPIError(f"Datadog API returned a statement with an invalid Resource field: {resource!r} ({context})")
+    condition = statement.get("Condition")
+    if condition is not None and not isinstance(condition, dict):
+        raise DatadogAPIError(f"Datadog API returned a statement with an invalid Condition field: {condition!r} ({context})")
+
+
 def resolve_instrumentation_statement_chunks(
     attributes, account_id, partition, standard_statements, max_size=MAX_MANAGED_POLICY_SIZE
 ):
     # Prefer structured policy_statements when the API provides them; policy_statements is
-    # always unchunked, so QuickStart renders, filters, and chunks it by size (Task 5).
+    # always unchunked, so QuickStart renders, filters, and chunks it by size.
     # permissions is pre-chunked server-side for the legacy fallback — filter each chunk but
     # preserve its boundaries rather than re-chunking, since dd-source already sized it to fit.
     policy_statements = attributes.get("policy_statements")
     if policy_statements is not None:
+        if not isinstance(policy_statements, list):
+            raise DatadogAPIError(
+                f"Datadog API returned non-list policy_statements: {type(policy_statements).__name__}"
+            )
+        for statement in policy_statements:
+            _validate_policy_statement(statement, "policy_statements")
         rendered = [render_placeholders(s, account_id, partition) for s in policy_statements]
         filtered = filter_instrumentation_statements(rendered, standard_statements)
         return chunk_statements(filtered, max_size=max_size)
 
-    legacy_statements = [legacy_chunk_to_statement(chunk) for chunk in attributes.get("permissions", [])]
+    # policy_statements missing or null falls back to the legacy permissions field, but that
+    # field must actually be present and a list — schema drift that drops both fields must raise
+    # here rather than silently resolving to an empty candidate set, which attach_instrumentation_permissions
+    # would treat as a valid "nothing to attach" and use to wipe previously-attached policies.
+    if "permissions" not in attributes:
+        raise DatadogAPIError("Datadog API response is missing both policy_statements and permissions")
+    legacy_permissions = attributes["permissions"]
+    if not isinstance(legacy_permissions, list):
+        raise DatadogAPIError(
+            f"Datadog API returned non-list permissions: {type(legacy_permissions).__name__}"
+        )
+    for chunk in legacy_permissions:
+        _validate_action_field(chunk if isinstance(chunk, list) else None, "permissions")
+    legacy_statements = [legacy_chunk_to_statement(chunk) for chunk in legacy_permissions]
     filtered = filter_instrumentation_statements(legacy_statements, standard_statements)
     return [[statement] for statement in filtered]
 
@@ -135,9 +193,25 @@ def _condition_covers(standard_condition, instrumentation_condition):
     return standard_condition == instrumentation_condition
 
 
+def _as_action_list(action):
+    # IAM's Action field may be a single string or a list; a bare string must not be iterated
+    # character-by-character.
+    if action is None:
+        return []
+    return action if isinstance(action, list) else [action]
+
+
+def _iam_action_matches(action, pattern):
+    # IAM action wildcards only recognize "*" and "?"; fnmatch also grants bracket
+    # character-class syntax that IAM does not, which could over-match an action pattern
+    # containing literal brackets.
+    regex = "".join(".*" if ch == "*" else "." if ch == "?" else re.escape(ch) for ch in pattern)
+    return re.fullmatch(regex, action) is not None
+
+
 def _action_covered_by_standard(action, resource, condition, standard_index):
-    for action_set, standard_resource, standard_condition in standard_index:
-        if action not in action_set:
+    for action_patterns, standard_resource, standard_condition in standard_index:
+        if not any(_iam_action_matches(action.lower(), pattern.lower()) for pattern in action_patterns):
             continue
         if not _resource_covers(standard_resource, resource):
             continue
@@ -152,10 +226,8 @@ def filter_instrumentation_statements(statements, standard_statements):
     # grants an equivalent-or-broader permission across Action, Resource, and Condition.
     # Action-name-only filtering would silently drop scoped statements just because the same
     # action name also appears, unconditioned, in the standard policy.
-    # Action lists are converted to sets once here so each instrumentation action's membership
-    # check is O(1) instead of a fresh linear scan per action, per standard statement.
     standard_index = [
-        (set(s.get("Action", [])), s.get("Resource", "*"), s.get("Condition")) for s in standard_statements
+        (_as_action_list(s.get("Action", [])), s.get("Resource", "*"), s.get("Condition")) for s in standard_statements
     ]
     filtered = []
     for statement in statements:
@@ -163,7 +235,7 @@ def filter_instrumentation_statements(statements, standard_statements):
         condition = statement.get("Condition")
         kept_actions = [
             action
-            for action in statement.get("Action", [])
+            for action in _as_action_list(statement.get("Action", []))
             if not _action_covered_by_standard(action, resource, condition, standard_index)
         ]
         if kept_actions:
@@ -171,17 +243,18 @@ def filter_instrumentation_statements(statements, standard_statements):
     return filtered
 
 
+def _policy_document_json(statements):
+    return json.dumps({"Version": "2012-10-17", "Statement": statements}, separators=(',', ':'))
+
+
 def _policy_document_size(statements):
-    return len(json.dumps({"Version": "2012-10-17", "Statement": statements}, separators=(',', ':')))
+    return len(_policy_document_json(statements))
 
 
 def chunk_statements(statements, max_size=MAX_MANAGED_POLICY_SIZE):
     # Chunk by measured minified-document size, never splitting a statement's Action away from
     # its Resource/Condition. If a single statement alone exceeds max_size it still gets its own
     # (oversize) chunk; validate_statement_chunks is what catches that before anything is deleted.
-    # Tracks the running document size incrementally (base wrapper + each statement's own
-    # serialized size + one comma per additional array element) instead of re-serializing the
-    # whole accumulated chunk on every statement, which would be O(n^2) in the chunk size.
     base_size = _policy_document_size([])
     chunks = []
     current = []
@@ -237,8 +310,7 @@ def _chunked_policy_names(role_name, prefix, max_policies=MAX_ATTACHED_MANAGED_P
 
 def _count_non_instrumentation_attached_policies(iam_client, role_name):
     # Exact-name match against what cleanup_instrumentation_policies actually deletes, not a
-    # substring check — a stale/unrelated policy whose name merely contains the instrumentation
-    # prefix isn't touched by cleanup and must still count against the attachment quota.
+    # substring check.
     instrumentation_names = set(_chunked_policy_names(role_name, BASE_POLICY_PREFIX_INSTRUMENTATION))
     count = 0
     paginator = iam_client.get_paginator("list_attached_role_policies")
@@ -303,23 +375,16 @@ def cleanup_legacy_base_policies(iam_client, role_name, account_id, partition, m
 
 def attach_standard_permissions(iam_client, role_name):
     permissions = fetch_permissions_from_datadog(STANDARD_PERMISSIONS_API_URL)
-    policy_document = {
-        "Version": "2012-10-17",
-        "Statement": [{"Effect": "Allow", "Action": permissions, "Resource": "*"}],
-    }
     iam_client.put_role_policy(
         RoleName=role_name,
         PolicyName=POLICY_NAME_STANDARD,
-        PolicyDocument=json.dumps(policy_document, separators=(',', ':')),
+        PolicyDocument=_policy_document_json([{"Effect": "Allow", "Action": permissions, "Resource": "*"}]),
     )
     return permissions
 
 
 def _create_and_attach_statement_policy(iam_client, role_name, policy_name, statements):
-    policy_json = json.dumps(
-        {"Version": "2012-10-17", "Statement": statements},
-        separators=(',', ':'),
-    )
+    policy_json = _policy_document_json(statements)
     LOGGER.info(f"Creating policy {policy_name} with {len(statements)} statements ({len(policy_json)} characters)")
     policy = iam_client.create_policy(PolicyName=policy_name, PolicyDocument=policy_json)
     iam_client.attach_role_policy(RoleName=role_name, PolicyArn=policy['Policy']['Arn'])
@@ -347,13 +412,11 @@ def attach_instrumentation_permissions(
     # Fetch before cleanup so that a transient API failure on an Update leaves the
     # previously-attached policies in place instead of silently revoking them.
     #
-    # standard_statements defaults to empty (no filtering): filtering is only safe when the caller
-    # knows what's actually attached to this role. The role-creation path passes the exact
-    # permissions it just attached via attach_standard_permissions in this same invocation; the
-    # add-on path (ManageBasePermissions=false) doesn't manage or verify the role's standard
-    # permissions, so it can't assume the *current* Datadog standard permission list reflects
-    # what's really on the role — filtering against that would risk dropping instrumentation
-    # actions the role doesn't actually have covered.
+    # standard_statements defaults to empty (no filtering): filtering is only safe when the
+    # caller knows what's actually attached. The role-creation path passes what it just attached
+    # via attach_standard_permissions in this same invocation; the add-on path
+    # (ManageBasePermissions=false) can't verify the role's real standard permissions, so it must
+    # skip filtering.
     if not resource_types:
         # Only clean up if the previous Update had instrumentation enabled — avoids running
         # delete calls on stacks that never opted in to instrumentation in the first place.
diff --git a/aws_quickstart/attach_integration_permissions_test.py b/aws_quickstart/attach_integration_permissions_test.py
index 366207b2..b60041a4 100644
--- a/aws_quickstart/attach_integration_permissions_test.py
+++ b/aws_quickstart/attach_integration_permissions_test.py
@@ -29,6 +29,7 @@
     chunk_statements,
     validate_statement_chunks,
     _count_non_instrumentation_attached_policies,
+    DatadogAPIError,
     InstrumentationPolicyLimitError,
     MAX_MANAGED_POLICY_SIZE,
     MAX_ATTACHED_MANAGED_POLICIES,
@@ -619,6 +620,138 @@ def test_structured_statements_are_size_chunked(self):
             size = len(json.dumps({"Version": "2012-10-17", "Statement": chunk}, separators=(',', ':')))
             self.assertLessEqual(size, 200)
 
+    def test_missing_both_fields_raises_instead_of_resolving_empty(self):
+        # Regression: schema drift that drops both policy_statements and permissions must not
+        # silently resolve to an empty candidate set — that would look like a valid "nothing to
+        # attach" result and cause the caller to wipe previously-attached instrumentation policies.
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks({}, "123456789012", "aws", [])
+
+    def test_null_policy_statements_falls_back_to_legacy(self):
+        attributes = {"policy_statements": None, "permissions": [["ec2:DescribeInstances"]]}
+        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+        self.assertEqual(chunks, [[legacy_chunk_to_statement(["ec2:DescribeInstances"])]])
+
+    def test_null_policy_statements_with_missing_permissions_raises(self):
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks({"policy_statements": None}, "123456789012", "aws", [])
+
+    def test_wrong_type_policy_statements_raises(self):
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(
+                {"policy_statements": "not-a-list"}, "123456789012", "aws", []
+            )
+
+    def test_wrong_type_permissions_raises(self):
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks({"permissions": "not-a-list"}, "123456789012", "aws", [])
+
+    def test_wrong_type_permissions_chunk_raises(self):
+        # permissions must be a list of chunks (each a list of action strings); a top-level list
+        # of bare strings would otherwise be silently iterated character-by-character downstream.
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(
+                {"permissions": ["ec2:DescribeInstances"]}, "123456789012", "aws", []
+            )
+
+    def test_policy_statement_missing_action_raises(self):
+        # A present statement with no Action would otherwise be silently dropped by
+        # filter_instrumentation_statements, collapsing the whole candidate set to [] and
+        # triggering the same wipe-previously-attached-policies failure mode as a missing
+        # top-level field.
+        attributes = {"policy_statements": [{"Effect": "Allow", "Resource": "*"}]}
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+
+    def test_policy_statement_empty_action_list_raises(self):
+        attributes = {"policy_statements": [{"Effect": "Allow", "Action": [], "Resource": "*"}]}
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+
+    def test_policy_statement_non_string_action_element_raises(self):
+        attributes = {"policy_statements": [{"Effect": "Allow", "Action": [123], "Resource": "*"}]}
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+
+    def test_legacy_chunk_with_empty_list_raises(self):
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks({"permissions": [[]]}, "123456789012", "aws", [])
+
+    def test_legacy_chunk_with_non_string_element_raises(self):
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks({"permissions": [[123]]}, "123456789012", "aws", [])
+
+    def test_policy_statement_non_dict_raises(self):
+        attributes = {"policy_statements": ["not-a-dict"]}
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+
+    def test_policy_statement_missing_effect_raises(self):
+        # A statement missing Effect/Resource would otherwise pass prep, reach cleanup, and
+        # only fail once create_policy rejects it as malformed — wiping previously-attached
+        # policies with nothing successfully re-attached in their place.
+        attributes = {"policy_statements": [{"Action": ["ec2:DescribeInstances"], "Resource": "*"}]}
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+
+    def test_policy_statement_deny_effect_raises(self):
+        attributes = {
+            "policy_statements": [{"Effect": "Deny", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]
+        }
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+
+    def test_policy_statement_missing_resource_raises(self):
+        attributes = {"policy_statements": [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"]}]}
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+
+    def test_policy_statement_non_dict_condition_raises(self):
+        attributes = {
+            "policy_statements": [
+                {"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*", "Condition": "not-a-dict"}
+            ]
+        }
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+
+    def test_policy_statement_whitespace_only_action_raises(self):
+        attributes = {"policy_statements": [{"Effect": "Allow", "Action": "   ", "Resource": "*"}]}
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+
+    def test_policy_statement_whitespace_only_action_element_raises(self):
+        attributes = {
+            "policy_statements": [
+                {"Effect": "Allow", "Action": ["ec2:DescribeInstances", "   "], "Resource": "*"}
+            ]
+        }
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+
+    def test_policy_statement_whitespace_only_resource_raises(self):
+        attributes = {"policy_statements": [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "  "}]}
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+
+    def test_legacy_chunk_with_whitespace_only_element_raises(self):
+        with self.assertRaises(DatadogAPIError):
+            resolve_instrumentation_statement_chunks({"permissions": [["   "]]}, "123456789012", "aws", [])
+
+    def test_strip_check_does_not_mutate_action_or_resource_values(self):
+        # .strip() is used only to detect blank strings; it must never rewrite the stored Action
+        # or Resource value, or a legitimately-formatted (if oddly spaced) statement would be
+        # silently altered on its way through.
+        attributes = {
+            "policy_statements": [
+                {"Effect": "Allow", "Action": [" ec2:DescribeInstances "], "Resource": " * "}
+            ]
+        }
+        chunks = resolve_instrumentation_statement_chunks(attributes, "123456789012", "aws", [])
+        self.assertEqual(
+            chunks, [[{"Effect": "Allow", "Action": [" ec2:DescribeInstances "], "Resource": " * "}]]
+        )
+
 
 class TestFilterInstrumentationStatements(unittest.TestCase):
     def test_removes_action_when_standard_is_equivalent_or_broader(self):
@@ -691,6 +824,53 @@ def test_absent_standard_condition_covers_conditioned_instrumentation(self):
         standard = [{"Action": ["iam:PassRole"], "Resource": "*"}]
         self.assertEqual(filter_instrumentation_statements(statements, standard), [])
 
+    def test_standard_wildcard_action_covers_concrete_instrumentation_action(self):
+        # The standard role policy commonly grants patterns like "ec2:Describe*"; a concrete
+        # instrumentation action matching that pattern must be treated as already covered,
+        # rather than kept as redundant and pushed into the size/quota preflight.
+        statements = [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]
+        standard = [{"Action": ["ec2:Describe*"], "Resource": "*"}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), [])
+
+    def test_standard_wildcard_action_match_is_case_insensitive(self):
+        statements = [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]
+        standard = [{"Action": ["EC2:describe*"], "Resource": "*"}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), [])
+
+    def test_standard_wildcard_action_does_not_match_unrelated_service(self):
+        statements = [{"Effect": "Allow", "Action": ["eks:DescribeCluster"], "Resource": "*"}]
+        standard = [{"Action": ["ec2:Describe*"], "Resource": "*"}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), statements)
+
+    def test_standard_action_as_bare_string_is_not_iterated_as_characters(self):
+        # IAM's Action field may be a single string rather than a list; treating it as an
+        # iterable of characters would falsely "match" via fnmatch against single-char patterns.
+        statements = [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}]
+        standard = [{"Action": "ec2:Describe*", "Resource": "*"}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), [])
+
+    def test_instrumentation_action_as_bare_string_is_not_iterated_as_characters(self):
+        # The instrumentation statement's own Action field may also be a bare string; it must be
+        # normalized to a single-element list, not iterated as characters.
+        statements = [{"Effect": "Allow", "Action": "ec2:DescribeInstances", "Resource": "*"}]
+        standard = [{"Action": ["ec2:DescribeInstances"], "Resource": "*"}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), [])
+
+    def test_instrumentation_action_as_bare_string_is_kept_when_uncovered(self):
+        statements = [{"Effect": "Allow", "Action": "ec2:DescribeInstances", "Resource": "*"}]
+        standard = [{"Action": ["iam:PassRole"], "Resource": "*"}]
+        self.assertEqual(
+            filter_instrumentation_statements(statements, standard),
+            [{"Effect": "Allow", "Action": ["ec2:DescribeInstances"], "Resource": "*"}],
+        )
+
+    def test_standard_action_bracket_characters_are_matched_literally(self):
+        # IAM action wildcards only recognize "*" and "?"; fnmatch's bracket character-class
+        # syntax must not be honored, or a literal "[" in an action pattern would over-match.
+        statements = [{"Effect": "Allow", "Action": ["ec2:Describe[Instances]"], "Resource": "*"}]
+        standard = [{"Action": ["ec2:Describe[a]"], "Resource": "*"}]
+        self.assertEqual(filter_instrumentation_statements(statements, standard), statements)
+
 
 class TestChunkAndValidateStatements(unittest.TestCase):
     def test_chunk_statements_never_exceeds_max_size(self):
diff --git a/aws_quickstart/datadog_integration_permissions.yaml b/aws_quickstart/datadog_integration_permissions.yaml
index 17e0b096..feb9562f 100644
--- a/aws_quickstart/datadog_integration_permissions.yaml
+++ b/aws_quickstart/datadog_integration_permissions.yaml
@@ -99,6 +99,7 @@ Resources:
         ZipFile: |
           import json
           import logging
+          import re
           from urllib.request import Request
           import urllib.error
           import urllib.parse
@@ -197,20 +198,77 @@ Resources:
               return {"Effect": "Allow", "Action": list(chunk), "Resource": "*"}
 
 
+          def _is_nonblank_string_or_list(value):
+              # .strip() only tests for blankness — it never rewrites the value, so a legitimately-formatted
+              # field passes through untouched.
+              if isinstance(value, str):
+                  return bool(value.strip())
+              if isinstance(value, list):
+                  return bool(value) and all(isinstance(v, str) and v.strip() for v in value)
+              return False
+
+
+          def _validate_action_field(action, context):
+              # A statement/chunk with a missing, null, or empty Action silently disappears in
+              # filter_instrumentation_statements (kept_actions ends up empty), which can collapse the
+              # entire candidate set to [] and trigger the same wipe-previously-attached-policies failure
+              # mode as a missing top-level field — just one level deeper.
+              if not _is_nonblank_string_or_list(action):
+                  raise DatadogAPIError(f"Datadog API returned a statement with an invalid Action field: {action!r} ({context})")
+
+
+          def _validate_policy_statement(statement, context):
+              # A statement missing Effect/Resource, or with a non-"Allow" Effect, passes prep and
+              # reaches cleanup_instrumentation_policies before create_policy rejects it as malformed —
+              # the same wipe-previously-attached-policies failure mode as an invalid Action.
+              if not isinstance(statement, dict):
+                  raise DatadogAPIError(f"Datadog API returned a non-dict statement: {statement!r} ({context})")
+              if statement.get("Effect") != "Allow":
+                  raise DatadogAPIError(
+                      f"Datadog API returned a statement with an invalid Effect field: {statement.get('Effect')!r} ({context})"
+                  )
+              _validate_action_field(statement.get("Action"), context)
+              resource = statement.get("Resource")
+              if not _is_nonblank_string_or_list(resource):
+                  raise DatadogAPIError(f"Datadog API returned a statement with an invalid Resource field: {resource!r} ({context})")
+              condition = statement.get("Condition")
+              if condition is not None and not isinstance(condition, dict):
+                  raise DatadogAPIError(f"Datadog API returned a statement with an invalid Condition field: {condition!r} ({context})")
+
+
           def resolve_instrumentation_statement_chunks(
               attributes, account_id, partition, standard_statements, max_size=MAX_MANAGED_POLICY_SIZE
           ):
               # Prefer structured policy_statements when the API provides them; policy_statements is
-              # always unchunked, so QuickStart renders, filters, and chunks it by size (Task 5).
+              # always unchunked, so QuickStart renders, filters, and chunks it by size.
               # permissions is pre-chunked server-side for the legacy fallback — filter each chunk but
               # preserve its boundaries rather than re-chunking, since dd-source already sized it to fit.
               policy_statements = attributes.get("policy_statements")
               if policy_statements is not None:
+                  if not isinstance(policy_statements, list):
+                      raise DatadogAPIError(
+                          f"Datadog API returned non-list policy_statements: {type(policy_statements).__name__}"
+                      )
+                  for statement in policy_statements:
+                      _validate_policy_statement(statement, "policy_statements")
                   rendered = [render_placeholders(s, account_id, partition) for s in policy_statements]
                   filtered = filter_instrumentation_statements(rendered, standard_statements)
                   return chunk_statements(filtered, max_size=max_size)
 
-              legacy_statements = [legacy_chunk_to_statement(chunk) for chunk in attributes.get("permissions", [])]
+              # policy_statements missing or null falls back to the legacy permissions field, but that
+              # field must actually be present and a list — schema drift that drops both fields must raise
+              # here rather than silently resolving to an empty candidate set, which attach_instrumentation_permissions
+              # would treat as a valid "nothing to attach" and use to wipe previously-attached policies.
+              if "permissions" not in attributes:
+                  raise DatadogAPIError("Datadog API response is missing both policy_statements and permissions")
+              legacy_permissions = attributes["permissions"]
+              if not isinstance(legacy_permissions, list):
+                  raise DatadogAPIError(
+                      f"Datadog API returned non-list permissions: {type(legacy_permissions).__name__}"
+                  )
+              for chunk in legacy_permissions:
+                  _validate_action_field(chunk if isinstance(chunk, list) else None, "permissions")
+              legacy_statements = [legacy_chunk_to_statement(chunk) for chunk in legacy_permissions]
               filtered = filter_instrumentation_statements(legacy_statements, standard_statements)
               return [[statement] for statement in filtered]
 
@@ -234,9 +292,25 @@ Resources:
               return standard_condition == instrumentation_condition
 
 
+          def _as_action_list(action):
+              # IAM's Action field may be a single string or a list; a bare string must not be iterated
+              # character-by-character.
+              if action is None:
+                  return []
+              return action if isinstance(action, list) else [action]
+
+
+          def _iam_action_matches(action, pattern):
+              # IAM action wildcards only recognize "*" and "?"; fnmatch also grants bracket
+              # character-class syntax that IAM does not, which could over-match an action pattern
+              # containing literal brackets.
+              regex = "".join(".*" if ch == "*" else "." if ch == "?" else re.escape(ch) for ch in pattern)
+              return re.fullmatch(regex, action) is not None
+
+
           def _action_covered_by_standard(action, resource, condition, standard_index):
-              for action_set, standard_resource, standard_condition in standard_index:
-                  if action not in action_set:
+              for action_patterns, standard_resource, standard_condition in standard_index:
+                  if not any(_iam_action_matches(action.lower(), pattern.lower()) for pattern in action_patterns):
                       continue
                   if not _resource_covers(standard_resource, resource):
                       continue
@@ -251,10 +325,8 @@ Resources:
               # grants an equivalent-or-broader permission across Action, Resource, and Condition.
               # Action-name-only filtering would silently drop scoped statements just because the same
               # action name also appears, unconditioned, in the standard policy.
-              # Action lists are converted to sets once here so each instrumentation action's membership
-              # check is O(1) instead of a fresh linear scan per action, per standard statement.
               standard_index = [
-                  (set(s.get("Action", [])), s.get("Resource", "*"), s.get("Condition")) for s in standard_statements
+                  (_as_action_list(s.get("Action", [])), s.get("Resource", "*"), s.get("Condition")) for s in standard_statements
               ]
               filtered = []
               for statement in statements:
@@ -262,7 +334,7 @@ Resources:
                   condition = statement.get("Condition")
                   kept_actions = [
                       action
-                      for action in statement.get("Action", [])
+                      for action in _as_action_list(statement.get("Action", []))
                       if not _action_covered_by_standard(action, resource, condition, standard_index)
                   ]
                   if kept_actions:
@@ -270,17 +342,18 @@ Resources:
               return filtered
 
 
+          def _policy_document_json(statements):
+              return json.dumps({"Version": "2012-10-17", "Statement": statements}, separators=(',', ':'))
+
+
           def _policy_document_size(statements):
-              return len(json.dumps({"Version": "2012-10-17", "Statement": statements}, separators=(',', ':')))
+              return len(_policy_document_json(statements))
 
 
           def chunk_statements(statements, max_size=MAX_MANAGED_POLICY_SIZE):
               # Chunk by measured minified-document size, never splitting a statement's Action away from
               # its Resource/Condition. If a single statement alone exceeds max_size it still gets its own
               # (oversize) chunk; validate_statement_chunks is what catches that before anything is deleted.
-              # Tracks the running document size incrementally (base wrapper + each statement's own
-              # serialized size + one comma per additional array element) instead of re-serializing the
-              # whole accumulated chunk on every statement, which would be O(n^2) in the chunk size.
               base_size = _policy_document_size([])
               chunks = []
               current = []
@@ -336,8 +409,7 @@ Resources:
 
           def _count_non_instrumentation_attached_policies(iam_client, role_name):
               # Exact-name match against what cleanup_instrumentation_policies actually deletes, not a
-              # substring check — a stale/unrelated policy whose name merely contains the instrumentation
-              # prefix isn't touched by cleanup and must still count against the attachment quota.
+              # substring check.
               instrumentation_names = set(_chunked_policy_names(role_name, BASE_POLICY_PREFIX_INSTRUMENTATION))
               count = 0
               paginator = iam_client.get_paginator("list_attached_role_policies")
@@ -402,23 +474,16 @@ Resources:
 
           def attach_standard_permissions(iam_client, role_name):
               permissions = fetch_permissions_from_datadog(STANDARD_PERMISSIONS_API_URL)
-              policy_document = {
-                  "Version": "2012-10-17",
-                  "Statement": [{"Effect": "Allow", "Action": permissions, "Resource": "*"}],
-              }
               iam_client.put_role_policy(
                   RoleName=role_name,
                   PolicyName=POLICY_NAME_STANDARD,
-                  PolicyDocument=json.dumps(policy_document, separators=(',', ':')),
+                  PolicyDocument=_policy_document_json([{"Effect": "Allow", "Action": permissions, "Resource": "*"}]),
               )
               return permissions
 
 
           def _create_and_attach_statement_policy(iam_client, role_name, policy_name, statements):
-              policy_json = json.dumps(
-                  {"Version": "2012-10-17", "Statement": statements},
-                  separators=(',', ':'),
-              )
+              policy_json = _policy_document_json(statements)
               LOGGER.info(f"Creating policy {policy_name} with {len(statements)} statements ({len(policy_json)} characters)")
               policy = iam_client.create_policy(PolicyName=policy_name, PolicyDocument=policy_json)
               iam_client.attach_role_policy(RoleName=role_name, PolicyArn=policy['Policy']['Arn'])
@@ -446,13 +511,11 @@ Resources:
               # Fetch before cleanup so that a transient API failure on an Update leaves the
               # previously-attached policies in place instead of silently revoking them.
               #
-              # standard_statements defaults to empty (no filtering): filtering is only safe when the caller
-              # knows what's actually attached to this role. The role-creation path passes the exact
-              # permissions it just attached via attach_standard_permissions in this same invocation; the
-              # add-on path (ManageBasePermissions=false) doesn't manage or verify the role's standard
-              # permissions, so it can't assume the *current* Datadog standard permission list reflects
-              # what's really on the role — filtering against that would risk dropping instrumentation
-              # actions the role doesn't actually have covered.
+              # standard_statements defaults to empty (no filtering): filtering is only safe when the
+              # caller knows what's actually attached. The role-creation path passes what it just attached
+              # via attach_standard_permissions in this same invocation; the add-on path
+              # (ManageBasePermissions=false) can't verify the role's real standard permissions, so it must
+              # skip filtering.
               if not resource_types:
                   # Only clean up if the previous Update had instrumentation enabled — avoids running
                   # delete calls on stacks that never opted in to instrumentation in the first place.