index.ts

import * as cdk from 'aws-cdk-lib';
import * as lambda_nodejs from "aws-cdk-lib/aws-lambda-nodejs";
import * as s3 from 'aws-cdk-lib/aws-s3';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as events from 'aws-cdk-lib/aws-events';
import * as targets from 'aws-cdk-lib/aws-events-targets';
import * as s3n from 'aws-cdk-lib/aws-s3-notifications';
import * as ssm from 'aws-cdk-lib/aws-ssm';
import { Construct } from 'constructs';

const DEFAULT_ATTACHMENTAV_ENDPOINT = 'https://eu.developer.attachmentav.com';

/**
 * Strategy for triggering the scanner Lambda function
 */
export enum TriggerStrategy {
  /**
   * Use S3 event notification to directly invoke the Lambda function
   * This is the most efficient approach with the lowest latency
   */
  S3_EVENT = 'S3_EVENT',

  /**
   * Use EventBridge to route S3 events to the Lambda function
   * This allows for more complex event filtering and routing
   */
  EVENTBRIDGE = 'EVENTBRIDGE',
}

/**
 * Properties for the AttachmentAV construct
 */
export interface AttachmentAVBucketScanProps {
  /**
   * The S3 bucket to monitor for new uploads
   */
  readonly bucket: s3.IBucket;

  /**
   * SSM Parameter Store parameter name containing the attachmentAV API key
   * Must be a SecureString parameter
   *
   * @example '/attachmentav/api-key'
   */
  readonly apiKeyParameterName: string;

  /**
   * Strategy for triggering the scanner Lambda
   *
   * @default TriggerStrategy.S3_EVENT
   */
  readonly triggerStrategy?: TriggerStrategy;

  /**
   * Whether to tag S3 objects with scan results
   *
   * @default true
   */
  readonly tagObjectWithScanResult?: boolean;

  /**
   * Whether to delete infected S3 objects
   * If true, infected files will be automatically deleted from S3
   *
   * @default false
   */
  readonly deleteInfectedObject?: boolean;

  /**
   * Optional S3 key prefix filter for triggering scans.
   * Only objects with keys starting with this prefix will trigger scans.
   *
   * @example 'uploads/'
   * @default undefined (all objects trigger scans)
   */
  readonly s3KeyPrefix?: string;

  /**
   * Optional S3 key suffix filter for triggering scans.
   * Only objects with keys ending with this suffix will trigger scans.
   *
   * @example '.pdf'
   * @default undefined (all objects trigger scans)
   */
  readonly s3KeySuffix?: string;

  /**
   * The attachmentAV API base URL
   *
   * @default 'https://eu.developer.attachmentav.com'
   * @example 'https://us.developer.attachmentav.com'
   */
  readonly apiUrl?: string;

  /**
   * The attachmentAV tenant id.
   * If provided, all incoming callback events are verified.
   */
  readonly tenantId?: string;
}

/**
 * CDK Construct for integrating attachmentAV malware scanning with S3 buckets
 *
 * This construct creates two Lambda functions that automatically scan files uploaded to an S3 bucket
 * using the attachmentAV API and tag these files with their scan result or even deletes infected files if configured.
 *
 * @example
 * ```typescript
 * const bucket = new s3.Bucket(this, 'MyBucket');
 *
 * new AttachmentAVBucketScan(this, 'AttachmentAVBucketScan', {
 *   bucket: bucket,
 *   apiKeyParameterName: '/attachmentav/api-key',
 *   triggerStrategy: TriggerStrategy.S3_EVENT, // default
 *   tagObjectWithScanResult: true, // default
 *   deleteInfectedObject: false, // default
 *   s3KeyPrefix: 'uploads/', // optional
 *   s3KeySuffix: '.pdf', // optional
 *   apiUrl: 'https://eu.developer.attachmentav.com', // default
 *   tenantId: '...', //optional
 * });
 * ```
 */
export class AttachmentAVBucketScan extends Construct {
  /**
   * The Lambda function that scans uploaded files
   */
  public readonly scannerFunction: lambda.Function;

  /**
   * The Lambda function that receives scan results via webhook
   */
  public readonly callbackFunction: lambda.Function;

  /**
   * The Function URL for the callback Lambda
   */
  public readonly callbackUrl: lambda.FunctionUrl;

  /**
   * The IAM User used for generating presigned URLs
   */
  public readonly downloadUser: iam.User;

  /**
   * The SSM Parameter containing the IAM User's Access Key ID
   */
  public readonly accessKeyIdParameter: ssm.StringParameter;

  /**
   * The SSM Parameter containing the IAM User's Secret Access Key
   */
  public readonly secretAccessKeyParameter: ssm.StringParameter;

