[AWSX] properly handle Transit gateway log source (#844)

DataDog · Sep 4, 2024 · 2cb0f09 · 2cb0f09
1 parent 5fdf26e
commit 2cb0f09
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 3 deletions.
diff --git a/aws/logs_monitoring/steps/enums.py b/aws/logs_monitoring/steps/enums.py
@@ -93,7 +93,7 @@ def __init__(self, string, event_source):
     REDSHIFT = ("_redshift_", AwsEventSource.REDSHIFT)
     # e.g. AWSLogs/123456779121/vpcdnsquerylogs/vpc-********/2021/05/11/vpc-********_vpcdnsquerylogs_********_20210511T0910Z_71584702.log.gz
     ROUTE53 = ("vpcdnsquerylogs", AwsEventSource.ROUTE53)
-    TRANSITAGATEWAY = ("transit-gateway", AwsEventSource.TRANSITGATEWAY)
+    TRANSITGATEWAY = ("transit-gateway", AwsEventSource.TRANSITGATEWAY)
     VERIFIED_ACCESS = ("verified-access", AwsEventSource.VERIFIED_ACCESS)
     # e.g. AWSLogs/123456779121/vpcflowlogs/us-east-1/2020/10/02/123456779121_vpcflowlogs_us-east-1_fl-xxxxx.log.gz
     VPC = ("vpcflowlogs", AwsEventSource.VPC)

diff --git a/aws/logs_monitoring/steps/handlers/s3_handler.py b/aws/logs_monitoring/steps/handlers/s3_handler.py
@@ -73,8 +73,8 @@ def _extract_event(self, event):
 
     def _set_source(self, event):
         self.data_store.source = parse_event_source(event, self.data_store.key)
-        if str(AwsS3EventSourceKeyword.TRANSITAGATEWAY) in self.data_store.bucket:
-            self.data_store.source = AwsEventSource.TRANSITGATEWAY
+        if str(AwsS3EventSourceKeyword.TRANSITGATEWAY) in self.data_store.bucket:
+            self.data_store.source = str(AwsEventSource.TRANSITGATEWAY)
         self.metadata[DD_SOURCE] = self.data_store.source
 
     def _set_host(self):

diff --git a/aws/logs_monitoring/tests/test_s3_handler.py b/aws/logs_monitoring/tests/test_s3_handler.py
@@ -267,6 +267,26 @@ def test_set_source_cloudfront(self):
             "cloudfront",
         )
 
+    def test_set_source_transit_gateway(self):
+        self.s3_handler.data_store.key = "AWSLogs/1234566312/vpcflowlogs/us-east-1/2024/08/09/11/123455660991_vpcflowlogs_us-east-1_fl-01fb37"
+        self.s3_handler.data_store.bucket = "my-bucket-transit-gateway"
+        self.s3_handler._set_source(
+            {
+                "Records": [
+                    {
+                        "s3": {
+                            "bucket": {"name": "my-bucket-transit-gateway"},
+                            "object": {"key": self.s3_handler.data_store.key},
+                        }
+                    }
+                ]
+            }
+        )
+        self.assertEqual(
+            self.s3_handler.data_store.source,
+            "transitgateway",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()