#
# 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: 'AWS IoT Security Baseline'

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - 
        Label:
          default: "Required Parameters"
        Parameters:
          - IoTSecurityNotificationsDL
          - CertificateRotationPolicy
          - AWSIoTCA
          - AWSIoTCoreLogging
    ParameterLabels:
      IoTSecurityNotificationsDL:
        default: "Please provide the email to receive security event notifications"
      CertificateRotationPolicy:
        default: "Please provide the certificate rotation policy requirements (in years)"
      AWSIoTCA:
        default: "Are you deploying X.509 certificates generated by AWS IoT?"
      AWSIoTCoreLogging:
        default: "Is AWS IoT Core logging enabled in the account?"

Parameters:
  IoTSecurityNotificationsDL:
    Type: String
    AllowedPattern: "^[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\\.[a-zA-Z0-9-.]+$"
  CertificateRotationPolicy:
    Type: Number
  AWSIoTCA:
    Type: String
    Default: No
    AllowedValues:
      - Yes
      - No
  AWSIoTCoreLogging:
    Type: String
    Default: No
    AllowedValues:
      - Yes
      - No

Conditions:
  CreateCertificateExpiryResources: !Equals [ !Ref AWSIoTCA, Yes ]
  CreateAWSIoTLogging: !Equals [ !Ref AWSIoTCoreLogging, No ]

Resources:

#===============================================================================================================================
# Enable / configure AWS IoT Logging
#===============================================================================================================================
  IoTCoreLoggingRole:
    Type: 'AWS::IAM::Role'
    Condition: CreateAWSIoTLogging
    Properties:
        AssumeRolePolicyDocument:
          Version: "2012-10-17"
          Statement:
              - 
                Action: 'sts:AssumeRole'
                Effect: Allow
                Principal: 
                  Service: 
                    "iot.amazonaws.com"

  IoTCoreLoggingRolePolicy:
    Condition: CreateAWSIoTLogging
    Type: 'AWS::IAM::Policy'
    Properties:
      PolicyName: "IoTCoreLoggingRolePolicy"
      Roles: 
        - !Ref IoTCoreLoggingRole
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
            - 
              Effect: Allow
              Action:
                  - 'logs:CreateLogGroup'
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
                  - 'logs:PutMetricFilter'
                  - 'logs:PutRetentionPolicy'
              Resource: !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:AWSIotLogsV2*"

  IoTCoreLogging:
    Type: AWS::IoT::Logging
    Condition: CreateAWSIoTLogging
    Properties: 
      AccountId: !Ref AWS::AccountId
      DefaultLogLevel: ERROR
      RoleArn: !GetAtt IoTCoreLoggingRole.Arn

#===============================================================================================================================
# SNS topic to send security notifications to customer security team
#===============================================================================================================================
  IoTSecurityNotificationsTopic:
    Type: 'AWS::SNS::Topic'
    Properties:
      TopicName: IoTSecurityNotificationsTopic
      KmsMasterKeyId: "alias/aws/sns"
      Subscription:
        - Endpoint: !Ref IoTSecurityNotificationsDL
          Protocol: email

#===============================================================================================================================
# IAM role to allow DD to assume role for audit and sending notifications
#===============================================================================================================================

  IoTDDRole:
    Type: 'AWS::IAM::Role'
    Properties:
        AssumeRolePolicyDocument:
          Version: "2012-10-17"
          Statement:
              - 
                Action: 'sts:AssumeRole'
                Effect: Allow
                Principal: 
                  Service: 
                    "iot.amazonaws.com"
        ManagedPolicyArns:
          - "arn:aws:iam::aws:policy/service-role/AWSIoTDeviceDefenderAudit"
          - "arn:aws:iam::aws:policy/AmazonSNSReadOnlyAccess"
          - "arn:aws:iam::aws:policy/service-role/AWSIoTDeviceDefenderPublishFindingsToSNSMitigationAction"

#===============================================================================================================================
# IoT DD Audit Configuration
#===============================================================================================================================
  
  IoTAuditConfiguration:
    Type: AWS::IoT::AccountAuditConfiguration
    Properties: 
      AccountId: !Sub "${AWS::AccountId}"
      AuditCheckConfigurations: 
          AuthenticatedCognitoRoleOverlyPermissiveCheck: 
            Enabled: true
          CaCertificateExpiringCheck: 
            Enabled: true
          CaCertificateKeyQualityCheck: 
            Enabled: true
          ConflictingClientIdsCheck: 
            Enabled: true
          DeviceCertificateExpiringCheck: 
            Enabled: true
          DeviceCertificateKeyQualityCheck: 
            Enabled: true
          DeviceCertificateSharedCheck: 
            Enabled: true
          IotPolicyOverlyPermissiveCheck: 
            Enabled: true
          IotRoleAliasAllowsAccessToUnusedServicesCheck: 
            Enabled: true
          IotRoleAliasOverlyPermissiveCheck: 
            Enabled: true
          LoggingDisabledCheck: 
            Enabled: true
          RevokedCaCertificateStillActiveCheck: 
            Enabled: true
          RevokedDeviceCertificateStillActiveCheck: 
            Enabled: true
          UnauthenticatedCognitoRoleOverlyPermissiveCheck: 
            Enabled: true
      AuditNotificationTargetConfigurations: 
        Sns: 
            Enabled: true
            RoleArn: !GetAtt IoTDDRole.Arn
            TargetArn: !Ref IoTSecurityNotificationsTopic
      RoleArn: !GetAtt IoTDDRole.Arn

