diff --git a/uncoder-core/app/routers/translate.py b/uncoder-core/app/routers/translate.py index 7acdaaee..009cab03 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 @@ -15,7 +16,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 +30,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..2fd36c45 --- /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""" diff --git a/uncoder-core/app/translator/core/render.py b/uncoder-core/app/translator/core/render.py index 81dec670..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 @@ -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/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.\-*]+)" 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)