fix(sdk): limit ECS task definitions by registration date (#11891)

Co-authored-by: Hugo Pereira Brito <101209179+HugoPBrito@users.noreply.github.com>

prowler/changelog.d/ecs-task-definitions-registration-date.fixed.md

@@ -0,0 +1 @@
+ECS task definition resource limits now select the latest task definitions by registration date instead of relying on ARN ordering

prowler/providers/aws/services/ecs/ecs_service.py

@@ -81,12 +81,13 @@ class ECS(AWSService):
 
         Resources already fetched are memoized in ``self.task_definitions`` and
         reused across checks (checks run sequentially, so no locking is needed).
-        The configured resource limit bounds ``describe_task_definition`` calls.
+        Task definitions are described before applying the configured resource
+        limit because AWS exposes ``registeredAt`` only through
+        ``describe_task_definition``. The limit bounds the task definitions
+        exposed to checks for analysis.
         """
         task_definitions = []
-        for arn, region in limit_resources(
-            self._list_task_definition_arns(), self.task_definition_limit
-        ):
+        for arn, region in self._list_task_definition_arns():
             task_definition = self.task_definitions.get(arn)
             if task_definition is None:
                 task_definition = TaskDefinition(
@@ -102,12 +103,42 @@ class ECS(AWSService):
 
         self.__threading_call__(self._describe_task_definition, task_definitions)
 
-        for arn, _ in limit_resources(
-            self._list_task_definition_arns(), self.task_definition_limit
-        ):
-            task_definition = self.task_definitions[arn]
+        selected_task_definitions = list(
+            limit_resources(
+                self._sort_task_definitions_by_registration_date(
+                    self.task_definitions.values()
+                ),
+                self.task_definition_limit,
+            )
+        )
+        self.task_definitions = {
+            task_definition.arn: task_definition
+            for task_definition in selected_task_definitions
+        }
+        for task_definition in selected_task_definitions:
             yield task_definition
 
+    @staticmethod
+    def _sort_task_definitions_by_registration_date(task_definitions):
+        task_definitions = list(task_definitions)
+        if not any(
+            task_definition.registered_at for task_definition in task_definitions
+        ):
+            return task_definitions
+
+        return sorted(
+            task_definitions,
+            key=lambda task_definition: (
+                task_definition.registered_at is None,
+                (
+                    -task_definition.registered_at.timestamp()
+                    if task_definition.registered_at
+                    else 0
+                ),
+                task_definition.arn,
+            ),
+        )
+
     def _describe_task_definition(self, task_definition):
         logger.info("ECS - Describing Task Definition...")
         try:

tests/providers/aws/services/ecs/ecs_service_test.py

@@ -1,3 +1,4 @@
+from datetime import datetime, timezone
 from unittest.mock import patch
 
 import botocore
@@ -136,6 +137,110 @@ def mock_generate_multi_region_clients(provider, service):
     }
 
 
+def mock_make_api_call_task_definitions_by_registration_date(
+    self, operation_name, kwarg
+):
+    task_definition_dates = {
+        f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:task-definition/eu-old:1": datetime(
+            2024, 1, 1, tzinfo=timezone.utc
+        ),
+        f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:task-definition/eu-second-newest:1": datetime(
+            2024, 5, 1, tzinfo=timezone.utc
+        ),
+        f"arn:aws:ecs:{AWS_REGION_US_EAST_1}:123456789012:task-definition/us-newest:1": datetime(
+            2024, 6, 1, tzinfo=timezone.utc
+        ),
+        f"arn:aws:ecs:{AWS_REGION_US_EAST_1}:123456789012:task-definition/us-old:1": datetime(
+            2024, 2, 1, tzinfo=timezone.utc
+        ),
+    }
+    task_definitions_by_region = {
+        AWS_REGION_EU_WEST_1: [
+            task_definition
+            for task_definition in task_definition_dates
+            if f":{AWS_REGION_EU_WEST_1}:" in task_definition
+        ],
+        AWS_REGION_US_EAST_1: [
+            task_definition
+            for task_definition in task_definition_dates
+            if f":{AWS_REGION_US_EAST_1}:" in task_definition
+        ],
+    }
+
+    if operation_name == "ListTaskDefinitions":
+        return {"taskDefinitionArns": task_definitions_by_region[self.region]}
+    if operation_name == "DescribeTaskDefinition":
+        return {
+            "taskDefinition": {
+                "containerDefinitions": [],
+                "registeredAt": task_definition_dates[kwarg["taskDefinition"]],
+            },
+            "tags": [],
+        }
+    if operation_name == "ListClusters":
+        return {"clusterArns": []}
+    return make_api_call(self, operation_name, kwarg)
+
+
+def mock_make_api_call_task_definitions_with_equal_registration_dates(
+    self, operation_name, kwarg
+):
+    registered_at = datetime(2024, 1, 1, tzinfo=timezone.utc)
+    task_definition_dates = {
+        f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:task-definition/zzz-task:1": registered_at,
+        f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:task-definition/aaa-task:1": registered_at,
+        f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:task-definition/mmm-task:1": registered_at,
+    }
+
+    if operation_name == "ListTaskDefinitions":
+        return {"taskDefinitionArns": list(task_definition_dates)}
+    if operation_name == "DescribeTaskDefinition":
+        return {
+            "taskDefinition": {
+                "containerDefinitions": [],
+                "registeredAt": task_definition_dates[kwarg["taskDefinition"]],
+            },
+            "tags": [],
+        }
+    if operation_name == "ListClusters":
+        return {"clusterArns": []}
+    return make_api_call(self, operation_name, kwarg)
+
+
+DESCRIBED_TASK_DEFINITIONS = []
+
+
+def mock_make_api_call_task_definitions_with_audit_resources(
+    self, operation_name, kwarg
+):
+    task_definition_dates = {
+        f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:task-definition/audited-older:1": datetime(
+            2024, 1, 1, tzinfo=timezone.utc
+        ),
+        f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:task-definition/audited-newer:1": datetime(
+            2024, 2, 1, tzinfo=timezone.utc
+        ),
+        f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:task-definition/unaudited-newest:1": datetime(
+            2024, 3, 1, tzinfo=timezone.utc
+        ),
+    }
+
+    if operation_name == "ListTaskDefinitions":
+        return {"taskDefinitionArns": list(task_definition_dates)}
+    if operation_name == "DescribeTaskDefinition":
+        DESCRIBED_TASK_DEFINITIONS.append(kwarg["taskDefinition"])
+        return {
+            "taskDefinition": {
+                "containerDefinitions": [],
+                "registeredAt": task_definition_dates[kwarg["taskDefinition"]],
+            },
+            "tags": [],
+        }
+    if operation_name == "ListClusters":
+        return {"clusterArns": []}
+    return make_api_call(self, operation_name, kwarg)
+
+
 @patch(
     "prowler.providers.aws.aws_provider.AwsProvider.generate_regional_clients",
     new=mock_generate_regional_clients,
@@ -292,9 +397,9 @@ class Test_ECS_Service:
             ecs = ECS(aws_provider)
 
             assert [td.revision for td in ecs.task_definitions.values()] == ["3", "2"]
-            assert len(describe_calls) == 2
+            assert len(describe_calls) == 3
 
-    def test_task_definition_limit_bounds_describe_calls(self):
+    def test_task_definition_limit_describes_candidates_before_exposing_limit(self):
         describe_calls = []
 
         def counting_make_api_call(self, operation_name, kwarg):
@@ -326,9 +431,82 @@ class Test_ECS_Service:
             ecs = ECS(aws_provider)
 
             assert [td.revision for td in ecs.task_definitions.values()] == ["3"]
-            assert describe_calls == [
-                "arn:aws:ecs:eu-west-1:123456789012:task-definition/fam:3"
-            ]
+            assert sorted(describe_calls) == sorted(
+                [
+                    "arn:aws:ecs:eu-west-1:123456789012:task-definition/fam:3",
+                    "arn:aws:ecs:eu-west-1:123456789012:task-definition/fam:2",
+                    "arn:aws:ecs:eu-west-1:123456789012:task-definition/fam:1",
+                ]
+            )
+
+    @patch(
+        "botocore.client.BaseClient._make_api_call",
+        new=mock_make_api_call_task_definitions_by_registration_date,
+    )
+    def test_task_definition_limit_uses_global_latest_registration_dates(self):
+        aws_provider = set_mocked_aws_provider(
+            [AWS_REGION_EU_WEST_1, AWS_REGION_US_EAST_1],
+            audit_config={"max_ecs_task_definitions": 2},
+        )
+        with patch(
+            "prowler.providers.aws.aws_provider.AwsProvider.generate_regional_clients",
+            new=mock_generate_multi_region_clients,
+        ):
+            ecs = ECS(aws_provider)
+
+        assert list(ecs.task_definitions) == [
+            f"arn:aws:ecs:{AWS_REGION_US_EAST_1}:123456789012:task-definition/us-newest:1",
+            f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:task-definition/eu-second-newest:1",
+        ]
+
+    @patch(
+        "botocore.client.BaseClient._make_api_call",
+        new=mock_make_api_call_task_definitions_with_equal_registration_dates,
+    )
+    def test_task_definition_limit_uses_arn_order_for_equal_registration_dates(self):
+        aws_provider = set_mocked_aws_provider(
+            [AWS_REGION_EU_WEST_1], audit_config={"max_ecs_task_definitions": 2}
+        )
+        ecs = ECS(aws_provider)
+
+        assert list(ecs.task_definitions) == [
+            f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:task-definition/aaa-task:1",
+            f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:task-definition/mmm-task:1",
+        ]
+
+    @patch(
+        "botocore.client.BaseClient._make_api_call",
+        new=mock_make_api_call_task_definitions_with_audit_resources,
+    )
+    def test_task_definition_limit_applies_after_audit_resources_filter(self):
+        DESCRIBED_TASK_DEFINITIONS.clear()
+        audited_older_task_definition = (
+            f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:"
+            "task-definition/audited-older:1"
+        )
+        audited_newer_task_definition = (
+            f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:"
+            "task-definition/audited-newer:1"
+        )
+        unaudited_newest_task_definition = (
+            f"arn:aws:ecs:{AWS_REGION_EU_WEST_1}:123456789012:"
+            "task-definition/unaudited-newest:1"
+        )
+        aws_provider = set_mocked_aws_provider(
+            [AWS_REGION_EU_WEST_1], audit_config={"max_ecs_task_definitions": 1}
+        )
+        aws_provider._audit_resources = [
+            audited_older_task_definition,
+            audited_newer_task_definition,
+        ]
+
+        ecs = ECS(aws_provider)
+
+        assert sorted(DESCRIBED_TASK_DEFINITIONS) == sorted(
+            [audited_older_task_definition, audited_newer_task_definition]
+        )
+        assert list(ecs.task_definitions) == [audited_newer_task_definition]
+        assert unaudited_newest_task_definition not in ecs.task_definitions
 
     def test_task_definition_limit_does_not_starve_later_regions(self):
         describe_calls = []
@@ -381,6 +559,8 @@ class Test_ECS_Service:
             ]
             assert set(describe_calls) == {
                 "arn:aws:ecs:eu-west-1:123456789012:task-definition/fam:3",
+                "arn:aws:ecs:eu-west-1:123456789012:task-definition/fam:2",
+                "arn:aws:ecs:eu-west-1:123456789012:task-definition/fam:1",
                 "arn:aws:ecs:us-east-1:123456789012:task-definition/fam:9",
             }