#===============================================================================================================================
# IoT DD Audit Schedule Configuration
#===============================================================================================================================

  IoTAuditSchedule:
    Type: AWS::IoT::ScheduledAudit
    DependsOn: IoTAuditConfiguration
    Properties:
      ScheduledAuditName: ScheduledIoTAudit
      TargetCheckNames:
        - AUTHENTICATED_COGNITO_ROLE_OVERLY_PERMISSIVE_CHECK
        - CA_CERTIFICATE_EXPIRING_CHECK
        - CA_CERTIFICATE_KEY_QUALITY_CHECK
        - CONFLICTING_CLIENT_IDS_CHECK
        - DEVICE_CERTIFICATE_EXPIRING_CHECK
        - DEVICE_CERTIFICATE_KEY_QUALITY_CHECK
        - DEVICE_CERTIFICATE_SHARED_CHECK
        - IOT_POLICY_OVERLY_PERMISSIVE_CHECK
        - IOT_ROLE_ALIAS_ALLOWS_ACCESS_TO_UNUSED_SERVICES_CHECK
        - IOT_ROLE_ALIAS_OVERLY_PERMISSIVE_CHECK
        - LOGGING_DISABLED_CHECK
        - REVOKED_CA_CERTIFICATE_STILL_ACTIVE_CHECK
        - REVOKED_DEVICE_CERTIFICATE_STILL_ACTIVE_CHECK
        - UNAUTHENTICATED_COGNITO_ROLE_OVERLY_PERMISSIVE_CHECK
      Frequency: DAILY

#===============================================================================================================================
# IoT DD Detect Configuration using ML Detect
#===============================================================================================================================

  IoTMLDetectSecurityProfile:
    Type: AWS::IoT::SecurityProfile
    Properties:
      SecurityProfileName: IoTMLDetectSecurityProfile
      SecurityProfileDescription: IoT Device Defender Detect Security Profile using ML Detect
      AlertTargets:
        SNS:
          AlertTargetArn: !Ref IoTSecurityNotificationsTopic
          RoleArn: !GetAtt IoTDDRole.Arn
      TargetArns:
        - !Sub "arn:aws:iot:${AWS::Region}:${AWS::AccountId}:all/things"
      Behaviors: 
        - Name: "Authorization_Failures"
          Metric: "aws:num-authorization-failures"
          Criteria: 
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false
        - Name: "Message_Size"
          Metric: "aws:message-byte-size"
          Criteria: 
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false
        - Name: "Message_Sent"
          Metric: "aws:num-messages-sent"
          Criteria: 
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false
        - Name: "Message_Received"
          Metric: "aws:num-messages-received"
          Criteria: 
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false
        - Name: "Connection_Attempts"
          Metric: "aws:num-connection-attempts"
          Criteria: 
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false
        - Name: "Disconnects"
          Metric: "aws:num-disconnects"
          Criteria: 
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false
        - Name: "Outbound_traffic"
          Metric: "aws:all-bytes-out"
          Criteria:
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false
        - Name: "TCP_inbound_traffic"
          Metric: "aws:all-bytes-in"
          Criteria:
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false
        - Name: "Max_TCP_port"
          Metric: "aws:num-listening-tcp-ports"
          Criteria:
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false
        - Name: "Max_UDP_ports"
          Metric: "aws:num-listening-tcp-ports"
          Criteria:
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false
        - Name: "Outbound_sent"
          Metric: "aws:all-packets-out"
          Criteria:
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false
        - Name: "Inbound_sent"
          Metric: "aws:all-packets-in"
          Criteria:
            ConsecutiveDatapointsToAlarm: 1
            ConsecutiveDatapointsToClear: 1
            MlDetectionConfig:
              ConfidenceLevel: "HIGH"
          SuppressAlerts: false

