Add missed AWS Sagemaker permission changes (#124)

Author: sureshcsdp

cleancloud/doctor/aws.py

@@ -645,13 +645,26 @@ def run_aws_ai_doctor(profile: Optional[str], region: Optional[str] = None) -> N
         # present since ListEndpoints already confirmed SageMaker access.
         endpoints = sagemaker.list_endpoints(MaxResults=1, StatusEquals="InService")
         endpoint_list = endpoints.get("Endpoints", [])
+        endpoint_config_name = None
         if endpoint_list:
-            sagemaker.describe_endpoint(EndpointName=endpoint_list[0]["EndpointName"])
+            ep = sagemaker.describe_endpoint(EndpointName=endpoint_list[0]["EndpointName"])
+            endpoint_config_name = ep.get("EndpointConfigName")
         permissions_tested.append("sagemaker:DescribeEndpoint")
         success("sagemaker:DescribeEndpoint")
     except Exception as e:
         permissions_failed.append(("sagemaker:DescribeEndpoint", str(e)))
         warn(f"sagemaker:DescribeEndpoint - {e}")
+        endpoint_config_name = None
+
+    try:
+        # DescribeEndpointConfig — needed to resolve instance type (not in DescribeEndpoint response)
+        if endpoint_config_name:
+            sagemaker.describe_endpoint_config(EndpointConfigName=endpoint_config_name)
+        permissions_tested.append("sagemaker:DescribeEndpointConfig")
+        success("sagemaker:DescribeEndpointConfig")
+    except Exception as e:
+        permissions_failed.append(("sagemaker:DescribeEndpointConfig", str(e)))
+        warn(f"sagemaker:DescribeEndpointConfig - {e}")
 
     try:
         cloudwatch = session.client("cloudwatch", region_name=region)

cleancloud/providers/aws/rules/sagemaker_endpoint_idle.py

@@ -199,7 +199,7 @@ def find_idle_sagemaker_endpoints(
             raise PermissionError(
                 "Missing required IAM permissions: "
                 "sagemaker:ListEndpoints, sagemaker:DescribeEndpoint, "
-                "cloudwatch:GetMetricStatistics"
+                "sagemaker:DescribeEndpointConfig, cloudwatch:GetMetricStatistics"
             ) from e
         raise
 
@@ -246,25 +246,42 @@ def _describe_endpoint(
 ) -> Tuple[float, bool, int, int, Optional[str]]:
     """Return (monthly_cost, is_gpu, variant_count, total_instances, primary_instance_type).
 
+    Instance type lives in the endpoint *config* (describe_endpoint_config), not in
+    the endpoint summary (describe_endpoint ProductionVariantSummary has no InstanceType
+    field — only DesiredInstanceCount). We make two calls: one for instance counts,
+    one for instance types, then pair them by VariantName.
+
     Cost is computed per-variant using DesiredInstanceCount × per-instance cost, summed
     across all variants. GPU flag is True if any variant uses an accelerator instance.
 
-    Returns (default_cost, False, 1, 1, None) on failure to ensure the endpoint is
-    still flagged conservatively rather than silently dropped.
+    Returns (0, False, 0, 0, None) on failure so the endpoint is skipped rather than
+    flagged with assumed values.
     """
     try:
-        response = sagemaker.describe_endpoint(EndpointName=endpoint_name)
-        variants = response.get("ProductionVariants", [])
+        endpoint = sagemaker.describe_endpoint(EndpointName=endpoint_name)
+        variants = endpoint.get("ProductionVariants", [])
         if not variants:
             return _DEFAULT_MONTHLY_COST, False, 0, 0, None
 
+        # Fetch instance types from the endpoint config
+        config_name = endpoint.get("EndpointConfigName", "")
+        instance_type_by_variant: dict = {}
+        try:
+            config = sagemaker.describe_endpoint_config(EndpointConfigName=config_name)
+            for cv in config.get("ProductionVariants", []):
+                itype = cv.get("InstanceType")
+                if itype:
+                    instance_type_by_variant[cv["VariantName"]] = itype
+        except ClientError:
+            pass  # config inaccessible — costs/GPU will use defaults
+
         total_monthly_cost = 0.0
         is_gpu = False
         total_instances = 0
         primary_instance_type: Optional[str] = None
 
         for i, v in enumerate(variants):
-            itype = v.get("CurrentInstanceType")
+            itype = instance_type_by_variant.get(v.get("VariantName", ""))
             count = v.get("DesiredInstanceCount") or 0
             total_instances += count
 

deploy/cloudformation/cleancloud-role.yaml

@@ -23,8 +23,18 @@ Parameters:
       Optional. If set, the hub account must provide this value when assuming the role
       (confused deputy protection). Pass the same value via --external-id in CleanCloud.
 
+  EnableAIScan:
+    Type: String
+    Default: "false"
+    AllowedValues: ["true", "false"]
+    Description: >
+      Set to true to attach the AI/ML policy (SageMaker idle endpoint detection).
+      Required for: cleancloud scan --category ai
+      See: security/aws/ai-readonly.json
+
 Conditions:
   UseExternalId: !Not [!Equals [!Ref ExternalId, ""]]
+  EnableAIScan: !Equals [!Ref EnableAIScan, "true"]
 
 Resources:
   CleanCloudRole:
@@ -97,6 +107,24 @@ Resources:
         - Key: Purpose
           Value: CrossAccountReadOnlyScanning
 
+  CleanCloudAIPolicy:
+    Type: AWS::IAM::Policy
+    Condition: EnableAIScan
+    Properties:
+      PolicyName: CleanCloudAIReadOnly
+      Roles:
+        - !Ref CleanCloudRole
+      PolicyDocument:
+        Version: "2012-10-17"
+        Statement:
+          - Sid: SageMakerReadOnly
+            Effect: Allow
+            Action:
+              - sagemaker:ListEndpoints
+              - sagemaker:DescribeEndpoint
+              - sagemaker:DescribeEndpointConfig
+            Resource: "*"
+
 Outputs:
   RoleArn:
     Description: ARN of the CleanCloud IAM role.

deploy/terraform/aws/main.tf

@@ -33,6 +33,29 @@ resource "aws_iam_role" "cleancloud" {
   })
 }
 