  constructor(scope: Construct, id: string, props: AttachmentAVBucketScanProps) {
    super(scope, id);

    const triggerStrategy = props.triggerStrategy ?? TriggerStrategy.S3_EVENT;
    const apiUrl = props.apiUrl ?? DEFAULT_ATTACHMENTAV_ENDPOINT;
    const tagObjectWithScanResult = props.tagObjectWithScanResult ?? true;
    const deleteInfectedObject = props.deleteInfectedObject ?? false;

    // Create IAM User for generating presigned URLs with permanent credentials
    this.downloadUser = new iam.User(this, 'DownloadUser');

    // Grant the IAM User read permissions to the S3 bucket
    props.bucket.grantRead(this.downloadUser);

    // Create Access Key for the IAM User
    const accessKey = new iam.CfnAccessKey(this, 'DownloadAccessKey', {
      userName: this.downloadUser.userName,
    });

    this.accessKeyIdParameter = new ssm.StringParameter(this, 'DownloadAccessKeyIdParameter', {
      parameterName: `/${cdk.Stack.of(this).stackName}/${id}/presign-download-user/access-key-id`,
      stringValue: accessKey.ref,
      description: 'Access Key ID for attachmentAV presigned URL generation',
    });

    this.secretAccessKeyParameter = new ssm.StringParameter(this, 'DownloadSecretAccessKeyParameter', {
      parameterName: `/${cdk.Stack.of(this).stackName}/${id}/presign-download-user/secret-access-key`,
      stringValue: accessKey.attrSecretAccessKey,
      description: 'Secret Access Key for attachmentAV presigned URL generation',
    });

    // Create callback Lambda function that receives scan results
    this.callbackFunction = new lambda_nodejs.NodejsFunction(this, 'CallbackFunction', {
      runtime: lambda.Runtime.NODEJS_24_X,
      timeout: cdk.Duration.seconds(10),
      memorySize: 256,
      description: 'Handles and logs attachmentAV scan results',
      environment: {
        TAG_OBJECT_WITH_SCAN_RESULT: tagObjectWithScanResult.toString(),
        DELETE_INFECTED_OBJECT: deleteInfectedObject.toString(),
        ...(props.tenantId ? { TENANT_ID: props.tenantId } : {}),
      },
    });

    // Create Function URL for callback Lambda (no auth required for webhook)
    this.callbackUrl = this.callbackFunction.addFunctionUrl({
      authType: lambda.FunctionUrlAuthType.NONE,
    });

    // Grant callback Lambda permissions based on configuration
    if (tagObjectWithScanResult) {
      // Grant permission to tag S3 objects
      this.callbackFunction.addToRolePolicy(
        new iam.PolicyStatement({
          actions: ['s3:ListBucket'],
          resources: [props.bucket.bucketArn],
        })
      );
      this.callbackFunction.addToRolePolicy(
        new iam.PolicyStatement({
          actions: ['s3:PutObjectTagging', 's3:PutObjectVersionTagging'],
          // TODO: apply prefix and suffix
          resources: [`${props.bucket.bucketArn}/*`],
        })
      );
    }

    if (deleteInfectedObject) {
      // Grant permission to delete S3 objects
      // TODO: apply prefix and suffix
      props.bucket.grantDelete(this.callbackFunction);
    }

    // Create scanner Lambda function that initiates scans
    this.scannerFunction = new lambda_nodejs.NodejsFunction(this, 'ScannerFunction', {
      runtime: lambda.Runtime.NODEJS_24_X,
      timeout: cdk.Duration.seconds(30),
      memorySize: 256,
      description: 'Initiates attachmentAV scans for S3 objects',
      environment: {
        ATTACHMENTAV_API_KEY_PARAMETER: props.apiKeyParameterName,
        ATTACHMENTAV_API_URL: apiUrl,
        CALLBACK_URL: this.callbackUrl.url,
        AWS_ACCESS_KEY_ID_PARAMETER: this.accessKeyIdParameter.parameterName,
        AWS_SECRET_ACCESS_KEY_PARAMETER: this.secretAccessKeyParameter.parameterName,
      },
    });

    // Grant scanner Lambda permission to read SSM parameters
    this.scannerFunction.addToRolePolicy(
      new iam.PolicyStatement({
        actions: ['ssm:GetParameter'],
        resources: [
          `arn:aws:ssm:${cdk.Stack.of(this).region}:${cdk.Stack.of(this).account}:parameter${props.apiKeyParameterName}`,
          this.accessKeyIdParameter.parameterArn,
          this.secretAccessKeyParameter.parameterArn,
        ],
      })
    );

    // Configure trigger based on strategy
    if (triggerStrategy === TriggerStrategy.S3_EVENT) {
      // Use S3 event notification (default)
      if (props.s3KeyPrefix || props.s3KeySuffix) {
        props.bucket.addEventNotification(
          s3.EventType.OBJECT_CREATED,
          new s3n.LambdaDestination(this.scannerFunction),
          { prefix: props.s3KeyPrefix, suffix: props.s3KeySuffix },
        );
      } else {
        props.bucket.addEventNotification(
          s3.EventType.OBJECT_CREATED,
          new s3n.LambdaDestination(this.scannerFunction),
        );
      }
    } else if (triggerStrategy === TriggerStrategy.EVENTBRIDGE) {
      // Use EventBridge for more complex routing

      // Enable EventBridge on the bucket if not already enabled
      // Note: This is a bucket-level setting and may affect other event consumers
      props.bucket.enableEventBridgeNotification();

      // Create EventBridge rule
      const rule = new events.Rule(this, 'ScannerRule', {
        description: 'Trigger attachmentAV scanner on S3 object creation',
        eventPattern: {
          source: ['aws.s3'],
          detailType: ['Object Created'],
          detail: {
            bucket: {
              name: [props.bucket.bucketName],
            },
            ...(props.s3KeyPrefix || props.s3KeySuffix
              ? {
                object: {
                  key: [
                    { wildcard: `${props.s3KeyPrefix ?? ''}*${props.s3KeySuffix ?? ''}` },
                  ],
                },
              }
              : {}),
          },
        },
      });

      rule.addTarget(new targets.LambdaFunction(this.scannerFunction));
    }
  }
}