template.yaml

AWSTemplateFormatVersion: '2010-09-09'
Transform:
  - GenerateNgwAlarms
  - AWS::Serverless-2016-10-31
Description: SAM template automatic VPC Flow Log recording

Parameters:
  VPCFlowLogRetentionTimeDays:
    Description: The time in days to retain the log group
    Type: Number
    Default: 7
    AllowedValues: [1,3,5,7,14,30,60,90,120,150,180,365,400,545,731,1827,3653]
    ConstraintDescription: Must be one of the allowed values
      [1,3,5,7,14,30,60,90,120,150,180,365,400,545,731,1827,3653]

  LambdaLogGroupRetentionTimeDays:
    Description: The time in days to retain the Lambda log group
    Type: Number
    Default: 7
    AllowedValues: [1,3,5,7,14,30,60,90,120,150,180,365,400,545,731,1827,3653]
    ConstraintDescription: Must be one of the allowed values
      [1,3,5,7,14,30,60,90,120,150,180,365,400,545,731,1827,3653]

  EmailAddress:
    Type: String
    Description: The email address to subscribe to the SNS Topic
  
  NatGatewayIDs:
    Type: String
    Description: A comma separated list of the IDs of the NAT Gateways to alarm on

  SamplingPeriod:
    Type: Number
    Default: 60
    MinValue: 10
    Description: The sampling period in seconds

  EvaluationPeriods:
    Type: Number
    Default: 3
    MinValue: 1
    Description: The number of evaluation periods

  Threshold:
    Type: Number
    Default: 100000000
    MinValue: 1000
    Description: The threshold for the alarm
  
  DatapointsToAlarm:
    Type: Number
    Default: 1
    MinValue: 1
    Description: The number of datapoints that must be breaching the threshold to trigger the alarm
  
  RecordingTime:
    Type: Number
    Default: 900
    MinValue: 0
    Description: Number of seconds to record traffic. 0 means to record until the CloudWatch Alarm returns to OK.

  NumOfRecordings:
    Type: Number
    Default: 0
    MinValue: 0
    Description: How many recordings to perform. 0 = whenever there is an Alarm. Above 0 will record only the specified number of Alarms.

