Merge 43ca1fb into 08bd015

lambda_functions/analyzer/binary_info.py

@@ -1,5 +1,6 @@
 """Keeps track of all information associated with and computed about a binary."""
 import os
+import subprocess
 import tempfile
 import time
 from typing import Any, Dict, List, Set
@@ -71,13 +72,9 @@ def __enter__(self):
         return self
 
     def __exit__(self, exception_type, exception_value, traceback):
-        """Remove the downloaded binary from local disk."""
-        # In Lambda, "os.remove" does not actually remove the file as expected.
-        # Thus, we first truncate the file to set its size to 0 before removing it.
-        if os.path.isfile(self.download_path):
-            with open(self.download_path, 'wb') as file:
-                file.truncate()
-            os.remove(self.download_path)
+        """Shred the downloaded binary and delete it from disk."""
+        # Note: This runs even during exception handling (it is the "with" context).
+        subprocess.check_call(['shred', '-u', self.download_path])
 
     @property
     def matched_rule_ids(self) -> Set[str]:

tests/lambda_functions/analyzer/main_test.py

@@ -53,6 +53,7 @@ def metadata(self):
         return GOOD_FILE_METADATA if self.key == GOOD_S3_OBJECT_KEY else EVIL_FILE_METADATA
 
 
+@mock.patch.object(subprocess, 'check_call')
 @mock.patch.object(subprocess, 'check_output', return_value=b'[{"yara_matches_found": false}]')
 class MainTest(fake_filesystem_unittest.TestCase):
     """Test end-to-end functionality of the analyzer."""
@@ -97,7 +98,7 @@ def setUp(self):
         # Mock S3 Object
         self.main.analyzer_aws_lib.S3.Object = MockS3Object
 
-    def test_analyze_lambda_handler(self, mock_suprocess: mock.MagicMock):
+    def test_analyze_lambda_handler(self, mock_output: mock.MagicMock, mock_call: mock.MagicMock):
         """Verify return value, logging, and boto3 calls when multiple files match YARA rules."""
         with mock.patch.object(self.main, 'LOGGER') as mock_logger:
             result = self.main.analyze_lambda_handler(self._test_event, TEST_CONTEXT)
@@ -119,11 +120,17 @@ def test_analyze_lambda_handler(self, mock_suprocess: mock.MagicMock):
             ])
 
             # Verify 2 subprocess calls (yextend over each binary)
-            mock_suprocess.assert_has_calls([
+            mock_output.assert_has_calls([
                 mock.call(['./yextend', '-r', COMPILED_RULES_FILEPATH, '-t', mock.ANY, '-j']),
                 mock.call(['./yextend', '-r', COMPILED_RULES_FILEPATH, '-t', mock.ANY, '-j'])
             ])
 
+            # Verify 2 shred calls
+            mock_call.assert_has_calls([
+                mock.call(['shred', '-u', mock.ANY]),
+                mock.call(['shred', '-u', mock.ANY])
+            ])
+
         # Verify return value.
         good_s3_id = 'S3:{}:{}'.format(MOCK_S3_BUCKET_NAME, GOOD_S3_OBJECT_KEY)
         evil_s3_id = 'S3:{}:{}'.format(MOCK_S3_BUCKET_NAME, EVIL_S3_OBJECT_KEY)
@@ -242,6 +249,3 @@ def test_analyze_lambda_handler(self, mock_suprocess: mock.MagicMock):
                 Namespace='BinaryAlert'
             )
         ])
-
-        # Verify that the downloaded file was removed from temp storage.
-        self.assertEqual([], os.listdir(tempfile.gettempdir()))