# Copyright 2022 Amazon.com and its affiliates; all rights reserved. 
# This file is Amazon Web Services Content and may not be duplicated or distributed without permission.
AWSTemplateFormatVersion: "2010-09-09"
Description: Export Sbom via Amazon Inspector

Parameters:
  EnvironmentName:
    Description: An environment name that is prefixed to resource names
    Type: String
    Default: sbom-inspector
    MinLength: 2
    MaxLength: 25

  KMSAdmin:
    Description: Arn of user/role that will have administrative rights over the KMS key
    Type: String
    AllowedPattern: ^arn:aws:iam::(\d{12})?:.*/[\w+=,.@-]{1,128}$

Resources:
  EcrRepo:
    Type: AWS::ECR::Repository
    Properties:
      RepositoryName: !Sub ${EnvironmentName}-respository

  SbomCMK:
    Type: AWS::KMS::Key
    DependsOn: SbomTransformationRole
    Properties:
      Description: CMS created for encryption
      EnableKeyRotation: true
      KeyPolicy:
        Version: '2012-10-17'
        Id: key-default-1
        Statement:
          - Sid: Enable IAM User Permissions
            Effect: Allow
            Principal:
              AWS: !Join [ ':', [ 'arn', !Ref "AWS::Partition", 'iam:', !Ref "AWS::AccountId", 'root' ]]
            Action: 'kms:*'
            Resource: '*'
          - Sid: Allow Amazon Inspector to use the key
            Effect: Allow
            Principal:
              Service: "inspector2.amazonaws.com"
            Action: 
              - 'kms:decrypt'
              - 'kms:GenerateDataKey'
            Resource: "*"
            Condition:
              StringEquals:
                "aws:SourceAccount": !Ref "AWS::AccountId"
              ArnLike:
                "aws:SourceArn": !Join [ ':', [ 'arn:aws:inspector2', !Ref "AWS::Region", !Ref "AWS::AccountId", '*/*' ]]
          - Sid: Allow access for Key Administrators
            Effect: Allow
            Principal:
              AWS: !Sub ${KMSAdmin}
            Action: 
              - "kms:Create"
              - "kms:Describe"
              - "kms:Enable"
              - "kms:List"
              - "kms:Put"
              - "kms:Update"
              - "kms:Revoke"
              - "kms:Disable"
              - "kms:Get"
              - "kms:Delete"
              - "kms:TagResource"
              - "kms:UntagResource"
              - "kms:ScheduleKeyDeletion"
              - "kms:CancelKeyDeletion"
            Resource: "*"
          - Sid: Allow Transform Lambda to use the key
            Effect: Allow
            Principal:
              AWS: !Join [ ':', [ 'arn', !Ref "AWS::Partition", 'iam:', !Ref "AWS::AccountId", !Sub 'role/${EnvironmentName}-transformation-role' ]]
            Action: 
              - "kms:Encrypt"
              - "kms:Decrypt"
              - "kms:ReEncrypt"
              - "kms:GenerateDataKey"
              - "kms:DescribeKey"
            Resource: '*'
          - Sid: Allow Trigger Lambda to use the key
            Effect: Allow
            Principal:
              AWS: !Join [ ':', [ 'arn', !Ref "AWS::Partition", 'iam:', !Ref "AWS::AccountId", !Sub 'role/${EnvironmentName}-trigger-role' ]]
            Action: 
              - "kms:Encrypt"
              - "kms:Decrypt"
              - "kms:ReEncrypt"
              - "kms:GenerateDataKey"
              - "kms:DescribeKey"
            Resource: '*'
          - Sid: Allow use of the key
            Effect: Allow
            Principal:
              AWS: !Sub ${KMSAdmin}
            Action: 
              - "kms:Encrypt"
              - "kms:Decrypt"
              - "kms:ReEncrypt"
              - "kms:GenerateDataKey"
              - "kms:DescribeKey"
            Resource: '*'
          - Sid: Allow attachment of persistent resources
            Effect: Allow
            Principal:
              AWS: !Sub ${KMSAdmin}
            Action: 
              - "kms:CreateGrant"
              - "kms:ListGrants"
              - "kms:RevokeGrant"
            Resource: "*"
            Condition:
              StringEquals:
                "aws:SourceAccount": !Ref "AWS::AccountId"
              Bool:
                 "kms:GrantIsForAWSResource": "true"
              ArnLike:
                "aws:SourceArn": !Join [ ':', [ 'arn:aws:inspector2', !Ref "AWS::Region", !Ref "AWS::AccountId", '*/*' ]]
            
  SbomCMKAlias:
    Type: 'AWS::KMS::Alias'
    Properties:
      AliasName: alias/sbom-kms
      TargetKeyId: !Ref SbomCMK
  
  SbomTriggerRole:
    Type: "AWS::IAM::Role"
    Properties:
      Path: "/"
      RoleName: !Sub ${EnvironmentName}-trigger-role
      AssumeRolePolicyDocument: '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"Service":"lambda.amazonaws.com"},"Action":"sts:AssumeRole"}]}'
      MaxSessionDuration: 43200
      Description: "Allows Lambda functions to call AWS Inspector."
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/AmazonInspector2FullAccess"
        - "arn:aws:iam::aws:policy/CloudWatchLogsFullAccess"

  SbomTransformationRole:
    Type: "AWS::IAM::Role"
    Properties:
      Path: "/"
      RoleName: !Sub ${EnvironmentName}-transformation-role
      AssumeRolePolicyDocument: '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"Service":"lambda.amazonaws.com"},"Action":"sts:AssumeRole"}]}'
      MaxSessionDuration: 43200
      Description: "Allows Lambda functions to call AWS Inspector."
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/CloudWatchLogsFullAccess"
      Policies:      
        - PolicyName: 's3-bucket-access-user-policy'
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
            - Sid: 'KMSList'
              Effect: 'Allow'
              Action:
                - 'kms:ListKeys'
              Resource: '*'
            - Sid: 'S3BucketAccess'
              Effect: 'Allow'
              Action:
                - 's3:PutObject'
                - 's3:GetObject'
                - 's3:ListBucket'
                - 'kms:GetPublicKey'
                - 'kms:GetKeyPolicy'
                - 'kms:DescribeKey'
                - 'kms:Decrypt'
                - 'kms:GenerateDataKey'
              Resource: "*"
  
  RawS3Bucket:
    Type: 'AWS::S3::Bucket'
    DependsOn: SbomTransformationLambdaPermission
    DeletionPolicy: Delete
    Properties:
      BucketName: !Join ['-', [!Ref EnvironmentName, !Ref "AWS::AccountId", "raw"]]
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: aws:kms
              KMSMasterKeyID: !Ref SbomCMK
            BucketKeyEnabled: true
      NotificationConfiguration:
        LambdaConfigurations:
          - Event: "s3:ObjectCreated:Put"
            Filter:
              S3Key:
                Rules:
                  - Name: Prefix
                    Value: "sbom/"
                  - Name: Suffix
                    Value: ".json"
            Function: !GetAtt SbomTransformationFunction.Arn

  RawS3BucketPolicy:
    Type: AWS::S3::BucketPolicy
    DependsOn: RawS3Bucket
    Properties:
      Bucket: !Ref RawS3Bucket
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Sid: "allow-inspector"
            Effect: "Allow"
            Action:
            - "s3:PutObject"
            - "s3:PutObjectAcl"
            - "s3:AbortMultipartUpload"
            Principal:
              Service: "inspector2.amazonaws.com"
            Resource: !Join ['/', [!GetAtt RawS3Bucket.Arn, "*"]]
          - Sid: "allow-lambda"
            Effect: "Allow"
            Action:
            - "s3:PutObject"
            - "s3:PutObjectAcl"
            - "s3:AbortMultipartUpload"
            - "s3:GetObject"
            Principal:
              Service: "lambda.amazonaws.com"
            Resource: !Join ['/', [!GetAtt RawS3Bucket.Arn, "*"]]

  TransformedS3Bucket:
    Type: 'AWS::S3::Bucket'
    DeletionPolicy: Delete
    Properties:
      BucketName: !Join ['-', [!Ref EnvironmentName, !Ref "AWS::AccountId", "transformed"]]

  AthenaS3Bucket:
    Type: 'AWS::S3::Bucket'
    DeletionPolicy: Delete
    Properties:
      BucketName: !Join ['-', [!Ref EnvironmentName, !Ref "AWS::AccountId", "athena"]]


  EventBridgeTriggerRule: 
    Type: AWS::Events::Rule
    DependsOn: SbomTriggerFunction
    Properties:
      Description: "Trigger SBOM export lambda once image successfully pushed into ECR repo"
      EventBusName: default
      EventPattern: |
        {
          "source": ["aws.ecr"],
          "detail-type": ["ECR Image Action"],
          "detail": {
            "action-type": ["PUSH"]
          }
        }
      Name: "trigger-sbom-export"
      State: ENABLED
      Targets: 
        - Arn: !GetAtt SbomTriggerFunction.Arn
          Id: !Join ['-', [!Ref EnvironmentName, "trigger" ,!Ref "AWS::AccountId", "rule"]]

  SbomTriggerLambdaPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: 'lambda:InvokeFunction'
      FunctionName: !Ref SbomTriggerFunction
      Principal: "events.amazonaws.com"
      SourceArn: !Sub arn:aws:events:${AWS::Region}:${AWS::AccountId}:rule/trigger-sbom-export
      SourceAccount: !Ref "AWS::AccountId"
  
  SbomTransformationLambdaPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: 'lambda:InvokeFunction'
      FunctionName: !Ref SbomTransformationFunction
      Principal: s3.amazonaws.com
      SourceArn: !Sub 'arn:aws:s3:::${EnvironmentName}-${AWS::AccountId}-raw'
      SourceAccount: !Ref "AWS::AccountId"

  SbomTriggerFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: sbom-trigger
      Handler: index.lambda_handler
      Environment:
        Variables:
          s3Bucket: !Ref RawS3Bucket
          kmsKey: !GetAtt SbomCMK.Arn
      Runtime: python3.11
      Role:
        Fn::GetAtt:
          - SbomTriggerRole
          - Arn
      Timeout: 60
      Code:
        ZipFile: |
          import json
          import boto3
          import os

          def lambda_handler(event, context):

              # Populate env variables to feed into inspector sbom function

              imagetag = event['detail']['image-tag']
              reponame = event['detail']['repository-name']
              s3bucket = os.environ["s3Bucket"]
              kmskeyarn = os.environ["kmsKey"]
              print(imagetag, reponame)
              
              client = boto3.client('inspector2')
              
              response = client.create_sbom_export(
                  reportFormat='SPDX_2_3',
                  resourceFilterCriteria={
                      'ecrImageTags': [
                          {
                              'comparison': 'EQUALS',
                              'value': imagetag
                          },
                      ],
                      'ecrRepositoryName': [
                      {
                              'comparison': 'EQUALS',
                              'value': reponame
                          },
                      ]
                  },
                  s3Destination={
                      'bucketName': s3bucket,
                      'keyPrefix': "sbom",
                      'kmsKeyArn': kmskeyarn
                  }
              )
              
              print(response)
              return {
                  'statusCode': 200,
                  'body': json.dumps('Exporting SBOM')
              }

  SbomTransformationFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: sbom-transformation
      Handler: index.lambda_handler
      Environment:
        Variables:
          s3Bucket: !Ref TransformedS3Bucket
      Runtime: python3.11
      Layers:
        - !Sub "arn:aws:lambda:${AWS::Region}:336392948345:layer:AWSSDKPandas-Python311:1"
      Role:
        Fn::GetAtt:
          - SbomTransformationRole
          - Arn
      Timeout: 60
      Code:
        ZipFile: |
          import json
          import boto3
          import pandas as pd
          import urllib.parse
          import os as ev

          def lambda_handler(event, context):
              
              print(event)
              
              bucket = event['Records'][0]['s3']['bucket']['name']
              object_key = urllib.parse.unquote(event['Records'][0]['s3']['object']['key'])
              filename = object_key.split("/")[-1]
              print(filename)
              s3bucket = ev.environ["s3Bucket"]
              
              # Download log file from s3
              s3 = boto3.resource('s3')
              s3.meta.client.download_file(bucket, object_key, "/tmp/" + filename)
              
              f = open('/tmp/'+filename)
              data = json.load(f)
              columns = ['ImageArn', 'CreatedOn', 'AccountNo', 'DocumentNamespace', 'Package', 'PackageVersion', 'OS', 'SHA']
              
              reponame = data['name'].rsplit("/", 1)[0].split("ECR/")[-1]
              createdon = data['creationInfo']['created']
              account = data['name'].split("/")[0]
              documentns = data['documentNamespace']
              os = data['documentNamespace'].split('/')[2]
              sha = data['documentNamespace'].split('/')[-1]
              
              sbom = []
              cve = []
              
              # Create SBOM dataframe
              for package in data['packages']:
                  packageName = package['name']
                  packageVersion = package['versionInfo']
                  sbom.append([reponame, createdon, account, documentns, packageName, packageVersion, os, sha])
              
              sbom_df = pd.DataFrame(sbom, columns=columns)
              
              sbom_filenm = filename.replace(".json", "") + ".parquet"
              
              sbom_df.to_parquet("/tmp/" + sbom_filenm)
                  
              # Create CVE dataframe
              for package in data['packages']:
                  for extref in package['externalRefs']:
                      if extref['referenceType'] == "vulnerability":
                          cve.append([sha, package['name'], package['versionInfo'], extref["referenceLocator"]])
                          
              cve_df = pd.DataFrame(cve, columns=["SHA", "Package", "PackageVersion", "CVE"])
              
              cve_filenm = "cve_" + filename.replace(".json", "") + ".parquet"
              
              cve_df.to_parquet("/tmp/" + cve_filenm)
              
              s3 = boto3.resource('s3')
              s3.meta.client.upload_file("/tmp/" + sbom_filenm, s3bucket, "sbom/" + sbom_filenm)
              s3.meta.client.upload_file("/tmp/" + cve_filenm, s3bucket, "cve/" + cve_filenm)
              return {
                  'statusCode': 200,
                  'body': json.dumps('Transferred SBOM parquert file to S3')
              }