fix: handle quoted alert JSON in workflow functions

keep/iohandler/iohandler.py

@@ -129,16 +129,27 @@ def extract_keep_functions(self, text):
                     escape_next = False
                     quote_char = ""
                     escapes = {}
+                    embedded_json_end = None
+                    string_content_started = False
                     while i < len(text) and (parent_count > 0 or in_string):
                         if text[i] == "\\" and in_string and not escape_next:
                             escape_next = True
                             i += 1
                             continue
+                        elif (
+                            in_string
+                            and embedded_json_end is not None
+                            and i < embedded_json_end
+                        ):
+                            if text[i] == quote_char and not escape_next:
+                                escapes[i - start] = text[i]
                         elif text[i] in ('"', "'"):
                             if not in_string:
                                 # Detecting the beginning of the string
                                 in_string = True
                                 quote_char = text[i]
+                                embedded_json_end = None
+                                string_content_started = False
                             elif (
                                 text[i] == quote_char
                                 and not escape_next
@@ -152,9 +163,21 @@ def extract_keep_functions(self, text):
                                 in_string = False
                                 quote_char = ""
                             elif text[i] == quote_char and not escape_next:
-                                escapes[i] = text[
+                                escapes[i - start] = text[
                                     i
                                 ]  # Save the quote character where we need to escape for valid ast parsing
+                        elif (
+                            in_string
+                            and not string_content_started
+                            and not text[i].isspace()
+                        ):
+                            string_content_started = True
+                            if text[i] == "{":
+                                try:
+                                    _, json_end = json.JSONDecoder().raw_decode(text[i:])
+                                    embedded_json_end = i + json_end
+                                except json.JSONDecodeError:
+                                    pass
                         elif text[i] == "(" and not in_string:
                             parent_count += 1
                         elif text[i] == ")" and not in_string:
@@ -173,6 +196,12 @@ def extract_keep_functions(self, text):
                 i += 1
         return matches
 
+    @staticmethod
+    def _escape_token_quotes(token, escapes):
+        for escape in sorted(escapes, reverse=True):
+            token = token[:escape] + "\\" + token[escape:]
+        return token
+
     def _trim_token_error(self, token):
         # trim too long tokens so that the error message will be readable
         if len(token) > 64:
@@ -230,15 +259,8 @@ def parse(self, string, safe=False, default="", additional_context=None):
             token, escapes = tokens[0]
             token_to_replace = token
             try:
-                escapes_counter = 0
                 if escapes:
-                    for escape in escapes:
-                        token = (
-                            token[: escape + escapes_counter]
-                            + "\\"
-                            + token[escape + escapes_counter :]
-                        )
-                        escapes_counter += 1  # we need to increment the counter because we added a character
+                    token = self._escape_token_quotes(token, escapes)
                 val = self._parse_token(token)
             except Exception as e:
                 # trim stacktrace since we have limitation on the error message
@@ -275,8 +297,7 @@ def parse(self, string, safe=False, default="", additional_context=None):
             token_to_replace = token
             try:
                 if escapes:
-                    for escape in escapes:
-                        token = token[:escape] + "\\" + token[escape:]
+                    token = self._escape_token_quotes(token, escapes)
                 val = self._parse_token(token)
 
             except Exception as e:

tests/test_iohandler.py

@@ -111,6 +111,47 @@ def test_with_function_2(context_manager):
     assert s == "hello 1"
 
 
+def test_dictget_with_quotes_and_parentheses_in_rendered_alert(context_manager):
+    description = "Hello ('Hola') everyone. How are you all doing?"
+    context_manager.event_context = AlertDto(
+        id="test",
+        name="test",
+        lastReceived="2024-03-20T00:00:00.000Z",
+        description=description,
+    )
+    iohandler = IOHandler(context_manager)
+
+    assert (
+        iohandler.render("keep.dictget('{{ alert }}', 'description', 'unknown')")
+        == description
+    )
+    assert (
+        iohandler.render('keep.dictget("{{ alert }}", "description", "unknown")')
+        == description
+    )
+    assert (
+        iohandler.render(
+            "Description: keep.dictget('{{ alert }}', 'description', 'unknown')"
+        )
+        == f"Description: {description}"
+    )
+    assert (
+        iohandler.render(
+            "keep.dictget('{{ alert }}', 'description', 'unknown') | "
+            "keep.dictget('{{ alert }}', 'some_field', 'unknown')"
+        )
+        == f"{description} | unknown"
+    )
+
+
+@pytest.mark.parametrize("value", ["(abc", "[WARN", "{name", "(don't)"])
+def test_extract_function_with_non_json_literal_prefix(context_manager, value):
+    iohandler = IOHandler(context_manager)
+    template = f'keep.split("{value}", "x")'
+
+    assert len(iohandler.extract_keep_functions(template)) == 1
+
+
 def test_with_json_dumps(context_manager):
     iohandler = IOHandler(context_manager)
     context_manager.steps_context = {