feat(sagemaker): add sagemaker_endpoint_config_kms_encryption_enabled check (#12118)

Co-authored-by: Daniel Barranquero <danielbo2001@gmail.com>
Co-authored-by: Alex Chen <l46983284@gmail.com>

prowler/changelog.d/sagemaker-endpoint-config-kms-encryption-enabled.added.md

@@ -0,0 +1 @@
+`sagemaker_endpoint_config_kms_encryption_enabled` check verifying SageMaker endpoint configurations use a KMS key for storage volume encryption

prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.metadata.json

@@ -0,0 +1,46 @@
+{
+  "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": "Other",
+  "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 on the endpoint configuration so inference data volumes and related storage use KMS encryption.",
+  "Risk": "Without **at-rest encryption** using a KMS key on endpoint configurations, model artifacts and inference-related data may be exposed through storage access or compromised hosts, reducing **confidentiality** and limiting **key rotation** and **revocation** controls.",
+  "RelatedUrl": "",
+  "AdditionalURLs": [
+    "https://docs.aws.amazon.com/sagemaker/latest/dg/key-management.html",
+    "https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateEndpointConfig.html",
+    "https://docs.aws.amazon.com/securityhub/latest/userguide/sagemaker-controls.html"
+  ],
+  "Remediation": {
+    "Code": {
+      "CLI": "aws sagemaker create-endpoint-config --endpoint-config-name <endpoint-config-name> --production-variants VariantName=AllTraffic,ModelName=<model-name>,InitialInstanceCount=1,InstanceType=ml.m5.large --kms-key-id <kms-key-id>",
+      "NativeIaC": "```yaml\n# CloudFormation: SageMaker EndpointConfig with KMS encryption\nResources:\n  <example_resource_name>:\n    Type: AWS::SageMaker::EndpointConfig\n    Properties:\n      ProductionVariants:\n        - VariantName: AllTraffic\n          ModelName: <example_resource_name>\n          InitialInstanceCount: 1\n          InstanceType: ml.m5.large\n      KmsKeyId: <example_resource_id>  # Critical: encrypts endpoint data at rest with KMS\n```",
+      "Other": "1. Open Amazon SageMaker > Inference > Endpoint configurations\n2. Create a new endpoint configuration (endpoint configs are immutable)\n3. Configure production variants as required\n4. Under Encryption, select a KMS key\n5. Create the configuration and update any endpoints to use the new encrypted configuration\n6. Delete the old unencrypted endpoint configuration when safe",
+      "Terraform": "```hcl\n# SageMaker endpoint configuration with KMS encryption\nresource \"aws_sagemaker_endpoint_configuration\" \"<example_resource_name>\" {\n  name        = \"<example_resource_name>\"\n  kms_key_arn = \"<example_resource_arn>\" # Critical: enables at-rest encryption with KMS\n\n  production_variants {\n    variant_name           = \"AllTraffic\"\n    model_name             = \"<example_resource_name>\"\n    instance_type          = \"ml.m5.large\"\n    initial_instance_count = 1\n  }\n}\n```"
+    },
+    "Recommendation": {
+      "Text": "Always set `KmsKeyId` on SageMaker endpoint configurations. Prefer a **customer-managed KMS key** with least-privilege key policies, enable **rotation**, and ensure endpoints are updated to use the encrypted configuration.",
+      "Url": "https://hub.prowler.com/check/sagemaker_endpoint_config_kms_encryption_enabled"
+    }
+  },
+  "Categories": [
+    "encryption",
+    "gen-ai"
+  ],
+  "DependsOn": [],
+  "RelatedTo": [
+    "sagemaker_notebook_instance_encryption_enabled",
+    "sagemaker_training_jobs_volume_and_output_encryption_enabled"
+  ],
+  "Notes": ""
+}

prowler/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled.py

@@ -0,0 +1,31 @@
+from typing import List
+
+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):
+    """Ensure SageMaker endpoint configurations encrypt data at rest with a KMS key."""
+
+    def execute(self) -> List[Check_Report_AWS]:
+        """Return PASS/FAIL findings for each SageMaker endpoint configuration."""
+        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 data encryption "
+                f"enabled with KMS key."
+            )
+            if not endpoint_config.kms_key_id:
+                report.status = "FAIL"
+                report.status_extended = (
+                    f"Sagemaker Endpoint Config {endpoint_config.name} does not have "
+                    f"data encryption enabled with a KMS key."
+                )
+
+            findings.append(report)
+
+        return findings

