From 98d4d9f7265d6197781c666d8969d1d4d0ebd978 Mon Sep 17 00:00:00 2001 From: Dmytro <47281757+tarnopolskyi@users.noreply.github.com> Date: Wed, 12 Jun 2024 13:40:38 +0200 Subject: [PATCH 1/3] added return only one query context --- uncoder-core/app/routers/translate.py | 9 +++++++-- uncoder-core/app/translator/core/context_vars.py | 4 ++++ uncoder-core/app/translator/core/render.py | 4 ++++ .../platforms/forti_siem/renders/forti_siem_rule.py | 3 +++ .../logrhythm_axon/renders/logrhythm_axon_query.py | 3 +++ 5 files changed, 21 insertions(+), 2 deletions(-) create mode 100644 uncoder-core/app/translator/core/context_vars.py diff --git a/uncoder-core/app/routers/translate.py b/uncoder-core/app/routers/translate.py index 7acdaaee..8ef9fdfc 100644 --- a/uncoder-core/app/routers/translate.py +++ b/uncoder-core/app/routers/translate.py @@ -1,6 +1,7 @@ from fastapi import APIRouter, Body from app.models.translation import InfoMessage, OneTranslationData, Platform, TranslatorPlatforms +from app.translator.core.context_vars import return_only_first_query_ctx_var from app.translator.cti_translator import CTITranslator from app.translator.translator import Translator @@ -14,8 +15,9 @@ def translate_one( source_platform_id: str = Body(..., embed=True), target_platform_id: str = Body(..., embed=True), - text: str = Body(..., embed=True), + return_only_first_query: bool = False, ) -> OneTranslationData: + return_only_first_query_ctx_var.set(return_only_first_query) status, data = translator.translate_one(text=text, source=source_platform_id, target=target_platform_id) if status: return OneTranslationData(status=status, translation=data, target_platform_id=target_platform_id) @@ -27,8 +29,11 @@ def translate_one( @st_router.post("/translate/all", tags=["translator"], description="Generate all translations") @st_router.post("/translate/all/", include_in_schema=False) def translate_all( - source_platform_id: str = Body(..., embed=True), text: str = Body(..., embed=True) + source_platform_id: str = Body(..., embed=True), + text: str = Body(..., embed=True), + return_only_first_query: bool = False, ) -> list[OneTranslationData]: + return_only_first_query_ctx_var.set(return_only_first_query) result = translator.translate_all(text=text, source=source_platform_id) translations = [] for platform_result in result: diff --git a/uncoder-core/app/translator/core/context_vars.py b/uncoder-core/app/translator/core/context_vars.py new file mode 100644 index 00000000..5ddc1ef9 --- /dev/null +++ b/uncoder-core/app/translator/core/context_vars.py @@ -0,0 +1,4 @@ +from contextvars import ContextVar + +return_only_first_query_ctx_var: ContextVar[bool] = ContextVar("return_only_first_query_ctx_var", default=False) +"""Set to True to return ony first query if rendered multiple options""" \ No newline at end of file diff --git a/uncoder-core/app/translator/core/render.py b/uncoder-core/app/translator/core/render.py index 81dec670..01f4e3e1 100644 --- a/uncoder-core/app/translator/core/render.py +++ b/uncoder-core/app/translator/core/render.py @@ -35,6 +35,7 @@ from app.translator.core.models.query_container import MetaInfoContainer, RawQueryContainer, TokenizedQueryContainer from app.translator.core.str_value_manager import StrValue, StrValueManager from app.translator.core.tokenizer import TOKEN_TYPE +from app.translator.core.context_vars import return_only_first_query_ctx_var class BaseQueryFieldValue(ABC): @@ -283,6 +284,7 @@ def finalize_query( **kwargs, # noqa: ARG002 ) -> str: query = self.query_pattern.format(prefix=prefix, query=query, functions=functions).strip() + query = self.wrap_query_with_meta_info(meta_info=meta_info, query=query) if not_supported_functions: rendered_not_supported = self.render_not_supported_functions(not_supported_functions) @@ -368,6 +370,8 @@ def _generate_from_tokenized_query_container(self, query_container: TokenizedQue meta_info=query_container.meta_info, source_mapping=source_mapping, ) + if return_only_first_query_ctx_var.get() is True: + return finalized_query queries_map[source_mapping.source_id] = finalized_query if not queries_map and errors: raise errors[0] diff --git a/uncoder-core/app/translator/platforms/forti_siem/renders/forti_siem_rule.py b/uncoder-core/app/translator/platforms/forti_siem/renders/forti_siem_rule.py index bef9392b..1c2d14c0 100644 --- a/uncoder-core/app/translator/platforms/forti_siem/renders/forti_siem_rule.py +++ b/uncoder-core/app/translator/platforms/forti_siem/renders/forti_siem_rule.py @@ -18,6 +18,7 @@ from typing import Optional, Union from app.translator.const import DEFAULT_VALUE_TYPE +from app.translator.core.context_vars import return_only_first_query_ctx_var from app.translator.core.custom_types.meta_info import SeverityType from app.translator.core.custom_types.tokens import GroupType, LogicalOperatorType, OperatorType from app.translator.core.custom_types.values import ValueType @@ -273,6 +274,8 @@ def _generate_from_tokenized_query_container(self, query_container: TokenizedQue source_mapping=source_mapping, fields=mapped_fields_set, ) + if return_only_first_query_ctx_var.get() is True: + return finalized_query queries_map[source_mapping.source_id] = finalized_query return self.finalize(queries_map) diff --git a/uncoder-core/app/translator/platforms/logrhythm_axon/renders/logrhythm_axon_query.py b/uncoder-core/app/translator/platforms/logrhythm_axon/renders/logrhythm_axon_query.py index 9be24b73..2bf38391 100644 --- a/uncoder-core/app/translator/platforms/logrhythm_axon/renders/logrhythm_axon_query.py +++ b/uncoder-core/app/translator/platforms/logrhythm_axon/renders/logrhythm_axon_query.py @@ -19,6 +19,7 @@ from typing import Union from app.translator.const import DEFAULT_VALUE_TYPE +from app.translator.core.context_vars import return_only_first_query_ctx_var from app.translator.core.custom_types.tokens import LogicalOperatorType from app.translator.core.custom_types.values import ValueType from app.translator.core.exceptions.core import StrictPlatformException @@ -262,6 +263,8 @@ def _generate_from_tokenized_query_container(self, query_container: TokenizedQue meta_info=query_container.meta_info, source_mapping=source_mapping, ) + if return_only_first_query_ctx_var.get() is True: + return finalized_query queries_map[source_mapping.source_id] = finalized_query return self.finalize(queries_map) From a8098c9836fc1c8a84929b55ad19e91bdcb95789 Mon Sep 17 00:00:00 2001 From: Dmytro <47281757+tarnopolskyi@users.noreply.github.com> Date: Wed, 12 Jun 2024 13:45:52 +0200 Subject: [PATCH 2/3] added deleted field --- uncoder-core/app/routers/translate.py | 1 + 1 file changed, 1 insertion(+) diff --git a/uncoder-core/app/routers/translate.py b/uncoder-core/app/routers/translate.py index 8ef9fdfc..009cab03 100644 --- a/uncoder-core/app/routers/translate.py +++ b/uncoder-core/app/routers/translate.py @@ -15,6 +15,7 @@ def translate_one( source_platform_id: str = Body(..., embed=True), target_platform_id: str = Body(..., embed=True), + text: str = Body(..., embed=True), return_only_first_query: bool = False, ) -> OneTranslationData: return_only_first_query_ctx_var.set(return_only_first_query) From e60bbfd9e0f6012f5cc7b56ad6f6fbe1f59a7e2c Mon Sep 17 00:00:00 2001 From: Dmytro <47281757+tarnopolskyi@users.noreply.github.com> Date: Wed, 12 Jun 2024 13:46:26 +0200 Subject: [PATCH 3/3] run linter --- uncoder-core/app/translator/core/context_vars.py | 2 +- uncoder-core/app/translator/core/render.py | 2 +- uncoder-core/app/translator/platforms/base/aql/const.py | 4 +++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/uncoder-core/app/translator/core/context_vars.py b/uncoder-core/app/translator/core/context_vars.py index 5ddc1ef9..2fd36c45 100644 --- a/uncoder-core/app/translator/core/context_vars.py +++ b/uncoder-core/app/translator/core/context_vars.py @@ -1,4 +1,4 @@ from contextvars import ContextVar return_only_first_query_ctx_var: ContextVar[bool] = ContextVar("return_only_first_query_ctx_var", default=False) -"""Set to True to return ony first query if rendered multiple options""" \ No newline at end of file +"""Set to True to return ony first query if rendered multiple options""" diff --git a/uncoder-core/app/translator/core/render.py b/uncoder-core/app/translator/core/render.py index 01f4e3e1..1120e784 100644 --- a/uncoder-core/app/translator/core/render.py +++ b/uncoder-core/app/translator/core/render.py @@ -21,6 +21,7 @@ from typing import Optional, Union from app.translator.const import DEFAULT_VALUE_TYPE +from app.translator.core.context_vars import return_only_first_query_ctx_var from app.translator.core.custom_types.tokens import LogicalOperatorType, OperatorType from app.translator.core.custom_types.values import ValueType from app.translator.core.escape_manager import EscapeManager @@ -35,7 +36,6 @@ from app.translator.core.models.query_container import MetaInfoContainer, RawQueryContainer, TokenizedQueryContainer from app.translator.core.str_value_manager import StrValue, StrValueManager from app.translator.core.tokenizer import TOKEN_TYPE -from app.translator.core.context_vars import return_only_first_query_ctx_var class BaseQueryFieldValue(ABC): diff --git a/uncoder-core/app/translator/platforms/base/aql/const.py b/uncoder-core/app/translator/platforms/base/aql/const.py index 063c6d78..7cab3078 100644 --- a/uncoder-core/app/translator/platforms/base/aql/const.py +++ b/uncoder-core/app/translator/platforms/base/aql/const.py @@ -1,5 +1,7 @@ UTF8_PAYLOAD_PATTERN = r"UTF8\(payload\)" NUM_VALUE_PATTERN = r"(?P\d+(?:\.\d+)*)" -SINGLE_QUOTES_VALUE_PATTERN = r"""'(?P(?:[:a-zA-Z\*0-9=+%#\-\/\\|,;_<>`~".$&^@!?\(\)\{\}\[\]\s]|'')*)'""" +SINGLE_QUOTES_VALUE_PATTERN = ( + r"""'(?P(?:[:a-zA-Zа-яА-Я\*0-9=+%#\-\/\\|,;_<>`~".$&^@!?\(\)\{\}\[\]\s]|'')*)'""" +) TABLE_PATTERN = r"\s+FROM\s+[a-zA-Z.\-*]+" TABLE_GROUP_PATTERN = r"\s+FROM\s+(?P[a-zA-Z.\-*]+)"