# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

AWSTemplateFormatVersion: 2010-09-09
Description: Create IAM Roles for SageMaker user profiles and SAML backend Lambda function

Outputs:

  SageMakerStudioExecutionRoleDefaultArn:
    Description: The ARN of the SageMaker default execution role
    Value: !GetAtt SageMakerStudioExecutionRoleDefault.Arn

  SageMakerStudioExecutionRoleTeam1Arn:
    Description: The ARN of the SageMaker Team1 execution role
    Value: !GetAtt SageMakerStudioExecutionRoleTeam1.Arn
  
  SageMakerStudioExecutionRoleTeam2Arn:
    Description: The ARN of the SageMaker Team2 execution role
    Value: !GetAtt SageMakerStudioExecutionRoleTeam2.Arn

  SetupLambdaExecutionRoleArn:
    Description: Lambda execution role for stack setup Lambda function
    Value: !GetAtt SetupLambdaExecutionRole.Arn
    
  SAMLBackendLambdaExecutionRoleArn:
    Description: Lambda execution role for SAML backend Lambda function
    Value: !GetAtt SAMLBackendLambdaExecutionRole.Arn

Parameters:
  EnvironmentName:
    Type: String
    AllowedPattern: '[a-z0-9\-]*'
    Description: Please specify your SageMaker environment name. 

  AllowedCIDR:
    Type: String
    Description: Allowed CIDR block for CreatePresignedDomainURL API call
  
Conditions:
  RestrictSageMakerToCIDRCondition: !Not [ !Equals [ !Ref AllowedCIDR, ''] ]

