Merge pull request #2 from blockmar/log-forwarder-eventbridge-s3-support-feedback

Log forwarder eventbridge s3 support feedback

Author: blockmar

aws/logs_monitoring/steps/handlers/s3_handler.py

@@ -58,42 +58,6 @@ def __init__(self, context, metadata, cache_layer):
         # a private data store for event attributes
         self.data_store = S3EventDataStore()
 
-    def _transform_eventbridge_to_s3_format(self, event):
-        """
-        Transform EventBridge S3 event to standard S3 event format.
-
-        EventBridge format:
-        {
-            "version": "0",
-            "detail-type": "Object Created",
-            "source": "aws.s3",
-            "detail": {
-                "bucket": {"name": "bucket-name"},
-                "object": {"key": "object-key", "size": 1234}
-            }
-        }
-
-        Standard S3 format:
-        {
-            "Records": [{
-                "s3": {
-                    "bucket": {"name": "bucket-name"},
-                    "object": {"key": "object-key"}
-                }
-            }]
-        }
-        """
-        return {
-            "Records": [{
-                "s3": {
-                    "bucket": event["detail"]["bucket"],
-                    "object": {
-                        "key": event["detail"]["object"]["key"]
-                    }
-                }
-            }]
-        }
-
     def handle(self, event):
         event = self._extract_event(event)
 
@@ -109,9 +73,6 @@ def handle(self, event):
         yield from self._get_structured_lines_for_s3_handler()
 
     def _extract_event(self, event):
-        # if this is an EventBridge S3 event, transform it to standard S3 format
-        if event.get("source") == "aws.s3" and "detail" in event:
-            event = self._transform_eventbridge_to_s3_format(event)
         # if this is a S3 event carried in a SNS message, extract it and override the event
         if "Sns" in event.get("Records")[0]:
             event = json.loads(event.get("Records")[0].get("Sns").get("Message"))

aws/logs_monitoring/steps/parsing.py

@@ -37,6 +37,7 @@ def parse(event, context, cache_layer):
                 events = aws_handler.handle(event)
                 return collect_and_count(events)
             case AwsEventType.S3:
+                event = _reformat_eventbridge_s3_event(event)
                 s3_handler = S3EventHandler(context, metadata, cache_layer)
                 events = s3_handler.handle(event)
             case AwsEventType.EVENTS:
@@ -86,6 +87,47 @@ def parse_event_type(event):
     raise Exception("Event type not supported (see #Event supported section)")
 
 
+def _reformat_eventbridge_s3_event(event):
+    """
+    Transform EventBridge S3 event to standard S3 event format.
+
+    EventBridge format:
+    {
+        "version": "0",
+        "detail-type": "Object Created",
+        "source": "aws.s3",
+        "detail": {
+            "bucket": {"name": "bucket-name"},
+            "object": {"key": "object-key", "size": 1234}
+        }
+    }
+
+    Standard S3 format:
+    {
+        "Records": [{
+            "s3": {
+                "bucket": {"name": "bucket-name"},
+                "object": {"key": "object-key"}
+            }
+        }]
+    }
+    """
+    if event.get("source") == "aws.s3" and "detail" in event:
+        return {
+            "Records": [
+                {
+                    "s3": {
+                        "bucket": event["detail"]["bucket"],
+                        "object": {
+                            "key": event["detail"]["object"]["key"],
+                        },
+                    }
+                }
+            ]
+        }
+    return event
+
+
 # Handle Cloudwatch Events
 def cwevent_handler(event, metadata):
     # Set the source on the log

aws/logs_monitoring/tests/test_parsing.py

@@ -1,8 +1,10 @@
 import unittest
+from unittest.mock import MagicMock, patch
 
 from settings import DD_CUSTOM_TAGS, DD_SOURCE
 from steps.common import get_service_from_tags_and_remove_duplicates, parse_event_source
 from steps.enums import AwsEventSource
+from steps.parsing import parse
 
 
 class TestParseEventSource(unittest.TestCase):
@@ -183,5 +185,46 @@ def test_get_service_from_tags_removing_duplicates(self):
         )
 
 
+class TestEventBridgeS3Parsing(unittest.TestCase):
+    class Context:
+        function_version = "$LATEST"
+        invoked_function_arn = "arn:aws:lambda:us-east-1:123456789012:function:datadog-forwarder"
+        function_name = "datadog-forwarder"
+        memory_limit_in_mb = "128"
+
+    @patch("steps.parsing.S3EventHandler")
+    def test_parse_normalizes_eventbridge_s3_event_before_s3_handler(self, mock_s3_handler_cls):
+        # Arrange: handler yields one log line; we only care about the input event it received
+        mock_s3_handler = mock_s3_handler_cls.return_value
+        mock_s3_handler.handle.return_value = iter([{"message": "ok"}])
+
+        eventbridge_event = {
+            "version": "0",
+            "detail-type": "Object Created",
+            "source": "aws.s3",
+            "detail": {
+                "bucket": {"name": "my-bucket"},
+                "object": {"key": "my-key.log", "size": 1234},
+            },
+        }
+
+        cache_layer = MagicMock()
+
+        # Act
+        _ = parse(eventbridge_event, self.Context(), cache_layer)
+
+        # Assert: parse() passed a canonical S3-shaped event into the S3 handler
+        mock_s3_handler.handle.assert_called_once()
+        (normalized_event,) = mock_s3_handler.handle.call_args.args
+
+        self.assertIn("Records", normalized_event)
+        self.assertEqual(
+            normalized_event["Records"][0]["s3"]["bucket"]["name"], "my-bucket"
+        )
+        self.assertEqual(
+            normalized_event["Records"][0]["s3"]["object"]["key"], "my-key.log"
+        )
+
+
 if __name__ == "__main__":
     unittest.main()

