chore: Set up CDK Nag (#25)

* add CDK nag for fmops

* Reduce the scope of CDK nag suppressions for mlflow modules

* Reduce the scope of CDK nag suppressions for endpoint module

* fix formatting

* add CDK nag to sagemaker studio module

* fix static checkers

* fix mypy error

* add cdk nag to the templates module

* Add CDK nag to sagemaker-notebook

* add CDK nag unit test checks

* fix formatting for fmops module

* fix sagemaker templates tests

* add CDK nag unit test for mlflow-farhate

* add CDK nag unit test for mlflow-image

* add CDK nag unit test for sagemaker-endpoint

* update tests for sagemaker-studio

* fix Lambda version suppression for sagemaker-studio

modules/fmops/sagemaker-jumpstart-fm-endpoint/index.ts

@@ -1,5 +1,6 @@
 import "source-map-support/register";
 import * as cdk from "aws-cdk-lib";
+import * as cdk_nag from "cdk-nag";
 import { SagemakerJumpStartFmEndpointStack } from "./lib/sagemaker-jumpstart-fm-endpoint-stack";
 
 const account = process.env.CDK_DEFAULT_ACCOUNT;
@@ -35,4 +36,6 @@ new cdk.CfnOutput(stack, "metadata", {
   }),
 });
 
+cdk.Aspects.of(app).add(new cdk_nag.AwsSolutionsChecks({ logIgnores: true }));
+
 app.synth();

modules/fmops/sagemaker-jumpstart-fm-endpoint/lib/sagemaker-jumpstart-fm-endpoint-stack.ts

@@ -8,6 +8,7 @@ import {
   SageMakerInstanceType,
   JumpStartSageMakerEndpoint,
 } from "@cdklabs/generative-ai-cdk-constructs";
+import * as cdk_nag from "cdk-nag";
 
 interface SagemakerJumpStartFmEndpointStackProps extends cdk.StackProps {
   projectName?: string;
@@ -91,5 +92,12 @@ export class SagemakerJumpStartFmEndpointStack extends cdk.Stack {
       role: this.role,
       vpcConfig: vpcConfig,
     });
+
+    cdk_nag.NagSuppressions.addResourceSuppressions(this.role, [
+      {
+        id: "AwsSolutions-IAM5",
+        reason: "Resource access restriced to S3 buckets (with a prefix) and ECR images",
+      },
+    ]);
   }
 }

modules/fmops/sagemaker-jumpstart-fm-endpoint/package.json

@@ -18,6 +18,7 @@
     "@types/jest": "^29.5.5",
     "@types/node": "20.7.1",
     "aws-cdk": "2.130.0",
+    "cdk-nag": "^2.28.55",
     "cypress": "^13.6.1",
     "jest": "^29.7.0",
     "prettier": "^3.1.1",

modules/fmops/sagemaker-jumpstart-fm-endpoint/test/sagemaker-fm-endpoint.test.ts

@@ -1,8 +1,8 @@
 import * as cdk from "aws-cdk-lib";
-import { Template } from "aws-cdk-lib/assertions";
+import { Annotations, Match, Template } from "aws-cdk-lib/assertions";
 import { SagemakerJumpStartFmEndpointStack } from "../lib/sagemaker-jumpstart-fm-endpoint-stack";
 
-test("Synth stack", () => {
+describe("Sagemaker JumpStart Fm Endpoint Stack", () => {
   const app = new cdk.App();
 
   const projectName = "mlops";
@@ -26,9 +26,16 @@ test("Synth stack", () => {
     env: { account, region },
   });
 
-  const template = Template.fromStack(stack);
+  test("Synth stack", () => {
+    const template = Template.fromStack(stack);
 
-  template.hasResource("AWS::SageMaker::Endpoint", {});
-  template.hasResource("AWS::IAM::Role", {});
-  template.hasResource("AWS::EC2::SecurityGroup", {});
+    template.hasResource("AWS::SageMaker::Endpoint", {});
+    template.hasResource("AWS::IAM::Role", {});
+    template.hasResource("AWS::EC2::SecurityGroup", {});
+  });
+
+  test("No CDK Nag Errors", () => {
+    const errors = Annotations.fromStack(stack).findError("*", Match.stringLikeRegexp("AwsSolutions-.*"));
+    expect(errors).toHaveLength(0);
+  });
 });

modules/mlflow/mlflow-fargate/app.py

@@ -6,6 +6,7 @@
 from typing import Optional, cast
 
 import aws_cdk
+import cdk_nag
 
 from stack import MlflowFargateStack, RDSSettings
 
@@ -116,4 +117,6 @@ def _param(name: str) -> str:
     ),
 )
 
+aws_cdk.Aspects.of(app).add(cdk_nag.AwsSolutionsChecks(log_ignores=True))
+
 app.synth()

modules/mlflow/mlflow-fargate/stack.py

