index.CallbackFunction.ts

import { S3Client, DeleteObjectCommand, PutObjectTaggingCommand } from '@aws-sdk/client-s3';
import { APIGatewayProxyEventV2, APIGatewayProxyResultV2 } from 'aws-lambda';
import { createPublicKey, createVerify } from 'node:crypto';

const s3Client = new S3Client({});

const TAG_OBJECT_WITH_SCAN_RESULT = process.env.TAG_OBJECT_WITH_SCAN_RESULT === 'true';
const DELETE_INFECTED_OBJECT = process.env.DELETE_INFECTED_OBJECT === 'true';
const TENANT_ID = process.env.TENANT_ID;
const SIGNATURE_TOLERANCE_IN_MILLIS = 5 * 60 * 1000;

const PUBLIC_KEY_PEM = `-----BEGIN RSA PUBLIC KEY-----
MIIBCgKCAQEAyLoZzjo1cQV9ZN2TH/alrxWiQ3u/ndT0HMrLMdBTVO3Tz1nUjLt6
SqKZsN8dQhvPoEjfyhCTEg7MPWopG3n0cf3NRxtoeXy/Z62b1zdUd426kMuKOQP8
Yy6cxa/RtK2tkHCnTGxjfvNmMK+m68sFmsilR88LnIN71my4cG8bIDGDftWublvK
AEOWhxSECYn1XEtyrQL5lm8HFnHdE9ys56xTJkdr5Mmkvanrnd/hXzTHzjruGcLv
bjciI82+Z335AzYgJcnmH4/zsBuyPL2FJSfQF9NsPaTJuQgkw1usAKBQcujcEriY
UDNWgTe1a+LOnCEMb+9mAYw8lMRYRd3CBwIDAQAB
-----END RSA PUBLIC KEY-----`;

interface AttachmentAVResult {
  status: 'clean' | 'infected' | 'no';
  finding?: string;
  size: number;
  scan_time: number;
  download_time: number;
  realfiletype?: string;
  custom_data?: string;
}

// custom data is the same data that has been provided to the endpoint in the Scanner Function
interface AttachmentAVCustomData {
  bucket: string;
  key: string;
  versionId?: string;
}

/**
 * Verifies the signature of a callback.
 *
 * @param {number} unixtimeInMillis - Current unixtime in milliseconds (Date.now())
 * @param {string} publicKeyPEM -  The public key in PEM PKCS1 format (available at https://attachmentav.com/help/virus-malware-scan-api/developer/publickey.html)
 * @param {string} timestamp - The value of the X-Timestamp header of the callback
 * @param {string} tenantId - Your tenant ID (get your tenant ID from https://attachmentav.com/help/virus-malware-scan-api/developer/definition.html#action-whoami)
 * @param {string} callbackUrl - The callback URL used when submitting the scan job
 * @param {string} body - The body of the callback
 * @param {string} signature - The value of the X-Signature header of the callback
 * @return {boolean} - true if the callback is valid
 */
function verify(unixtimeInMillis: number, publicKeyPEM: string, timestamp: string, tenantId: string, callbackUrl: string, body: string, signature: string) {
  const publicKey = createPublicKey({
    key: publicKeyPEM,
    format: 'pem',
    type: 'pkcs1'
  });
  const verify = createVerify('sha256');
  verify.update(timestamp);
  verify.update('.');
  verify.update(tenantId);
  verify.update('.');
  verify.update(callbackUrl);
  verify.update('.');
  verify.update(body);
  verify.end();
  const valid = verify.verify(publicKey, signature, 'hex');
  return valid && Math.abs(unixtimeInMillis-parseInt(timestamp, 10)) <= SIGNATURE_TOLERANCE_IN_MILLIS;
}

function logResult(bucket: string, key: string, versionId: string | undefined, result: AttachmentAVResult) {
  // Log structured scan results
  console.log('Scan Results:', JSON.stringify({
    bucket,
    key,
    versionId,
    status: result.status,
    finding: result.finding || 'none',
    size: result.size,
    realFileType: result.realfiletype,
    downloadTime: result.download_time,
    scanTime: result.scan_time,
    timestamp: new Date().toISOString(),
  }, null, 2));

  // Log specific messages based on scan status
  if (result.status === 'clean') {
    console.log(`✓ File is CLEAN - No threats detected`);
    console.log(`  File type: ${result.realfiletype}`);
    console.log(`  File size: ${result.size} bytes`);
  } else if (result.status === 'infected') {
    console.warn(`⚠ File is INFECTED - Malware detected!`);
    console.warn(`  Threat: ${result.finding || 'Unknown'}`);
    console.warn(`  File type: ${result.realfiletype}`);
    console.warn(`  File size: ${result.size} bytes`);
  } else {
    console.warn(`⚠ Scan status UNKNOWN - Unable to determine file status`);
    console.warn(`  File type: ${result.realfiletype}`);
    console.warn(`  File size: ${result.size} bytes`);
  }
}

