diff --git a/keep/iohandler/iohandler.py b/keep/iohandler/iohandler.py index 947b3b4f01..3779064ad4 100644 --- a/keep/iohandler/iohandler.py +++ b/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: diff --git a/tests/test_iohandler.py b/tests/test_iohandler.py index 0f4b63a76a..822420668e 100644 --- a/tests/test_iohandler.py +++ b/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 = {