Resources:
  NGWFlowlogsTable:
    Type: AWS::DynamoDB::Table
    Properties:
      TableName: !Sub
        - NGW-Flowlogs-${Suffix}
        - Suffix: !Select
            - 4
            - !Split
              - '-'
              - !Select
                - 2
                - !Split
                  - /
                  - !Ref AWS::StackId
      AttributeDefinitions:
        - AttributeName: ngwid
          AttributeType: S
      KeySchema:
        - AttributeName: ngwid
          KeyType: HASH
      BillingMode: PAY_PER_REQUEST

  VPCFlowLogsCloudWatchLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub
        - Auto-VPC-FlowLogs-${Suffix}
        - Suffix: !Select
            - 4
            - !Split
              - '-'
              - !Select
                - 2
                - !Split
                  - /
                  - !Ref AWS::StackId
      RetentionInDays: !Ref VPCFlowLogRetentionTimeDays
      LogGroupClass: INFREQUENT_ACCESS
    DeletionPolicy: Retain
    UpdateReplacePolicy: Retain

  ActivateVPCFlowLogFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !Sub
        - ActivateVPCFlowLog-${Suffix}
        - Suffix: !Select
            - 4
            - !Split
              - '-'
              - !Select
                - 2
                - !Split
                  - /
                  - !Ref AWS::StackId
      Handler: activate_vpc_flow_log.lambda_handler
      Runtime: python3.12
      Architectures:
        - arm64
      CodeUri: ./src-activate
      Policies:
        - AWSLambdaExecute
        - CloudWatchLogsFullAccess
        - Version: '2012-10-17'
          Statement:
            - Effect: Allow
              Action:
                - ec2:CreateFlowLogs
                - ec2:DescribeFlowLogs
                - ec2:DeleteFlowLogs
                - ec2:DescribeNatGateways
              Resource: '*'
            - Effect: Allow
              Action:
                - iam:PassRole
              Resource: !GetAtt VPCFlowLogRole.Arn
      Environment:
        Variables:
          LOG_GROUP_NAME: !Ref VPCFlowLogsCloudWatchLogGroup
          VPC_FLOW_LOG_ROLE: !GetAtt VPCFlowLogRole.Arn

  ActivateVPCFlowLogFunctionCloudWatchLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${ActivateVPCFlowLogFunction}
      RetentionInDays: !Ref LambdaLogGroupRetentionTimeDays
      LogGroupClass: INFREQUENT_ACCESS

  StoreWaitTokenFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !Sub
        - StoreWaitToken-${Suffix}
        - Suffix: !Select
            - 4
            - !Split
              - '-'
              - !Select
                - 2
                - !Split
                  - /
                  - !Ref AWS::StackId
      Handler: store_wait_token.lambda_handler
      Runtime: python3.12
      Architectures:
        - arm64
      CodeUri: ./src-store_wait_token
      Policies:
        - AWSLambdaExecute
        - CloudWatchLogsFullAccess
        - Version: '2012-10-17'
          Statement:
            - Effect: Allow
              Action:
                - dynamodb:UpdateItem
              Resource: !GetAtt NGWFlowlogsTable.Arn
            - Effect: Allow
              Action:
                - iam:PassRole
              Resource: !GetAtt VPCFlowLogRole.Arn

  StoreWaitTokenFunctionCloudWatchLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${StoreWaitTokenFunction}
      RetentionInDays: !Ref LambdaLogGroupRetentionTimeDays
      LogGroupClass: INFREQUENT_ACCESS

  AlarmEndedFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !Sub
        - VPCFlowLogs-AlarmEnded-${Suffix}
        - Suffix: !Select
            - 4
            - !Split
              - '-'
              - !Select
                - 2
                - !Split
                  - /
                  - !Ref AWS::StackId
      Handler: notify_alarm_ended.lambda_handler
      Runtime: python3.12
      Architectures:
        - arm64
      CodeUri: ./src-notify_alarm_ended
      Policies:
        - AWSLambdaExecute
        - CloudWatchLogsFullAccess
        - Version: '2012-10-17'
          Statement:
            - Effect: Allow
              Action:
                - dynamodb:GetItem
              Resource: !GetAtt NGWFlowlogsTable.Arn
            - Effect: Allow
              Action:
                - states:SendTaskSuccess
              Resource: !GetAtt StepFunction.Arn
            - Effect: Allow
              Action:
                - iam:PassRole
              Resource: !GetAtt VPCFlowLogRole.Arn
      Environment:
        Variables:
          TABLE_NAME: !Ref NGWFlowlogsTable         

  # Add a permission for the EventBridge rule "CloudWatchEventRuleAlarmOff" to invoke the Alarm Ended Lambda function
  AlarmEndedFunctionPermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt AlarmEndedFunction.Arn
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt CloudWatchEventRuleAlarmOff.Arn 

  AlarmEndedFunctionCloudWatchLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${AlarmEndedFunction}
      RetentionInDays: !Ref LambdaLogGroupRetentionTimeDays
      LogGroupClass: INFREQUENT_ACCESS

  DeleteVPCFlowLogFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !Sub
        - DeleteVPCFlowLog-${Suffix}
        - Suffix: !Select
            - 4
            - !Split
              - '-'
              - !Select
                - 2
                - !Split
                  - /
                  - !Ref AWS::StackId
      Handler: delete_vpc_flow_log.lambda_handler
      Runtime: python3.12
      Architectures:
        - arm64
      CodeUri: ./src-delete
      Policies:
        - AWSLambdaExecute
        - CloudWatchLogsFullAccess
        - Version: '2012-10-17'
          Statement:
            - Effect: Allow
              Action:
                - ec2:CreateFlowLogs
                - ec2:DescribeFlowLogs
                - ec2:DeleteFlowLogs
              Resource: '*'

  DeleteVPCFlowLogFunctionCloudWatchLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${DeleteVPCFlowLogFunction}
      RetentionInDays: !Ref LambdaLogGroupRetentionTimeDays
      LogGroupClass: INFREQUENT_ACCESS

  StepFunction:
    Type: AWS::Serverless::StateMachine
    Properties:
      Name: !Sub
        - enable-vpc-fl-for-spikes-${Suffix}
        - Suffix: !Select
            - 4
            - !Split
              - '-'
              - !Select
                - 2
                - !Split
                  - /
                  - !Ref AWS::StackId
      Definition:
        Comment: "State Machine to record Nat Gateway spikes in a VPC Flow Log"
        StartAt: "Add Number of Recordings To State Machine"
        States:
          "Add Number of Recordings To State Machine":
            Type: "Pass"
            Next: "Check-NGW-ID"
            Result: 
              NumOfRecordings:
              Fn::Transform:
                Name: "StringToInt-F39F5A878973"
                Parameters:
                  Value: !Ref NumOfRecordings
            ResultPath: $.NumOfRecordings
          Check-NGW-ID:
            InputPath: "$.detail.configuration.metrics[1].metricStat.metric.dimensions.NatGatewayId"
            Next: "Were Flow Logs recorded yet?"
            Parameters:
              Key:
                ngwid.$: "$"
              TableName: !Ref NGWFlowlogsTable
            Resource: "arn:aws:states:::dynamodb:getItem"
            ResultPath: $.taskresult
            Type: "Task"
          "Were Flow Logs recorded yet?":
            Type: "Choice"
            Choices:
              - IsPresent: false
                Variable: "$.taskresult.Item.ngwid.S"
                Next: "AddNGW-ID"
                Comment: If not recorded yet - go to add the Item to dynamoDB
            Default: "Convert Numbers"
          "Convert Numbers":
            Type: "Pass"
            Next: "ShouldWeRecord?"
            Parameters:
              val.$: "States.StringToJson($.taskresult.Item.NumOfRecordings.N)"
            ResultPath: $.RecordedSoFar
          ShouldWeRecord?:
            Type: Choice
            Choices:
              - And:
                  - Variable: "$.taskresult.Item.ngwid.S"
                    IsPresent: true
                  - Variable: "$.NumOfRecordings"
                    NumericGreaterThan: 0
                  - Variable: "$.RecordedSoFar.val"
                    NumericGreaterThanEqualsPath: "$.NumOfRecordings"
                Comment: "Recorded Max Times"
                Next: VPCflAlreadyRecorded
            Default: "Increment Number Of Recordings"
          "Increment Number Of Recordings":
            Type: Task
            Resource: arn:aws:states:::dynamodb:updateItem
            Parameters:
              TableName: !Ref NGWFlowlogsTable
              Key:
                ngwid:
                  S.$: $.detail.configuration.metrics[1].metricStat.metric.dimensions.NatGatewayId
              UpdateExpression: "ADD NumOfRecordings :inc"
              ExpressionAttributeValues:
                ":inc":
                  N: "1"
            Next: Extract-NGW-ID
            ResultPath: null            
          AddNGW-ID:
            Next: Extract-NGW-ID
            Parameters:
              Item:
                ngwid:
                  S.$: "$.detail.configuration.metrics[1].metricStat.metric.dimensions.NatGatewayId"
                NumOfRecordings:
                  N: "1"
              TableName: !Ref NGWFlowlogsTable
            Resource: "arn:aws:states:::dynamodb:putItem"
            ResultPath: null
            Type: "Task"
          Extract-NGW-ID:
            Next: "ActivateVPCFlowLog"
            Type: "Pass"
            Parameters:
              ngwid.$: "$.detail.configuration.metrics[1].metricStat.metric.dimensions.NatGatewayId"
          ActivateVPCFlowLog:
            Next: "WaitForAlarmOrTimeout"
            Parameters:
              ngwid.$: "$.ngwid"
            Resource: !GetAtt ActivateVPCFlowLogFunction.Arn
            ResultPath: "$.taskresult"
            Type: "Task"
          WaitForAlarmOrTimeout:
            Type: Task
            Resource: arn:aws:states:::lambda:invoke.waitForTaskToken
            TimeoutSeconds:
              Fn::Transform:
                Name: "StringToInt-F39F5A878973"
                Parameters:
                  Value: !Ref RecordingTime
            Parameters:
              FunctionName: !GetAtt StoreWaitTokenFunction.Arn
              Payload:
                ngwid.$: $.ngwid
                token.$: $$.Task.Token
                TableName: !Ref NGWFlowlogsTable
            Catch:
              - ErrorEquals:
                  - States.Timeout
                ResultPath: null
                Next: RemoveToken
                Comment: "Timeout"
            ResultPath: null
            Next: RemoveToken     
          RemoveToken:
            Type: "Task"                
            Parameters:
              Key:
                ngwid:
                  S.$: "$.ngwid"
              UpdateExpression: REMOVE tokenID
              TableName: !Ref NGWFlowlogsTable
            Resource: "arn:aws:states:::dynamodb:updateItem"
            ResultPath: null
            Next: "DeleteVPCFlowLog"
          DeleteVPCFlowLog:
            Next: "NotifySNS"
            Parameters:
              FlowLogIds.$: "$.taskresult.body.FlowLogIds"
            Resource: !GetAtt DeleteVPCFlowLogFunction.Arn
            Type: "Task"
            ResultPath: null
          NotifySNS:
            End: true
            Parameters:
              Message.$: !Sub States.Format('The VPC Flow Logs has been recorded for NAT Gateway {}\nCloudWatch Logs Group {}\nStream Prefix {}\n', $.ngwid, 'https://console.aws.amazon.com/cloudwatch/home?region=${AWS::Region}#logsV2:log-groups/log-group/${VPCFlowLogsCloudWatchLogGroup}', $.taskresult.body.ENI)                         
              Subject: "NAT Gateway VPC Flow Log Recorded"
              TopicArn: !Ref NatGatewayFlowLogsTopic
            Resource: "arn:aws:states:::sns:publish"
            Type: "Task"
          VPCflAlreadyRecorded:
            End: true
            Parameters:
              Message: "Flow Log was already recorded for NatGateway"
              ngwid.$: "$.taskresult.Item.ngwid.S"
            Type: "Pass"
      Role: !GetAtt StepFunctionExecutionRole.Arn

  StepFunctionExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: states.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: StepFunctionPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - dynamodb:GetItem
                  - dynamodb:PutItem
                  - dynamodb:UpdateItem
                  - lambda:InvokeFunction
                Resource: '*'
              - Effect: Allow
                Action:
                  - sns:publish
                Resource: !Ref NatGatewayFlowLogsTopic

  # Define the CloudWatch Event Rule for Nat Gateway Alarm ON
  CloudWatchEventRuleAlarmOn:
    Type: AWS::Events::Rule
    Properties:
      Name: !Sub 
        - NatGatewayAlarmON-${Suffix}
        - Suffix: !Select
            - 4
            - !Split
              - '-'
              - !Select
                - 2
                - !Split
                  - /
                  - !Ref AWS::StackId
      Description: EventBridge rule when a Nat Gateway is triggered by an Alarm on BytesInFromDestination or BytesInFromSource
      EventPattern:
        detail-type:
          - CloudWatch Alarm State Change
        source:
          - aws.cloudwatch
        detail:
          state:
            value:
              - ALARM
          configuration:
            metrics:
              metricStat:
                metric:
                  namespace:
                    - AWS/NATGateway
                  name:
                    - BytesInFromDestination
                    - BytesInFromSource
                  dimensions:
                    NatGatewayId:
                      !Split [",", !Ref NatGatewayIDs]
             
      State: ENABLED
      Targets:
        - Arn: !Ref StepFunction
          Id: StepFunctionTarget
          RoleArn: !GetAtt CloudWatchEventAlarmOnRole.Arn

  # Define the Role for the CloudWatch Event Rule for Nat Gateway Alarm ON
  CloudWatchEventAlarmOnRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: events.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: AllowInvokeStepFunction
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: states:StartExecution
                Resource: !Ref StepFunction

  # Define the CloudWatch Event Rule for Nat Gateway Alarm Off
  CloudWatchEventRuleAlarmOff:
    Type: AWS::Events::Rule
    Properties:
      Name: !Sub 
        - NatGatewayAlarmOFF-${Suffix}
        - Suffix: !Select
            - 4
            - !Split
              - '-'
              - !Select
                - 2
                - !Split
                  - /
                  - !Ref AWS::StackId
      Description: EventBridge rule when a Nat Gateway is triggered by the end of an Alarm on BytesInFromDestination or BytesInFromD
      EventPattern:
        detail-type:
          - CloudWatch Alarm State Change
        source:
          - aws.cloudwatch
        detail:
          state:
            value:
              - OK
          configuration:
            metrics:
              metricStat:
                metric:
                  namespace:
                    - AWS/NATGateway
                  name:
                    - BytesInFromDestination
                    - BytesInFromSource
                  dimensions:
                    NatGatewayId:
                      !Split [",", !Ref NatGatewayIDs]
      State: ENABLED
      Targets:
        - Arn: !GetAtt AlarmEndedFunction.Arn
          Id: LambdaTarget

  # Define the Role for the CloudWatch Event Rule for Nat Gateway Alarm OFF
  CloudWatchEventRuleAlarmOffRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: events.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: AllowInvokeLambdaFunction
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: lambda:InvokeFunction
                Resource: !GetAtt AlarmEndedFunction.Arn

  VPCFlowLogRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: vpc-flow-logs.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: VPCFlowLogPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                  - logs:DescribeLogGroups
                  - logs:DescribeLogStreams
                Resource: '*'

  NatGatewayFlowLogsTopic:
    Type: 'AWS::SNS::Topic'
    Properties:
      TopicName: 'NATGatewayFlowLogsRecorded'

  EmailSubscription:
    Type: 'AWS::SNS::Subscription'
    Properties: 
      Protocol: 'email'
      Endpoint: !Ref EmailAddress
      TopicArn: !Ref NatGatewayFlowLogsTopic

