#
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0
#
# 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.
#
# 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.
#

AWSTemplateFormatVersion: "2010-09-09"
Description: Creates a function and schedule to monitor IP utilization
Parameters:
  ScheduleInterval:
    Type: String
    Default: '15 minutes'
    Description: 'The interval that schedules functions run.'
    AllowedValues:
      - '5 minutes'
      - '10 minutes'
      - '15 minutes'
      - '30 minutes'
      - '45 minutes'
      - '1 hour'
      - '2 hours'
      - '4 hours'
      - '6 hours'
      - '8 hours'
      - '12 hours'
      - '1 day'
Resources:
  IpAllocationMonitorIAMRole:
    Type: AWS::IAM::Role
    Properties:
      Path: /
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
      Description: IAM Role for the IP Allocation Monitor function
      Policies:
        - PolicyName: root
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Sid: CreateLogGroup
                Effect: Allow
                Action: logs:CreateLogGroup
                Resource: !Sub 'arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:*'
              - Sid: LambdaLogPermissions
                Effect: Allow
                Action:
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: !Sub 'arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${AWS::StackName}-IP-Allocation-Monitor:*'
              - Sid: DescribeSubnetsPermissions
                Effect: Allow
                Action: ec2:DescribeSubnets
                Resource: '*'
              - Sid: PutMetricsPermissions
                Effect: Allow
                Action: cloudwatch:PutMetricData
                Resource: '*'
                Condition:
                  StringEquals:
                    "cloudwatch:namespace": "IPUsage"

  IpAllocationMonitorFunction:
    DependsOn: IpAllocationMonitorIAMRole
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          import logging
          import boto3
          from ipaddress import ip_network
          from datetime import datetime
          ec2_client = boto3.client('ec2')
          cloudwatch_client = boto3.client('cloudwatch')
          def get_usage():
              subnet_data = dict()
              paginator = ec2_client.get_paginator('describe_subnets')
              iterator = paginator.paginate()
              for page in iterator:
                  for subnet in page['Subnets']:
                      logging.info(str(subnet))
                      if not subnet['Ipv6Native']:
                          try:
                              number_of_addresses = ip_network(subnet['CidrBlock']).num_addresses
                          except:
                              logging.warning("Could not decipher Subnet CIDR.")
                              number_of_addresses = 0
                          try:
                              available_addresses = int(subnet['AvailableIpAddressCount'])
                          except:
                              logging.warning("Could not decipher available IPs.")
                              available_addresses = 0
                          used_ips = number_of_addresses - available_addresses
                          if number_of_addresses:
                              utilization = (used_ips / number_of_addresses) * 100
                          else:
                              utilization = 1
                          subnet_data[subnet['SubnetId']] = {'VpcId': subnet['VpcId'], 'AvailableIPs': available_addresses, 'TotalIPs': number_of_addresses, 'UsedIPs': used_ips, 'Utilization': utilization}
              return subnet_data
          def post_metric(metric_name: str, subnet_id: str, vpc_id: str, value: float, unit: str):
              cloudwatch_client.put_metric_data(
                  Namespace='IPUsage',
                  MetricData=[{'MetricName': metric_name,'Dimensions': [{'Name': 'SubnetId', 'Value': subnet_id}, {'Name': 'VpcId', 'Value': vpc_id}],'Timestamp': datetime.utcnow(), 'Value': value, 'Unit': unit}])
          def post_usage(subnet_data):
              for subnet in subnet_data.keys():
                  try:
                      post_metric('AvailableIPs', str(subnet), str(subnet_data[subnet]['VpcId']), float(subnet_data[subnet]['AvailableIPs']), 'Count')
                  except:
                      logging.critical("Could not post Available IPs metric for "+str(subnet))
                  try:
                      post_metric('TotalIPs', str(subnet), str(subnet_data[subnet]['VpcId']), float(subnet_data[subnet]['TotalIPs']), 'Count')
                  except:
                      logging.critical("Could not post Total IPs metric for "+str(subnet))
                  try:
                      post_metric('UsedIPs', str(subnet), str(subnet_data[subnet]['VpcId']), float(subnet_data[subnet]['UsedIPs']), 'Count')
                  except:
                      logging.critical("Could not post Used IPs metric for "+str(subnet))
                  try:
                      post_metric('Utilization', str(subnet), str(subnet_data[subnet]['VpcId']), float(subnet_data[subnet]['Utilization']), 'Percent')
                  except:
                      logging.critical("Could not post Utilization metric for "+str(subnet))
          def main(event, context):
              subnet_data = get_usage()
              post_usage(subnet_data)
      Description: This function monitors subnet IPv4 utilization and publishes it to CloudWatch
      FunctionName: !Sub '${AWS::StackName}-IP-Allocation-Monitor'
      Handler: index.main
      MemorySize: 128
      Role: !GetAtt IpAllocationMonitorIAMRole.Arn
      Runtime: 'python3.9'
      Timeout: 300

  IPUsageMonitorEventRulePermission:
    DependsOn:
      - IPUsageMonitorEventRule
      - IpAllocationMonitorFunction
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !Ref IpAllocationMonitorFunction
      Action: 'lambda:InvokeFunction'
      Principal: 'events.amazonaws.com'
      SourceArn: !GetAtt IPUsageMonitorEventRule.Arn

  IPUsageMonitorEventRule:
    DependsOn: IpAllocationMonitorFunction
    Type: AWS::Events::Rule
    Properties:
      Description: 'Schedule to run the IP Usage Monitoring function on a regular interval.'
      ScheduleExpression: !Sub 'rate(${ScheduleInterval})'
      State: 'ENABLED'
      Targets:
        - Arn: !GetAtt IpAllocationMonitorFunction.Arn
          Id: !Ref IpAllocationMonitorFunction