From 28c6b148280fa72fa06d9eca164a6181393fdd29 Mon Sep 17 00:00:00 2001
From: "chenguang.xu@datadoghq.com" <chenguang.xu@datadoghq.com>
Date: Mon, 30 Oct 2023 15:52:02 -0400
Subject: [PATCH] To support new delimiters for stepFunction execution_arn

---
 aws/logs_monitoring/parsing.py            | 21 ++++++++++-----
 aws/logs_monitoring/tests/test_parsing.py | 31 +++++++++++++++++++++++
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/aws/logs_monitoring/parsing.py b/aws/logs_monitoring/parsing.py
index d6c726ec3..ede21cd20 100644
--- a/aws/logs_monitoring/parsing.py
+++ b/aws/logs_monitoring/parsing.py
@@ -536,13 +536,11 @@ def awslogs_handler(event, context, metadata):
     ):
         state_machine_arn = ""
         try:
-            message = json.loads(logs["logEvents"][0]["message"])
-            if message.get("execution_arn") is not None:
-                execution_arn = message["execution_arn"]
-                arn_tokens = execution_arn.split(":")
-                arn_tokens[5] = "stateMachine"
-                metadata[DD_HOST] = ":".join(arn_tokens[:-1])
-                state_machine_arn = ":".join(arn_tokens[:7])
+            state_machine_arn = get_state_machine_arn(
+                json.loads(logs["logEvents"][0]["message"])
+            )
+            if state_machine_arn:  # not empty
+                metadata[DD_HOST] = state_machine_arn
         except Exception as e:
             logger.debug(
                 "Unable to set stepfunction host or get state_machine_arn: %s" % e
@@ -856,3 +854,12 @@ def normalize_events(events, metadata):
     )
 
     return normalized
+
+
+def get_state_machine_arn(message):
+    if message.get("execution_arn") is not None:
+        execution_arn = message["execution_arn"]
+        arn_tokens = re.split(r"[:/\\]", execution_arn)
+        arn_tokens[5] = "stateMachine"
+        return ":".join(arn_tokens[:7])
+    return ""
diff --git a/aws/logs_monitoring/tests/test_parsing.py b/aws/logs_monitoring/tests/test_parsing.py
index fe6fce435..02328050a 100644
--- a/aws/logs_monitoring/tests/test_parsing.py
+++ b/aws/logs_monitoring/tests/test_parsing.py
@@ -27,6 +27,7 @@
     separate_security_hub_findings,
     parse_aws_waf_logs,
     get_service_from_tags,
+    get_state_machine_arn,
 )
 from settings import (
     DD_CUSTOM_TAGS,
@@ -966,5 +967,35 @@ def test_get_service_from_tags_default_to_source(self):
         self.assertEqual(get_service_from_tags(metadata), "ecs")
 
 
+class TestParsingStepFunctionLogs(unittest.TestCase):
+    def test_get_state_machine_arn(self):
+        invalid_sf_log_message = {"no_execution_arn": "xxxx/yyy"}
+        self.assertEqual(get_state_machine_arn(invalid_sf_log_message), "")
+
+        normal_sf_log_message = {
+            "execution_arn": "arn:aws:states:sa-east-1:425362996713:express:my-Various-States:7f653fda-c79a-430b-91e2-3f97eb87cabb:862e5d40-a457-4ca2-a3c1-78485bd94d3f"
+        }
+        self.assertEqual(
+            get_state_machine_arn(normal_sf_log_message),
+            "arn:aws:states:sa-east-1:425362996713:stateMachine:my-Various-States",
+        )
+
+        forward_slash_sf_log_message = {
+            "execution_arn": "arn:aws:states:sa-east-1:425362996713:express:my-Various-States/7f653fda-c79a-430b-91e2-3f97eb87cabb:862e5d40-a457-4ca2-a3c1-78485bd94d3f"
+        }
+        self.assertEqual(
+            get_state_machine_arn(forward_slash_sf_log_message),
+            "arn:aws:states:sa-east-1:425362996713:stateMachine:my-Various-States",
+        )
+
+        back_slash_sf_log_message = {
+            "execution_arn": "arn:aws:states:sa-east-1:425362996713:express:my-Various-States\\7f653fda-c79a-430b-91e2-3f97eb87cabb:862e5d40-a457-4ca2-a3c1-78485bd94d3f"
+        }
+        self.assertEqual(
+            get_state_machine_arn(back_slash_sf_log_message),
+            "arn:aws:states:sa-east-1:425362996713:stateMachine:my-Various-States",
+        )
+
+
 if __name__ == "__main__":
     unittest.main()