AWSTemplateFormatVersion: 2010-09-09
Description: Create IAM Roles for SageMaker Notebook Instance and KMS Key for Encrypting Notebook EBS Volume.

Parameters:
  ProjectName:
    Type: String
  DataBucket:
    Type: String
  KMSAlias:
    Type: String
    Default: sagemaker-kms-example

Resources:
  SageMakerNotebookInstancePolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Description: SageMaker Notebook restricted access
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Sid: CloudWatchLogsAccess
            Effect: Allow
            Action:
              - logs:CreateLogGroup
              - logs:CreateLogStream
              - logs:DescribeLogStreams
              - logs:PutLogEvents
              - logs:GetLogEvents
            Resource:
              - !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/sagemaker/*
              - !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-stream:*
          - Sid: CloudWatchMetricsAccess
            Effect: Allow
            Action:
              - cloudwatch:GetMetricStatistics
              - cloudwatch:ListMetrics
              - cloudwatch:PutMetricData
              - cloudwatch:GetMetricData
              - cloudwatch:PutMetricAlarm
              - cloudwatch:DescribeAlarms
            Resource: '*'
          - Sid: KMSAccess
            Effect: Allow
            Action:
              - 'kms:CreateGrant'
              - 'kms:Decrypt'
              - 'kms:DescribeKey'
              - 'kms:Encrypt'
              - 'kms:GenerateDataKey'
              - 'kms:ListAliases'
            Resource: '*'
          - Sid: S3Access
            Effect: Allow
            Action:
              - s3:GetObject
              - s3:PutObject
              - s3:DeleteObject
              - s3:ListBucket
            Resource:
              - !Sub arn:aws:s3:::${DataBucket}/*
          - Sid: ECRAccess
            Effect: Allow
            Action:
              - ecr:GetAuthorizationToken
              - ecr:GetDownloadUrlForLayer
              - ecr:BatchGetImage
              - ecr:BatchCheckLayerAvailability
              - ecr:SetRepositoryPolicy
              - ecr:CompleteLayerUpload
              - ecr:BatchDeleteImage
              - ecr:UploadLayerPart
              - ecr:InitiateLayerUpload
              - ecr:PutImage
            Resource: 'arn:aws:ecr:*:*:repository/*sagemaker*'
          - Sid: PassRole
            Effect: Allow
            Action:
              - iam:PassRole
            Resource: '*'
            Condition:
              StringEquals:
                iam:PassedToService: sagemaker.amazonaws.com
          - Sid: SageMakerAccess
            Action:
              - sagemaker:CreateTrainingJob
              - sagemaker:CreateProcessingJob
              - sagemaker:CreateModel
              - sagemaker:CreateHyperParameterTuningJob
            Resource: '*'
            Effect: Deny
            Condition:
              'Null':
                'sagemaker:VpcSubnets': 'true'
          - Sid: SageMakerList
            Action:
              - sagemaker:Describe*
              - sagemaker:List*
            Resource: '*'
            Effect: Allow
          - Sid: SageMakerStudioSignedURLCreation
            Action:
              - sagemaker:CreateApp
            Resource: '*'
            Effect: Allow
          - Sid: EC2Access
            Effect: Allow
            Action:
              - ec2:CreateNetworkInterface
              - ec2:CreateNetworkInterfacePermission
              - ec2:DeleteNetworkInterface
              - ec2:DeleteNetworkInterfacePermission
              - ec2:DescribeNetworkInterfaces
              - ec2:DescribeVpcs
              - ec2:DescribeDhcpOptions
              - ec2:DescribeSubnets
              - ec2:DescribeSecurityGroups
              - ec2:DescribeVpcEndpoints
            Resource: '*'

  SageMakerNotebookInstanceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service: sagemaker.amazonaws.com
            Action: sts:AssumeRole
      Path: /
      RoleName: !Sub ${ProjectName}-notebook-role
      ManagedPolicyArns:
        - !Ref SageMakerNotebookInstancePolicy



  KMSKey:
    Type: AWS::KMS::Key
    Properties:
      Description: "Generated KMS Key for sagemaker Notebook's EBS encryption"
      EnableKeyRotation: true
      Enabled: true
      KeyPolicy:
        Version: 2012-10-17
        Id: allow-root-access-to-key
        Statement:
          - Sid: allow-root-to-delegate-actions
            Effect: Allow
            Principal:
              AWS: !Sub arn:aws:iam::${AWS::AccountId}:root
            Action:
              - kms:DeleteAlias
              - kms:DescribeKey
              - kms:EnableKey
              - kms:GetKeyPolicy
              - kms:UpdateAlias
              - kms:CreateAlias
              - kms:GetKeyPolicy
              - kms:CreateGrant
              - kms:DisableKey
              - kms:Revoke*
              - kms:Disable*
              - kms:CancelKeyDeletion
              - kms:ScheduleKeyDeletion
              - kms:PutKeyPolicy
              - kms:RevokeGrant
              - kms:TagResource
              - kms:UnTagResource
              - kms:EnableKeyRotation
              - kms:ListResourceTags
            Resource: '*'
  KeyAlias:
    Type: AWS::KMS::Alias
    Properties:
      AliasName: !Sub alias/${KMSAlias}
      TargetKeyId: !Ref KMSKey  


Outputs:
  ExecutionRoleARN:
    Value: !GetAtt SageMakerNotebookInstanceRole.Arn
  KMSKeyArn:
    Value: !GetAtt KMSKey.Arn