AWSTemplateFormatVersion: 2010-09-09
Description: Invoking Amazon Rekognition Content Moderation via API Gateway

Parameters:
  apiGatewayName:
    Type: String
    Default: ContentModeration
  apiGatewayStageName:
    Type: String
    AllowedPattern: "[a-z0-9]+"
    Default: production
  apiGatewayHTTPMethod:
    Type: String
    Default: POST
  lambdaFunctionName:
    Type: String
    AllowedPattern: "[a-zA-Z0-9]+[a-zA-Z0-9-]+[a-zA-Z0-9]+"
    Default: invokeContentModeration

Resources:
  apiGateway:
    Type: AWS::ApiGateway::RestApi
    Properties:
      Description: API Gateway
      EndpointConfiguration:
        Types:
          - REGIONAL
      Name: !Ref apiGatewayName

  apiGatewayRootMethod:
    Type: AWS::ApiGateway::Method
    Properties:
      AuthorizationType: NONE
      HttpMethod: !Ref apiGatewayHTTPMethod
      Integration:
        IntegrationHttpMethod: POST
        Type: AWS_PROXY
        Uri: !Sub
          - arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${lambdaArn}/invocations
          - lambdaArn: !GetAtt lambdaFunction.Arn
      ResourceId: !GetAtt apiGateway.RootResourceId
      RestApiId: !Ref apiGateway

  apiGatewayDeployment:
    Type: AWS::ApiGateway::Deployment
    DependsOn:
      - apiGatewayRootMethod
    Properties:
      RestApiId: !Ref apiGateway
      StageName: !Ref apiGatewayStageName

  lambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          import json
          import boto3
          import urllib3
          import io


          client = boto3.client('rekognition')
          manager = urllib3.PoolManager()


          def getModerationForUrl(url):
              try:
                  extensions = ['jpg', 'jpeg', 'png']
                  if not any(url.lower().endswith(ext) for ext in extensions):
                      return 400, "Amazon Rekognition supports only the following image formats: jpg, jpeg, png"


                  response = manager.request('GET', url, preload_content=False)
                  if response.status == 404:
                      return 404, "Image not found"

                  try:
                      reader = io.BufferedReader(response, 8)
                      readBytes = reader.read()
                  finally:
                      if(reader is not None):
                          reader.close()

                  if(len(readBytes) > 5242880):
                      return 400, "Amazon Rekognition does not support images more than 5MB in this implementation. Use images stored on Amazon S3. See here: https://docs.aws.amazon.com/rekognition/latest/dg/limits.html"

                  response = client.detect_moderation_labels(Image={'Bytes': readBytes}, MinConfidence=60)
                  return 200, response['ModerationLabels']

              except Exception as e:
                  return 503, "Unexpected error: " + str(e)

          def lambda_handler(event, context):
              print(f'event: {json.dumps(event)}')

              body = event.get('body')
              if body is None:
                  raise KeyError("payload is missing")

              url = json.loads(body)['url']

              if url is None:
                  raise KeyError("url is missing from the payload")

              moderationResponse = getModerationForUrl(url)
              print(f'returning moderationResponse: {json.dumps(moderationResponse)}')

              return {
                  'statusCode': moderationResponse[0],
                  'body': json.dumps(moderationResponse[1])
              }


      Description: Lambda function
      FunctionName: !Ref lambdaFunctionName
      Handler: index.lambda_handler
      Timeout: 35
      MemorySize: 128
      Role: !GetAtt lambdaIAMRole.Arn
      Runtime: python3.11

  lambdaApiGatewayInvoke:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !GetAtt lambdaFunction.Arn
      Principal: apigateway.amazonaws.com
      SourceArn: !Sub arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${apiGateway}/${apiGatewayStageName}/${apiGatewayHTTPMethod}/

  lambdaIAMRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Action:
              - sts:AssumeRole
            Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
      Policies:
        - PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Effect: Allow
                Resource:
                  - !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${lambdaFunctionName}:*
          PolicyName: lambdaLogs
        - PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Action:
                  - rekognition:DetectModerationLabels
                Effect: Allow
                Resource:
                  - "*"
          PolicyName: lambdaRekognition

  lambdaLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${lambdaFunctionName}
      RetentionInDays: 90

Outputs:
  apiGatewayInvokeURL:
    Value: !Sub https://${apiGateway}.execute-api.${AWS::Region}.amazonaws.com/${apiGatewayStageName}
