fix(aws): avoid full ec2 inventory in dlm check (#11850)

This commit is contained in:
Hugo Pereira Brito
2026-07-08 15:59:16 +01:00
committed by GitHub
parent 52875b5c7c
commit 20aad80a78
5 changed files with 185 additions and 42 deletions
@@ -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
@@ -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(),
@@ -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
@@ -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"
@@ -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"},
]