Change analyzer to read bucket name from event

lambda_functions/analyzer/main.py

@@ -1,6 +1,5 @@
 """AWS Lambda function for testing a binary against a list of YARA rules."""
 # Expects the following environment variables:
-#   S3_BUCKET_NAME: Name of the S3 bucket from which to download binaries.
 #   SQS_QUEUE_URL: URL of the queue from which messages originated (needed for message deletion).
 #   YARA_MATCHES_DYNAMO_TABLE_NAME: Name of the Dynamo table which stores YARA match results.
 #   YARA_ALERTS_SNS_TOPIC_ARN: ARN of the SNS topic which should be alerted on a YARA match.
@@ -35,10 +34,23 @@ def analyze_lambda_handler(event_data: Dict[str, Any], lambda_context) -> Dict[s
 
     Args:
         event_data: [dict] of the form: {
-            'S3Objects': [...],  # S3 object keys.
+            'Records': [
+                {
+                    "s3": {
+                        "object": {
+                            "key": "FileName.txt"
+                        },
+                        "bucket": {
+                            "name": "mybucket"
+                        }
+                    }
+                }
+            ],
             'SQSReceipts': [...]  # SQS receipt handles (to be deleted after processing).
         }
             There can be any number of S3objects, but no more than 10 SQS receipts.
+            The Records are the same format as the S3 Put event, which means the analyzer could be
+            directly linked to an S3 bucket notification if needed.
         lambda_context: LambdaContext object (with .function_version).
 
     Returns:
@@ -60,13 +72,13 @@ def analyze_lambda_handler(event_data: Dict[str, Any], lambda_context) -> Dict[s
     except ValueError:
         lambda_version = -1
 
-    LOGGER.info('Processing %d record(s)', len(event_data['S3Objects']))
-    for s3_key in event_data['S3Objects']:
-        # S3 keys in event notifications are url-encoded.
-        s3_key = urllib.parse.unquote_plus(s3_key)
-        LOGGER.info('Analyzing "%s"', s3_key)
+    LOGGER.info('Processing %d record(s)', len(event_data['Records']))
+    for record in event_data['Records']:
+        bucket_name = record['s3']['bucket']['name']
+        s3_key = urllib.parse.unquote_plus(record['s3']['object']['key'])
+        LOGGER.info('Analyzing "%s:%s"', bucket_name, s3_key)
 
-        with binary_info.BinaryInfo(os.environ['S3_BUCKET_NAME'], s3_key, ANALYZER) as binary:
+        with binary_info.BinaryInfo(bucket_name, s3_key, ANALYZER) as binary:
             result[binary.s3_identifier] = binary.summary()
             binaries.append(binary)
 
@@ -79,7 +91,10 @@ def analyze_lambda_handler(event_data: Dict[str, Any], lambda_context) -> Dict[s
                 LOGGER.info('%s did not match any YARA rules', binary)
 
     # Delete all of the SQS receipts (mark them as completed).
-    analyzer_aws_lib.delete_sqs_messages(os.environ['SQS_QUEUE_URL'], event_data['SQSReceipts'])
+    analyzer_aws_lib.delete_sqs_messages(
+        os.environ['SQS_QUEUE_URL'],
+        event_data.get('SQSReceipts', [])
+    )
 
     # Publish metrics.
     try:

terraform/lambda.tf

@@ -98,7 +98,6 @@ module "binaryalert_analyzer" {
   filename        = "lambda_analyzer.zip"
 
   environment_variables = {
-    S3_BUCKET_NAME                 = "${aws_s3_bucket.binaryalert_binaries.id}"
     SQS_QUEUE_URL                  = "${aws_sqs_queue.s3_object_queue.id}"
     YARA_MATCHES_DYNAMO_TABLE_NAME = "${aws_dynamodb_table.binaryalert_yara_matches.name}"
     YARA_ALERTS_SNS_TOPIC_ARN      = "${aws_sns_topic.yara_match_alerts.arn}"

tests/lambda_functions/analyzer/main_test.py

@@ -52,6 +52,12 @@ def metadata(self):
         return GOOD_FILE_METADATA if self.key == GOOD_S3_OBJECT_KEY else EVIL_FILE_METADATA
 
 
+@mock.patch.dict(os.environ, {
+    'LAMBDA_TASK_ROOT': '/var/task',
+    'SQS_QUEUE_URL': MOCK_SQS_URL,
+    'YARA_MATCHES_DYNAMO_TABLE_NAME': MOCK_DYNAMO_TABLE_NAME,
+    'YARA_ALERTS_SNS_TOPIC_ARN': MOCK_SNS_TOPIC_ARN
+})
 @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):
@@ -66,17 +72,23 @@ def setUp(self):
         os.makedirs(os.path.dirname(COMPILED_RULES_FILEPATH))
         yara_mocks.save_test_yara_rules(COMPILED_RULES_FILEPATH)
 
-        # Set environment variables.
-        os.environ['LAMBDA_TASK_ROOT'] = '/var/task'
-        os.environ['S3_BUCKET_NAME'] = MOCK_S3_BUCKET_NAME
-        os.environ['SQS_QUEUE_URL'] = MOCK_SQS_URL
-        os.environ['YARA_MATCHES_DYNAMO_TABLE_NAME'] = MOCK_DYNAMO_TABLE_NAME
-        os.environ['YARA_ALERTS_SNS_TOPIC_ARN'] = MOCK_SNS_TOPIC_ARN
-
         # Create test event.
         self._test_event = {
-            # Two objects, which match different YARA rules.
-            'S3Objects': [urllib.parse.quote_plus(GOOD_S3_OBJECT_KEY), EVIL_S3_OBJECT_KEY],
+            # Two objects which match different YARA rules.
+            'Records': [
+                {
+                    's3': {
+                        'bucket': {'name': MOCK_S3_BUCKET_NAME},
+                        'object': {'key': urllib.parse.quote_plus(GOOD_S3_OBJECT_KEY)}
+                    }
+                },
+                {
+                    's3': {
+                        'bucket': {'name': MOCK_S3_BUCKET_NAME},
+                        'object': {'key': urllib.parse.quote_plus(EVIL_S3_OBJECT_KEY)}
+                    }
+                }
+            ],
             'SQSReceipts': MOCK_SQS_RECEIPTS
         }
 
@@ -104,13 +116,13 @@ def test_analyze_lambda_handler(self, mock_output: mock.MagicMock, mock_call: mo
             # Verify logging statements.
             mock_logger.assert_has_calls([
                 mock.call.info('Processing %d record(s)', 2),
-                mock.call.info('Analyzing "%s"', GOOD_S3_OBJECT_KEY),
+                mock.call.info('Analyzing "%s:%s"', MOCK_S3_BUCKET_NAME, GOOD_S3_OBJECT_KEY),
                 mock.call.warning(
                     '%s matched YARA rules: %s',
                     mock.ANY,
                     {'externals.yar:filename_contains_win32'}
                 ),
-                mock.call.info('Analyzing "%s"', EVIL_S3_OBJECT_KEY),
+                mock.call.info('Analyzing "%s:%s"', MOCK_S3_BUCKET_NAME, EVIL_S3_OBJECT_KEY),
                 mock.call.warning(
                     '%s matched YARA rules: %s',
                     mock.ANY,