@@ -47,7 +47,7 @@ def __init__(
 
         cdk.Tags.of(scope=cast(IConstruct, self)).add(key="Deployment", value=app_prefix[:64])
 
-        role = iam.Role(
+        task_role = iam.Role(
             self,
             "TaskRole",
             assumed_by=iam.ServicePrincipal(service="ecs-tasks.amazonaws.com"),
@@ -58,7 +58,7 @@ def __init__(
 
         # Grant artifacts bucket read-write permissions
         model_bucket = s3.Bucket.from_bucket_name(self, "ArtifactsBucket", bucket_name=artifacts_bucket_name)
-        model_bucket.grant_read_write(role)
+        model_bucket.grant_read_write(task_role)
 
         vpc = ec2.Vpc.from_lookup(self, "Vpc", vpc_id=vpc_id)
         subnets = [ec2.Subnet.from_subnet_id(self, f"Subnet {subnet_id}", subnet_id) for subnet_id in subnet_ids]
@@ -76,7 +76,7 @@ def __init__(
         task_definition = ecs.FargateTaskDefinition(
             self,
             "MlflowTask",
-            task_role=role,
+            task_role=task_role,
             cpu=task_cpu_units,
             memory_limit_mib=task_memory_limit_mb,
         )
@@ -172,8 +172,8 @@ def __init__(
         fs.connections.allow_default_port_from(fargate_service.service.connections)
 
         # Setup security group
-        fargate_securigy_group = fargate_service.service.connections.security_groups[0]
-        fargate_securigy_group.add_ingress_rule(
+        fargate_security_group = fargate_service.service.connections.security_groups[0]
+        fargate_security_group.add_ingress_rule(
             peer=ec2.Peer.ipv4(vpc.vpc_cidr_block),
             connection=ec2.Port.tcp(5000),
             description="Allow inbound from VPC for mlflow",
@@ -185,7 +185,7 @@ def __init__(
                 self, "RDS Security Group", rds_settings["security_group_id"]
             )
             rds_security_group.add_ingress_rule(
-                peer=ec2.Peer.security_group_id(fargate_securigy_group.security_group_id),
+                peer=ec2.Peer.security_group_id(fargate_security_group.security_group_id),
                 connection=ec2.Port.tcp(rds_settings["port"]),
                 description="Allow inbound access to RDS for mlflow",
             )
@@ -199,12 +199,10 @@ def __init__(
             scale_out_cooldown=cdk.Duration.seconds(60),
         )
 
-        # Add CDK nag solutions checks
-        cdk.Aspects.of(self).add(cdk_nag.AwsSolutionsChecks(log_ignores=True))
-
-        cdk_nag.NagSuppressions.add_stack_suppressions(
-            self,
-            apply_to_nested_stacks=True,
+        # Add CDK nag suppressions
+        cdk_nag.NagSuppressions.add_resource_suppressions(
+            [task_role, task_definition.execution_role],  # type: ignore[list-item]
+            apply_to_children=True,
             suppressions=[
                 cdk_nag.NagPackSuppression(
                     id="AwsSolutions-IAM4",
@@ -214,10 +212,20 @@ def __init__(
                     id="AwsSolutions-IAM5",
                     reason="Resource access restricted to resources",
                 ),
+            ],
+        )
+        cdk_nag.NagSuppressions.add_resource_suppressions(
+            task_definition,
+            suppressions=[
                 cdk_nag.NagPackSuppression(
                     id="AwsSolutions-ECS2",
                     reason="Not passing secrets via env variables",
                 ),
+            ],
+        )
+        cdk_nag.NagSuppressions.add_resource_suppressions(
+            lb_access_logs_bucket,
+            suppressions=[
                 cdk_nag.NagPackSuppression(
                     id="AwsSolutions-S1",
                     reason="Access logs not required for access logs bucket",

modules/mlflow/mlflow-fargate/tests/test_stack.py

@@ -3,8 +3,9 @@
 from unittest import mock
 
 import aws_cdk as cdk
+import cdk_nag
 import pytest
-from aws_cdk.assertions import Template
+from aws_cdk.assertions import Annotations, Match, Template
 
 
 @pytest.fixture(scope="function")
@@ -20,8 +21,8 @@ def stack_defaults():
         yield
 
 
-@pytest.mark.parametrize("use_rds", [False, True])
-def test_synthesize_stack(stack_defaults, use_rds) -> None:
+@pytest.fixture(scope="function")
+def stack(stack_defaults, use_rds: bool) -> cdk.Stack:
     import stack
 
     app = cdk.App()
@@ -49,7 +50,7 @@ def test_synthesize_stack(stack_defaults, use_rds) -> None:
     else:
         rds_settings = None
 
-    stack = stack.MlflowFargateStack(
+    return stack.MlflowFargateStack(
         scope=app,
         id=app_prefix,
         app_prefix=app_prefix,
@@ -71,6 +72,20 @@ def test_synthesize_stack(stack_defaults, use_rds) -> None:
         ),
     )
 
+
+@pytest.mark.parametrize("use_rds", [False, True])
+def test_synthesize_stack(stack: cdk.Stack, use_rds: bool) -> None:
     template = Template.from_stack(stack)
     template.resource_count_is("AWS::ECS::Cluster", 1)
     template.resource_count_is("AWS::ECS::TaskDefinition", 1)
+
+
+@pytest.mark.parametrize("use_rds", [False, True])
+def test_no_cdk_nag_errors(stack: cdk.Stack, use_rds: bool) -> None:
+    cdk.Aspects.of(stack).add(cdk_nag.AwsSolutionsChecks())
+
+    nag_errors = Annotations.from_stack(stack).find_error(
+        "*",
+        Match.string_like_regexp(r"AwsSolutions-.*"),
+    )
+    assert not nag_errors, f"Found {len(nag_errors)} CDK nag errors"

modules/mlflow/mlflow-image/app.py

@@ -4,6 +4,7 @@
 import os
 
 import aws_cdk
+import cdk_nag
 
 from stack import MlflowImagePublishingStack
 
@@ -51,4 +52,6 @@ def _param(name: str) -> str:
     ),
 )
 
+aws_cdk.Aspects.of(app).add(cdk_nag.AwsSolutionsChecks(log_ignores=True))
+
 app.synth()

modules/mlflow/mlflow-image/stack.py

@@ -41,16 +41,16 @@ def __init__(
             dest=ecr_deployment.DockerImageName(self.image_uri),
         )
 
-        # Add CDK nag solutions checks
-        cdk.Aspects.of(self).add(cdk_nag.AwsSolutionsChecks(log_ignores=True))
-
+        # Add CDK nag suppressions
         cdk_nag.NagSuppressions.add_stack_suppressions(
             self,
-            apply_to_nested_stacks=True,
             suppressions=[
                 cdk_nag.NagPackSuppression(
                     id="AwsSolutions-IAM4",
                     reason="Managed Policies are for src account roles only",
+                    applies_to=[
+                        "Policy::arn:<AWS::Partition>:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole",
+                    ],
                 ),
                 cdk_nag.NagPackSuppression(
                     id="AwsSolutions-IAM5",

modules/mlflow/mlflow-image/tests/test_stack.py

@@ -2,8 +2,9 @@
 import sys
 
 import aws_cdk as cdk
+import cdk_nag
 import pytest
-from aws_cdk.assertions import Template
+from aws_cdk.assertions import Annotations, Match, Template
 
 
 @pytest.fixture(scope="function")
@@ -16,7 +17,8 @@ def stack_defaults() -> None:
         del sys.modules["stack"]
 
 
-def test_synthesize_stack() -> None:
+@pytest.fixture(scope="function")
+def stack(stack_defaults) -> cdk.Stack:
     import stack
 
     app = cdk.App()
@@ -28,7 +30,7 @@ def test_synthesize_stack() -> None:
 
     ecr_repo_name = "repo"
 
-    stack = stack.MlflowImagePublishingStack(
+    return stack.MlflowImagePublishingStack(
         scope=app,
         id=app_prefix,
         app_prefix=app_prefix,
@@ -39,5 +41,17 @@ def test_synthesize_stack() -> None:
         ),
     )
 
+
+def test_synthesize_stack(stack: cdk.Stack) -> None:
     template = Template.from_stack(stack)
     template.resource_count_is("Custom::CDKBucketDeployment", 1)
+
+
+def test_no_cdk_nag_errors(stack: cdk.Stack) -> None:
+    cdk.Aspects.of(stack).add(cdk_nag.AwsSolutionsChecks())
+
+    nag_errors = Annotations.from_stack(stack).find_error(
+        "*",
+        Match.string_like_regexp(r"AwsSolutions-.*"),
+    )
+    assert not nag_errors, f"Found {len(nag_errors)} CDK nag errors"

modules/sagemaker/sagemaker-endpoint/app.py

@@ -5,6 +5,7 @@
 import os
 
 import aws_cdk
+import cdk_nag
 
 from stack import DeployEndpointStack
 
@@ -84,7 +85,7 @@ def _param(name: str) -> str:
     id="metadata",
     value=stack.to_json_string(
         {
-            "ModelExecutionRoleArn": stack.model_execution_role_arn,
+            "ModelExecutionRoleArn": stack.model_execution_role.role_arn,
             "ModelName": stack.model.model_name,
             "ModelPackageArn": stack.model_package_arn,
             "EndpointName": stack.endpoint.attr_endpoint_name,
@@ -93,5 +94,6 @@ def _param(name: str) -> str:
     ),
 )
 
+aws_cdk.Aspects.of(app).add(cdk_nag.AwsSolutionsChecks(log_ignores=True))
 
 app.synth()