#===============================================================================================================================
# Certificate expiry notifications
#===============================================================================================================================

  CertRotationLambdaExecutionRole:
    Type: 'AWS::IAM::Role'
    Condition: CreateCertificateExpiryResources
    Properties:
      AssumeRolePolicyDocument:
          Version: '2012-10-17'
          Statement:
              - 
                  Effect: Allow
                  Principal:
                      Service:
                          - lambda.amazonaws.com
                  Action:
                      - 'sts:AssumeRole'
  CertRotationLambdaExecutionRolePolicy:
    Condition: CreateCertificateExpiryResources
    Type: 'AWS::IAM::Policy'
    Properties:
      PolicyName: CertRotationLambdaExecutionRolePolicy
      Roles: 
        - !Ref CertRotationLambdaExecutionRole
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
            - 
              Effect: Allow
              Action:
                  - 'logs:CreateLogGroup'
              Resource: !Sub "arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:*"
            - 
              Effect: Allow
              Action:
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
              Resource: !Sub "arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:*"
            - 
              Action:
                  - 'iot:DescribeCertificate'
              Effect: Allow
              Resource: !Sub "arn:${AWS::Partition}:iot:${AWS::Region}:${AWS::AccountId}:*"
            - 
              Action:
                  - 'iot:ListCertificates'
                  - 'iot:ListPrincipalThings'
              Effect: Allow
              Resource: "*"
            -
              Action:
                  - 'sns:Publish'
                  - 'sns:ListTopics'
              Effect: Allow
              Resource: !Ref IoTSecurityNotificationsTopic

  CertRotationLambdaFunction:
    Type: "AWS::Lambda::Function"
    Condition: CreateCertificateExpiryResources
    Properties:
      Code:
        ZipFile: |
          import boto3
          import json
          import os
          from datetime import date
          from dateutil.relativedelta import relativedelta

          iotclient = boto3.client('iot')
          snsclient = boto3.client('sns')

          def list_certificates_paginator():
            # Create a reusable Paginator
            paginator = iotclient.get_paginator('list_certificates')

            # Create a PageIterator from the Paginator
            page_iterator = paginator.paginate()
            
            certificates = []

            for page in page_iterator:
              for certificate in page['certificates']:
                certificates.append(certificate)
            return(certificates)

          def listPrincipalThings(certificateArn):
            return iotclient.list_principal_things(principal=certificateArn)

          def handler(event, context):
              response = list_certificates_paginator()
              topic = os.environ['sns_topic']
              cert_rotation_policy = os.environ['certificate_rotation_policy']
              notification = []
              notification_overdue = []
              for r in response:
                  msg = {
                  "certificate": '',
                  "device": '',
                  "status": ''
                  }
                  msg_overdue = {
                  "certificate": '',
                  "device": '',
                  "status": '',
                  "account_no": ''
                  }
                  expiry_date = r['creationDate'].date() + relativedelta(years=+int(cert_rotation_policy))
                  if ((expiry_date - date.today()).days < 30 and (expiry_date - date.today()).days > 0 and r['status'] == 'ACTIVE'):
                      response = listPrincipalThings(r['certificateArn'])
                      msg['device'] = response['things']
                      msg['certificate'] = r['certificateId']
                      msg['status'] = f"This certificate will expire in {(expiry_date - date.today()).days} days. Please renew"
                      notification.append(msg)
                  elif ((expiry_date - date.today()).days < 0 and r['status'] == 'ACTIVE'):
                      response = listPrincipalThings(r['certificateArn'])
                      msg_overdue['device'] = response['things']
                      msg_overdue['certificate'] = r['certificateId']
                      msg_overdue['status'] = f"This certificate expired {(date.today() - expiry_date).days} days ago. Please renew immidiately"
                      msg_overdue['account_no'] = context.invoked_function_arn.split(":")[4]
                      notification_overdue.append(msg_overdue)
                  else:
                      print("Valid certificate: ", r['certificateId'])
              print("notification", notification, "\n")
              print("notification_overdue: ", notification_overdue, "\n")
              
              if notification:
                response = snsclient.publish(
                  TopicArn = topic,
                  Message = json.dumps(notification, indent=4),
                  Subject = 'IoT Security Notification - Certificates Expiring in 30 days',
                  MessageStructure = 'raw'
                )

              if notification_overdue:
                response_overdue = snsclient.publish(
                  TopicArn = topic,
                  Message = json.dumps(notification_overdue, indent=4),
                  Subject = 'IoT Security Notification - Expired Certificates',
                  MessageStructure = 'raw'
                )
      Description: Lambda function to encode / decode OS logs
      FunctionName: CertRotationLambdaFunction
      Handler: index.handler
      Role: !GetAtt CertRotationLambdaExecutionRole.Arn
      ReservedConcurrentExecutions: 10
      Environment:
        Variables: 
          sns_topic: !Ref IoTSecurityNotificationsTopic
          certificate_rotation_policy: !Ref CertificateRotationPolicy
      Runtime: python3.8
      Timeout: 300
  
  CWEPermissionToCallLambda:
    Type: "AWS::Lambda::Permission"
    Condition: CreateCertificateExpiryResources
    Properties:
      Action: 'lambda:InvokeFunction'
      FunctionName: !GetAtt CertRotationLambdaFunction.Arn
      Principal: 'events.amazonaws.com'
      SourceArn: !GetAtt CertRotationCloudWatchEventRule.Arn

  CertRotationCloudWatchEventRule:
    Type: AWS::Events::Rule
    Condition: CreateCertificateExpiryResources
    Properties:
      Name: CertRotationCloudWatchEventRule
      Description: Publishes Certificate Rotation Notifications to an SNS topic
      ScheduleExpression: "rate(30 days)"
      State: ENABLED
      Targets:
        -
          Arn: !GetAtt CertRotationLambdaFunction.Arn
          Id: CertRotationLambdaFunction