aws/logs_monitoring/tests/test_s3_handler.py

@@ -325,105 +325,6 @@ def test_non_vpc_flowlog_includes_first_line(self):
         self.assertEqual(len(structured_lines), 3)
         self.assertEqual(structured_lines[0]["message"], "first line of data")
 
-    @patch("steps.handlers.s3_handler.boto3")
-    def test_eventbridge_s3_object_created_event(self, mock_boto3):
-        """Test EventBridge S3 Object Created event is transformed and processed correctly"""
-        mock_s3_client = MagicMock()
-        mock_s3_client.get_object.return_value = {
-            "Body": MagicMock(read=MagicMock(return_value=b"test log line"))
-        }
-        mock_boto3.client.return_value = mock_s3_client
-
-        eventbridge_event = {
-            "version": "0",
-            "id": "test-event-id",
-            "detail-type": "Object Created",
-            "source": "aws.s3",
-            "account": "123456789012",
-            "time": "2024-01-15T12:00:00Z",
-            "region": "us-east-1",
-            "resources": ["arn:aws:s3:::my-bucket"],
-            "detail": {
-                "version": "0",
-                "bucket": {"name": "my-bucket"},
-                "object": {
-                    "key": "my-key.log",
-                    "size": 1234,
-                    "etag": "abc123",
-                    "sequencer": "xyz789",
-                },
-                "request-id": "request-id",
-                "requester": "123456789012",
-                "source-ip-address": "1.2.3.4",
-                "reason": "PutObject",
-            },
-        }
-
-        logs = list(self.s3_handler.handle(eventbridge_event))
-
-        self.assertEqual(len(logs), 1)
-        self.assertEqual(logs[0]["message"], "test log line")
-        self.assertEqual(logs[0]["aws"]["s3"]["bucket"], "my-bucket")
-        self.assertEqual(logs[0]["aws"]["s3"]["key"], "my-key.log")
-        mock_s3_client.get_object.assert_called_once_with(
-            Bucket="my-bucket", Key="my-key.log"
-        )
-
-    @patch("steps.handlers.s3_handler.boto3")
-    def test_eventbridge_s3_with_gzip_file(self, mock_boto3):
-        """Test EventBridge S3 event with gzipped file"""
-        mock_s3_client = MagicMock()
-        gzip_data = gzip.compress(b"gzipped log line")
-        mock_s3_client.get_object.return_value = {
-            "Body": MagicMock(read=MagicMock(return_value=gzip_data))
-        }
-        mock_boto3.client.return_value = mock_s3_client
-
-        eventbridge_event = {
-            "version": "0",
-            "detail-type": "Object Created",
-            "source": "aws.s3",
-            "detail": {
-                "bucket": {"name": "my-bucket"},
-                "object": {"key": "my-key.log.gz"},
-            },
-        }
-
-        logs = list(self.s3_handler.handle(eventbridge_event))
-
-        self.assertEqual(len(logs), 1)
-        self.assertEqual(logs[0]["message"], "gzipped log line")
-
-    @patch("steps.handlers.s3_handler.boto3")
-    def test_eventbridge_s3_cloudtrail(self, mock_boto3):
-        """Test EventBridge S3 event with CloudTrail logs"""
-        mock_s3_client = MagicMock()
-        cloudtrail_data = {
-            "Records": [
-                {"eventVersion": "1.05", "eventName": "AssumeRole", "eventSource": "sts.amazonaws.com"}
-            ]
-        }
-        mock_s3_client.get_object.return_value = {
-            "Body": MagicMock(read=MagicMock(return_value=gzip.compress(bytes(str(cloudtrail_data).replace("'", '"'), "utf-8"))))
-        }
-        mock_boto3.client.return_value = mock_s3_client
-
-        eventbridge_event = {
-            "version": "0",
-            "detail-type": "Object Created",
-            "source": "aws.s3",
-            "detail": {
-                "bucket": {"name": "my-bucket"},
-                "object": {"key": "123456779121_CloudTrail_eu-west-3_20180707T1735Z_abcdefghi0MCRL2O.json.gz"},
-            },
-        }
-
-        logs = list(self.s3_handler.handle(eventbridge_event))
-
-        self.assertEqual(len(logs), 1)
-        self.assertEqual(logs[0]["eventName"], "AssumeRole")
-        self.assertEqual(logs[0]["aws"]["s3"]["bucket"], "my-bucket")
-
     @patch("steps.handlers.s3_handler.boto3")
     def test_backward_compatibility_direct_s3(self, mock_boto3):
         """Test that direct S3 events still work after EventBridge changes"""