CloudFront Based Token Revocation System (only token_revoked may be called by CloudFront, responses will be cached)

token_revoked.py

from typing import Dict
import boto3


def lambda_handler(event: Dict, context):
    try:
        path: str = event['path']
        if path.startswith('/api/token/revoked/'):
            token_id: str = path.split('/')[4]

            response: Dict = get_table().get_item(
                Key={'Id': token_id},
                AttributesToGet=['Id'],
                ConsistentRead=False
            )

            if "Item" in response:
                result = "TRUE"
            else:
                result = "FALSE"

            return {
                "statusCode": 200,
                "headers": {
                    "Content-Type": "text/plain",
                    "Cache-Control": f"max-age=31536000"
                },
                "body": result
            }
        else:
            return {
                "statusCode": 400,
                "headers": {
                    "Content-Type": "text/plain",
                    "Cache-Control": "no-cache"
                },
                "body": f"Invalid path '{path}'"
            }
    except Exception as e:
        print(e)
        return {
            "statusCode": 500,
            "headers": {
                "Content-Type": "text/plain",
                "Cache-Control": "no-cache"
            },
            "body": f"Unexpected error."
        }


def get_table():
    return boto3.resource('dynamodb').Table("Table-APITokenRevoked")

token_revoked_admin.py

import time
from typing import Dict
import boto3


def lambda_handler(event: Dict, context):
    try:
        path: str = event['path']
        if path.startswith('/api/token/admin/'):
            path_elements = path.split('/')
            action: str = path_elements[4]
            token_id: str = path_elements[5]

            if action == "revoke":
                get_table().put_item(
                    Item={'Id': token_id}
                )
            elif action == "unrevoke":
                get_table().delete_item(
                    Key={'Id': token_id}
                )
            else:
                return {
                    "statusCode": 400,
                    "headers": {
                        "Content-Type": "text/plain",
                        "Cache-Control": "no-cache"
                    },
                    "body": f"Invalid action '{action}'. Only revoke or unrevoke are supported actions."
                }

            try:
                # Invalidate the CloudFront cache for this item
                epoch_time = f'{int(time.time())}'
                cloudfront_client = boto3.client('cloudfront')
                cloudfront_client.create_invalidation(
                    DistributionId='E109A3SUUOZZRV',
                    InvalidationBatch={
                        'Paths': {
                            'Quantity': 1,
                            'Items': [
                                f'/{token_id}',
                            ]
                        },
                        'CallerReference': epoch_time
                    }
                )
                cloud_front_failed = ""
            except Exception as e:
                print(e)
                cloud_front_failed = "-CF"

            return {
                "statusCode": 200,
                "headers": {
                    "Content-Type": "text/plain",
                    "Cache-Control": "no-cache"
                },
                "body": f"DONE{cloud_front_failed}"
            }

        else:
            return {
                "statusCode": 400,
                "headers": {
                    "Content-Type": "text/plain",
                    "Cache-Control": "no-cache"
                },
                "body": f"Invalid path '{path}'"
            }
    except Exception as e:
        print(e)
        return {
            "statusCode": 500,
            "headers": {
                "Content-Type": "text/plain",
                "Cache-Control": "no-cache"
            },
            "body": f"Unexpected error."
        }


def get_table():
    return boto3.resource('dynamodb').Table("Table-APITokenRevoked")

upload_cloudformation.yml

AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: API Revocation Service (fronted by ALB)

Globals:
  Function:
    Timeout: 3

Resources:

## Public API to Check for Revocation ##
  ApiFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: api_token_revoked/app.py
      Handler: app.lambda_handler
      Runtime: python3.7
      Policies:
        - DynamoDBReadPolicy:
            TableName: !Ref DynamoTable

  InvokePermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt ApiFunction.Arn
      Action: 'lambda:InvokeFunction'
      Principal: elasticloadbalancing.amazonaws.com

  TargetGroup:
    Type: AWS::ElasticLoadBalancingV2::TargetGroup
    DependsOn: InvokePermission
    Properties:
      TargetType: lambda
      Targets:
        - Id: !GetAtt ApiFunction.Arn

  ListenerRule:
    Type: AWS::ElasticLoadBalancingV2::ListenerRule
    Properties:
      Actions:
        - TargetGroupArn: !Ref TargetGroup
          Type: forward
      Conditions:
        - Field: path-pattern
          Values: ['/api/token/revoked/*']
      ListenerArn: !ImportValue ALBHTTPSListernArn
      Priority: 8272

## Private Admin Lambda to Revoke Tokens ##
  AdminApiFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: api_token_revoked/app_admin.py
      Handler: app_admin.lambda_handler
      Runtime: python3.7
      Policies:
        - DynamoDBCrudPolicy:
            TableName: !Ref DynamoTable
        - Statement:
          - Sid: CloudFrontInvalidatePolicy
            Effect: Allow
            Action:
            - cloudfront:CreateInvalidation
            - cloudfront:GetInvalidation
            - cloudfront:ListInvalidations
            Resource: '*'

  AdminTargetGroup:
    Type: AWS::ElasticLoadBalancingV2::TargetGroup
    DependsOn: AdminInvokePermission
    Properties:
      TargetType: lambda
      Targets:
        - Id: !GetAtt AdminApiFunction.Arn

  AdminInvokePermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt AdminApiFunction.Arn
      Action: 'lambda:InvokeFunction'
      Principal: elasticloadbalancing.amazonaws.com

  AdminListenerRule:
    Type: AWS::ElasticLoadBalancingV2::ListenerRule
    Properties:
      Actions:
        - TargetGroupArn: !Ref AdminTargetGroup
          Type: forward
      Conditions:
        - Field: path-pattern
          Values: ['/api/token/admin/*']
      ListenerArn: !ImportValue ALBHTTPSListernArn
      Priority: 8273

## Token Revocation Database ##
  DynamoTable:
    Type: AWS::DynamoDB::Table
    Properties:
      TableName: !Sub Table-${AWS::StackName}
      AttributeDefinitions:
        - AttributeName: Id
          AttributeType: S
      KeySchema:
        - AttributeName: Id
          KeyType: HASH
      ProvisionedThroughput:
        ReadCapacityUnits: 5
        WriteCapacityUnits: 5