Resources:

  SageMakerDeniedServicesPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Description: Explicit deny for specific SageMaker services
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Sid: AmazonSageMakerDeniedServices
            Action:
              - sagemaker:CreatePresignedNotebookInstanceUrl
              - sagemaker:*NotebookInstance
              - sagemaker:*NotebookInstanceLifecycleConfig
              - sagemaker:CreateUserProfile
              - sagemaker:DeleteDomain
              - sagemaker:DeleteUserProfile
            Resource: 
              - '*'
            Effect: Deny

  SageMakerReadOnlyPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Description: Read-only baseline policy for SageMaker execution role
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Sid: AmazonSageMakerDescribeReadyOnlyPolicy
            Action:
              - sagemaker:Describe*
              - sagemaker:GetSearchSuggestions
            Resource:
              - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*'
            Effect: Allow
          - Sid: AmazonSageMakerListOnlyPolicy
            Action:
              - 'sagemaker:List*'
            Resource:
              - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*'
            Effect: Allow
          - Sid: AmazonSageMakerUIandMetricsOnlyPolicy
            Action:
              - sagemaker:*App
              - sagemaker:Search
              - sagemaker:RenderUiTemplate
              - sagemaker:BatchGetMetrics
            Resource:
              - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*' 
            Effect: Allow
          - Sid: AmazonSageMakerEC2ReadOnlyPolicy
            Action:
              - ec2:DescribeDhcpOptions
              - ec2:DescribeNetworkInterfaces
              - ec2:DescribeRouteTables
              - ec2:DescribeSecurityGroups
              - ec2:DescribeSubnets
              - ec2:DescribeVpcEndpoints
              - ec2:DescribeVpcs
            Resource:
              - !Sub 'arn:aws:ec2:*:${AWS::AccountId}:*'
            Effect: Allow
          - Sid: AmazonSageMakerIAMReadOnlyPolicy
            Action:
              - iam:ListRoles
            Resource:
              - !Sub 'arn:aws:iam::${AWS::AccountId}:*'
            Effect: Allow

  SageMakerAccessSupportingServicesPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Description: Read-only baseline policy for SageMaker execution role
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Sid: AmazonSageMakerCRUDAccessS3Policy
            Action:
              - s3:PutObject
              - s3:GetObject
              - s3:AbortMultipartUpload
              - s3:DeleteObject
              - s3:CreateBucket
              - s3:ListBucket
              - s3:PutBucketCORS
              - s3:ListAllMyBuckets
              - s3:GetBucketCORS
              - s3:GetBucketLocation
            Resource:
              - arn:aws:s3:::*SageMaker*
              - arn:aws:s3:::*Sagemaker*
              - arn:aws:s3:::*sagemaker*
            Effect: Allow
          - Sid: AmazonSageMakerReadOnlyAccessKMSPolicy
            Action:
              - kms:DescribeKey
              - kms:ListAliases
            Resource:
              - !Sub 'arn:aws:kms:*:${AWS::AccountId}:*'
            Effect: Allow
          - Sid: AmazonSageMakerCRUDAccessECRPolicy
            Action:
              - ecr:Set*
              - ecr:CompleteLayerUpload
              - ecr:Batch*
              - ecr:Upload*
              - ecr:InitiateLayerUpload
              - ecr:Put*
              - ecr:Describe*
              - ecr:CreateRepository
              - ecr:Get*
              - ecr:StartImageScan
            Resource:
              - '*'
            Effect: Allow
          - Sid: AmazonSageMakerCRUDAccessCloudWatchPolicy
            Action:
              - cloudwatch:Put*
              - cloudwatch:Get*
              - cloudwatch:List*
              - cloudwatch:DescribeAlarms
              - logs:Put*
              - logs:Get*
              - logs:List*
              - logs:CreateLogGroup
              - logs:CreateLogStream
              - logs:ListLogDeliveries
              - logs:Describe*
              - logs:CreateLogDelivery
              - logs:PutResourcePolicy
              - logs:UpdateLogDelivery
            Resource:
              - '*'
            Effect: Allow

  SageMakerStudioDeveloperAccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Description: Read-only baseline policy for SageMaker execution role
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Sid: AmazonSageMakerStudioCreateApp
            Action:
              - sagemaker:CreateApp
            Resource:
              - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*'
            Effect: Allow
          - Sid: AmazonSageMakerStudioIAMPassRole
            Action:
              - iam:PassRole
            Resource:
              - !Sub 'arn:aws:iam::${AWS::AccountId}:role/*AmazonSageMaker*'
            Effect: Allow
            Condition:
              StringEquals:
                iam:PassedToService: sagemaker.amazonaws.com
          - Sid: AmazonSageMakerInvokeEndPointRole
            Action:
              - sagemaker:InvokeEndpoint
            Resource:
              - '*'
            Effect: Allow
          - Sid: AmazonSageMakerTags
            Action:
              - sagemaker:AddTags
              - sagemaker:ListTags
            Resource:
              - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*'
            Effect: Allow
          - Sid: AmazonSageMakerCreate
            Action:
              - sagemaker:Create*
            Resource:
              - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*'
            Effect: Allow
            Condition:
              ForAnyValue:StringEquals:
                aws:TagKeys:
                  - Team
              StringEqualsIfExists:
                aws:RequestTag/Team: ${aws:PrincipalTag/Team}
          - Sid: AmazonSageMakerUpdateDeleteExecutePolicy
            Action:
              - sagemaker:Delete*
              - sagemaker:Stop*
              - sagemaker:Update*
              - sagemaker:Start*
              - sagemaker:DisassociateTrialComponent
              - sagemaker:AssociateTrialComponent
              - sagemaker:BatchPutMetrics
            Resource:
              - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:*'
            Effect: Allow
            Condition:
              StringEquals:
                aws:PrincipalTag/Team: ${sagemaker:ResourceTag/Team}

  SageMakerStudioExecutionRoleDefault:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - sagemaker.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      ManagedPolicyArns:
        - !Ref SageMakerReadOnlyPolicy
        - !Ref SageMakerDeniedServicesPolicy
      Tags:
        - Key: EnvironmentName
          Value: !Ref EnvironmentName

  SageMakerStudioExecutionRoleTeam1:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - sagemaker.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      ManagedPolicyArns:
        - !Ref SageMakerReadOnlyPolicy
        - !Ref SageMakerAccessSupportingServicesPolicy
        - !Ref SageMakerStudioDeveloperAccessPolicy
        - !Ref SageMakerDeniedServicesPolicy
      Tags:
        - Key: EnvironmentName
          Value: !Ref EnvironmentName
        - Key: Team
          Value: Team1

  SageMakerStudioExecutionRoleTeam2:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - sagemaker.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      ManagedPolicyArns:
        - !Ref SageMakerReadOnlyPolicy
        - !Ref SageMakerAccessSupportingServicesPolicy
        - !Ref SageMakerStudioDeveloperAccessPolicy
        - !Ref SageMakerDeniedServicesPolicy
      Tags:
        - Key: EnvironmentName
          Value: !Ref EnvironmentName
        - Key: Team
          Value: Team2

  RestrictSageMakerToCIDRPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Path: /
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Deny
            Action:
              - 'sagemaker:*'
            Resource: !Sub 'arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:user-profile/*/*'
            Condition:
              NotIpAddress:
                  aws:VpcSourceIp: !If [RestrictSageMakerToCIDRCondition, !Ref AllowedCIDR, '0.0.0.0/0']

  SageMakerPermissionsPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Path: /
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Action:
              - 'sagemaker:CreatePresignedDomainUrl'
              - 'sagemaker:ListUserProfiles'
              - 'sagemaker:CreateUserProfile'
              - 'sagemaker:DescribeUserProfile'
              - 'sagemaker:AddTags'
              - 'sagemaker:ListTags'
            Resource: !Sub 'arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:user-profile/*/*'

  SAMLBackendLambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/CloudWatchLambdaInsightsExecutionRolePolicy'
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
        - !Ref SageMakerPermissionsPolicy
        - !If [RestrictSageMakerToCIDRCondition, !Ref RestrictSageMakerToCIDRPolicy, !Ref AWS::NoValue]
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          -
            Effect: Allow
            Principal:
              Service:
                - 'lambda.amazonaws.com'
            Action:
              - 'sts:AssumeRole'
      Policies:
        - PolicyName: PassExecutionRole
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - 'iam:PassRole'
                Resource:
                  - !GetAtt SageMakerStudioExecutionRoleTeam1.Arn
                  - !GetAtt SageMakerStudioExecutionRoleTeam2.Arn
        - PolicyName: AccessVPCResources
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - 'ec2:CreateNetworkInterface'
                  - 'ec2:DescribeNetworkInterfaces'
                  - 'ec2:DeleteNetworkInterface'
                  - 'ec2:DescribeSecurityGroups'
                  - 'ec2:DescribeSubnets'
                  - 'ec2:DescribeVpcs'
                Resource:
                  - '*'  
      Tags:
        - Key: EnvironmentName
          Value: !Ref EnvironmentName

  SetupLambdaExecutionPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Path: /
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Sid: SageMakerDomainPermission
            Effect: Allow
            Action:
              - sagemaker:ListDomains
              - sagemaker:DescribeDomain
              - sagemaker:UpdateDomain
              - sagemaker:ListUserProfiles
              - sagemaker:DescribeUserProfile
              - sagemaker:ListApps
              - sagemaker:DescribeApp
              - sagemaker:DeleteApp
              - sagemaker:UpdateApp
              - sagemaker:DeleteUserProfile
            Resource:
              - !Sub "arn:${AWS::Partition}:sagemaker:*:*:domain/*"
              - !Sub "arn:${AWS::Partition}:sagemaker:*:*:user-profile/*"
              - !Sub "arn:${AWS::Partition}:sagemaker:*:*:app/*"
          - Sid: SCPermissions
            Effect: Allow
            Action:
              - servicecatalog:Associate*
              - servicecatalog:Accept*
              - servicecatalog:Enable*
              - servicecatalog:Get*
              - servicecatalog:List*
              - servicecatalog:Describe*
            Resource:
              - !Sub 'arn:aws:servicecatalog:*:${AWS::AccountId}:*'
              - !Sub 'arn:aws:catalog:*:${AWS::AccountId}:*'
          - # Authorization strategy is ActionOnly for these two operations and require * in resource field
            Sid: SageMakerEnableSCPortfolio
            Effect: Allow
            Action:
              - sagemaker:EnableSagemakerServicecatalogPortfolio
              - sagemaker:DisableSagemakerServicecatalogPortfolio
            Resource:
              - '*'
          - Sid: SageMakerExecPassRole
            Effect: Allow
            Action:
              - iam:PassRole
              - iam:GetRole
            Resource: 
              - !GetAtt SageMakerStudioExecutionRoleDefault.Arn
              - !GetAtt SageMakerStudioExecutionRoleTeam1.Arn
              - !GetAtt SageMakerStudioExecutionRoleTeam2.Arn

  SetupLambdaExecutionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      ManagedPolicyArns:
        - !Ref SetupLambdaExecutionPolicy
        - 'arn:aws:iam::aws:policy/AmazonEC2ReadOnlyAccess'
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
      Tags:
      - Key: EnvironmentName
        Value: !Ref EnvironmentName