Optional EFS mount usage to scan large files

README.md

@@ -103,6 +103,8 @@ the table below for reference.
 | AV_PROCESS_ORIGINAL_VERSION_ONLY | Controls that only original version of an S3 key is processed (if bucket versioning is enabled) | False | No |
 | AV_DELETE_INFECTED_FILES | Controls whether infected files should be automatically deleted | False | No |
 | EVENT_SOURCE | The source of antivirus scan event "S3" or "SNS" (optional) | S3 | No |
+| AV_EFS_MOUNT_POINT | EFS mount point that used to scan larger files (optional) | | No |
+| AV_EFS_LARGE_FILE_SIZE_THRESHOLD | Threshold for a file size in bytes to be copied to EFS-mounted system instead of lambda internal storage. Only useful if AV_EFS_MOUNT_POINT is set | 314572800 | No |
 | S3_ENDPOINT | The Endpoint to use when interacting wth S3 | None | No |
 | SNS_ENDPOINT | The Endpoint to use when interacting wth SNS | None | No |
 | LAMBDA_ENDPOINT | The Endpoint to use when interacting wth Lambda | None | No |

common.py

@@ -41,6 +41,11 @@
 )
 AV_DELETE_INFECTED_FILES = os.getenv("AV_DELETE_INFECTED_FILES", "False")
 
+AV_EFS_MOUNT_POINT = os.getenv("AV_EFS_MOUNT_POINT")
+AV_EFS_LARGE_FILE_SIZE_THRESHOLD = os.getenv(
+    "AV_EFS_LARGE_FILE_SIZE_THRESHOLD", 314572800
+)  # 300MB
+
 AV_DEFINITION_FILE_PREFIXES = ["main", "daily", "bytecode"]
 AV_DEFINITION_FILE_SUFFIXES = ["cld", "cvd"]
 SNS_ENDPOINT = os.getenv("SNS_ENDPOINT", None)

scan.py

@@ -39,9 +39,13 @@
 from common import AV_TIMESTAMP_METADATA
 from common import SNS_ENDPOINT
 from common import S3_ENDPOINT
+from common import AV_EFS_MOUNT_POINT
+from common import AV_EFS_LARGE_FILE_SIZE_THRESHOLD
 from common import create_dir
 from common import get_timestamp
 
+DEFAULT_SCAN_DIR = "/tmp"
+
 
 def event_object(event, event_source="s3"):
 
@@ -100,8 +104,21 @@ def verify_s3_object_version(s3, s3_object):
         )
 
 
-def get_local_path(s3_object, local_prefix):
-    return os.path.join(local_prefix, s3_object.bucket_name, s3_object.key)
+def get_local_path(s3_object):
+    return get_local_path_internal(
+        s3_object,
+        DEFAULT_SCAN_DIR,
+        AV_EFS_MOUNT_POINT,
+        int(AV_EFS_LARGE_FILE_SIZE_THRESHOLD),
+    )
+
+
+def get_local_path_internal(s3_object, local_prefix, efs_prefix, efs_threshold):
+    if efs_prefix and s3_object.content_length > efs_threshold:
+        prefix = efs_prefix
+    else:
+        prefix = local_prefix
+    return os.path.join(prefix, s3_object.bucket_name, s3_object.key)
 
 
 def delete_s3_object(s3_object):
@@ -221,7 +238,7 @@ def lambda_handler(event, context):
         start_scan_time = get_timestamp()
         sns_start_scan(sns_client, s3_object, AV_SCAN_START_SNS_ARN, start_scan_time)
 
-    file_path = get_local_path(s3_object, "/tmp")
+    file_path = get_local_path(s3_object)
     create_dir(os.path.dirname(file_path))
     s3_object.download_file(file_path)
 

scan_test.py

@@ -29,7 +29,7 @@
 from common import get_timestamp
 from scan import delete_s3_object
 from scan import event_object
-from scan import get_local_path
+from scan import get_local_path_internal
 from scan import set_av_metadata
 from scan import set_av_tags
 from scan import sns_start_scan
@@ -261,13 +261,44 @@ def test_sns_start_scan(self):
             s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
             sns_start_scan(self.sns_client, s3_obj, sns_arn, timestamp)
 
-    def test_get_local_path(self):
-        local_prefix = "/tmp"
-
+    def test_get_local_path_internal(self):
         s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
-        file_path = get_local_path(s3_obj, local_prefix)
-        expected_file_path = "/tmp/test_bucket/test_key"
-        self.assertEquals(file_path, expected_file_path)
+        s3_stubber_resource = Stubber(self.s3.meta.client)
+        content_length = 200
+        head_object_response = {
+            "ContentType": "content",
+            "Metadata": {},
+            "ContentLength": content_length,
+        }
+        head_object_expected_params = {
+            "Bucket": self.s3_bucket_name,
+            "Key": self.s3_key_name,
+        }
+        s3_stubber_resource.add_response(
+            "head_object", head_object_response, head_object_expected_params
+        )
+
+        with s3_stubber_resource:
+            file_path = get_local_path_internal(
+                s3_obj, "/tmp", "/mnt", content_length - 1
+            )
+
+            expected_file_path = "/mnt/test_bucket/test_key"
+            self.assertEquals(file_path, expected_file_path)
+
+        with s3_stubber_resource:
+            file_path = get_local_path_internal(
+                s3_obj, "/tmp", "/mnt", content_length + 1
+            )
+
+            expected_file_path = "/tmp/test_bucket/test_key"
+            self.assertEquals(file_path, expected_file_path)
+
+        with s3_stubber_resource:
+            file_path = get_local_path_internal(s3_obj, "/tmp", None, None)
+
+            expected_file_path = "/tmp/test_bucket/test_key"
+            self.assertEquals(file_path, expected_file_path)
 
     def test_set_av_metadata(self):
         scan_result = "CLEAN"