Added File Size Check to Guard API and Trigger

.gitignore

@@ -2,3 +2,4 @@ lambda.zip
 s3-antivirus/node_modules
 .idea
 .DS_store
+.vscode

s3-antivirus/antivirus.js

@@ -10,6 +10,17 @@ const s3 = new AWS.S3();
 const utils = require('./utils');
 const constants = require('./constants');
 
+async function sizeOf(key, bucket) {
+    let res = await s3.headObject({ Key: key, Bucket: bucket }).promise();
+    return res.ContentLength;
+}
+
+async function isS3FileToBig(s3ObjectKey, s3ObjectBucket)
+{
+    let fileSize = await sizeOf(s3ObjectKey, s3ObjectBucket);
+    return (fileSize > constants.MAX_FILE_SIZE);
+}
+
 function downloadFileFromS3(s3ObjectKey, s3ObjectBucket) {
     const downloadDir = `/tmp/download`;
     if (!fs.existsSync(downloadDir)){
@@ -46,9 +57,18 @@ async function lambdaHandleEvent(event, context) {
     //Here to end of function can be functioned out for reuse.
     await clamav.downloadAVDefinitions(constants.CLAMAV_BUCKET_NAME, constants.PATH_TO_AV_DEFINITIONS);
 
-    await downloadFileFromS3(s3ObjectKey, s3ObjectBucket);
+    let virusScanStatus;
+
+    //You need to verify that you are not getting to git of a file
+    //at this point in time lambdas max out at 500MB storage.
+    if(await isS3FileToBig(s3ObjectKey, s3ObjectBucket)){
+        virusScanStatus = constants.STATUS_SKIPPED_FILE;
+    }
+    else{
+        await downloadFileFromS3(s3ObjectKey, s3ObjectBucket);
+        clamav.scanLocalFile(path.basename(s3ObjectKey));
+    }
 
-    let virusScanStatus = clamav.scanLocalFile(path.basename(s3ObjectKey));
 
     var taggingParams = {
         Bucket: s3ObjectBucket,
@@ -91,6 +111,8 @@ async function scanS3Object(s3ObjectKey, s3ObjectBucket){
 
 module.exports = {
     lambdaHandleEvent:  lambdaHandleEvent,
-    scanS3Object:       scanS3Object
+    scanS3Object:       scanS3Object,
+    isS3FileToBig:      isS3FileToBig,
+    sizeOf:             sizeOf
 };
 

s3-antivirus/api.js

@@ -2,15 +2,15 @@ const AWS = require('aws-sdk');
 const path = require('path');
 const utils = require('./utils');
 const av = require('./antivirus');
+const constants = require('./constants');
 
 async function lambdaHandleEvent(event, context) {
 
-    let jsonBody = JSON.parse(event.body);
+    let s3ObjectKey = utils.extractKeyFromApiEvent(event);
+    let s3ObjectBucket = utils.extractBucketFromApiEvent(event);
 
-    let s3ObjectKey = utils.extractKeyFromApiEvent(jsonBody);
-    let s3ObjectBucket = utils.extractBucketFromApiEvent(jsonBody);
-
-    let virusScanStatus = await av.scanS3Object(s3ObjectKey, s3ObjectBucket);
+    let virusScanStatus =
+     await av.isS3FileToBig(s3ObjectKey, s3ObjectBucket) ? constants.STATUS_SKIPPED_FILE : await av.scanS3Object(s3ObjectKey, s3ObjectBucket);
 
     return virusScanStatus;
 }

s3-antivirus/constants.js

@@ -27,8 +27,10 @@ const FRESHCLAM_WORK_DIR = '/tmp/';
 const STATUS_CLEAN_FILE = process.env.STATUS_CLEAN_FILE || 'CLEAN';
 const STATUS_INFECTED_FILE = process.env.STATUS_INFECTED_FILE || 'INFECTED';
 const STATUS_ERROR_PROCESSING_FILE = process.env.STATUS_ERROR_PROCESSING_FILE || 'ERROR';
+const STATUS_SKIPPED_FILE = process.env.STATUS_SKIPPED_FILE || 'SKIPPED';
 const VIRUS_SCAN_STATUS_KEY = process.env.VIRUS_SCAN_STATUS_KEY || 'virusScanStatus';
 const VIRUS_SCAN_TIMESTAMP_KEY = process.env.VIRUS_SCAN_TIMESTAMP_KEY || 'virusScanTimestamp';
+const MAX_FILE_SIZE = process.env.MAX_FILE_SIZE || '314572800';
 
 // List of CLAMAV definition files. These are the compressed files.
 const CLAMAV_DEFINITIONS_FILES = ['main.cvd', 'daily.cvd', 'bytecode.cvd'];
@@ -45,6 +47,8 @@ module.exports = {
     STATUS_CLEAN_FILE           : STATUS_CLEAN_FILE,
     STATUS_INFECTED_FILE        : STATUS_INFECTED_FILE,
     STATUS_ERROR_PROCESSING_FILE: STATUS_ERROR_PROCESSING_FILE,
+    STATUS_SKIPPED_FILE         : STATUS_SKIPPED_FILE,
     VIRUS_STATUS_STATUS_KEY: VIRUS_SCAN_STATUS_KEY,
-    VIRUS_SCAN_TIMESTAMP_KEY: VIRUS_SCAN_TIMESTAMP_KEY
+    VIRUS_SCAN_TIMESTAMP_KEY: VIRUS_SCAN_TIMESTAMP_KEY,
+    MAX_FILE_SIZE               : MAX_FILE_SIZE
 };