# License: Copyright 2021 Amazon Web Services, Inc. or its affiliates.
# All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#    http://aws.amazon.com/apache2.0/
# or in the "license" file accompanying this file.

# This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
# CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.

AWSTemplateFormatVersion: "2010-09-09"
Description: "Security Hub disable controls"

Resources:
  SHLambdaFunction:
    Type: 'AWS::Lambda::Function'
    Properties:
      Code:
        ZipFile: |
          import sys,json,boto3,time,os
          import cfnresponse
          from botocore.exceptions import ClientError

          def lambda_handler(event,context):
            print('this is the event ' + json.dumps(event))
            StackName = event['StackId']
            LogicalResourceId = event['LogicalResourceId']
            UniqueId = event['RequestId']
            props = event['ResourceProperties']
            region = os.environ['AWS_REGION']
            account_id = context.invoked_function_arn.split(":")[4]
            sh = boto3.client('securityhub')
            controls = f"arn:aws:securityhub:{region}:{account_id}:control/aws-foundational-security-best-practices/v/1.0.0/IAM.2,arn:aws:securityhub:{region}:{account_id}:control/aws-foundational-security-best-practices/v/1.0.0/IAM.6,arn:aws:securityhub:{region}:{account_id}:control/aws-foundational-security-best-practices/v/1.0.0/RDS.5,arn:aws:securityhub:{region}:{account_id}:control/aws-foundational-security-best-practices/v/1.0.0/RDS.6,arn:aws:securityhub:{region}:{account_id}:control/aws-foundational-security-best-practices/v/1.0.0/Lambda.4,arn:aws:securityhub:{region}:{account_id}:control/aws-foundational-security-best-practices/v/1.0.0/KMS.3,arn:aws:securityhub:{region}:{account_id}:control/aws-foundational-security-best-practices/v/1.0.0/DynamoDB.1,arn:aws:securityhub:{region}:{account_id}:control/pci-dss/v/3.2.1/PCI.IAM.2,arn:aws:securityhub:{region}:{account_id}:control/pci-dss/v/3.2.1/PCI.IAM.4,arn:aws:securityhub:{region}:{account_id}:control/pci-dss/v/3.2.1/PCI.IAM.5,arn:aws:securityhub:{region}:{account_id}:control/pci-dss/v/3.2.1/PCI.IAM.6,arn:aws:securityhub:{region}:{account_id}:control/pci-dss/v/3.2.1/PCI.S3.3,arn:aws:securityhub:{region}:{account_id}:control/pci-dss/v/3.2.1/PCI.Lambda.2,arn:aws:securityhub:{region}:{account_id}:control/cis-aws-foundations-benchmark/v/1.2.0/1.13,arn:aws:securityhub:{region}:{account_id}:control/cis-aws-foundations-benchmark/v/1.2.0/1.14,arn:aws:securityhub:{region}:{account_id}:control/cis-aws-foundations-benchmark/v/1.2.0/3.2,arn:aws:securityhub:{region}:{account_id}:control/cis-aws-foundations-benchmark/v/1.2.0/3.9,arn:aws:securityhub:{region}:{account_id}:control/cis-aws-foundations-benchmark/v/1.2.0/3.7,arn:aws:securityhub:{region}:{account_id}:control/cis-aws-foundations-benchmark/v/1.2.0/3.13,arn:aws:securityhub:{region}:{account_id}:control/cis-aws-foundations-benchmark/v/1.2.0/3.8"
            controls_list = controls
            if (event['RequestType'] == 'Create' or event['RequestType'] == 'Update'):
              try:                
                for control in controls_list.split(","):
                  print(f"Trying to disable {control}")
                  try:
                    sh.update_standards_control(StandardsControlArn=f'{control.strip()}',ControlStatus='DISABLED',DisabledReason='Disabled by Lambda')                
                    print(f"Disabled: {control}")
                  except Exception as ex:
                    print(f"Cannot disable {control}")
                print("Respond: SUCCESS")
                cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
              except Exception as ex:
                print("failed executing actions: " + str(ex))
                cfnresponse.send(event, context, cfnresponse.FAILED, ex)
            else:
              try:
                for control in controls_list.split(","):
                  print(f"Trying to enable {control}")
                  sh.update_standards_control(StandardsControlArn=f'{control}',ControlStatus='ENABLED')                
                  print(f"Enable {control}")
                print("Respond: SUCCESS")
                cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
              except Exception as ex:
                print("failed executing actions: " + str(ex))
                print("Respond: SUCCESS")
                cfnresponse.send(event, context, cfnresponse.SUCCESS, ex)
            
      Handler: index.lambda_handler
      MemorySize: 128
      Role:
        'Fn::GetAtt':
          - SHLambdaRole
          - Arn
      Runtime: "python3.8"
      Timeout: 300

  SHConfiguration:
    Type: 'Custom::SHConfiguration'
    Properties:
      ServiceToken:
        'Fn::GetAtt':
          - SHLambdaFunction
          - Arn

  SHLambdaRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: 
                - lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      Path: /
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
        - 'arn:aws:iam::aws:policy/AWSSecurityHubReadOnlyAccess'
      Policies:
        - PolicyName: iam-sh-update-controls
          PolicyDocument:
            Statement:
              - Effect: Allow
                Action:
                  - securityhub:UpdateStandardsControl
                Resource: "*"