Change prio on setting s3 log source WAF > Cloudfront

DataDog · Jun 4, 2024 · b9e0b0b · b9e0b0b
1 parent 83c4e68
commit b9e0b0b
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 3 deletions.
diff --git a/aws/logs_monitoring/steps/enums.py b/aws/logs_monitoring/steps/enums.py
@@ -71,6 +71,10 @@ def __init__(self, string, event_source):
         self.string = string
         self.event_source = event_source
 
+    # e.g. 2020/10/02/21/aws-waf-logs-testing-1-2020-10-02-21-25-30-x123x-x456x or AWSLogs/123456779121/WAFLogs/us-east-1/xxxxxx-waf/2022/10/11/14/10/123456779121_waflogs_us-east-1_xxxxx-waf_20221011T1410Z_12756524.log.gz
+    WAF_0 = ("aws-waf-logs", AwsEventSource.WAF)
+    WAF_1 = ("waflogs", AwsEventSource.WAF)
+
     BEDROCK = ("bedrock", AwsEventSource.BEDROCK)
     # e.g. carbon-black-cloud-forwarder/alerts/org_key=*****/year=2021/month=7/day=19/hour=18/minute=15/second=41/8436e850-7e78-40e4-b3cd-6ebbc854d0a2.jsonl.gz
     CARBONBLACK = ("carbon-black", AwsEventSource.CARBONBLACK)
@@ -91,9 +95,6 @@ def __init__(self, string, event_source):
     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)
-    # e.g. 2020/10/02/21/aws-waf-logs-testing-1-2020-10-02-21-25-30-x123x-x456x or AWSLogs/123456779121/WAFLogs/us-east-1/xxxxxx-waf/2022/10/11/14/10/123456779121_waflogs_us-east-1_xxxxx-waf_20221011T1410Z_12756524.log.gz
-    WAF_0 = ("aws-waf-logs", AwsEventSource.WAF)
-    WAF_1 = ("waflogs", AwsEventSource.WAF)
 
     def __str__(self):
         return f"{self.string}"

diff --git a/aws/logs_monitoring/tests/test_s3_handler.py b/aws/logs_monitoring/tests/test_s3_handler.py
@@ -219,6 +219,46 @@ def test_elb_s3_key_multi_prefix_gov(self):
             "/123456789aabcdef",
         )
 
+    def test_set_source_waf_cloudfront(self):
+        self.s3_handler.data_store.key = "AWSLogs/123456779121/WAFLogs/cloudfront/this/is/a/prio/test.log.gz"
+        self.s3_handler.data_store.bucket = "my-bucket"
+        self.s3_handler._set_source(
+            {
+                "Records": [
+                    {
+                        "s3": {
+                            "bucket": {"name": "my-bucket"},
+                            "object": {"key": self.s3_handler.data_store.key},
+                        }
+                    }
+                ]
+            }
+        )
+        self.assertEqual(
+            self.s3_handler.data_store.source,
+            "waf",
+        )
+
+    def test_set_source_cloudfront(self):
+        self.s3_handler.data_store.key = "AWSLogs/123456779121/CloudFront/us-east-1/2020/10/02/21/123456779121_CloudFront_us-east-1_20201002T2100Z_abcdef.log.gz"
+        self.s3_handler.data_store.bucket = "my-bucket"
+        self.s3_handler._set_source(
+            {
+                "Records": [
+                    {
+                        "s3": {
+                            "bucket": {"name": "my-bucket"},
+                            "object": {"key": self.s3_handler.data_store.key},
+                        }
+                    }
+                ]
+            }
+        )
+        self.assertEqual(
+            self.s3_handler.data_store.source,
+            "cloudfront",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()