ds_admin_detective.yaml

# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0
Description: |
  Deploy a Lambda function and CloudWatch trigger to inspect SageMaker resources and ensure they are attached to a VPC.

Parameters:
  StackSetName:
    Type: String 
    Description: A name to be used across nested stacks 

Resources:
  SageMakerDetectiveControlExecutionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
      RoleName:
        !Join
          - ''
          -
            - !Sub '${StackSetName}-SageMaker-DetectiveControl-'
            - !Select
              - 4
              - !Split
                - '-'
                - !Select
                  - 2
                  - !Split
                    - /
                    - !Ref 'AWS::StackId'
      Policies:
        - PolicyName: LambdaInlineForSageMaker
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Sid: VisualEditor0
                Effect: Allow
                Action:
                  - 'sagemaker:DeleteTags'
                  - 'sagemaker:DeleteEndpointConfig'
                  - 'sagemaker:ListTags'
                  - 'sagemaker:ListTransformJobs'
                  - 'sagemaker:StopTrainingJob'
                  - 'sagemaker:DeleteModel'
                  - 'sagemaker:ListTrainingJobs'
                  - 'sagemaker:ListHyperParameterTuningJobs'
                  - 'sagemaker:DeleteEndpoint'
                  - 'sagemaker:ListModels'
                  - 'sagemaker:StopTransformJob'
                  - 'sagemaker:AddTags'
                  - 'sagemaker:ListEndpoints'
                Resource: '*'
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'

  SageMakerVPCDetectiveControl: 
    Type: 'AWS::Lambda::Function'
    Properties:
      FunctionName: SageMakerVPCEnforcer
      Description: Detective control to enforce VPC attachment of SageMaker resources
      Runtime: python3.7
      Code: vpc_detective_control.zip
      Handler: inspect_sagemaker_resource.lambda_handler
      MemorySize: 320
      Timeout: 180
      Role: !GetAtt SageMakerDetectiveControlExecutionRole.Arn

  SageMakerVPCEnforcementRule:
    Type: 'AWS::Events::Rule'
    DependsOn: SageMakerVPCDetectiveControl
    Properties:
      Description: The Cloudwatch Rule checking VPC enablement of SageMaker resources
      EventPattern:
        source:
          - aws.sagemaker
        detail-type:
          - AWS API Call via CloudTrail
        detail:
          eventSource:
            - sagemaker.amazonaws.com
          eventName:
            - CreateTrainingJob
            - CreateModel
      Name: SageMakerVPCEnforcementRule
      State: ENABLED
      Targets:
        - Arn: !GetAtt SageMakerVPCDetectiveControl.Arn
          Id: SagemakerVPCEnforcementLambda

  InvokeLambdaPermission:
    Type: 'AWS::Lambda::Permission'
    DependsOn: SageMakerVPCEnforcementRule
    Properties:
      FunctionName: !GetAtt SageMakerVPCDetectiveControl.Arn
      Action: 'lambda:InvokeFunction'
      Principal: events.amazonaws.com
      SourceArn: !GetAtt SageMakerVPCEnforcementRule.Arn