Outputs:
  ActivateVPCFlowLogFunction:
    Description: Activate VPC Flow Log Lambda Function Console Link
    Value: !Sub https://console.aws.amazon.com/lambda/home?region=${AWS::Region}#/functions/${ActivateVPCFlowLogFunction}

  StoreWaitTokenFunction:
    Description: Store Wait Token Lambda Function Console Link
    Value: !Sub https://console.aws.amazon.com/lambda/home?region=${AWS::Region}#/functions/${StoreWaitTokenFunction}

  AlarmEndedFunction:
    Description: Cloudwatch Nat Gateway Alarm Ended Function Console Link
    Value: !Sub https://console.aws.amazon.com/lambda/home?region=${AWS::Region}#/functions/${AlarmEndedFunction}    

  DeleteVPCFlowLogFunction:
    Description: Delete VPC Flow Log Lambda Function Console Link
    Value: !Sub https://console.aws.amazon.com/lambda/home?region=${AWS::Region}#/functions/${DeleteVPCFlowLogFunction}

  StepFunction:
    Description: ENI Flowlogs Step Function Console Link
    Value: !Sub https://console.aws.amazon.com/states/home?region=${AWS::Region}#/statemachines/view/${StepFunction}

  NGWFlowlogsTable:
    Description: NGW Flowlogs DynamoDB Table Console Link
    Value: !Sub https://console.aws.amazon.com/dynamodb/home?region=${AWS::Region}#tables:selected=${NGWFlowlogsTable}

  NatGatewayFlowLogsTopicArn:
    Description: 'SNS Topic in the AWS Management Console'
    Value: !Sub 'https://console.aws.amazon.com/sns/v3/home?region=${AWS::Region}#/topic/${NatGatewayFlowLogsTopic}'

  VPCFlowLogsCloudWatchLogGroupLink:
    Description: "VPC Flow Logs CloudWatch Group"
    Value: !Sub
      - "https://console.aws.amazon.com/cloudwatch/home?region=${AWS::Region}#logStream:group=${LogGroupName}"
      - LogGroupName: !Ref VPCFlowLogsCloudWatchLogGroup