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

AWSTemplateFormatVersion: 2010-09-09
Description: Creates necessary IAM roles in target accounts (staging, prod) for StackSet operations and multi-account model deployment

Parameters:
  EnvName:
    Type: String
    AllowedPattern: '[a-z0-9\-]*'
    Description: Please specify your data science environment/team name.  Used as a suffix for environment resource names.

  EnvType:
    Description: System Environment (e.g. dev, staging, prod). Used as a suffix for environment resource names.
    Type: String
    Default: dev

  SageMakerModelExecutionRoleName:
    Type: String
    Description: Name of the SageMaker model execution role. If empty, the name will be generated
    Default: ''
  
  StackSetExecutionRoleName:
    Type: String
    Description: Name of the CloudFormation StackSet execution role. If empty, the name will be generated
    Default: ''

  PipelineExecutionRoleArn:
    Type: String
    Description: CI/CD execution role in the development account. If empty, the name AmazonSageMakerServiceCatalogProductsUseRole is used
    Default: ''
  
  AdministratorAccountId:
    Type: String
    Description: AWS Account Id of the account in which StackSets will be created (for multi-account deployment)
    MaxLength: 12
    MinLength: 12
  
  ModelS3KMSKeyArn:
    Type: String
    Description: KMS key ARN for model encryption in the S3 bucket

  ModelBucketName:
    Type: String
    Description: Name of the S3 bucket where models are stored    

Outputs:
  SageMakerModelExecutionRoleArn:
    Description: The ARN of the SageMaker model deployment and execution role 
    Value: !GetAtt SageMakerModelExecutionRole.Arn

  SageMakerModelExecutionRoleName:
    Description: The ARN of the SageMaker model deployment and execution role 
    Value: !Select ['1', !Split [ '/', !Select ['5', !Split [':', !GetAtt SageMakerModelExecutionRole.Arn ]]]]

  StackSetExecutionRoleArn:
    Description: The ARN of the AWS CloudFormation StackSet execution role
    Value: !GetAtt StackSetExecutionRole.Arn

  StackSetExecutionRoleName:
    Description: The name of the AWS CloudFormation StackSet execution role
    Value: !Select ['1', !Split [ '/', !Select ['5', !Split [':', !GetAtt StackSetExecutionRole.Arn ]]]]

Conditions:
  GenerateSageMakerModelExecutionRoleName: !Equals [ !Ref SageMakerModelExecutionRoleName, '']
  GenerateStackSetExecutionRoleName: !Equals [ !Ref StackSetExecutionRoleName, '']
  PipelineExecutionRoleCondition: !Equals [ !Ref PipelineExecutionRoleArn, '']

