CFNAccountCreditTracking.yaml

AWSTemplateFormatVersion: 2010-09-09
Description: >-
  This template builds an AWS account credit tracking stack with a DynamoDB
  table, SNS notifications, and a scheduled Lambda function. Customizable
  parameters include the credit award name, effective and expiry dates, credit
  amount, and alert threshold.
Parameters:
  SNSNotificationTopicName:
    Type: String
    Description: >-
      Name of the SNS topic that will be used to send notifications when the
      utilization threshold is reached.
    Default: aws-credit-status
  SNSNotificationEmail:
    Type: String
    Description: >-
      Email address for receiving alerts when the credit balance falls below the
      utilization threshold.
  CreditsTableName:
    Type: String
    Description: Name of the DynamoDB table that will store credit status information.
    Default: aws-credit-status
  CreditAwardName:
    Type: String
    Description: Name of the credit award to track.
  CreditEffectiveDate:
    Type: String
    Description: >-
      Date on which the credit award becomes effective. This parameter should be
      a string in the format 'YYYY-MM-DD'.
  CreditExpiry:
    Type: String
    Description: >-
      The date on which the credit award expires. This parameter should be a
      string in the format 'YYYY-MM-DD'.
  CreditAmount:
    Type: Number
    Description: Amount of credit awarded.
    Default: '500.0'
  AlertThreshold:
    Type: Number
    Description: >-
      The threshold at which an alert is triggered when the credit balance falls
      below this amount.
    Default: '400'
  LambdaEventTriggerCron:
    Type: String
    Description: >-
      Cron expression for scheduling the Lambda function to be triggered. This
      parameter should be a valid cron expression that specifies the frequency
      at which the function should be triggered. The default value is set to
      trigger the function every day at midnight (12:00 AM UTC).
    Default: cron(0 0 * * ? *)
Metadata:
  'AWS::CloudFormation::Interface':
    ParameterGroups:
      - Label:
          default: AWS Credits Table
        Parameters:
          - CreditsTableName
          - CreditAwardName
          - CreditEffectiveDate
          - CreditExpiry
          - CreditAmount
      - Label:
          default: Alert Parameters
        Parameters:
          - AlertThreshold
          - LambdaEventTriggerCron
      - Label:
          default: Notification Parameters
        Parameters:
          - SNSNotificationTopicName
          - SNSNotificationEmail