+resource "aws_iam_role_policy" "cleancloud_ai" {
+  count = var.enable_ai ? 1 : 0
+
+  name = "CleanCloudAIReadOnly"
+  role = aws_iam_role.cleancloud.id
+
+  policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Sid    = "SageMakerReadOnly"
+        Effect = "Allow"
+        Action = [
+          "sagemaker:ListEndpoints",
+          "sagemaker:DescribeEndpoint",
+          "sagemaker:DescribeEndpointConfig",
+        ]
+        Resource = "*"
+      },
+    ]
+  })
+}
+
 resource "aws_iam_role_policy" "cleancloud" {
   name = "CleanCloudReadOnly"
   role = aws_iam_role.cleancloud.id

deploy/terraform/aws/variables.tf

@@ -20,6 +20,12 @@ variable "external_id" {
   description = "Optional. If set, adds an ExternalId condition to the trust policy (confused deputy protection). Pass the same value via --external-id in CleanCloud."
 }
 
+variable "enable_ai" {
+  type        = bool
+  default     = false
+  description = "Attach the AI/ML policy (SageMaker idle endpoint detection). Required for: cleancloud scan --category ai. See: security/aws/ai-readonly.json"
+}
+
 variable "tags" {
   type        = map(string)
   default     = {}

docs/rules.md

@@ -731,6 +731,7 @@ Confidence thresholds and signal weighting are documented in [confidence.md](con
 **Required permissions:**
 - `sagemaker:ListEndpoints`
 - `sagemaker:DescribeEndpoint`
+- `sagemaker:DescribeEndpointConfig`
 - `cloudwatch:GetMetricStatistics`
 
 > **Not run by default.** AI/ML rules are opt-in to avoid surprising users who don't use these services. Run with `cleancloud scan --provider aws --category ai` (or `--category all` to combine with hygiene rules). If the permissions above are not granted, the rule is gracefully skipped and reported in the skipped rules section — it will not fail the scan. Attach [`security/aws/ai-readonly.json`](../security/aws/ai-readonly.json) to your IAM role to enable this rule.

security/aws/ai-readonly.json

@@ -6,7 +6,8 @@
       "Effect": "Allow",
       "Action": [
         "sagemaker:ListEndpoints",
-        "sagemaker:DescribeEndpoint"
+        "sagemaker:DescribeEndpoint",
+        "sagemaker:DescribeEndpointConfig"
       ],
       "Resource": "*"
     }

tests/cleancloud/providers/aws/test_aws_sagemaker_endpoint_idle.py

@@ -34,16 +34,27 @@ def _make_endpoint(name="test-endpoint", age_days=30):
 def _make_describe_response(
     instance_type="ml.m5.xlarge", variant_count=1, desired_instance_count=1
 ):
-    """Build a describe_endpoint response with DesiredInstanceCount on each variant."""
+    """Build a describe_endpoint response.
+
+    ProductionVariantSummary does NOT include InstanceType — that lives in the
+    endpoint config. We include EndpointConfigName so the rule can fetch it.
+    """
     variants = [
         {
             "VariantName": f"variant-{i}",
-            "CurrentInstanceType": instance_type,
             "CurrentInstanceCount": desired_instance_count,
             "DesiredInstanceCount": desired_instance_count,
         }
         for i in range(variant_count)
     ]
+    return {"ProductionVariants": variants, "EndpointConfigName": "test-config"}
+
+
+def _make_describe_config_response(instance_type="ml.m5.xlarge", variant_count=1):
+    """Build a describe_endpoint_config response with InstanceType per variant."""
+    variants = [
+        {"VariantName": f"variant-{i}", "InstanceType": instance_type} for i in range(variant_count)
+    ]
     return {"ProductionVariants": variants}
 
 
@@ -73,6 +84,7 @@ def test_idle_cpu_endpoint_detected():
     paginator = sagemaker.get_paginator.return_value
     paginator.paginate.return_value = [{"Endpoints": [_make_endpoint(age_days=30)]}]
     sagemaker.describe_endpoint.return_value = _make_describe_response("ml.m5.xlarge")
+    sagemaker.describe_endpoint_config.return_value = _make_describe_config_response("ml.m5.xlarge")
     cloudwatch.get_metric_statistics.return_value = _no_invocations()
 
     session = _make_session(sagemaker, cloudwatch)
@@ -99,6 +111,9 @@ def test_idle_gpu_endpoint_detected_high_risk():
     paginator = sagemaker.get_paginator.return_value
     paginator.paginate.return_value = [{"Endpoints": [_make_endpoint(age_days=30)]}]
     sagemaker.describe_endpoint.return_value = _make_describe_response("ml.p3.2xlarge")
+    sagemaker.describe_endpoint_config.return_value = _make_describe_config_response(
+        "ml.p3.2xlarge"
+    )
     cloudwatch.get_metric_statistics.return_value = _no_invocations()
 
     session = _make_session(sagemaker, cloudwatch)
@@ -156,6 +171,7 @@ def test_timezone_naive_creation_time_handled():
 
     paginator.paginate.return_value = [{"Endpoints": [endpoint]}]
     sagemaker.describe_endpoint.return_value = _make_describe_response("ml.m5.xlarge")
+    sagemaker.describe_endpoint_config.return_value = _make_describe_config_response("ml.m5.xlarge")
     cloudwatch.get_metric_statistics.return_value = _no_invocations()
 
     session = _make_session(sagemaker, cloudwatch)
@@ -212,7 +228,8 @@ def test_missing_desired_instance_count_treated_as_zero():
     paginator.paginate.return_value = [{"Endpoints": [_make_endpoint(age_days=30)]}]
     # No DesiredInstanceCount key — AWS response omits it
     sagemaker.describe_endpoint.return_value = {
-        "ProductionVariants": [{"VariantName": "v1", "CurrentInstanceType": "ml.m5.xlarge"}]
+        "ProductionVariants": [{"VariantName": "v1"}],
+        "EndpointConfigName": "test-config",
     }
     cloudwatch.get_metric_statistics.return_value = _no_invocations()
 
@@ -233,9 +250,10 @@ def test_partial_scaled_to_zero_still_flagged():
     # 2 variants: one with 1 instance, one with 0
     sagemaker.describe_endpoint.return_value = {
         "ProductionVariants": [
-            {"VariantName": "v1", "CurrentInstanceType": "ml.m5.xlarge", "DesiredInstanceCount": 1},
-            {"VariantName": "v2", "CurrentInstanceType": "ml.m5.xlarge", "DesiredInstanceCount": 0},
-        ]
+            {"VariantName": "v1", "DesiredInstanceCount": 1},
+            {"VariantName": "v2", "DesiredInstanceCount": 0},
+        ],
+        "EndpointConfigName": "test-config",
     }
     cloudwatch.get_metric_statistics.return_value = _no_invocations()
 
@@ -386,6 +404,9 @@ def test_g4dn_instance_detected_as_gpu():
     paginator = sagemaker.get_paginator.return_value
     paginator.paginate.return_value = [{"Endpoints": [_make_endpoint(age_days=30)]}]
     sagemaker.describe_endpoint.return_value = _make_describe_response("ml.g4dn.xlarge")
+    sagemaker.describe_endpoint_config.return_value = _make_describe_config_response(
+        "ml.g4dn.xlarge"
+    )
     cloudwatch.get_metric_statistics.return_value = _no_invocations()
 
     session = _make_session(sagemaker, cloudwatch)
@@ -404,6 +425,9 @@ def test_inf1_instance_detected_as_gpu():
     paginator = sagemaker.get_paginator.return_value
     paginator.paginate.return_value = [{"Endpoints": [_make_endpoint(age_days=30)]}]
     sagemaker.describe_endpoint.return_value = _make_describe_response("ml.inf1.xlarge")
+    sagemaker.describe_endpoint_config.return_value = _make_describe_config_response(
+        "ml.inf1.xlarge"
+    )
     cloudwatch.get_metric_statistics.return_value = _no_invocations()
 
     session = _make_session(sagemaker, cloudwatch)
@@ -427,6 +451,9 @@ def test_multi_variant_cost_scaled():
     sagemaker.describe_endpoint.return_value = _make_describe_response(
         "ml.m5.xlarge", variant_count=3
     )
+    sagemaker.describe_endpoint_config.return_value = _make_describe_config_response(
+        "ml.m5.xlarge", variant_count=3
+    )
     cloudwatch.get_metric_statistics.return_value = _no_invocations()
 
     session = _make_session(sagemaker, cloudwatch)
@@ -448,16 +475,15 @@ def test_multi_variant_mixed_instance_types_cost():
     paginator.paginate.return_value = [{"Endpoints": [_make_endpoint(age_days=30)]}]
     sagemaker.describe_endpoint.return_value = {
         "ProductionVariants": [
-            {
-                "VariantName": "cpu",
-                "CurrentInstanceType": "ml.m5.xlarge",
-                "DesiredInstanceCount": 2,
-            },
-            {
-                "VariantName": "gpu",
-                "CurrentInstanceType": "ml.g4dn.xlarge",
-                "DesiredInstanceCount": 1,
-            },
+            {"VariantName": "cpu", "DesiredInstanceCount": 2},
+            {"VariantName": "gpu", "DesiredInstanceCount": 1},
+        ],
+        "EndpointConfigName": "test-config",
+    }
+    sagemaker.describe_endpoint_config.return_value = {
+        "ProductionVariants": [
+            {"VariantName": "cpu", "InstanceType": "ml.m5.xlarge"},
+            {"VariantName": "gpu", "InstanceType": "ml.g4dn.xlarge"},
         ]
     }
     cloudwatch.get_metric_statistics.return_value = _no_invocations()