---
AWSTemplateFormatVersion: "2010-09-09"
Description: Amazon CloudFormation template for setting up resources for centralized logging for AWS Security Hub

Resources:
  AWSSecurityHubBucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Sub aws-securityhub-blog-${AWS::AccountId}-${AWS::Region}
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        BlockPublicPolicy: true
        IgnorePublicAcls: true
        RestrictPublicBuckets: true

  SecurityHubGlueCrawler:
    Type: AWS::Glue::Crawler
    Properties:
      Name: aws-securityhub
      Role: !Ref SecurityHubGlueCrawlerRole
      DatabaseName: aws-securityhub-db
      Targets:
        S3Targets:
          - Path: !Ref AWSSecurityHubBucket

  SecurityHubGlueCrawlerRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - glue.amazonaws.com
            Action:
              - sts:AssumeRole
        Version: "2012-10-17"
      ManagedPolicyArns:
        - !Sub arn:${AWS::Partition}:iam::aws:policy/service-role/AWSGlueServiceRole
      Path: /
      Policies:
        - PolicyName: sechub-GlueCrawlerRolePolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - s3:ListBucket
                  - s3:GetObject
                Resource:
                  - !Sub arn:${AWS::Partition}:s3:::${AWSSecurityHubBucket}*
              - Effect: Allow
                Action:
                  - s3:PutObject
                  - s3:GetObject
                Resource:
                  - !Sub arn:${AWS::Partition}:s3:::${AWSSecurityHubBucket}*
                  - "*"

  SecurityHubLambdaFunctionLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      RetentionInDays: 30

  SecurityHubLambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      LoggingConfig:
        LogGroup: !Ref SecurityHubLambdaFunctionLogGroup
      FunctionName: !Sub
        - aws-securityhub-${Id}
        - Id: !Select
            - 0
            - !Split
              - "-"
              - !Ref AWS::StackName
      Handler: index.lambda_handler
      Runtime: python3.10
      MemorySize: 128
      Timeout: 900
      Role: !GetAtt SecurityHubLambdaFunctionRole.Arn
      Environment:
        Variables:
          S3_BUCKET: !Ref AWSSecurityHubBucket
          CRAWLER_NAME: !Ref SecurityHubGlueCrawler
          PREFIX: aws-securityhub
          ManagementAccountId: !Ref AWS::AccountId
      Code:
        ZipFile: |
          import json
          import boto3
          import os
          import logging
          import re
          from datetime import date, timedelta, datetime
          from boto3.s3.transfer import S3Transfer
          from botocore.exceptions import ClientError

          BUCKET = os.environ['S3_BUCKET']
          PREFIX = os.environ['PREFIX']
          CRAWLER_NAME = os.environ['CRAWLER_NAME']
          MANAGEMENT_ACCOUNT_ID = os.environ['ManagementAccountId']
          def start_crawler():
              glue = boto3.client('glue')
              try:
                  glue.start_crawler(Name=CRAWLER_NAME)
              except Exception as exc:
                  logging.warning(exc)
          def store_data_to_s3(flattened_data, path):
              today = date.today()
              year = today.year
              month = today.strftime('%m')
              day = today.day
              hour = datetime.now()
              key = today.strftime(f"{path}/year={year}/month={month}/day={day}/hour={hour.hour}/{year}-{month}-{day}-{hour.hour}.json")
              local_file = '/tmp/tmp.json'
              s3client = boto3.resource('s3')
              try:
                  obj=s3client.Bucket(BUCKET).download_file(key,local_file)
              except ClientError as e:
                  if e.response['Error']['Code'] == "404":
                      print("The object does not exist.")
                      with open(local_file, 'w') as f:
                          f.write(flattened_data)
                  else:
                      raise
              with open(local_file, 'a') as f:
                  f.write('\n')
                  f.write(flattened_data)
              s3client.meta.client.upload_file(local_file, BUCKET, key)
              return "Successful"
          def rename_keys(payload):
              if isinstance(payload, dict):
                  renamed_payload = {}
                  for key, value in payload.items():
                      renamed_key = re.sub(r'\W+', '_', key)
                      renamed_value = rename_keys(value)
                      renamed_payload[renamed_key] = renamed_value
                  return renamed_payload
              elif isinstance(payload, list):
                  renamed_list=[]
                  for item in payload:
                      renamed_item = rename_keys(item)
                      renamed_list.append(renamed_item)
                      return renamed_list
              elif isinstance(payload,(str,int,float,bool,type(None))):
                  return payload
              else:
                  if hasattr(payload, "__dict__"):
                      renamed_payload = {}
                      for key, value in payload.__dict__.items():
                          renamed_key = re.sub(r'\W+', '_', key)
                          renamed_value = rename_keys(value)
                          renamed_payload[renamed_key] = renamed_value
                          return renamed_payload
                      else:
                          return payload
          def lambda_handler(event, context):
              logger = logging.getLogger()
              sts = boto3.client('sts')
              payload = event
              renamed_payload=rename_keys(payload)
              flattened_data=json.dumps(renamed_payload, separators=(',',':'))
              store_data_to_s3(flattened_data, f'{PREFIX}/securityhub-data/management_id={MANAGEMENT_ACCOUNT_ID}')
              total_count = 0
              total_count += len(flattened_data)
              if total_count:
                   start_crawler()
              return "Successful"

  SecurityHubLambdaFunctionRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: AWS-securityhub-Execute-Lambda
      AssumeRolePolicyDocument:
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole

  S3BucketAccessPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: securityhub-S3BucketAccessPolicy
      Roles:
        - !Ref SecurityHubLambdaFunctionRole
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Action:
              - s3:ListBucket
              - s3:PutObject
              - s3:GetObject
            Resource:
              - !Sub arn:${AWS::Partition}:s3:::${AWSSecurityHubBucket}
              - !Sub arn:${AWS::Partition}:s3:::${AWSSecurityHubBucket}/*

  SecurityHubGlueCrawlerPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: securityhub-GlueCrawlerPolicy
      Roles:
        - !Ref SecurityHubLambdaFunctionRole
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Action:
              - glue:StartCrawler
            Resource:
              - !Sub arn:${AWS::Partition}:glue:${AWS::Region}:${AWS::AccountId}:crawler/${SecurityHubGlueCrawler}

  SecurityHubCloudWatchLogsPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: securityhub-CloudWatchLogsPolicy
      Roles:
        - !Ref SecurityHubLambdaFunctionRole
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Action:
              - logs:CreateLogGroup
              - logs:CreateLogStream
              - logs:PutLogEvents
            Resource: !GetAtt SecurityHubLambdaFunctionLogGroup.Arn

  SecHubDatabase:
    Type: AWS::Glue::Database
    Properties:
      CatalogId: !Ref AWS::AccountId
      DatabaseInput:
        Name: aws-securityhub-db

  SecurityHubAthenaView:
    Type: AWS::Athena::NamedQuery
    Properties:
      Name: AWS-securityhub
      Database: aws-securityhub-db
      QueryString: !Sub |
        CREATE OR REPLACE VIEW "security_hub_findings" AS SELECT DISTINCT management_id,
          year Year,
          month Month,
          day Day,
          detail.findings [ 1 ].AwsAccountId LinkedAccountId,
          detail.findings [ 1 ].region Region,
          detail.findings [ 1 ].Id Finding_Id,
          detail.findings [ 1 ].GeneratorId GeneratorId,
          detail.findings [ 1 ].ProductName ProductName,
          CAST(
            "from_iso8601_timestamp"(detail.findings [ 1 ].CreatedAt) AS date
          ) Created_At,
          CAST(
            "from_iso8601_timestamp"(detail.findings [ 1 ].UpdatedAt) AS date
          ) UpdatedAt,
          CAST(
            "from_iso8601_timestamp"(detail.findings [ 1 ].lastobservedat) AS date
          ) lastobservedat,
          detail.findings [ 1 ].resources [ 1 ].type Resource_Type,
          detail.findings [ 1 ].resources [ 1 ].id Resource_Id,
          detail.findings [ 1 ].resources [ 1 ].region Resource_Region,
          detail.findings [ 1 ].Severity.Normalized Normalized_Severity,
          detail.findings [ 1 ].Severity.Label Normalized_Severity_Label,
          detail.findings [ 1 ].Severity.Original Product_Severity_Label,
          COALESCE(
            regexp_extract(detail.findings [ 1 ].Types [ 1 ], '.*/(.*)$', 1),
            detail.findings [ 1 ].Types [ 1 ]
          ) Finding_Types,
          detail.findings [ 1 ].Compliance.Status Compliance_Status,
          detail.findings [ 1 ].compliance.securitycontrolid SecurityControlId,
          upper(
            split_part(detail.findings [ 1 ].resources [ 1 ].id, ':', 3)
          ) Finding_Service
        FROM aws_securityhub_${AWS::AccountId}_${AWS::Region} WHERE (
            CAST(
              "from_iso8601_timestamp"(detail.findings [ 1 ].CreatedAt) AS date
            ) >= date_parse(
              date_format(date_add('month', -3, current_date), '%Y-%m-%d'),
              '%Y-%m-%d'
            )
          )
        GROUP BY 1,
          2,
          3,
          4,
          5,
          6,
          7,
          8,
          9,
          10,
          11,
          12,
          13,
          14,
          15,
          16,
          17,
          18,
          19,
          20,
          21,
          22

  SecurityHubEventRule:
    Type: AWS::Events::Rule
    Properties:
      Description: EventBridge Rule to ingest security hub findings
      EventPattern:
        source:
          - aws.securityhub
        detail-type:
          - Security Hub Findings - Imported
      Targets:
        - Arn: !GetAtt SecurityHubLambdaFunction.Arn
          Id: TargetFunctionV1
      Name: securityhub-eventrule
      State: ENABLED

  PermissionForEventsToInvokeLambda:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !Ref SecurityHubLambdaFunction
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt SecurityHubEventRule.Arn

Outputs:
  SecurityHubLambdaRoleARN:
    Description: Role for Lambda execution of AWS SecurityHub extraction.
    Value: !GetAtt SecurityHubLambdaFunctionRole.Arn

  LambdaSecurityHubARN:
    Description: Lambda function ARN for retrieving AWS SecurityHub extraction
    Value: !GetAtt SecurityHubLambdaFunction.Arn

  GlueCrawler:
    Value: !Ref SecurityHubGlueCrawler