prowler/providers/aws/services/sagemaker/sagemaker_service.py

@@ -499,6 +499,8 @@ class SageMaker(AWSService):
                     )
                 )
             endpoint_config.production_variants = production_variants
+            if "KmsKeyId" in describe_endpoint_config:
+                endpoint_config.kms_key_id = describe_endpoint_config["KmsKeyId"]
         except Exception as error:
             logger.error(
                 f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
@@ -623,6 +625,7 @@ class EndpointConfig(BaseModel):
     region: str
     arn: str
     production_variants: list[ProductionVariant] = []
+    kms_key_id: Optional[str] = None
     tags: Optional[list] = []
 
 

tests/providers/aws/services/sagemaker/sagemaker_endpoint_config_kms_encryption_enabled/sagemaker_endpoint_config_kms_encryption_enabled_test.py

@@ -0,0 +1,138 @@
+from unittest import mock
+
+from boto3 import client
+from moto import mock_aws
+
+from tests.providers.aws.utils import (
+    AWS_REGION_EU_WEST_1,
+    AWS_REGION_US_EAST_1,
+    set_mocked_aws_provider,
+)
+
+
+class Test_sagemaker_endpoint_config_kms_encryption_enabled:
+    @mock_aws
+    def test_no_endpoint_configs(self):
+        from prowler.providers.aws.services.sagemaker.sagemaker_service import SageMaker
+
+        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.sagemaker.sagemaker_endpoint_config_kms_encryption_enabled.sagemaker_endpoint_config_kms_encryption_enabled.sagemaker_client",
+                new=SageMaker(aws_provider),
+            ),
+        ):
+            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
+
+    @mock_aws
+    def test_endpoint_config_without_kms(self):
+        sagemaker_client = client("sagemaker", region_name=AWS_REGION_EU_WEST_1)
+        endpoint_config_name = "endpoint-config-no-kms"
+        model_name = "model-v1"
+        sagemaker_client.create_model(ModelName=model_name)
+        endpoint_config = sagemaker_client.create_endpoint_config(
+            EndpointConfigName=endpoint_config_name,
+            ProductionVariants=[
+                {
+                    "VariantName": "AllTraffic",
+                    "ModelName": model_name,
+                    "InitialInstanceCount": 1,
+                    "InstanceType": "ml.m5.large",
+                    "InitialVariantWeight": 1.0,
+                }
+            ],
+        )
+
+        from prowler.providers.aws.services.sagemaker.sagemaker_service import SageMaker
+
+        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",
+                new=SageMaker(aws_provider),
+            ),
+        ):
+            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 {endpoint_config_name} does not have data encryption enabled with a KMS key."
+            )
+            assert result[0].resource_id == endpoint_config_name
+            assert result[0].resource_arn == endpoint_config["EndpointConfigArn"]
+
+    @mock_aws
+    def test_endpoint_config_with_kms(self):
+        kms_client = client("kms", region_name=AWS_REGION_EU_WEST_1)
+        key = kms_client.create_key()["KeyMetadata"]["KeyId"]
+
+        sagemaker_client = client("sagemaker", region_name=AWS_REGION_EU_WEST_1)
+        endpoint_config_name = "endpoint-config-with-kms"
+        model_name = "model-v1"
+        sagemaker_client.create_model(ModelName=model_name)
+        endpoint_config = sagemaker_client.create_endpoint_config(
+            EndpointConfigName=endpoint_config_name,
+            KmsKeyId=key,
+            ProductionVariants=[
+                {
+                    "VariantName": "AllTraffic",
+                    "ModelName": model_name,
+                    "InitialInstanceCount": 1,
+                    "InstanceType": "ml.m5.large",
+                    "InitialVariantWeight": 1.0,
+                }
+            ],
+        )
+
+        from prowler.providers.aws.services.sagemaker.sagemaker_service import SageMaker
+
+        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",
+                new=SageMaker(aws_provider),
+            ),
+        ):
+            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 {endpoint_config_name} has data encryption enabled with KMS key."
+            )
+            assert result[0].resource_id == endpoint_config_name
+            assert result[0].resource_arn == endpoint_config["EndpointConfigArn"]