---
AWSTemplateFormatVersion: "2010-09-09"
Transform: AWS::Serverless-2016-10-31

Parameters:
  CognitoUserPoolID:
    Type: String
    Default: ''
    Description: Cognito User Pool ID

Resources:
  CognitoEBQueue:
    Type: AWS::SQS::Queue
    Properties:
      QueueName: CognitoEBQueue
      KmsMasterKeyId: alias/aws/sqs
      ReceiveMessageWaitTimeSeconds: 20 
      VisibilityTimeout: 105  
  
  CognitoEBQueuePolicy:
    Type: AWS::SQS::QueuePolicy
    Properties:
      Queues:
        - !Ref CognitoEBQueue
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Sid: cognitopolicy
          Effect: Allow
          Principal: 
            Service: 
            - events.amazonaws.com
          Action:
          - 'sqs:SendMessage'
          - 'sqs:ReceiveMessage'
          Resource: !GetAtt CognitoEBQueue.Arn
            
  CognitoSignInFunctionInvokeConfig:
    Type: AWS::Lambda::EventInvokeConfig
    Properties:
      FunctionName: !Ref CognitoAuthEvents
      Qualifier: "$LATEST"
      MaximumRetryAttempts: 0
      MaximumEventAgeInSeconds: 21600

  LambdaFunctionEventSourceMapping:
    Type: AWS::Lambda::EventSourceMapping
    Properties:
      Enabled: true
      EventSourceArn: !GetAtt CognitoEBQueue.Arn
      FunctionName: !GetAtt CognitoAuthEvents.Arn

  CognitoSignInRule:
    Type: AWS::Events::Rule
    Properties:
      Name: CognitoSignInRule
      
      EventPattern:
       source: 
        - aws.cognito-idp 
       detail-type: 
        - 'AWS API Call via CloudTrail'
       detail:
          eventSource:
            - cognito-idp.amazonaws.com
          eventName: 
            - InitiateAuth
      Targets:
        - Id: "CognitoEBQueue"
          Arn: !GetAtt CognitoEBQueue.Arn

  CognitoLambdaExecutionRole:
    Type: "AWS::IAM::Role"
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: "sts:AssumeRole"
      Policies:
        - PolicyName: CognitoAuthEventsRole
          PolicyDocument: 
            Version: "2012-10-17"
            Statement:
              - Action:
                  - 'logs:CreateLogGroup'
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
                  - 's3:ListBucket'
                  - 's3:GetBucketLocation'
                  - 's3:GetObject'
                  - 's3:PutObject'
                  
                Resource: "*"  
                Effect: Allow
              - Action:
                  - 'sqs:DeleteMessage'
                  - 'sqs:GetQueueAttributes'
                  - 'sqs:ReceiveMessage'
                Resource: !GetAtt CognitoEBQueue.Arn
                Effect: Allow 
              - Action:
                  - 's3:ListBucket'
                  - 's3:GetBucketLocation'
                  - 's3:GetObject'
                  - 's3:PutObject'
                Resource: !GetAtt LogExtensionsBucket.Arn
                Effect: Allow 
    Metadata:
      cfn_nag:
        rules_to_suppress:
          - id: W11
            reason: Resource * is required due to the administrative nature of the solution.
  
  S3LogExtensionsLayer:
    Type: AWS::Lambda::LayerVersion
    Properties:
      Description: Layer containing extension(s)
      Content:
        S3Bucket: awsiammedia
        S3Key: public/sample/1209-how-to-automate-AWS-WAF/S3LogExtensionsLayer.zip
      CompatibleRuntimes:
      - python3.8
      LicenseInfo: Available under the MIT-0 license
  
  LogExtensionsBucket:
    Type: AWS::S3::Bucket
    Properties:
      AccessControl: LogDeliveryWrite 
      BucketEncryption: 
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
      LifecycleConfiguration:
        Rules:
        - Id: DeleteAfterSevenDays
          Status: Enabled
          ExpirationInDays: 7
    Metadata:
      cfn_nag:
        rules_to_suppress:
          - id: W51
            reason: Bucket is a logging bucket
          - id: W35
            reason: Bucket has acl set for log delivery

  CognitoAuthEvents:
    Type: AWS::Lambda::Function
    Properties:
      Description: Cognito authentication events function
      Runtime: python3.8
      Role: !GetAtt CognitoLambdaExecutionRole.Arn
      Handler: index.lambda_handler
      MemorySize: 128
      ReservedConcurrentExecutions: 10
      Timeout: 100
      Code: 
        ZipFile: |
          import boto3 
          import json
          import os
          import logging
         

          cognito_records = boto3.client('cognito-idp')
          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
        


          def lambda_handler(event, context):
          
              PoolID = os.environ['COGNITO_POOL']

              #For every cloudtrail event extract the sub identifier
              try:
                  for record in event['Records']:
                      body = json.loads(record["body"])
                      #get sub from cloudtrail event to get the username with API call
                      userid = body['detail']['additionalEventData']['sub']
                    
              except Exception as e:
                  # Send some context about this error to Lambda Logs
                  print(e)
                  # throw exception, do not handle. Lambda will make message visible again.
                  raise e
                  
              #Get username using the sub
            
              authusers = cognito_records.list_users(
                  UserPoolId=PoolID,
                  AttributesToGet=['sub',],
                  Filter=f"sub^=\"{userid}\"")
                  
              for authuser in authusers['Users']:
                  authusername = authuser['Username']
                  auth_events = cognito_records.admin_list_user_auth_events(       #Get authentication events for the users, limit to 5
                      UserPoolId=PoolID, 
                      Username=authusername, 
                      MaxResults=5)
                  logging.info("auth events, %s", auth_events) #Send the user events to the logs
                  
              return {
                  'statusCode': 200
              }
      Environment:
        Variables:
          S3_BUCKET_NAME: !Ref LogExtensionsBucket
          COGNITO_POOL: !Ref CognitoUserPoolID
      Layers: 
        - !Ref S3LogExtensionsLayer
    Metadata:
      cfn_nag:
        rules_to_suppress:
          - id: W89
            reason: Lambda is a target for SQS

   # CloudWatch Logs Insights query - CompromisedCredentialsDetected
  compromisedCredentialDetectedQueryDefinition:
    Type: AWS::Logs::QueryDefinition
    Properties:
      Name: "compromisedCredentialQueryDefinition"
      LogGroupNames:
          - yourlambdaloggroup
      QueryString: 'fields @message
      | filter @message like /INFO/
      | filter AuthEvents.0.EventType like "SignIn"
      | filter AuthEvents.0.EventRisk.RiskDecision like "AccountTakeover" and AuthEvents.0.EventRisk.CompromisedCredentialsDetected =! "false"
      | stats count(*) as RequestsperIP by AuthEvents.0.EventContextData.IpAddress as IP
      | sort desc
      | limit 10'

  # CloudWatch Logs Insights query - RiskLevelHigh
  riskLevelHighQueryDefinition:
    Type: AWS::Logs::QueryDefinition
    Properties:
      Name: "riskLevelHighQueryDefinition"
      LogGroupNames:
          - yourlambdaloggroup
      QueryString: 'fields @message
      | filter @message like /INFO/
      | filter AuthEvents.0.EventRisk.RiskLevel like "High"
      | stats count(*) as RequestsperIP by AuthEvents.0.EventContextData.IpAddress as IP, AuthEvents.0.EventType as EventType
      | sort desc
      | limit 10'

  # CloudWatch Logs Insights query - RiskLevelContext
  riskLevelContextQueryDefinition:
    Type: AWS::Logs::QueryDefinition
    Properties:
      Name: "riskLevelContextQueryDefinition"
      LogGroupNames:
          - yourlambdaloggroup
      QueryString: 'fields @message
      | filter @message like /INFO/
      | filter AuthEvents.0.EventRisk.RiskLevel like "High"
      |stats count(*) as RequestsperCountry by AuthEvents.0.EventContextData.Country as Country
      | sort desc'


  
Outputs:
  ExtensionsLayer:
    Value:
      Ref: S3LogExtensionsLayer
    Description: Log Extension Layer ARN
  Function:
    Value:
      Ref: CognitoAuthEvents
    Description: Lambda Function
  LogExtensionsBucketName:
    Value:
      Ref: LogExtensionsBucket