AWSTemplateFormatVersion: 2010-09-09
Description: >-
  Deploy Lambda Function to Enrich Amazon Comprehend Medical Output with an Amazon Comprehend Custom Classification and Custom Code.

Parameters:
  ComprehendCustomEndpoint:
    Type: String
  StaticCodeBucket:
    Type: String
    Default: "aws-hcls-ml"
  StaticCodeKey:
    Type: String
    Default: 'blog_post_support_materials/comprehend_medical_custom_lambda_code/functions/packages/MyFunction/lambda.zip'


Resources:
#Create the buckets that we need
  NotificationBucket:
    Type: 'AWS::S3::Bucket'
  OutputBucket:
    Type: 'AWS::S3::Bucket'


  LambdaTrigger:
    Type: 'Custom::LambdaTrigger'
    DependsOn: LambdaInvokePermission
    Properties:
      ServiceToken: !GetAtt CustomResourceLambdaFunction.Arn
      LambdaArn: !GetAtt CallCMandCustomCode.Arn
      Bucket: !Ref NotificationBucket

  #this is the lambda function that call comprehend medical and the custom models
  CallCMandCustomCode: 
    Type: 'AWS::Lambda::Function'
    Properties:
      Code:
        S3Bucket: !Ref StaticCodeBucket
        S3Key: !Ref StaticCodeKey
      Handler: call_cm_other_models_lambda.lambda_handler
      Role: !GetAtt LambdaIAMRole.Arn
      Runtime: python3.6
      Timeout: 300
      MemorySize: 128
      Environment:
        Variables:
          NOTIFICATION_BUCKET_NAME: !Ref NotificationBucket
          OUTPUT_BUCKET_NAME: !Ref OutputBucket
          COMPREHEND_CUSTOM_ENDPOINT: !Ref ComprehendCustomEndpoint

  LambdaInvokePermission:
      Type: 'AWS::Lambda::Permission'
      Properties:
        FunctionName: !GetAtt CallCMandCustomCode.Arn
        Action: 'lambda:InvokeFunction'
        Principal: s3.amazonaws.com
        SourceAccount: !Ref 'AWS::AccountId'
        SourceArn: !Sub 'arn:aws:s3:::${NotificationBucket}'

  CustomResourceLambdaFunction:
    Type: 'AWS::Lambda::Function'
    Properties:
      Handler: index.lambda_handler
      Role: !GetAtt LambdaIAMRole.Arn
      Code:
        ZipFile: 
            |
            from __future__ import print_function
            import json
            import boto3
            import cfnresponse
            import os
            
            SUCCESS = "SUCCESS"
            FAILED = "FAILED"
            
            print('Loading function')
            s3 = boto3.resource('s3')
           
            def lambda_handler(event, context):
                print("Received event: " + json.dumps(event, indent=2))
                responseData={}
                try:
                    if event['RequestType'] == 'Delete':
                        print("Request Type:",event['RequestType'])
                        Bucket=event['ResourceProperties']['Bucket']
                        delete_notification(Bucket)
                        print("Sending response to custom resource after Delete")
                    elif event['RequestType'] == 'Create' or event['RequestType'] == 'Update':
                        print("Request Type:",event['RequestType'])
                        LambdaArn=event['ResourceProperties']['LambdaArn']
                        Bucket=event['ResourceProperties']['Bucket']
                        add_notification(LambdaArn, Bucket)
                        responseData={'Bucket':Bucket}
                        print("Sending response to custom resource")
                    responseStatus = 'SUCCESS'
                except Exception as e:
                    print('Failed to process:', e)
                    responseStatus = 'FAILURE'
                    responseData = {'Failure': 'Something bad happened.'}
                cfnresponse.send(event, context, responseStatus, responseData)

            def add_notification(LambdaArn, Bucket):
                bucket_notification = s3.BucketNotification(Bucket)
                response = bucket_notification.put(
                  NotificationConfiguration={
                    'LambdaFunctionConfigurations': [
                      {
                          'LambdaFunctionArn': LambdaArn,
                          'Events': [
                              's3:ObjectCreated:*'
                          ]
                      }
                    ]
                  }
                )
                print("Put request completed....")
              
            def delete_notification(Bucket):
                bucket_notification = s3.BucketNotification(Bucket)
                response = bucket_notification.put(
                    NotificationConfiguration={}
                )
                print("Delete request completed....")


      Runtime: python3.6
      Timeout: 500
      MemorySize: 1000


  LambdaIAMRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
                - comprehendmedical.amazonaws.com
                - comprehend.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      Policies:
        - PolicyName: root
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - 's3:GetBucketNotification'
                  - 's3:PutBucketNotification'
                Resource: !Sub 'arn:aws:s3:::${NotificationBucket}'
              - Effect: Allow
                Action:
                  - 's3:*'
                Resource: !Sub 'arn:aws:s3:::${NotificationBucket}/*'

              - Effect: Allow
                Action:
                  - 's3:*'
                Resource: !Sub 'arn:aws:s3:::${OutputBucket}/*'
              - Effect: Allow
                Action:
                  - 's3:*'
                Resource: !Sub 'arn:aws:s3:::${StaticCodeBucket}/*'
              - Effect: Allow
                Action:
                  - 'logs:CreateLogGroup'
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
                Resource: 'arn:aws:logs:*:*:*'
              - Effect: Allow
                Action: 
                  - "comprehendmedical:*"
                Resource: '*'
              - Effect: Allow
                Action: 
                  - "comprehend:*"
                Resource: '*'



Outputs:
  InputBucket:
    Description: S3 bucket where raw data is placed by the user.
    Value: !Ref NotificationBucket
  ResultsBucket:
    Description: S3 bucket where processed data is placed after being processed by Lambda.
    Value: !Ref OutputBucket
  #ZipCodeBucket:  #for debugging
  #  Description: S3 bucket where Lambda zipfile code is located.
  #  Value: !Ref StaticDataBucket
  #LambdaFunctionCreated:
  #  Value: !Ref CallCMandCustomCode
  #IamRoleCreated:
  #  Value: !GetAtt LambdaIAMRole.Arn
  #ResourceLambdaFunctionCreated:
  #  Value: !Ref CustomResourceLambdaFunction