From 4cf0c7adfe6468d82b10e2138ab452336cca40b5 Mon Sep 17 00:00:00 2001 From: goutham-hari Date: Sat, 4 Jul 2026 07:24:27 +0530 Subject: [PATCH 1/2] feat(sagemaker): add sagemaker_endpoint_config_kms_encryption_enabled check This adds a check to ensure SageMaker Endpoint Configurations are encrypted with a KMS key. --- .../__init__.py | 1 + ...onfig_kms_encryption_enabled.metadata.json | 41 +++++++ ..._endpoint_config_kms_encryption_enabled.py | 20 ++++ .../services/sagemaker/sagemaker_service.py | 2 + ...oint_config_kms_encryption_enabled_test.py | 112 ++++++++++++++++++ 5 files changed, 176 insertions(+) create mode 100644 prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/__init__.py create mode 100644 prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.metadata.json create mode 100644 prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.py create mode 100644 tests/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled_test.py diff --git a/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/__init__.py b/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/__init__.py new file mode 100644 index 00000000000..912256a340a --- /dev/null +++ b/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/__init__.py @@ -0,0 +1 @@ +# empty init diff --git a/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.metadata.json b/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.metadata.json new file mode 100644 index 00000000000..c5d17990aa3 --- /dev/null +++ b/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.metadata.json @@ -0,0 +1,41 @@ +{ + "Provider": "aws", + "CheckID": "sagemaker_endpoint_config_kms_encryption_enabled", + "CheckTitle": "SageMaker endpoint configuration is encrypted with a KMS key", + "CheckType": [ + "Software and Configuration Checks/AWS Security Best Practices", + "Software and Configuration Checks/Industry and Regulatory Standards/AWS Foundational Security Best Practices", + "Effects/Data Exposure" + ], + "ServiceName": "sagemaker", + "SubServiceName": "", + "ResourceIdTemplate": "", + "Severity": "medium", + "ResourceType": "AwsSageMakerEndpointConfig", + "ResourceGroup": "ai_ml", + "Description": "**Amazon SageMaker endpoint configurations** are assessed for **at-rest encryption** using an AWS KMS key. The finding reflects whether a `KmsKeyId` is configured for the endpoint's ML volume encryption.", + "Risk": "Without **at-rest encryption** using a customer-managed KMS key, data on endpoint EBS volumes and snapshots can be exposed via storage access, copied backups, or host compromise, reducing **confidentiality** and limiting **key rotation** and **revocation** controls.", + "RelatedUrl": "", + "AdditionalURLs": [ + "https://docs.aws.amazon.com/sagemaker/latest/dg/key-management.html" + ], + "Remediation": { + "Code": { + "CLI": "", + "NativeIaC": "```yaml\n# CloudFormation: SageMaker endpoint config with KMS encryption\nResources:\n :\n Type: AWS::SageMaker::EndpointConfig\n Properties:\n KmsKeyId: # Critical: encrypts the endpoint config's EBS volume with the specified KMS key\n```", + "Other": "1. Open the AWS Console > Amazon SageMaker > Endpoint configurations\n2. Note the details of the failing endpoint configuration\n3. Click Create endpoint configuration\n4. Enter name and details identical to the previous one\n5. In Encryption key, select your KMS key\n6. Create the endpoint configuration\n7. Update any endpoints using the old config to use the new config\n8. Delete the old unencrypted endpoint configuration", + "Terraform": "```hcl\n# SageMaker endpoint config with KMS encryption\nresource \"aws_sagemaker_endpoint_configuration\" \"\" {\n name = \"\"\n kms_key_arn = \"\" # Critical: enables EBS encryption using this KMS key\n}\n```" + }, + "Recommendation": { + "Text": "Use a **customer-managed KMS key** for endpoint ML volumes by setting `KmsKeyId`, and apply KMS to related S3 inputs/outputs. Enforce **least privilege** on key usage, enable **rotation**, and align key access with **defense in depth** to protect data at rest.", + "Url": "https://hub.prowler.com/check/sagemaker_endpoint_config_kms_encryption_enabled" + } + }, + "Categories": [ + "encryption", + "gen-ai" + ], + "DependsOn": [], + "RelatedTo": [], + "Notes": "" +} diff --git a/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.py b/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.py new file mode 100644 index 00000000000..ce4276ab748 --- /dev/null +++ b/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.py @@ -0,0 +1,20 @@ +from prowler.lib.check.models import Check, Check_Report_AWS +from prowler.providers.aws.services.sagemaker.sagemaker_client import sagemaker_client + + +class sagemaker_endpoint_config_kms_encryption_enabled(Check): + def execute(self): + findings = [] + for endpoint_config in sagemaker_client.endpoint_configs.values(): + report = Check_Report_AWS( + metadata=self.metadata(), resource=endpoint_config + ) + report.status = "PASS" + report.status_extended = f"Sagemaker Endpoint Config {endpoint_config.name} has KMS encryption enabled." + if not endpoint_config.kms_key_id: + report.status = "FAIL" + report.status_extended = f"Sagemaker Endpoint Config {endpoint_config.name} does not have KMS encryption enabled." + + findings.append(report) + + return findings diff --git a/prowler/providers/aws/services/sagemaker/sagemaker_service.py b/prowler/providers/aws/services/sagemaker/sagemaker_service.py index 20ea4c0280c..2f32a56e804 100644 --- a/prowler/providers/aws/services/sagemaker/sagemaker_service.py +++ b/prowler/providers/aws/services/sagemaker/sagemaker_service.py @@ -443,6 +443,7 @@ def _describe_endpoint_config(self, endpoint_config): ) ) endpoint_config.production_variants = production_variants + endpoint_config.kms_key_id = describe_endpoint_config.get("KmsKeyId") except Exception as error: logger.error( f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" @@ -560,6 +561,7 @@ class EndpointConfig(BaseModel): region: str arn: str production_variants: list[ProductionVariant] = [] + kms_key_id: Optional[str] = None tags: Optional[list] = [] diff --git a/tests/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled_test.py b/tests/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled_test.py new file mode 100644 index 00000000000..0571a4b8d40 --- /dev/null +++ b/tests/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled_test.py @@ -0,0 +1,112 @@ +from unittest import mock +from uuid import uuid4 + +from prowler.providers.aws.services.sagemaker.sagemaker_service import EndpointConfig +from tests.providers.aws.utils import ( + AWS_ACCOUNT_NUMBER, + AWS_REGION_EU_WEST_1, + set_mocked_aws_provider, +) + +test_endpoint_config = "test-endpoint-config" +endpoint_config_arn = f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:endpoint-config/{test_endpoint_config}" +kms_key = str(uuid4()) + + +class Test_sagemaker_endpoint_config_kms_encryption_enabled: + def test_no_endpoint_configs(self): + sagemaker_client = mock.MagicMock + sagemaker_client.endpoint_configs = {} + + aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1]) + + with ( + mock.patch( + "prowler.providers.common.provider.Provider.get_global_provider", + return_value=aws_provider, + ), + mock.patch( + "prowler.providers.aws.services.sagemaker.sagemaker_endpoint_config_kms_encryption_enabled.sagemaker_endpoint_config_kms_encryption_enabled.sagemaker_client", + sagemaker_client, + ), + ): + from prowler.providers.aws.services.sagemaker.sagemaker_endpoint_config_kms_encryption_enabled.sagemaker_endpoint_config_kms_encryption_enabled import ( + sagemaker_endpoint_config_kms_encryption_enabled, + ) + + check = sagemaker_endpoint_config_kms_encryption_enabled() + result = check.execute() + assert len(result) == 0 + + def test_endpoint_config_with_kms_key(self): + sagemaker_client = mock.MagicMock + sagemaker_client.endpoint_configs = {} + sagemaker_client.endpoint_configs[endpoint_config_arn] = EndpointConfig( + name=test_endpoint_config, + arn=endpoint_config_arn, + region=AWS_REGION_EU_WEST_1, + kms_key_id=kms_key, + ) + + aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1]) + + with ( + mock.patch( + "prowler.providers.common.provider.Provider.get_global_provider", + return_value=aws_provider, + ), + mock.patch( + "prowler.providers.aws.services.sagemaker.sagemaker_endpoint_config_kms_encryption_enabled.sagemaker_endpoint_config_kms_encryption_enabled.sagemaker_client", + sagemaker_client, + ), + ): + from prowler.providers.aws.services.sagemaker.sagemaker_endpoint_config_kms_encryption_enabled.sagemaker_endpoint_config_kms_encryption_enabled import ( + sagemaker_endpoint_config_kms_encryption_enabled, + ) + + check = sagemaker_endpoint_config_kms_encryption_enabled() + result = check.execute() + assert len(result) == 1 + assert result[0].status == "PASS" + assert ( + result[0].status_extended + == f"Sagemaker Endpoint Config {test_endpoint_config} has KMS encryption enabled." + ) + assert result[0].resource_id == test_endpoint_config + assert result[0].resource_arn == endpoint_config_arn + + def test_endpoint_config_no_kms_key(self): + sagemaker_client = mock.MagicMock + sagemaker_client.endpoint_configs = {} + sagemaker_client.endpoint_configs[endpoint_config_arn] = EndpointConfig( + name=test_endpoint_config, + arn=endpoint_config_arn, + region=AWS_REGION_EU_WEST_1, + ) + + aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1]) + + with ( + mock.patch( + "prowler.providers.common.provider.Provider.get_global_provider", + return_value=aws_provider, + ), + mock.patch( + "prowler.providers.aws.services.sagemaker.sagemaker_endpoint_config_kms_encryption_enabled.sagemaker_endpoint_config_kms_encryption_enabled.sagemaker_client", + sagemaker_client, + ), + ): + from prowler.providers.aws.services.sagemaker.sagemaker_endpoint_config_kms_encryption_enabled.sagemaker_endpoint_config_kms_encryption_enabled import ( + sagemaker_endpoint_config_kms_encryption_enabled, + ) + + check = sagemaker_endpoint_config_kms_encryption_enabled() + result = check.execute() + assert len(result) == 1 + assert result[0].status == "FAIL" + assert ( + result[0].status_extended + == f"Sagemaker Endpoint Config {test_endpoint_config} does not have KMS encryption enabled." + ) + assert result[0].resource_id == test_endpoint_config + assert result[0].resource_arn == endpoint_config_arn From 4a2fe78560e8c5f5cc149ea985da25acf302f911 Mon Sep 17 00:00:00 2001 From: goutham-hari Date: Sun, 5 Jul 2026 07:51:23 +0530 Subject: [PATCH 2/2] fix(sagemaker): add CLI remediation to endpoint config KMS metadata --- ...gemaker_endpoint_config_kms_encryption_enabled.metadata.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.metadata.json b/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.metadata.json index c5d17990aa3..fc205deb844 100644 --- a/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.metadata.json +++ b/prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.metadata.json @@ -21,7 +21,7 @@ ], "Remediation": { "Code": { - "CLI": "", + "CLI": "aws sagemaker create-endpoint-config --endpoint-config-name --kms-key-id --production-variants file://production-variants.json", "NativeIaC": "```yaml\n# CloudFormation: SageMaker endpoint config with KMS encryption\nResources:\n :\n Type: AWS::SageMaker::EndpointConfig\n Properties:\n KmsKeyId: # Critical: encrypts the endpoint config's EBS volume with the specified KMS key\n```", "Other": "1. Open the AWS Console > Amazon SageMaker > Endpoint configurations\n2. Note the details of the failing endpoint configuration\n3. Click Create endpoint configuration\n4. Enter name and details identical to the previous one\n5. In Encryption key, select your KMS key\n6. Create the endpoint configuration\n7. Update any endpoints using the old config to use the new config\n8. Delete the old unencrypted endpoint configuration", "Terraform": "```hcl\n# SageMaker endpoint config with KMS encryption\nresource \"aws_sagemaker_endpoint_configuration\" \"\" {\n name = \"\"\n kms_key_arn = \"\" # Critical: enables EBS encryption using this KMS key\n}\n```"