fix(aws): target EC2 Amazon AMI loading (#11958)

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

prowler/changelog.d/ec2-targeted-ami-loading.fixed.md

@@ -0,0 +1 @@
+EC2 AMI loading now targets Amazon-owned AMIs used by audited instances, reducing AWS API calls during EC2 scans

prowler/providers/aws/services/ec2/ec2_ami_public/ec2_ami_public.py

@@ -18,4 +18,4 @@ class ec2_ami_public(Check):
 
                 findings.append(report)
 
-            return findings
+        return findings

prowler/providers/aws/services/ec2/ec2_instance_with_outdated_ami/ec2_instance_with_outdated_ami.py

@@ -26,11 +26,12 @@ class ec2_instance_with_outdated_ami(Check):
             List[Check_Report_AWS]: A list containing the results of the check for each instance.
         """
         findings = []
+        images_by_id = getattr(ec2_client, "images_by_id", None)
+        if images_by_id is None:
+            images_by_id = {image.id: image for image in ec2_client.images}
+
         for instance in ec2_client.instances:
-            ami = next(
-                (image for image in ec2_client.images if image.id == instance.image_id),
-                None,
-            )
+            ami = images_by_id.get(instance.image_id)
             if ami and ami.owner == "amazon":
                 report = Check_Report_AWS(metadata=self.metadata(), resource=instance)
                 report.status = "PASS"

prowler/providers/aws/services/ec2/ec2_service.py

@@ -13,6 +13,8 @@ from prowler.lib.resource_limit import (
 from prowler.lib.scan_filters.scan_filters import is_resource_filtered
 from prowler.providers.aws.lib.service.service import AWSService
 
+DESCRIBE_IMAGES_IMAGE_IDS_BATCH_SIZE = 200
+
 
 class EC2(AWSService):
     def __init__(self, provider):
@@ -39,6 +41,7 @@ class EC2(AWSService):
         self.network_interfaces = {}
         self.__threading_call__(self._describe_network_interfaces)
         self.images = []
+        self.images_by_id = {}
         self.__threading_call__(self._describe_images)
         self.volumes = []
         self.__threading_call__(self._describe_volumes)
@@ -372,36 +375,90 @@ class EC2(AWSService):
 
     def _describe_images(self, regional_client):
         try:
-            for owner in ["self", "amazon"]:
-                try:
-                    for image in regional_client.describe_images(
-                        Owners=[owner], IncludeDeprecated=True
-                    )["Images"]:
-                        arn = f"arn:{self.audited_partition}:ec2:{regional_client.region}:{self.audited_account}:image/{image['ImageId']}"
-                        if not self.audit_resources or (
-                            is_resource_filtered(arn, self.audit_resources)
-                        ):
-                            self.images.append(
-                                Image(
-                                    id=image["ImageId"],
-                                    arn=arn,
-                                    name=image.get("Name", ""),
-                                    public=image.get("Public", False),
-                                    region=regional_client.region,
-                                    tags=image.get("Tags"),
-                                    deprecation_time=image.get("DeprecationTime"),
-                                    owner=owner,
-                                )
-                            )
-                except Exception as error:
-                    logger.error(
-                        f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
-                    )
+            try:
+                for image in regional_client.describe_images(
+                    Owners=["self"], IncludeDeprecated=True
+                )["Images"]:
+                    self._add_image(image, regional_client.region, "self")
+            except Exception as error:
+                logger.error(
+                    f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+                )
+
+            amazon_image_ids = sorted(
+                {
+                    instance.image_id
+                    for instance in self.instances
+                    if instance.region == regional_client.region
+                    and instance.image_id
+                    and instance.image_id not in self.images_by_id
+                }
+            )
+
+            for image_batch in self._get_image_id_batches(amazon_image_ids):
+                for image in self._describe_images_by_id(regional_client, image_batch):
+                    if self._is_amazon_image(image):
+                        self._add_image(image, regional_client.region, "amazon")
         except Exception as error:
             logger.error(
                 f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
             )
 
+    def _add_image(self, image, region, owner):
+        arn = f"arn:{self.audited_partition}:ec2:{region}:{self.audited_account}:image/{image['ImageId']}"
+        if not self.audit_resources or (
+            is_resource_filtered(arn, self.audit_resources)
+        ):
+            ec2_image = Image(
+                id=image["ImageId"],
+                arn=arn,
+                name=image.get("Name", ""),
+                public=image.get("Public", False),
+                region=region,
+                tags=image.get("Tags"),
+                deprecation_time=image.get("DeprecationTime"),
+                owner=owner,
+            )
+            self.images.append(ec2_image)
+            self.images_by_id[ec2_image.id] = ec2_image
+
+    def _describe_images_by_id(self, regional_client, image_ids):
+        try:
+            return regional_client.describe_images(
+                ImageIds=image_ids, IncludeDeprecated=True
+            )["Images"]
+        except ClientError as error:
+            if error.response["Error"]["Code"] == "InvalidAMIID.NotFound":
+                if len(image_ids) == 1:
+                    logger.warning(
+                        f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+                    )
+                    return []
+
+                midpoint = len(image_ids) // 2
+                return self._describe_images_by_id(
+                    regional_client, image_ids[:midpoint]
+                ) + self._describe_images_by_id(regional_client, image_ids[midpoint:])
+
+            logger.error(
+                f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+            )
+            return []
+        except Exception as error:
+            logger.error(
+                f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+            )
+            return []
+
+    @staticmethod
+    def _get_image_id_batches(image_ids):
+        for index in range(0, len(image_ids), DESCRIBE_IMAGES_IMAGE_IDS_BATCH_SIZE):
+            yield image_ids[index : index + DESCRIBE_IMAGES_IMAGE_IDS_BATCH_SIZE]
+
+    @staticmethod
+    def _is_amazon_image(image):
+        return image.get("ImageOwnerAlias") == "amazon"
+
     def _describe_volumes(self, regional_client):
         try:
             describe_volumes_paginator = regional_client.get_paginator(

tests/providers/aws/services/ec2/ec2_ami_public/ec2_ami_public_test.py

@@ -141,3 +141,86 @@ class Test_ec2_ami_public:
             )
             assert result[0].region == AWS_REGION_US_EAST_1
             assert result[0].resource_tags == []
+
+    @mock_aws
+    def test_multiple_self_owned_amis_mixed_public_and_private(self):
+        ec2 = client("ec2", region_name=AWS_REGION_US_EAST_1)
+
+        reservation = ec2.run_instances(ImageId=EXAMPLE_AMI_ID, MinCount=1, MaxCount=1)
+        instance = reservation["Instances"][0]
+        instance_id = instance["InstanceId"]
+
+        private_image_id = ec2.create_image(
+            InstanceId=instance_id,
+            Name="test-private-ami",
+            Description="this is a private test ami",
+        )["ImageId"]
+        public_image_id = ec2.create_image(
+            InstanceId=instance_id,
+            Name="test-public-ami",
+            Description="this is a public test ami",
+        )["ImageId"]
+
+        image = resource("ec2", region_name=AWS_REGION_US_EAST_1).Image(public_image_id)
+        image.modify_attribute(
+            ImageId=public_image_id,
+            Attribute="launchPermission",
+            OperationType="add",
+            UserGroups=["all"],
+        )
+
+        from prowler.providers.aws.services.ec2.ec2_service import EC2
+
+        aws_provider = set_mocked_aws_provider(
+            [AWS_REGION_EU_WEST_1, AWS_REGION_US_EAST_1]
+        )
+
+        with (
+            mock.patch(
+                "prowler.providers.common.provider.Provider.get_global_provider",
+                return_value=aws_provider,
+            ),
+            mock.patch(
+                "prowler.providers.aws.services.ec2.ec2_ami_public.ec2_ami_public.ec2_client",
+                new=EC2(aws_provider),
+            ),
+        ):
+            from prowler.providers.aws.services.ec2.ec2_ami_public.ec2_ami_public import (
+                ec2_ami_public,
+            )
+
+            check = ec2_ami_public()
+            result = check.execute()
+
+            findings_by_resource_id = {
+                finding.resource_id: finding for finding in result
+            }
+
+            assert len(result) == 2
+            assert set(findings_by_resource_id) == {private_image_id, public_image_id}
+
+            private_finding = findings_by_resource_id[private_image_id]
+            assert private_finding.status == "PASS"
+            assert (
+                private_finding.status_extended
+                == "EC2 AMI test-private-ami is not public."
+            )
+            assert (
+                private_finding.resource_arn
+                == f"arn:{aws_provider.identity.partition}:ec2:{AWS_REGION_US_EAST_1}:{aws_provider.identity.account}:image/{private_image_id}"
+            )
+            assert private_finding.region == AWS_REGION_US_EAST_1
+            assert private_finding.resource_tags == []
+
+            public_finding = findings_by_resource_id[public_image_id]
+            assert public_finding.status == "FAIL"
+            assert (
+                public_finding.status_extended
+                == "EC2 AMI test-public-ami is currently public."
+            )
+            assert (
+                public_finding.resource_arn
+                == f"arn:{aws_provider.identity.partition}:ec2:{AWS_REGION_US_EAST_1}:{aws_provider.identity.account}:image/{public_image_id}"
+            )
+            assert public_finding.region == AWS_REGION_US_EAST_1
+            assert public_finding.resource_tags == []

tests/providers/aws/services/ec2/ec2_instance_with_outdated_ami/ec2_instance_with_outdated_ami_test.py

@@ -1,11 +1,17 @@
 from unittest import mock
 
 import botocore
-from moto import mock_aws
+import pytest
 
 from tests.providers.aws.utils import AWS_REGION_US_EAST_1, set_mocked_aws_provider
 
 make_api_call = botocore.client.BaseClient._make_api_call
+describe_images_calls = []
+
+
+@pytest.fixture(autouse=True)
+def clear_describe_images_calls():
+    describe_images_calls.clear()
 
 
 def mock_make_api_call(self, operation_name, kwarg):
@@ -27,16 +33,25 @@ def mock_make_api_call(self, operation_name, kwarg):
             ]
         }
     elif operation_name == "DescribeImages":
-        if "Owners" in kwarg and kwarg["Owners"] == ["amazon"]:
+        describe_images_calls.append(kwarg)
+        if kwarg.get("Owners") == ["self"]:
+            return {"Images": []}
+        if kwarg.get("Owners") == ["amazon"]:
+            raise AssertionError(
+                "Amazon AMIs must not be fetched with a broad owner lookup"
+            )
+        if kwarg.get("ImageIds") == ["ami-12345678"]:
             return {
                 "Images": [
                     {
                         "ImageId": "ami-12345678",
                         "DeprecationTime": "2050-01-01T00:00:00.000Z",
                         "Public": True,
+                        "ImageOwnerAlias": "amazon",
                     }
                 ]
             }
+        return {"Images": []}
     return make_api_call(self, operation_name, kwarg)
 
 
@@ -59,6 +74,13 @@ def mock_make_api_call_private(self, operation_name, kwarg):
             ]
         }
     elif operation_name == "DescribeImages":
+        describe_images_calls.append(kwarg)
+        if kwarg.get("Owners") == ["amazon"]:
+            raise AssertionError(
+                "Amazon AMIs must not be fetched with a broad owner lookup"
+            )
+        if kwarg.get("Owners") == ["self"]:
+            return {"Images": []}
         return {
             "Images": [
                 {
@@ -90,16 +112,25 @@ def mock_make_api_call_outdated_ami(self, operation_name, kwarg):
             ]
         }
     elif operation_name == "DescribeImages":
-        if "Owners" in kwarg and kwarg["Owners"] == ["amazon"]:
+        describe_images_calls.append(kwarg)
+        if kwarg.get("Owners") == ["self"]:
+            return {"Images": []}
+        if kwarg.get("Owners") == ["amazon"]:
+            raise AssertionError(
+                "Amazon AMIs must not be fetched with a broad owner lookup"
+            )
+        if kwarg.get("ImageIds") == ["ami-87654321"]:
             return {
                 "Images": [
                     {
                         "ImageId": "ami-87654321",
                         "DeprecationTime": "2022-01-01T00:00:00.000Z",
                         "Public": True,
+                        "ImageOwnerAlias": "amazon",
                     }
                 ]
             }
+        return {"Images": []}
     return make_api_call(self, operation_name, kwarg)
 
 
@@ -122,12 +153,32 @@ def mock_make_api_call_missing_ami(self, operation_name, kwarg):
             ]
         }
     elif operation_name == "DescribeImages":
+        describe_images_calls.append(kwarg)
+        if kwarg.get("Owners") == ["amazon"]:
+            raise AssertionError(
+                "Amazon AMIs must not be fetched with a broad owner lookup"
+            )
+        return {"Images": []}
+    return make_api_call(self, operation_name, kwarg)
+
+
+def mock_make_api_call_no_instances(self, operation_name, kwarg):
+    if operation_name == "DescribeInstances":
+        return {"Reservations": []}
+    elif operation_name == "DescribeImages":
+        describe_images_calls.append(kwarg)
+        if kwarg.get("Owners") == ["amazon"]:
+            raise AssertionError(
+                "Amazon AMIs must not be fetched with a broad owner lookup"
+            )
         return {"Images": []}
     return make_api_call(self, operation_name, kwarg)
 
 
 class Test_ec2_instance_with_outdated_ami:
-    @mock_aws
+    @mock.patch(
+        "botocore.client.BaseClient._make_api_call", new=mock_make_api_call_no_instances
+    )
     def test_ec2_no_instances(self):
         from prowler.providers.aws.services.ec2.ec2_service import EC2
 
@@ -151,6 +202,7 @@ class Test_ec2_instance_with_outdated_ami:
             result = check.execute()
 
             assert len(result) == 0
+            assert not any("ImageIds" in call for call in describe_images_calls)
 
     @mock.patch(
         "botocore.client.BaseClient._make_api_call", new=mock_make_api_call_private
@@ -178,6 +230,13 @@ class Test_ec2_instance_with_outdated_ami:
             result = check.execute()
 
             assert len(result) == 0
+            assert not any(
+                call.get("Owners") == ["amazon"] for call in describe_images_calls
+            )
+            assert any(
+                call.get("ImageIds") == ["ami-12345678"]
+                for call in describe_images_calls
+            )
 
     @mock.patch("botocore.client.BaseClient._make_api_call", new=mock_make_api_call)
     def test_instance_ami_not_outdated(self):
@@ -209,6 +268,13 @@ class Test_ec2_instance_with_outdated_ami:
                 result[0].status_extended
                 == "EC2 Instance i-0123456789abcdef0 is not using an outdated AMI."
             )
+            assert not any(
+                call.get("Owners") == ["amazon"] for call in describe_images_calls
+            )
+            assert any(
+                call.get("ImageIds") == ["ami-12345678"]
+                for call in describe_images_calls
+            )
 
     @mock.patch(
         "botocore.client.BaseClient._make_api_call", new=mock_make_api_call_outdated_ami
@@ -242,6 +308,13 @@ class Test_ec2_instance_with_outdated_ami:
                 result[0].status_extended
                 == "EC2 Instance i-0123456789abcdef0 is using outdated AMI ami-87654321."
             )
+            assert not any(
+                call.get("Owners") == ["amazon"] for call in describe_images_calls
+            )
+            assert any(
+                call.get("ImageIds") == ["ami-87654321"]
+                for call in describe_images_calls
+            )
 
     @mock.patch(
         "botocore.client.BaseClient._make_api_call", new=mock_make_api_call_missing_ami
@@ -269,3 +342,10 @@ class Test_ec2_instance_with_outdated_ami:
             result = check.execute()
 
             assert result == []
+            assert not any(
+                call.get("Owners") == ["amazon"] for call in describe_images_calls
+            )
+            assert any(
+                call.get("ImageIds") == ["ami-missing"]
+                for call in describe_images_calls
+            )

tests/providers/aws/services/ec2/ec2_launch_template_no_public_ip/ec2_launch_template_no_public_ip_test.py

@@ -39,6 +39,12 @@ def mock_make_api_call(self, operation_name, kwarg):
     return make_api_call(self, operation_name, kwarg)
 
 
+def get_mocked_ec2_client():
+    ec2_client = mock.MagicMock()
+    ec2_client.audit_config = {}
+    return ec2_client
+
+
 class Test_ec2_launch_template_no_public_ip:
     @mock_aws
     def test_no_launch_templates(self):
@@ -124,7 +130,7 @@ class Test_ec2_launch_template_no_public_ip:
             assert result[0].resource_tags == []
 
     def test_launch_template_public_ip_auto_assign(self):
-        ec2_client = mock.MagicMock()
+        ec2_client = get_mocked_ec2_client()
         launch_template_name = "tester"
         launch_template_id = "lt-1234567890"
         launch_template_arn = (
@@ -190,7 +196,7 @@ class Test_ec2_launch_template_no_public_ip:
     def test_network_interface_with_public_ipv4_network_interface_autoassign_true_and_false(
         self,
     ):
-        ec2_client = mock.MagicMock()
+        ec2_client = get_mocked_ec2_client()
         launch_template_name = "tester"
         launch_template_id = "lt-1234567890"
         launch_template_arn = (
@@ -290,7 +296,7 @@ class Test_ec2_launch_template_no_public_ip:
     def test_network_interface_with_public_ipv6_network_interface_autoassign_true_and_false(
         self,
     ):
-        ec2_client = mock.MagicMock()
+        ec2_client = get_mocked_ec2_client()
         launch_template_name = "tester"
         launch_template_id = "lt-1234567890"
         launch_template_arn = (

tests/providers/aws/services/ec2/ec2_service_test.py

@@ -6,12 +6,17 @@ from datetime import datetime
 import botocore
 import mock
 from boto3 import client, resource
+from botocore.exceptions import ClientError
 from dateutil.tz import tzutc
 from freezegun import freeze_time
 from moto import mock_aws
 
 from prowler.config.config import encoding_format_utf_8
-from prowler.providers.aws.services.ec2.ec2_service import EC2, Snapshot
+from prowler.providers.aws.services.ec2.ec2_service import (
+    DESCRIBE_IMAGES_IMAGE_IDS_BATCH_SIZE,
+    EC2,
+    Snapshot,
+)
 from tests.providers.aws.utils import (
     AWS_ACCOUNT_NUMBER,
     AWS_REGION_EU_WEST_1,
@@ -196,6 +201,123 @@ class Test_EC2_Service:
         assert ec2.volumes_with_snapshots == {"vol-old": True, "vol-new": True}
         assert [snapshot.id for snapshot in ec2.snapshots] == ["snap-new"]
 
+    def test_describe_images_by_id_preserves_valid_amis_when_batch_has_missing_ids(
+        self,
+    ):
+        missing_image_id = "ami-0000-missing"
+        valid_image_ids = [
+            f"ami-{index:04d}"
+            for index in range(1, DESCRIBE_IMAGES_IMAGE_IDS_BATCH_SIZE + 1)
+        ]
+        instance_image_ids = valid_image_ids + [missing_image_id]
+
+        class FakeInstance:
+            def __init__(self, image_id):
+                self.region = AWS_REGION_US_EAST_1
+                self.image_id = image_id
+
+        class FakeEC2Client:
+            region = AWS_REGION_US_EAST_1
+
+            def __init__(self):
+                self.image_id_calls = []
+
+            def describe_images(self, **kwargs):
+                if kwargs.get("Owners") == ["self"]:
+                    return {"Images": []}
+
+                image_ids = kwargs["ImageIds"]
+                self.image_id_calls.append(image_ids)
+                if missing_image_id in image_ids:
+                    raise ClientError(
+                        {
+                            "Error": {
+                                "Code": "InvalidAMIID.NotFound",
+                                "Message": "The image id does not exist",
+                            }
+                        },
+                        "DescribeImages",
+                    )
+
+                return {
+                    "Images": [
+                        {
+                            "ImageId": image_id,
+                            "Public": True,
+                            "ImageOwnerAlias": "amazon",
+                        }
+                        for image_id in image_ids
+                    ]
+                }
+
+        regional_client = FakeEC2Client()
+        ec2 = EC2.__new__(EC2)
+        ec2.instances = [FakeInstance(image_id) for image_id in instance_image_ids]
+        ec2.images = []
+        ec2.images_by_id = {}
+        ec2.audit_resources = []
+        ec2.audited_partition = "aws"
+        ec2.audited_account = AWS_ACCOUNT_NUMBER
+
+        ec2._describe_images(regional_client)
+
+        assert (
+            len(regional_client.image_id_calls[0])
+            == DESCRIBE_IMAGES_IMAGE_IDS_BATCH_SIZE
+        )
+        assert regional_client.image_id_calls[-1] == [valid_image_ids[-1]]
+        assert missing_image_id not in ec2.images_by_id
+        assert set(ec2.images_by_id) == set(valid_image_ids)
+        assert {image.owner for image in ec2.images} == {"amazon"}
+
+    def test_describe_images_ignores_non_amazon_public_images_returned_by_image_id(
+        self,
+    ):
+        marketplace_image_id = "ami-marketplace-public"
+        vendor_image_id = "ami-vendor-public"
+
+        class FakeInstance:
+            def __init__(self, image_id):
+                self.region = AWS_REGION_US_EAST_1
+                self.image_id = image_id
+
+        class FakeEC2Client:
+            region = AWS_REGION_US_EAST_1
+
+            def describe_images(self, **kwargs):
+                if kwargs.get("Owners") == ["self"]:
+                    return {"Images": []}
+
+                return {
+                    "Images": [
+                        {
+                            "ImageId": marketplace_image_id,
+                            "Public": True,
+                            "ImageOwnerAlias": "aws-marketplace",
+                        },
+                        {
+                            "ImageId": vendor_image_id,
+                            "Public": True,
+                        },
+                    ]
+                }
+
+        ec2 = EC2.__new__(EC2)
+        ec2.instances = [
+            FakeInstance(marketplace_image_id),
+            FakeInstance(vendor_image_id),
+        ]
+        ec2.images = []
+        ec2.images_by_id = {}
+        ec2.audit_resources = []
+        ec2.audited_partition = "aws"
+        ec2.audited_account = AWS_ACCOUNT_NUMBER
+
+        ec2._describe_images(FakeEC2Client())
+
+        assert ec2.images == []
+        assert ec2.images_by_id == {}
+
     # Test EC2 Describe Instances
     @mock_aws
     @freeze_time(MOCK_DATETIME)
@@ -743,9 +865,10 @@ class Test_EC2_Service:
             {"Key": "OS_Version", "Value": "AWS Linux 2"},
         ]
 
-        # Verify that Amazon images are also present
-        amazon_images = [img for img in ec2.images if img.owner == "amazon"]
-        assert len(amazon_images) > 0  # Should have Amazon AMIs
+        # Amazon public AMIs are fetched by targeted instance ImageIds only when
+        # AWS identifies them as Amazon-owned. Moto does not expose that owner
+        # alias for its fixture AMIs, so this service test only verifies the
+        # self-owned AMI behavior used by ec2_ami_public.
 
     # Test EC2 Describe Volumes
     @mock_aws