Resources:
  CreditsTable:
    Type: 'AWS::DynamoDB::Table'
    Properties:
      AttributeDefinitions:
        - AttributeName: credit_name
          AttributeType: S
      BillingMode: PAY_PER_REQUEST
      KeySchema:
        - AttributeName: credit_name
          KeyType: HASH
      PointInTimeRecoverySpecification:
        PointInTimeRecoveryEnabled: true
      TableName: !Ref CreditsTableName
      Tags:
        - Key: Project
          Value: AWS Sandbox
  AlarmAlertTargetSNS:
    Type: 'AWS::SNS::Topic'
    Properties:
      DisplayName: AWS Credits Utilization Notification
      TopicName: !Ref SNSNotificationTopicName
  AlarmSNSSubscription:
    Type: 'AWS::SNS::Subscription'
    Properties:
      Endpoint: !Ref SNSNotificationEmail
      Protocol: email
      TopicArn: !Ref AlarmAlertTargetSNS
  InitializeCreditTable:
    Type: 'AWS::CloudFormation::CustomResource'
    DependsOn: CreditsTable
    Properties:
      TableName: !Ref CreditsTableName
      ServiceToken: !GetAtt 
        - CreditsTableLambdaInitializer
        - Arn
      NameField: !Ref CreditAwardName
      EffectiveDate: !Ref CreditEffectiveDate
      ExpiryDate: !Ref CreditExpiry
      CreditAmount: !Ref CreditAmount
      AlertThreshold: !Ref AlertThreshold
      AlertNotificationTopic: !Ref AlarmAlertTargetSNS
  AWSCreditsTrackingCustomPolicy:
    Type: 'AWS::IAM::ManagedPolicy'
    Properties:
      Description: >-
        Policy allowing Lambda function to call Cost Explorer and access
        DynamoDB
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Sid: VisualEditor0
            Effect: Allow
            Action:
              - 'sns:Publish'
              - 'dynamodb:GetItem'
              - 'ce:GetCostAndUsage'
              - 'dynamodb:UpdateItem'
            Resource: '*'
  CreditTrackingLambdaRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: CreditTrackingTLambdaRole
      Description: >-
        Role to execute Lambda function to query Cost Explorer Server, calculate
        credits used, and send notification
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
        - !Ref AWSCreditsTrackingCustomPolicy
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
  CreditTrackingTableInitializerLambdaRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: CreditTrackingTableInitializerLambdaRole
      Description: Role to execute Lambda function to initialize credit tracking table
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
        - 'arn:aws:iam::aws:policy/AmazonDynamoDBFullAccess'
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
  CreditsTableLambdaInitializer:
    Type: 'AWS::Lambda::Function'
    Properties:
      Description: Lambda function to populate credits table with one record
      Role: !GetAtt 
        - CreditTrackingTableInitializerLambdaRole
        - Arn
      FunctionName: aws_credit_table_initializer
      Handler: index.lambda_handler
      Runtime: python3.8
      Code:
        ZipFile: !Join 
          - '\n'
          - - | 
              # MIT License
              #
              # Copyright (c) 2022 Allen Brodjeski
              #
              # Permission is hereby granted, free of charge, to any person obtaining a copy
              # of this software and associated documentation files (the "Software"), to deal
              # in the Software without restriction, including without limitation the rights
              # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
              # copies of the Software, and to permit persons to whom the Software is
              # furnished to do so, subject to the following conditions:
              #
              # The above copyright notice and this permission notice shall be included in all
              # copies or substantial portions of the Software.
              #
              # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
              # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
              # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
              # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
              # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
              # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
              # SOFTWARE.
              import json
              import urllib.request
              import boto3
              from botocore.exceptions import ClientError


              def lambda_handler(event, context):

                  # Log some information about the event and context
                  print('Initializing DynamoDB table')
                  print(json.dumps(event))
                  print(context)

                  # Get the request type and response URL from the event
                  request_type = event['RequestType']
                  response_url = event['ResponseURL']

                  # Get the resource properties and logical resource ID from the event
                  properties = event['ResourceProperties']
                  logical_resource_id = event['LogicalResourceId']

                  # Initialize an empty dictionary to store the response data
                  responseData = {}

                  if request_type == 'Create':
                      # Set default values for the response data
                      responseData['Status'] = 'SUCCESS'
                      responseData['Reason'] = 'SUCCESS'
                      responseData['PhysicalResourceId'] = 'InitializedTable'
                      responseData['StackId'] = event['StackId']
                      responseData['RequestId'] = event['RequestId']
                      responseData['LogicalResourceId'] = logical_resource_id

                      # Try to create the DynamoDB table
                      try:
                          # Get the DynamoDB client
                          ddb_client = boto3.client('dynamodb')
                          # Log the resource properties
                          print(properties['NameField'])
                          print(properties['EffectiveDate'])
                          print(properties['ExpiryDate'])
                          print(properties['CreditAmount'])
                          print(properties['AlertThreshold'])
                          print(properties['AlertNotificationTopic'])

                          # Put the item in the DynamoDB table
                          response = ddb_client.put_item(
                              TableName=properties['TableName'],
                              Item={
                                  'credit_name': {'S': properties['NameField']},
                                  'credit_start': {'S': properties['EffectiveDate']},
                                  'credit_expiry': {'S': properties['ExpiryDate']},
                                  'credit_limit': {'N': str(properties['CreditAmount'])},
                                  'credit_threshold': {'N': str(properties['AlertThreshold'])},
                                  'credit_notification_topic': {'S': properties['AlertNotificationTopic']}
                              }
                          )
                          print(response)

                      # Catch any errors that occur
                      except ClientError as e:
                          # Log the error
                          logging.error(e)

                          # Set the status and reason in the response data
                          responseData['Status'] = 'FAILED'
                          responseData['Reason'] = e.response['Error']['Message']

                      # Convert the response data to JSON
                      json_response = json.dumps(responseData)

                      print('SENDING RESPONSE')
                      print(json_response)

                      # Send the response
                      try:
                          print('Response URL = ' + response_url)
                          response = urllib.request.urlopen(urllib.request.Request(
                              url=response_url,
                              data=bytes(json_response, 'utf-8'),
                              method='PUT'
                          ),
                              timeout=5
                          )
                      except Exception as e:
                          print('send failed executing urllib.request: {e}')

                  # If the request type is 'Update', then update the DynamoDB table
                  elif request_type == 'Update':
                      # Set default values for the response data
                      responseData['Status'] = 'SUCCESS'
                      responseData['Reason'] = 'SUCCESS'
                      responseData['PhysicalResourceId'] = 'InitializedTable'
                      responseData['StackId'] = event['StackId']
                      responseData['RequestId'] = event['RequestId']
                      responseData['LogicalResourceId'] = logical_resource_id
                      try:
                          # Get the DynamoDB client
                          ddb_client = boto3.client('dynamodb')

                          # Log the resource properties
                          print(properties['NameField'])
                          print(properties['EffectiveDate'])
                          print(properties['ExpiryDate'])
                          print(properties['CreditAmount'])
                          print(properties['AlertThreshold'])
                          print(properties['AlertNotificationTopic'])

                          # Put the item in the DynamoDB table
                          response = ddb_client.put_item(
                              TableName=properties['TableName'],
                              Item={
                                  'credit_name': {'S': properties['NameField']},
                                  'credit_start': {'S': properties['EffectiveDate']},
                                  'credit_expiry': {'S': properties['ExpiryDate']},
                                  'credit_limit': {'N': str(properties['CreditAmount'])},
                                  'credit_threshold': {'N': str(properties['AlertThreshold'])},
                                  'credit_notification_topic': {'S': properties['AlertNotificationTopic']}
                              }
                          )
                          print(response)
                      except ClientError as e:
                          # Log the error
                          logging.error(e)

                          # Set the status and reason in the response data
                          responseData['Status'] = 'FAILED'
                          responseData['Reason'] = e.response['Error']['Message']

                      # Convert the response data to JSON
                      json_response = json.dumps(responseData)

                      # Send the response
                      try:
                          response = urllib.request.urlopen(urllib.request.Request(
                              url=response_url,
                              data=bytes(json_response, 'utf-8'),
                              method='PUT'
                          ),
                              timeout=5
                          )
                      except Exception as e:
                          logging.error('send failed executing urllib.request: {e}')

                  # If the request type is 'Delete'
                  elif request_type == 'Delete':
                      # Set default values for the response data
                      responseData['Status'] = 'SUCCESS'
                      responseData['Reason'] = 'SUCCESS'
                      responseData['PhysicalResourceId'] = 'InitializedTable'
                      responseData['StackId'] = event['StackId']
                      responseData['RequestId'] = event['RequestId']
                      responseData['LogicalResourceId'] = logical_resource_id

                      json_response = json.dumps(responseData)

                      # Send the response
                      try:
                          response = urllib.request.urlopen(urllib.request.Request(
                              url=response_url,
                              data=bytes(json_response, 'utf-8'),
                              method='PUT'
                          ),
                              timeout=5
                          )
                      except Exception as e:
                          logging.error('send failed executing urllib.request: {e}')
                  return responseData

  CreditsTableLambdaTracker:
    Type: 'AWS::Lambda::Function'
    Properties:
      Description: >-
        Lambda function to query Cost Explorer API, sum up credits used and send
        notification if threshold is reached
      Role: !GetAtt 
        - CreditTrackingLambdaRole
        - Arn
      FunctionName: aws-account-credits-tracking
      Handler: index.lambda_handler
      Runtime: python3.8
      Environment:
        Variables:
          AWS_CREDITS_NAME: !Ref CreditAwardName
          AWS_CREDITS_TABLE_NAME: !Ref CreditsTableName
      Code:
        ZipFile: !Join 
          - '\n'
          - - | 
              #MIT License
              #
              #Copyright (c) 2022 Allen Brodjeski
              #
              #Permission is hereby granted, free of charge, to any person obtaining a copy
              #of this software and associated documentation files (the "Software"), to deal
              #in the Software without restriction, including without limitation the rights
              #to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
              #copies of the Software, and to permit persons to whom the Software is
              #furnished to do so, subject to the following conditions:
              #
              #The above copyright notice and this permission notice shall be included in all
              #copies or substantial portions of the Software.
              #
              #THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
              #IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
              #FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
              #AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
              #LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
              #OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
              #SOFTWARE.
              import os
              import boto3
              from datetime import datetime
              from botocore.exceptions import ClientError

              # Get the required environment variables
              AWS_CREDITS_NAME = os.environ.get('AWS_CREDITS_NAME')
              AWS_CREDITS_TABLE_NAME = os.environ.get('AWS_CREDITS_TABLE_NAME')

              # Check if the required environment variables are set
              if not AWS_CREDITS_NAME:
                  raise ValueError("Missing required environment variable 'AWS_CREDITS_NAME'")

              if not AWS_CREDITS_TABLE_NAME:
                  raise ValueError("Missing required environment variable 'AWS_CREDITS_TABLE_NAME'")

              class Credit:

                  def __init__(self, credit_name: str) -> None:
                      print('Building Credit from DynamoDB')
                      self.ddb_client = boto3.client('dynamodb')
                      self.credit_start = ''
                      self.credit_expiry = ''
                      self.credit_threshold = 0
                      self.credit_limit = 0
                      self.credit_notification_topic = ''

                      print(f'Credit Name: {credit_name}')
                      try:
                          # Get the credit data from the DynamoDB table
                          response = self.ddb_client.get_item(
                              TableName=AWS_CREDITS_TABLE_NAME,
                              Key={'credit_name': {'S': credit_name}},
                              AttributesToGet=[
                                  'credit_start',
                                  'credit_expiry',
                                  'credit_limit',
                                  'credit_threshold',
                                  'credit_notification_topic'
                              ]
                          )
                          # Extract the credit data from the response
                          self.credit_start = response['Item']['credit_start']['S']
                          self.credit_expiry = response['Item']['credit_expiry']['S']
                          self.credit_limit = float(response['Item']['credit_limit']['N'])
                          self.credit_threshold = float(response['Item']['credit_threshold']['N'])
                          self.credit_notification_topic = response['Item']['credit_notification_topic']['S']

                      except ClientError as e:
                          if e.response['Error']['Code'] == 'ResourceNotFoundException':
                              # If the DynamoDB table does not exist, raise a ValueError
                              raise ValueError(f"DynamoDB table '{AWS_CREDITS_TABLE_NAME}' does not exist")
                          else:
                              raise
                      except ValueError as e:
                          logging.error(e)
                          raise

              def lambda_handler(event, context):
                  total_cost = 0.0

                  credit_table = Credit(AWS_CREDITS_NAME)

                  cost_client = boto3.client('ce')

                  account_id = boto3.client('sts').get_caller_identity().get('Account')
                  today = datetime.today().strftime('%Y-%m-%d')
                  print(f'Today: {today}')
                  print(f'Start time: {credit_table.credit_start}')
                  response = cost_client.get_cost_and_usage(
                      TimePeriod={
                          'Start': credit_table.credit_start,
                          'End': today
                      },
                      Granularity='DAILY',
                      Filter={
                          'Dimensions': {
                              'Key': 'RECORD_TYPE',
                              'Values': ['Credit'],
                              'MatchOptions': ['EQUALS']
                          }
                      },
                      Metrics=['UnblendedCost']
                  )

                  print(response)

                  for credit in response['ResultsByTime']:
                      total_cost += float(credit['Total']['UnblendedCost']['Amount'])

                  print(f'Total cost: {total_cost}')
                  credit_remaining = credit_table.credit_limit + total_cost

                  print(f'Credit threshold: {credit_table.credit_threshold}')
                  utilization = total_cost / credit_table.credit_limit * -100.0

                  if credit_remaining < credit_table.credit_threshold:
                      print('Sending utilization alert')
                      sns_client = boto3.client('sns')

                      response = sns_client.publish(
                          TopicArn=credit_table.credit_notification_topic,
                          Subject='AWS Account Credits Utilization Alert. Please see information below:',
                          Message=f'Account ID: {account_id}'
                                  f"\n\nCredits remaining: {credit_remaining:.2f}"
                                  f"\n\nUtilization: {utilization:.2f}%"
                      )
                      print(f'SNS response: {response}')
                  else:
                      print(f'Credit remaining: {credit_remaining}')

                  return {
                      'utilization': utilization,
                      'credit_remainning' : credit_remaining
                  }
  CreditTrackingLambdaEventPermission:
    Type: 'AWS::Lambda::Permission'
    Properties:
      Action: 'lambda:InvokeFunction'
      FunctionName: aws-account-credits-tracking
      Principal: events.amazonaws.com
      SourceArn: !GetAtt 
        - CreditThresholdLambdaTimedTriggerEvent
        - Arn
  CreditThresholdLambdaTimedTriggerEvent:
    Type: 'AWS::Events::Rule'
    Properties:
      Description: Event that triggers the credit threshold Lambda function
      Name: aws-account-credits-tracking-trigger
      ScheduleExpression: !Ref LambdaEventTriggerCron
      State: ENABLED
      Targets:
        - Arn: !Sub >-
            arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:aws-account-credits-tracking
          Id: CreditThresholdLambdaFunction