Resources:

  # This role is used to run a model endpoint and to test the model in SageMaker MLOps CI/CD pipeline
  # Assumed by AmazonSageMakerServiceCatalogProductsUseRole
  SageMakerModelExecutionRole:
    Type: AWS::IAM::Role
    Properties:
        AssumeRolePolicyDocument:
          Version: "2012-10-17"
          Statement:
            -
              Effect: "Allow"
              Principal:
                Service:
                  - "sagemaker.amazonaws.com"
              Action:
                - "sts:AssumeRole"
            -
              Effect: "Allow"
              Principal:
                AWS: !If 
                  - PipelineExecutionRoleCondition
                  - !Sub 'arn:aws:iam::${AdministratorAccountId}:role/service-role/AmazonSageMakerServiceCatalogProductsUseRole'
                  - !Sub ${PipelineExecutionRoleArn}
              Action:
                - "sts:AssumeRole"
        Path: "/"
        RoleName: !If
          - GenerateSageMakerModelExecutionRoleName
          - !Ref 'AWS::NoValue'
          - !Ref SageMakerModelExecutionRoleName
        Policies: 
          -
            PolicyName: !Sub "${EnvName}-${EnvType}-${AWS::Region}-sagemaker-model-execution-policy"
            PolicyDocument:
              Version: "2012-10-17"
              Statement:
                - # permission to create a model, endpoint config, and endpoint
                  Effect: "Allow"
                  Action:
                    - "sagemaker:*Model"
                    - "sagemaker:*EndpointConfig"
                    - "sagemaker:*Endpoint"
                  Resource:
                    - !Sub "arn:aws:sagemaker:*:${AWS::AccountId}:endpoint/*"
                    - !Sub "arn:aws:sagemaker:*:${AWS::AccountId}:endpoint-config/*"
                    - !Sub "arn:aws:sagemaker:*:${AWS::AccountId}:model/*"
                - # permission to create an ENI in the provided VPC
                  Effect: "Allow"
                  Action:
                    - "ec2:DeleteNetworkInterfacePermission"
                    - "ec2:DeleteNetworkInterface"
                    - "ec2:CreateNetworkInterfacePermission"
                    - "ec2:CreateNetworkInterface"
                  Resource: 
                    - !Sub 'arn:aws:ec2:*:${AWS::AccountId}:*'
                -
                  Effect: 'Allow'
                  Action:
                    - "ec2:DescribeVpcEndpoints"
                    - "ec2:DescribeDhcpOptions"
                    - "ec2:DescribeVpcs"
                    - "ec2:DescribeSubnets"
                    - "ec2:DescribeSecurityGroups"
                    - "ec2:DescribeNetworkInterfaces"
                  Resource:
                    - '*'
                - # permission to download the model artifact from the S3 bucket in model account
                  Effect: 'Allow'
                  Action:
                    - s3:GetObject
                    - s3:ListBucket
                  Resource:
                    - !Sub "arn:aws:s3:::${ModelBucketName}/*"
                    - !Sub "arn:aws:s3:::${ModelBucketName}"
                -  # permission to access the KMS key for S3 bucket with the model
                  Effect: 'Allow'
                  Action:
                    - kms:Encrypt
                    - kms:Decrypt
                    - kms:ReEncrypt*
                    - kms:GenerateDataKey*
                    - kms:DescribeKey
                  Resource: !Ref ModelS3KMSKeyArn
        Tags:
          - Key: EnvironmentName
            Value: !Ref EnvName
          - Key: EnvironmentType
            Value: !Ref EnvType

  # This role is used by StackSet operations in the target account
  StackSetExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
          - Effect: Allow
            Principal:
              AWS:
                - !Ref AdministratorAccountId
            Action:
              - sts:AssumeRole
      Path: /
      RoleName: !If
          - GenerateStackSetExecutionRoleName
          - !Ref 'AWS::NoValue'
          - !Ref StackSetExecutionRoleName
      Policies: 
        -
          PolicyName: !Sub "${EnvName}-${EnvType}-${AWS::Region}-stackset-execution-policy"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              -
                Effect: Allow
                Action:
                  - 'cloudformation:Cancel*'
                  - 'cloudformation:Continue*'
                  - 'cloudformation:Create*'
                  - 'cloudformation:Deactivate*'
                  - 'cloudformation:Delete*'
                  - 'cloudformation:Describe*'
                  - 'cloudformation:Get*'
                  - 'cloudformation:Detect*'
                  - 'cloudformation:Execute*'
                  - 'cloudformation:Estimate*'
                  - 'cloudformation:List*'
                  - 'cloudformation:Publish*'
                  - 'cloudformation:Set*'
                  - 'cloudformation:Register*'
                  - 'cloudformation:Signal*'
                  - 'cloudformation:Stop*'
                  - 'cloudformation:Update*'
                  - 'cloudformation:Validate*'
                  - 'cloudformation:Test*'
                Resource: 
                - !Sub 'arn:aws:cloudformation:*:${AWS::AccountId}:*'
              -
                Effect: Allow
                Action:
                  - 's3:Get*'
                  - 's3:List*'
                  - 's3:Create*'
                  - 's3:Head*'
                  - 's3:Put*'
                  - 's3:Upload*'
                Resource:
                  - 'arn:aws:s3:::*'
              -
                Effect: Allow
                Action:
                  - 'sns:Create*'
                  - 'sns:Get*'
                  - 'sns:List*'
                  - 'sns:Publish'
                  - 'sns:Set*'
                  - 'sns:Subscribe'
                  - 'sns:Tag*'
                  - 'sns:Unsubscribe'
                  - 'sns:Untag*'
                  - 'sns:Add*'
                Resource: 
                  - 'arn:aws:sns:*:*:*'
              -
                Effect: "Allow"
                Action:
                  - "sagemaker:*Model"
                  - "sagemaker:*EndpointConfig"
                  - "sagemaker:*Endpoint"
                  - "sagemaker:AddTags"
                Resource:
                  - !Sub "arn:aws:sagemaker:*:${AWS::AccountId}:endpoint/*"
                  - !Sub "arn:aws:sagemaker:*:${AWS::AccountId}:endpoint-config/*"
                  - !Sub "arn:aws:sagemaker:*:${AWS::AccountId}:model/*"
                  - !Sub "arn:aws:sagemaker:*:${AdministratorAccountId}:model-package/*"
              - 
                Effect: 'Allow'
                Action:
                  - 'ssm:GetParameter*'
                Resource:
                  - !Sub 'arn:aws:ssm:${AWS::Region}:${AWS::AccountId}:parameter/${EnvName}*'
              - 
                Effect: 'Allow'
                Action:
                  - 'iam:PassRole'
                Resource:
                  - !GetAtt SageMakerModelExecutionRole.Arn
      Tags:
        - Key: EnvironmentName
          Value: !Ref EnvName
        - Key: EnvironmentType
          Value: !Ref EnvType