async function tagObjectWithResult(bucket: string, key: string, versionId: string | undefined, result: AttachmentAVResult) {
  const objectRef = versionId
    ? `s3://${bucket}/${key} (version: ${versionId})`
    : `s3://${bucket}/${key}`;
  console.log(`Tagging object with scan results: ${objectRef}`);

  try {
    await s3Client.send(
      new PutObjectTaggingCommand({
        Bucket: bucket,
        Key: key,
        ...(versionId ? { VersionId: versionId } : {}),
        Tagging: {
          TagSet: [
            { Key: 'ScanStatus', Value: result.status },
            { Key: 'ScanFinding', Value: result.finding || 'none' },
            { Key: 'ScanTimestamp', Value: new Date().toISOString() },
          ],
        },
      })
    );
    console.log(`✓ Successfully tagged object with scan results`);
  } catch (tagError) {
    console.error(`Failed to tag object:`, tagError);
  }
}

async function deleteObject(bucket: string, key: string, versionId: string | undefined) {
  console.warn(`Deleting infected file: s3://${bucket}/${key} (version: ${versionId || ''}`);

  try {
    await s3Client.send(
      new DeleteObjectCommand({
        Bucket: bucket,
        Key: key,
        ...(versionId ? { VersionId: versionId } : {}),
      })
    );
    console.warn(`✓ Successfully deleted infected file`);
  } catch (deleteError) {
    console.error(`Failed to delete infected file:`, deleteError);
    // TODO: notify someone if deletion fails
  }
}

export const handler = async (event: APIGatewayProxyEventV2): Promise<APIGatewayProxyResultV2> => {
  console.debug(`Received callback from attachmentAV: ${JSON.stringify(event, null, 2)}`);

  if (!event.body) {
    console.error('No body in callback request');
    return {
      statusCode: 400,
      body: JSON.stringify({ error: 'Missing request body' }),
    };
  }

  if (TENANT_ID) {
    const attachmentAVTimestamp = event.headers['x-timestamp'];
    const attachmentAVSignature = event.headers['x-signature'];
    if (!attachmentAVTimestamp || !attachmentAVSignature) {
      console.error('can not validate signature as timestamp or signature as missing');
      return {
        statusCode: 432,
        body: JSON.stringify({ error: 'callback signature validation failed: no timestamp or signature provided' }),
      };
    }

    const callbackUrl = `https://${event.requestContext.domainName}/`;
    const isValidCallback = verify(Date.now(), PUBLIC_KEY_PEM, attachmentAVTimestamp, TENANT_ID, callbackUrl, event.body, attachmentAVSignature);
    console.log(`signature verification result: ${isValidCallback ? 'valid' : 'invalid'}`);

    if (!isValidCallback) {
      console.error('callback signature validation failed');
      return {
        statusCode: 432,
        body: JSON.stringify({ error: 'callback signature validation failed' }),
      };
    }
  }

  const result: AttachmentAVResult = JSON.parse(event.body);

  // Extract bucket, key, and versionId from custom_data
  const customData = result.custom_data ? JSON.parse(result.custom_data) as AttachmentAVCustomData : undefined;
  const bucket = customData?.bucket;
  const key = customData?.key;
  const versionId = customData?.versionId;

  if (!bucket || !key) {
    console.error('Missing bucket or key in custom_data');
    return {
      statusCode: 400,
      body: JSON.stringify({ error: 'Missing bucket or key in custom_data' }),
    };
  }

  if (versionId) {
    console.log(`Scan results for s3://${bucket}/${key} (version: ${versionId})`);
  } else {
    console.log(`Scan results for s3://${bucket}/${key}`);
  }

  logResult(bucket, key, versionId, result);

  // Tag object with scan results if configured
  if (TAG_OBJECT_WITH_SCAN_RESULT) {
    await tagObjectWithResult(bucket, key, versionId, result);
  }

  // Delete infected files if configured
  if (result.status === 'infected' && DELETE_INFECTED_OBJECT) {
    await deleteObject(bucket, key, versionId);
  }

  // Return success response
  return {
    statusCode: 200,
    body: JSON.stringify({ message: 'Scan results received and logged' }),
  };
};