index.ScannerFunction.ts

import { S3Client, GetObjectCommand } from '@aws-sdk/client-s3';
import { SSMClient, GetParameterCommand } from '@aws-sdk/client-ssm';
import { getSignedUrl } from '@aws-sdk/s3-request-presigner';
import { S3Event, S3ObjectCreatedNotificationEvent } from 'aws-lambda';

const PRESIGNED_URL_EXPIRATION_SECONDS = 3600;
const API_URL = process.env.ATTACHMENTAV_API_URL || 'https://eu.developer.attachmentav.com';
const API_KEY_PARAMETER = process.env.ATTACHMENTAV_API_KEY_PARAMETER;
const CALLBACK_URL = process.env.CALLBACK_URL;
const AWS_ACCESS_KEY_ID_PARAMETER = process.env.AWS_ACCESS_KEY_ID_PARAMETER;
const AWS_SECRET_ACCESS_KEY_PARAMETER = process.env.AWS_SECRET_ACCESS_KEY_PARAMETER;

const ssmClient = new SSMClient({});

let attachmentAVApiKey: string;
let awsCredentials: { accessKeyId: string, secretAccessKey: string };

interface AttachmentAVCustomData {
  bucket: string;
  key: string;
  versionId?: string;
}

interface AttachmentAVRequest {
  download_url: string;
  callback_url: string;
  custom_data?: string;
}

async function getAttachmentAVApiKey() {
  const parameterResponse = await ssmClient.send(
    new GetParameterCommand({
      Name: API_KEY_PARAMETER,
      WithDecryption: true,
    })
  );

  const apiKey = parameterResponse.Parameter?.Value;
  if (!apiKey) {
    throw new Error(`Failed to retrieve API key from parameter: ${API_KEY_PARAMETER}`);
  }

  return apiKey;
}

async function getAWSCredentials() {
  if (!AWS_ACCESS_KEY_ID_PARAMETER || !AWS_SECRET_ACCESS_KEY_PARAMETER) {
    throw new Error('AWS credential parameter names are not set');
  }

  // Retrieve Access Key ID
  const accessKeyIdResponse = await ssmClient.send(
    new GetParameterCommand({
      Name: AWS_ACCESS_KEY_ID_PARAMETER,
      WithDecryption: true,
    })
  );

  const accessKeyId = accessKeyIdResponse.Parameter?.Value;
  if (!accessKeyId) {
    throw new Error(`Failed to retrieve Access Key ID from parameter: ${AWS_ACCESS_KEY_ID_PARAMETER}`);
  }

  // Retrieve Secret Access Key
  const secretAccessKeyResponse = await ssmClient.send(
    new GetParameterCommand({
      Name: AWS_SECRET_ACCESS_KEY_PARAMETER,
      WithDecryption: true,
    })
  );

  const secretAccessKey = secretAccessKeyResponse.Parameter?.Value;
  if (!secretAccessKey) {
    throw new Error(`Failed to retrieve Secret Access Key from parameter: ${AWS_SECRET_ACCESS_KEY_PARAMETER}`);
  }

  return {
    accessKeyId,
    secretAccessKey,
  };
}

async function startFileScan(s3Client: S3Client, callbackUrl: string, bucket: string, key: string, versionId?: string) {
  // Generate presigned URL for the S3 object (valid for 1 hour)
  // Using IAM User credentials ensures the URL remains valid even after Lambda credentials expire
  const getObjectCommand = new GetObjectCommand({
    Bucket: bucket,
    Key: key,
    ...(versionId && { VersionId: versionId }),
  });
  const presignedUrl = await getSignedUrl(s3Client, getObjectCommand, {
    expiresIn: PRESIGNED_URL_EXPIRATION_SECONDS,
  });

  console.log('Generated presigned URL using IAM User credentials');

  // Call attachmentAV API with custom_data containing bucket, key, and versionId
  // This allows the callback Lambda to know which S3 object was scanned
  const customData: AttachmentAVCustomData = {
    bucket: bucket,
    key: key,
    ...(versionId && { versionId: versionId }),
  };
  const requestBody: AttachmentAVRequest = {
    download_url: presignedUrl,
    callback_url: callbackUrl,
    custom_data: JSON.stringify(customData),
  };

  const apiEndpoint = `${API_URL}/v1/scan/async/download`;
  console.log(`Calling attachmentAV API: ${apiEndpoint}`);

  const response = await fetch(apiEndpoint, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'x-api-key': attachmentAVApiKey,
    },
    body: JSON.stringify(requestBody),
  });

  if (response.status === 204) {
    console.log(`Successfully submitted scan request for s3://${bucket}/${key} (version: ${versionId || ''})`);
  } else {
    const responseText = await response.text();
    console.error(`Unexpected response status: ${response.status}`);
    console.error(`Response body: ${responseText}`);
    throw new Error(`attachmentAV API returned status ${response.status}: ${responseText}`);
  }
}

export const handler = async (event: S3Event | S3ObjectCreatedNotificationEvent): Promise<void> => {
  console.log('Received S3 event:', JSON.stringify(event, null, 2));

  // Validate environment variables
  if (!API_KEY_PARAMETER) {
    throw new Error('ATTACHMENTAV_API_KEY_PARAMETER environment variable is not set');
  }
  if (!CALLBACK_URL) {
    throw new Error('CALLBACK_URL environment variable is not set');
  }
  if (!AWS_ACCESS_KEY_ID_PARAMETER || !AWS_SECRET_ACCESS_KEY_PARAMETER) {
    throw new Error('AWS credential parameter environment variables are not set');
  }

  // Retrieve API key and AWS credentials from SSM Parameter Store
  if (!attachmentAVApiKey) {
    attachmentAVApiKey = await getAttachmentAVApiKey();
  }
  if (!awsCredentials) {
    awsCredentials = await getAWSCredentials();
  }

  // Create S3 client with IAM User credentials for presigned URL generation
  const s3Client = new S3Client({
    credentials: {
      accessKeyId: awsCredentials.accessKeyId,
      secretAccessKey: awsCredentials.secretAccessKey,
    },
  });

  if ((event as S3Event).Records) {
    console.log('event is S3 event notification - processing records');

    for (const record of (event as S3Event).Records) {
      const bucket = record.s3.bucket.name;
      const key = record.s3.object.key;
      const versionId = record.s3.object.versionId;

      console.log(`Processing file: s3://${bucket}/${key} (version: ${versionId || ''})`);

      try {
        await startFileScan(s3Client, CALLBACK_URL, bucket, key, versionId);
      } catch (error) {
        console.error(`Error processing file s3://${bucket}/${key} (version: ${versionId || ''}):`, error);
        throw error;
      }
    }
  } else {
    console.log('event is EventBridge S3 event');

    const eventBridgeEvent = event as S3ObjectCreatedNotificationEvent;
    const bucket = eventBridgeEvent.detail.bucket.name;
    const key = eventBridgeEvent.detail.object.key;
    const versionId = eventBridgeEvent.detail.object["version-id"];

    console.log(`Processing file: s3://${bucket}/${key} (version: ${versionId || ''})`);

    try {
      await startFileScan(s3Client, CALLBACK_URL, bucket, key, versionId);
    } catch (error) {
      console.error(`Error processing file s3://${bucket}/${key} (version: ${versionId || ''}):`, error);
      throw error;
    }
  }
};