AWSTemplateFormatVersion: "2010-09-09"
Description: Provisions Amazon SageMaker Studio domain.
Parameters:
  S3Bucket:
    Type: String
    Description: "S3 Bucket"
  VPCId:
    Type: AWS::EC2::VPC::Id
    Description: VPC Id to launch SageMaker Studio domain
  SubnetIds:
    Type: AWS::EC2::Subnet::Id
    Description: Public subnet id to launch SageMaker Studio domain
  S3KeyLambdaLayer:
    Type: String
    Description: S3Key for lambda layer, cfnResponse-layer.zip
    Default: "code/cfnResponse_layer.zip"
  S3KeyLambdaFunction:
    Type: String
    Description: S3Key for lambda function code, create_domain.zip
    Default: "code/create_domain.zip"

Resources:
  LambdaExecutionRole:
    Type: "AWS::IAM::Role"
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - "sts:AssumeRole"
      Path: /
      Policies:
        - PolicyName: root
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Sid: CloudWatchLogsPermissions
                Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: 
                  Fn::Sub: "arn:${AWS::Partition}:logs:*:*:*"
              - Sid: SageMakerDomainPermission
                Effect: Allow
                Action:
                  - sagemaker:CreateDomain
                  - sagemaker:DescribeDomain
                  - sagemaker:DeleteDomain
                  - sagemaker:UpdateDomain
                Resource: 
                  Fn::Sub: "arn:${AWS::Partition}:sagemaker:*:*:domain/*"
              - Sid: SageMakerExecPassRole
                Effect: Allow
                Action:
                  - iam:PassRole
                Resource: 
                 Fn::GetAtt: SageMakerExecutionRole.Arn
              - Sid: ServiceLinkedRole
                Effect: Allow
                Action:
                  - iam:CreateServiceLinkedRole
                Resource:
                  - !Sub "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/aws-service-role/sagemaker.amazonaws.com/AWSServiceRoleForAmazonSageMakerNotebooks"

  SageMakerExecutionRole:
    Type: "AWS::IAM::Role"
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - sagemaker.amazonaws.com
            Action:
              - "sts:AssumeRole"
      Path: /
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess

  StudioDomainFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: create_domain.lambda_handler
      Role: 
        Fn::GetAtt: LambdaExecutionRole.Arn
      Code:
        S3Bucket: 
         Ref: S3Bucket
        S3Key: 
          Ref: S3KeyLambdaFunction
      Runtime: python3.8
      Timeout: 900
      Layers:
        - Ref: CfnResponseLayer

  CfnResponseLayer:
    Type: AWS::Lambda::LayerVersion
    Properties:
      CompatibleRuntimes:
        - python3.8
      Content:
        S3Bucket: 
          Ref: S3Bucket
        S3Key: 
          Ref: S3KeyLambdaLayer
      Description: cfn-response layer
      LayerName: cfn-response

  StudioDomain:
    Type: Custom::StudioDomain
    Properties:
      ServiceToken: 
        Fn::GetAtt: StudioDomainFunction.Arn
      VPC: 
        Ref: VPCId
      SubnetIds: 
        Ref: SubnetIds
      DomainName: "MyDomainName"
      DefaultUserSettings:
        ExecutionRole: 
          Fn::GetAtt: SageMakerExecutionRole.Arn