AWSTemplateFormatVersion: '2010-09-09'
Description: Deploy Resources for enabling Tag Propagation for SageMaker
Resources:
  sagemakerTagPropagatorRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: sagemakerTagPropagatorRole
      Description: This role gives permission to Lambda function for creating tags and log groups resources
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Policies:
        - PolicyName: sagemakerTaggingPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Sid: PermissionToCreateTags
              - Effect: Allow
                Resource: "arn:aws:ec2:*:::security-group/*"
                Action:
                  - "ec2:CreateTags"     
              - Sid: PermissionToCreateLogsAndLogGroups         
              - Effect: Allow
                Resource: '*'
                Action:
                  - "sagemaker:ListTags"
                  - "efs:CreateTags"
                  - "logs:CreateLogGroup"
                  - "logs:CreateLogStream"
                  - "logs:PutLogEvents"  
      Tags:
        - Key: Name
          Value: sagemakerTagPropagatorRole        

  taggingSageMakerResources:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: propagateSagemakerTags
      Description: This Lambda function's scope is to propagate tags to resources that SageMaker creates like EFS and Security Groups
      Handler: propagateSagemakerTags.lambda_handler
      Role: !GetAtt sagemakerTagPropagatorRole.Arn
      Code:
        S3Bucket: tagpropagatorsagemaker
        S3Key: tagPropagatorforSageMaker.zip
      Runtime: python3.8
      Timeout: 60
      TracingConfig:
        Mode: Active
      Tags: 
        - Key: Name
          Value: propagateSagemakerTags

  propagateSageMakerTagsRuleForEFS:
    Type: AWS::Events::Rule
    Properties:
      Description: Invoke Lambda for tagging SageMaker Resources
      Name: tagPropagatorForSageMakerEFS
      EventPattern:
        source:
          - aws.elasticfilesystem
        detail-type:
          - AWS API Call via CloudTrail
        detail:
          eventSource:
            - elasticfilesystem.amazonaws.com
          eventName:
            - CreateFileSystem
      State: ENABLED
      Targets:
        - Arn: !GetAtt taggingSageMakerResources.Arn
          Id: LambdaFucntion

  PermissionForEFSEventsToInvokeLambda: 
    Type: AWS::Lambda::Permission
    Properties: 
      FunctionName: !Ref taggingSageMakerResources
      Action: "lambda:InvokeFunction"
      Principal: "events.amazonaws.com"
      SourceArn: !GetAtt propagateSageMakerTagsRuleForEFS.Arn

  propagateSageMakerTagsRuleForSG:
    Type: AWS::Events::Rule
    Properties:
      Description: Invoke Lambda for tagging SageMaker Resources
      Name: tagPropagatorForSageMakerSG
      EventPattern:
        source:
          - aws.ec2
        detail-type:
          - AWS API Call via CloudTrail
        detail:
          eventSource:
            - ec2.amazonaws.com
          eventName:
            - CreateTags
      State: ENABLED
      Targets:
        - Arn: !GetAtt taggingSageMakerResources.Arn
          Id: LambdaFucntion

  PermissionForSGEventsToInvokeLambda: 
    Type: AWS::Lambda::Permission
    Properties: 
      FunctionName: !Ref taggingSageMakerResources
      Action: "lambda:InvokeFunction"
      Principal: "events.amazonaws.com"
      SourceArn: !GetAtt propagateSageMakerTagsRuleForSG.Arn  