# 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 PCI Enabler"

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))
            region = os.environ['AWS_REGION']
            account_id = context.invoked_function_arn.split(":")[4]
            sh = boto3.client('securityhub')
            #Change it to the desired standard to enable
            standardArn=f"arn:aws:securityhub:{region}::standards/pci-dss/v/3.2.1"
            print(standardArn)
            #Change it to match the same standard, but adding account_id variable
            standardArnSubs=f"arn:aws:securityhub:{region}:{account_id}:subscription/pci-dss/v/3.2.1"
            print(standardArnSubs)
            standards_subscription_request = [{'StandardsArn': standardArn}]
            if (event['RequestType'] == 'Create' or event['RequestType'] == 'Update'):
              try:                
                sh.batch_enable_standards(StandardsSubscriptionRequests=standards_subscription_request)
                print("Security Hub standard enabled")
                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:                
                sh.batch_disable_standards(StandardsSubscriptionArns=[standardArnSubs,])
                print("Security Hub with standard disabled")
                cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
              except Exception as ex:
                print("failed executing actions: " + str(ex))
                cfnresponse.send(event, context, cfnresponse.SUCCESS, ex)
      Handler: index.lambda_handler
      MemorySize: 128
      Role:
        'Fn::GetAtt':
          - SHLambdaRole
          - Arn
      Runtime: "python3.7"
      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-enable-pci-standard
          PolicyDocument:
            Statement:
              - Effect: Allow
                Action:
                  - securityhub:BatchEnableStandards
                  - securityhub:BatchDisableStandards
                Resource: "*"