From 20aad80a78a9fb2c2f30504daa5ecc4b1d570958 Mon Sep 17 00:00:00 2001 From: Hugo Pereira Brito <101209179+HugoPBrito@users.noreply.github.com> Date: Wed, 8 Jul 2026 15:59:16 +0100 Subject: [PATCH] fix(aws): avoid full ec2 inventory in dlm check (#11850) --- ...dlm-ebs-snapshot-lifecycle-policy.fixed.md | 1 + ...lm_ebs_snapshot_lifecycle_policy_exists.py | 5 +- .../providers/aws/services/dlm/dlm_service.py | 34 +++++ ...s_snapshot_lifecycle_policy_exists_test.py | 119 ++++++++++++------ .../aws/services/dlm/dlm_service_test.py | 68 ++++++++++ 5 files changed, 185 insertions(+), 42 deletions(-) create mode 100644 prowler/changelog.d/dlm-ebs-snapshot-lifecycle-policy.fixed.md diff --git a/prowler/changelog.d/dlm-ebs-snapshot-lifecycle-policy.fixed.md b/prowler/changelog.d/dlm-ebs-snapshot-lifecycle-policy.fixed.md new file mode 100644 index 0000000000..2a07052351 --- /dev/null +++ b/prowler/changelog.d/dlm-ebs-snapshot-lifecycle-policy.fixed.md @@ -0,0 +1 @@ +`dlm_ebs_snapshot_lifecycle_policy_exists` no longer initializes the full EC2 inventory just to detect EBS snapshots, avoiding slow scans when checking DLM lifecycle policies diff --git a/prowler/providers/aws/services/dlm/dlm_ebs_snapshot_lifecycle_policy_exists/dlm_ebs_snapshot_lifecycle_policy_exists.py b/prowler/providers/aws/services/dlm/dlm_ebs_snapshot_lifecycle_policy_exists/dlm_ebs_snapshot_lifecycle_policy_exists.py index 3ee815fb23..c7eedc7f8f 100644 --- a/prowler/providers/aws/services/dlm/dlm_ebs_snapshot_lifecycle_policy_exists/dlm_ebs_snapshot_lifecycle_policy_exists.py +++ b/prowler/providers/aws/services/dlm/dlm_ebs_snapshot_lifecycle_policy_exists/dlm_ebs_snapshot_lifecycle_policy_exists.py @@ -1,6 +1,5 @@ from prowler.lib.check.models import Check, Check_Report_AWS from prowler.providers.aws.services.dlm.dlm_client import dlm_client -from prowler.providers.aws.services.ec2.ec2_client import ec2_client class dlm_ebs_snapshot_lifecycle_policy_exists(Check): @@ -8,8 +7,8 @@ class dlm_ebs_snapshot_lifecycle_policy_exists(Check): findings = [] for region in dlm_client.lifecycle_policies: if ( - region in ec2_client.regions_with_snapshots - and ec2_client.regions_with_snapshots[region] + region in dlm_client.regions_with_snapshots + and dlm_client.regions_with_snapshots[region] ): report = Check_Report_AWS( metadata=self.metadata(), diff --git a/prowler/providers/aws/services/dlm/dlm_service.py b/prowler/providers/aws/services/dlm/dlm_service.py index 1d6fff9b5a..f0f67e631c 100644 --- a/prowler/providers/aws/services/dlm/dlm_service.py +++ b/prowler/providers/aws/services/dlm/dlm_service.py @@ -9,7 +9,13 @@ class DLM(AWSService): # Call AWSService's __init__ super().__init__(__class__.__name__, provider) self.lifecycle_policies = {} + self.regions_with_snapshots = {} self.__threading_call__(self._get_lifecycle_policies) + ec2_regional_clients = provider.generate_regional_clients("ec2") or {} + self.__threading_call__( + self._get_regions_with_snapshots, + iterator=ec2_regional_clients.values(), + ) def _get_lifecycle_policy_arn_template(self, region): return ( @@ -35,6 +41,34 @@ class DLM(AWSService): f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" ) + def _get_regions_with_snapshots(self, regional_client): + logger.info("DLM - Checking regions with self-owned EBS snapshots...") + try: + self.regions_with_snapshots[regional_client.region] = False + next_token = None + while True: + describe_snapshots_args = { + "OwnerIds": ["self"], + "MaxResults": 5, + } + if next_token: + describe_snapshots_args["NextToken"] = next_token + + snapshots = regional_client.describe_snapshots( + **describe_snapshots_args + ) + if snapshots.get("Snapshots"): + self.regions_with_snapshots[regional_client.region] = True + break + + next_token = snapshots.get("NextToken") + if not next_token: + break + except Exception as error: + logger.error( + f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" + ) + class LifecyclePolicy(BaseModel): id: str diff --git a/tests/providers/aws/services/dlm/dlm_ebs_snapshot_lifecycle_policy_exists/dlm_ebs_snapshot_lifecycle_policy_exists_test.py b/tests/providers/aws/services/dlm/dlm_ebs_snapshot_lifecycle_policy_exists/dlm_ebs_snapshot_lifecycle_policy_exists_test.py index 3412aead68..47950ab33f 100644 --- a/tests/providers/aws/services/dlm/dlm_ebs_snapshot_lifecycle_policy_exists/dlm_ebs_snapshot_lifecycle_policy_exists_test.py +++ b/tests/providers/aws/services/dlm/dlm_ebs_snapshot_lifecycle_policy_exists/dlm_ebs_snapshot_lifecycle_policy_exists_test.py @@ -1,3 +1,5 @@ +import builtins +import sys from unittest import mock from boto3 import client, resource @@ -12,38 +14,42 @@ from tests.providers.aws.utils import ( ) LIFECYCLE_POLICY_ID = "policy-XXXXXXXXXXXX" +CHECK_MODULE = ( + "prowler.providers.aws.services.dlm." + "dlm_ebs_snapshot_lifecycle_policy_exists." + "dlm_ebs_snapshot_lifecycle_policy_exists" +) +DLM_CLIENT_MODULE = "prowler.providers.aws.services.dlm.dlm_client" +EC2_CLIENT_MODULE = "prowler.providers.aws.services.ec2.ec2_client" + + +def unload_dlm_check_modules(): + sys.modules.pop(CHECK_MODULE, None) + sys.modules.pop(DLM_CLIENT_MODULE, None) class Test_dlm_ebs_snapshot_lifecycle_policy_exists: @mock_aws def test_no_ebs_snapshot_no_lifecycle_policies(self): # DLM Mock Client - dlm_client = mock.MagicMock + dlm_client = mock.MagicMock() dlm_client.audited_account = AWS_ACCOUNT_NUMBER dlm_client.audited_account_arn = AWS_ACCOUNT_ARN dlm_client.lifecycle_policies = {} + dlm_client.regions_with_snapshots = {} aws_provider = set_mocked_aws_provider([AWS_REGION_US_EAST_1]) - - from prowler.providers.aws.services.ec2.ec2_service import EC2 + unload_dlm_check_modules() with ( mock.patch( "prowler.providers.aws.services.dlm.dlm_service.DLM", - new=dlm_client, + new=mock.MagicMock(return_value=dlm_client), ), mock.patch( "prowler.providers.common.provider.Provider.get_global_provider", return_value=aws_provider, ), - mock.patch( - "prowler.providers.aws.services.ec2.ec2_service.EC2", - return_value=EC2(aws_provider), - ) as ec2_client, - mock.patch( - "prowler.providers.aws.services.ec2.ec2_client.ec2_client", - new=ec2_client, - ), ): from prowler.providers.aws.services.dlm.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_ebs_snapshot_lifecycle_policy_exists import ( dlm_ebs_snapshot_lifecycle_policy_exists, @@ -92,22 +98,22 @@ class Test_dlm_ebs_snapshot_lifecycle_policy_exists: ) } } + dlm_client.regions_with_snapshots = {AWS_REGION_US_EAST_1: True} dlm_client.lifecycle_policy_arn_template = f"arn:{dlm_client.audited_partition}:dlm:{dlm_client.region}:{dlm_client.audited_account}:policy" dlm_client._get_lifecycle_policy_arn_template = mock.MagicMock( return_value=dlm_client.lifecycle_policy_arn_template ) aws_provider = set_mocked_aws_provider([AWS_REGION_US_EAST_1]) - - from prowler.providers.aws.services.ec2.ec2_service import EC2 + unload_dlm_check_modules() with ( mock.patch( - "prowler.providers.common.provider.Provider.get_global_provider", - return_value=aws_provider, + "prowler.providers.aws.services.dlm.dlm_service.DLM", + new=mock.MagicMock(return_value=dlm_client), ), mock.patch( - "prowler.providers.aws.services.dlm.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_ebs_snapshot_lifecycle_policy_exists.ec2_client", - new=EC2(aws_provider), + "prowler.providers.common.provider.Provider.get_global_provider", + return_value=aws_provider, ), mock.patch( "prowler.providers.aws.services.dlm.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_client", @@ -154,25 +160,23 @@ class Test_dlm_ebs_snapshot_lifecycle_policy_exists: )["SnapshotId"] # DLM Mock Client - dlm_client = mock.MagicMock + dlm_client = mock.MagicMock() dlm_client.audited_account = AWS_ACCOUNT_NUMBER dlm_client.audited_account_arn = AWS_ACCOUNT_ARN dlm_client.lifecycle_policies = {} - - # from prowler.providers.aws.services.ec2.ec2_service import EC2 + dlm_client.regions_with_snapshots = {AWS_REGION_US_EAST_1: True} aws_provider = set_mocked_aws_provider([AWS_REGION_US_EAST_1]) - - from prowler.providers.aws.services.ec2.ec2_service import EC2 + unload_dlm_check_modules() with ( mock.patch( - "prowler.providers.common.provider.Provider.get_global_provider", - return_value=aws_provider, + "prowler.providers.aws.services.dlm.dlm_service.DLM", + new=mock.MagicMock(return_value=dlm_client), ), mock.patch( - "prowler.providers.aws.services.dlm.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_ebs_snapshot_lifecycle_policy_exists.ec2_client", - new=EC2(aws_provider), + "prowler.providers.common.provider.Provider.get_global_provider", + return_value=aws_provider, ), mock.patch( "prowler.providers.aws.services.dlm.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_client", @@ -190,7 +194,7 @@ class Test_dlm_ebs_snapshot_lifecycle_policy_exists: @mock_aws def test_no_ebs_snapshot_and_dlm_lifecycle_policy(self): # DLM Mock Client - dlm_client = mock.MagicMock + dlm_client = mock.MagicMock() dlm_client.audited_account = AWS_ACCOUNT_NUMBER dlm_client.audited_account_arn = AWS_ACCOUNT_ARN dlm_client.lifecycle_policies = { @@ -203,30 +207,25 @@ class Test_dlm_ebs_snapshot_lifecycle_policy_exists: ) } } - - # from prowler.providers.aws.services.ec2.ec2_service import EC2 + dlm_client.regions_with_snapshots = {} aws_provider = set_mocked_aws_provider([AWS_REGION_US_EAST_1]) - - from prowler.providers.aws.services.ec2.ec2_service import EC2 + unload_dlm_check_modules() with ( + mock.patch( + "prowler.providers.aws.services.dlm.dlm_service.DLM", + new=mock.MagicMock(return_value=dlm_client), + ), mock.patch( "prowler.providers.common.provider.Provider.get_global_provider", return_value=aws_provider, ), - mock.patch( - "prowler.providers.aws.services.dlm.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_ebs_snapshot_lifecycle_policy_exists.ec2_client", - new=EC2(aws_provider), - ) as ec2_client, mock.patch( "prowler.providers.aws.services.dlm.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_client", new=dlm_client, ), ): - # Remove all snapshots - ec2_client.regions_with_snapshots = {} - from prowler.providers.aws.services.dlm.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_ebs_snapshot_lifecycle_policy_exists import ( dlm_ebs_snapshot_lifecycle_policy_exists, ) @@ -234,3 +233,45 @@ class Test_dlm_ebs_snapshot_lifecycle_policy_exists: check = dlm_ebs_snapshot_lifecycle_policy_exists() result = check.execute() assert len(result) == 0 + + @mock_aws + def test_check_does_not_import_ec2_service_client(self): + dlm_client = mock.MagicMock() + dlm_client.audited_account = AWS_ACCOUNT_NUMBER + dlm_client.audited_account_arn = AWS_ACCOUNT_ARN + dlm_client.audited_partition = "aws" + dlm_client.lifecycle_policies = {AWS_REGION_US_EAST_1: {}} + dlm_client.regions_with_snapshots = {AWS_REGION_US_EAST_1: True} + dlm_client._get_lifecycle_policy_arn_template = mock.MagicMock( + return_value=f"arn:aws:dlm:{AWS_REGION_US_EAST_1}:{AWS_ACCOUNT_NUMBER}:policy" + ) + + aws_provider = set_mocked_aws_provider([AWS_REGION_US_EAST_1]) + unload_dlm_check_modules() + sys.modules.pop(EC2_CLIENT_MODULE, None) + real_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + if name == EC2_CLIENT_MODULE: + raise AssertionError("DLM check must not import the EC2 service client") + return real_import(name, *args, **kwargs) + + with ( + mock.patch( + "prowler.providers.aws.services.dlm.dlm_service.DLM", + new=mock.MagicMock(return_value=dlm_client), + ), + mock.patch( + "prowler.providers.common.provider.Provider.get_global_provider", + return_value=aws_provider, + ), + mock.patch("builtins.__import__", side_effect=guarded_import), + ): + from prowler.providers.aws.services.dlm.dlm_ebs_snapshot_lifecycle_policy_exists.dlm_ebs_snapshot_lifecycle_policy_exists import ( + dlm_ebs_snapshot_lifecycle_policy_exists, + ) + + result = dlm_ebs_snapshot_lifecycle_policy_exists().execute() + + assert len(result) == 1 + assert result[0].status == "FAIL" diff --git a/tests/providers/aws/services/dlm/dlm_service_test.py b/tests/providers/aws/services/dlm/dlm_service_test.py index 38307c508c..b5b64322de 100644 --- a/tests/providers/aws/services/dlm/dlm_service_test.py +++ b/tests/providers/aws/services/dlm/dlm_service_test.py @@ -34,6 +34,8 @@ def mock_make_api_call(self, operation_name, kwargs): } ] } + if operation_name == "DescribeSnapshots": + return {"Snapshots": [{"SnapshotId": "snap-1234567890abcdef0"}]} return make_api_call(self, operation_name, kwargs) @@ -46,6 +48,13 @@ def mock_generate_regional_clients(provider, service): return {AWS_REGION_US_EAST_1: regional_client} +def mock_generate_regional_clients_without_ec2(provider, service): + if service == "ec2": + return None + + return mock_generate_regional_clients(provider, service) + + @patch( "prowler.providers.aws.aws_provider.AwsProvider.generate_regional_clients", new=mock_generate_regional_clients, @@ -92,3 +101,62 @@ class Test_DLM_Service: ) } } + + def test_get_regions_with_snapshots(self): + aws_provider = set_mocked_aws_provider() + dlm = DLM(aws_provider) + assert dlm.regions_with_snapshots == {AWS_REGION_US_EAST_1: True} + + +@patch( + "prowler.providers.aws.aws_provider.AwsProvider.generate_regional_clients", + new=mock_generate_regional_clients_without_ec2, +) +@patch("botocore.client.BaseClient._make_api_call", new=mock_make_api_call) +class Test_DLM_Service_Without_EC2_Regional_Clients: + def test_service_handles_missing_ec2_regional_clients(self): + aws_provider = set_mocked_aws_provider() + dlm = DLM(aws_provider) + + assert dlm.regions_with_snapshots == {} + + +class FakeEC2RegionalClient: + def __init__(self, responses): + self.region = AWS_REGION_US_EAST_1 + self.requests = [] + self.responses = list(responses) + + def describe_snapshots(self, **kwargs): + self.requests.append(kwargs) + return self.responses.pop(0) + + +class Test_DLM_Regions_With_Snapshots: + def test_get_regions_without_snapshots(self): + dlm = DLM.__new__(DLM) + dlm.regions_with_snapshots = {} + regional_client = FakeEC2RegionalClient([{"Snapshots": []}]) + + dlm._get_regions_with_snapshots(regional_client) + + assert dlm.regions_with_snapshots == {AWS_REGION_US_EAST_1: False} + assert regional_client.requests == [{"OwnerIds": ["self"], "MaxResults": 5}] + + def test_get_regions_with_snapshots_after_pagination(self): + dlm = DLM.__new__(DLM) + dlm.regions_with_snapshots = {} + regional_client = FakeEC2RegionalClient( + [ + {"Snapshots": [], "NextToken": "next-page"}, + {"Snapshots": [{"SnapshotId": "snap-1234567890abcdef0"}]}, + ] + ) + + dlm._get_regions_with_snapshots(regional_client) + + assert dlm.regions_with_snapshots == {AWS_REGION_US_EAST_1: True} + assert regional_client.requests == [ + {"OwnerIds": ["self"], "MaxResults": 5}, + {"OwnerIds": ["self"], "MaxResults": 5, "NextToken": "next-page"}, + ]