AWSTemplateFormatVersion: "2010-09-09"
Description: A template for creating the stack from "Redacting sensitive information in near-real time" Blog. Which creates a Lambda function, Kinesis Firehose, and all required IAM roles.
Resources:
  rLambdaRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Join
      - ""
      - - "lambda-role-"
        - !Ref AWS::StackName
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      Path: "/"        
      Policies:
        - PolicyName: !Sub ${AWS::StackName}-${AWS::Region}-AWSLambda-Firehose
          PolicyDocument:
            Version: '2012-10-17'  
            Statement:
              - Effect: Allow
                Action: 
                  - firehose:PutRecord
                  - firehose:PutRecordBatch
                Resource: "*"
        - PolicyDocument:
           Statement:
              - Effect: Allow
                Action:
                  - comprehend:DetectPiiEntities
                  - comprehend:ContainsPiiEntities
                Resource: "*"
           Version: '2012-10-17'
          PolicyName: !Sub ${AWS::StackName}-${AWS::Region}-AWSLambda-Comprehend
        - PolicyDocument:
            Statement:
              - Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Effect: Allow
                Resource: arn:aws:logs:*:*:*
            Version: '2012-10-17'
          PolicyName: !Sub ${AWS::StackName}-${AWS::Region}-AWSLambda-CW
  rLambdaLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub "/aws/lambda/${rLambdaTransformation}"
      RetentionInDays: 30
  rFirehoseRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: !Join
      - ""
      - - "firehose-role-"
        - !Ref AWS::StackName
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - firehose.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /        
      Policies:
        - PolicyName: !Join
          - ""
          - - "firehose-policy-"
            - !Ref AWS::StackName
          PolicyDocument:
           Version: '2012-10-17'
           Statement:
           - Effect: Allow
             Action:
               - s3:AbortMultipartUpload 
               - s3:GetBucketLocation
               - s3:GetObject
               - s3:ListBucket
               - s3:ListBucketMultipartUploads
               - s3:PutObject
             Resource:
             - !Sub arn:aws:s3:::${rS3BucketRaw}/*
             - !Sub arn:aws:s3:::${rS3BucketRaw}
             - !Sub arn:aws:s3:::${rS3BucketRedacted}/*
             - !Sub arn:aws:s3:::${rS3BucketRedacted}
        - PolicyDocument:
            Statement:
              - Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Effect: Allow
                Resource: arn:aws:logs:*:*:*
          PolicyName: !Sub ${AWS::StackName}-${AWS::Region}-AWSLambda-Firehose-Logs
        - PolicyDocument:
            Statement:
              - Action:
                  - lambda:InvokeFunction
                  - logs:CreateLogStream
                Effect: Allow
                Resource: !GetAtt rLambdaTransformation.Arn
          PolicyName: !Sub ${AWS::StackName}-${AWS::Region}-AWSLambda-Firehose-Lambda                
  rS3BucketRaw:
    Type: AWS::S3::Bucket
    Properties:
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: 'AES256'
      PublicAccessBlockConfiguration:
        BlockPublicAcls: "True"
        BlockPublicPolicy: "True"
        IgnorePublicAcls: "True"
        RestrictPublicBuckets: "True"
      BucketName: !Join
      - ""
      - - !Ref AWS::AccountId
        - "-"
        - "awscomprehendrealtimeredactionblog-raw"
  rS3BucketRedacted:
    Type: AWS::S3::Bucket
    Properties:
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: 'AES256'
      PublicAccessBlockConfiguration:
        BlockPublicAcls: True
        BlockPublicPolicy: True
        IgnorePublicAcls: True
        RestrictPublicBuckets: True
      BucketName: !Join
      - ""
      - - !Ref AWS::AccountId
        - "-"
        - "awscomprehendrealtimeredactionblog-redacted"
  rLambdaTransformation:
    Type: AWS::Lambda::Function
    Properties:
     FunctionName: AmazonComprehendPII-Redact
     Environment:
      Variables:
       keys: keys
     Runtime: python3.9
     Role: !GetAtt rLambdaRole.Arn
     Timeout: 90
     Handler: index.lambda_handler
     Code: 
       ZipFile: |
          import json
          import boto3
          import os
          import base64
          import sys

          def lambda_handler(event, context):
              
              output = []
              
              for record in event['records']:
                  
                  # Gathers keys from enviroment variables and makes a list of desired keys to check for PII
                  rawkeys = os.environ['keys']
                  splitkeys = rawkeys.split(", ")
                  print(splitkeys)
                  #decode base64
                  #Kinesis data is base64 encoded so decode here
                  payloadraw=base64.b64decode(record["data"]).decode('utf-8')
                  #Loads decoded payload into json
                  payloadjsonraw = json.loads(payloadraw)
                  
                  # Creates Comprehend client
                  comprehend_client = boto3.client('comprehend')
                  
                  
                  # This codes handles the logic to check for keys, identify if PII exists, and redact PII if available. 
                  for i in payloadjsonraw:
                      # checks if the key found in the message matches a redact
                      if i in splitkeys:
                          print("Redact key found, checking for PII")
                          payload = str(payloadjsonraw[i])
                          # check if payload size is less than 100KB
                          if sys.getsizeof(payload) < 99999:
                              print('Size is less than 100KB checking if value contains PII')
                              # Runs Comprehend ContainsPiiEntities API call to see if key value contains PII
                              pii_identified = comprehend_client.contains_pii_entities(Text=payload, LanguageCode='en')
                              
                              # If PII is not found, skip over key
                              if (pii_identified['Labels']) == []:
                                  print('No PII found')
                              else:
                              # if PII is found, run through redaction logic
                                  print('PII found redacting')
                                  # Runs Comprehend DetectPiiEntities call to find exact location of PII
                                  response = comprehend_client.detect_pii_entities(Text=payload, LanguageCode='en')
                                  entities = response['Entities']
                                  # creates redacted_payload which will be redacted
                                  redacted_payload = payload
                                  # runs through a loop that gathers necessary values from Comprehend API response and redacts values
                                  for entity in entities:
                                      char_offset_begin = entity['BeginOffset']
                                      char_offset_end = entity['EndOffset']
                                      redacted_payload = redacted_payload[:char_offset_begin] + '*'*(char_offset_end-char_offset_begin) + redacted_payload[char_offset_end:]
                                  # replaces original value with redacted value
                                  payloadjsonraw[i] = redacted_payload
                                  print(str(payloadjsonraw[i]))
                          else:
                              print ('Size is more than 100KB, skipping inspection')
                      else:
                          print("Key value not found in redaction list")
                  
                  redacteddata = json.dumps(payloadjsonraw)
                  
                  # adds inspected record to record
                  output_record = {
                      'recordId': record['recordId'],
                      'result': 'Ok',
                      'data' : base64.b64encode(redacteddata.encode('utf-8'))
                  }
                  output.append(output_record)
                  print(output_record)
                  
              print('Successfully processed {} records.'.format(len(event['records'])))
              
              return {'records': output}
         
  MyFirehoseLogGroup:
    Type: AWS::Logs::LogGroup
    Properties: 
      LogGroupName: !Join
      - ""
      - - !Ref AWS::StackName
        - "-"
        - "awscomprehendblog-loggroup"
      RetentionInDays: 30
  MyFirehoseLogStream:      
    Type: AWS::Logs::LogStream
    Properties: 
      LogGroupName: !Ref MyFirehoseLogGroup
      LogStreamName: !Join
      - ""
      - - "/aws/firehose"
        - !Ref AWS::StackName
        - "-"
        - "awscomprehendblog-stream"
  rKinesisFirehoseDeliveryStream:
    Type: AWS::KinesisFirehose::DeliveryStream
    Properties:
      DeliveryStreamName: !Join
      - ""
      - - !Ref AWS::AccountId
        - "-"
        - "awscomprehendblog-stream"
      DeliveryStreamType: DirectPut
      ExtendedS3DestinationConfiguration:
        RoleARN: !GetAtt rFirehoseRole.Arn
        BucketARN: !GetAtt rS3BucketRedacted.Arn
        BufferingHints:
          IntervalInSeconds: 300
          SizeInMBs: 128
        CloudWatchLoggingOptions: 
          Enabled: true
          LogGroupName: !Ref MyFirehoseLogGroup
          LogStreamName: !Ref MyFirehoseLogStream
        CompressionFormat: "UNCOMPRESSED"
        Prefix: >-
          redacted/year=!{timestamp:yyyy}/month=!{timestamp:MM}/day=!{timestamp:dd}/hour=!{timestamp:HH}/
        ErrorOutputPrefix: >-
          fherroroutputbase/!{firehose:random-string}/!{firehose:error-output-type}/!{timestamp:yyyy/MM/dd}/
        ProcessingConfiguration:
          Enabled: true 
          Processors:
            - Type: Lambda
              Parameters:
                - ParameterName: LambdaArn
                  ParameterValue: !GetAtt rLambdaTransformation.Arn
                - ParameterName: NumberOfRetries
                  ParameterValue: '3'
                - ParameterName: RoleArn
                  ParameterValue: !GetAtt rFirehoseRole.Arn
                - ParameterName: BufferSizeInMBs
                  ParameterValue: '1'
                - ParameterName: BufferIntervalInSeconds
                  ParameterValue: '60'
        S3BackupMode: Enabled
        S3BackupConfiguration:
            RoleARN: !GetAtt rFirehoseRole.Arn
            BucketARN: !GetAtt rS3BucketRaw.Arn
            BufferingHints: 
              SizeInMBs: 1
              IntervalInSeconds: 240
            CompressionFormat: GZIP
            CloudWatchLoggingOptions: 
              Enabled: true
              LogGroupName: !Ref MyFirehoseLogGroup
              LogStreamName: !Ref MyFirehoseLogStream
Outputs:
  KinesisFirehoseStream:
    Description: Firehose stream that will receive the real time data being generated by KDG.
    Value: !GetAtt rFirehoseRole.Arn
  LambdaFunction:
    Description: Lambda function that handles the 
    Value: !GetAtt rLambdaTransformation.Arn
  LambdaIAMRole:
    Description: IAM role created for the Lambda function
    Value: !GetAtt rLambdaRole.Arn
  FirehoseDeliveryRole:
    Description: IAM role created for the Firehose stream
    Value: !GetAtt rFirehoseRole.Arn
  S3BucketRaw:
    Description: Bucket that will receive the unredacted data.
    Value: !Ref rS3BucketRaw
  S3BucketRedacted:
    Description: Bucket that will receive the redacted data.
    Value: !Ref rS3BucketRedacted
