sync

Author: alexvolha

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:

uncoder-core/app/translator/core/custom_types/time.py

@@ -0,0 +1,9 @@
+from app.translator.tools.custom_enum import CustomEnum
+
+
+class TimeFrameType(CustomEnum):
+    years = "years"
+    months = "months"
+    days = "days"
+    hours = "hours"
+    minutes = "minutes"

uncoder-core/app/translator/core/exceptions/core.py

@@ -1,9 +1,7 @@
-class NotImplementedException(BaseException):
-    ...
+class NotImplementedException(BaseException): ...
 
 
-class BasePlatformException(BaseException):
-    ...
+class BasePlatformException(BaseException): ...
 
 
 class StrictPlatformException(BasePlatformException):

uncoder-core/app/translator/core/exceptions/iocs.py

@@ -1,9 +1,7 @@
-class BaseIOCsException(BaseException):
-    ...
+class BaseIOCsException(BaseException): ...
 
 
-class IocsLimitExceededException(BaseIOCsException):
-    ...
+class IocsLimitExceededException(BaseIOCsException): ...
 
 
 class EmptyIOCSException(BaseIOCsException):

uncoder-core/app/translator/core/exceptions/parser.py

@@ -1,5 +1,4 @@
-class BaseParserException(BaseException):
-    ...
+class BaseParserException(BaseException): ...
 
 
 class TokenizerGeneralException(BaseParserException):

uncoder-core/app/translator/core/exceptions/render.py

@@ -1,5 +1,4 @@
-class BaseRenderException(BaseException):
-    ...
+class BaseRenderException(BaseException): ...
 
 
 class UnexpectedLogsourceException(BaseRenderException):
@@ -8,8 +7,7 @@ def __init__(self, platform_name: str, log_source: str):
         super().__init__(message)
 
 
-class FunctionRenderException(BaseRenderException):
-    ...
+class FunctionRenderException(BaseRenderException): ...
 
 
 class UnsupportedRenderMethod(BaseRenderException):

uncoder-core/app/translator/core/mapping.py

@@ -20,6 +20,10 @@ def is_suitable(self, *args, **kwargs) -> bool:
     def __str__(self) -> str:
         raise NotImplementedError("Abstract method")
 
+    @property
+    def default_source(self) -> dict:
+        return self._default_source
+
 
 class FieldMapping:
     def __init__(self, generic_field_name: str, platform_field_name: str):

uncoder-core/app/translator/core/parser.py

@@ -15,6 +15,7 @@
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 -----------------------------------------------------------------
 """
+
 import re
 from abc import ABC, abstractmethod
 from typing import Union

uncoder-core/app/translator/core/render.py

@@ -16,11 +16,13 @@
 limitations under the License.
 -----------------------------------------------------------------
 """
+
 from abc import ABC, abstractmethod
 from collections.abc import Callable
-from typing import Optional, Union
+from typing import ClassVar, 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
@@ -164,7 +166,14 @@ class QueryRender(ABC):
     is_single_line_comment: bool = False
     unsupported_functions_text = "Unsupported functions were excluded from the result query:"
 
-    platform_functions: PlatformFunctions = PlatformFunctions()
+    platform_functions: PlatformFunctions = None
+
+    def __init__(self):
+        self.init_platform_functions()
+
+    def init_platform_functions(self) -> None:
+        self.platform_functions = PlatformFunctions()
+        self.platform_functions.platform_query_render = self
 
     def render_not_supported_functions(self, not_supported_functions: list) -> str:
         line_template = f"{self.comment_symbol} " if self.comment_symbol and self.is_single_line_comment else ""
@@ -192,9 +201,10 @@ class PlatformQueryRender(QueryRender):
     field_value_map = BaseQueryFieldValue(or_token=or_token)
 
     query_pattern = "{prefix}{query}{functions}"
-    raw_log_field_pattern: str = None
+    raw_log_field_pattern_map: ClassVar[dict[str, str]] = None
 
     def __init__(self):
+        super().__init__()
         self.operator_map = {
             LogicalOperatorType.AND: f" {self.and_token} ",
             LogicalOperatorType.OR: f" {self.or_token} ",
@@ -288,7 +298,8 @@ def finalize_query(
         *args,  # noqa: ARG002
         **kwargs,  # noqa: ARG002
     ) -> str:
-        query = self.query_concatenation(prefix=prefix, search=query, functions=functions)
+        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)
@@ -329,7 +340,23 @@ def _generate_from_raw_query_container(self, query_container: RawQueryContainer)
             prefix="", query=query_container.query, functions="", meta_info=query_container.meta_info
         )
 
+    def process_raw_log_field(self, field: str, field_type: str) -> Optional[str]:
+        if raw_log_field_pattern := self.raw_log_field_pattern_map.get(field_type):
+            return raw_log_field_pattern.pattern.format(field=field)
+
+    def process_raw_log_field_prefix(self, field: str, source_mapping: SourceMapping) -> Optional[list]:
+        if isinstance(field, list):
+            list_of_prefix = []
+            for f in field:
+                if prepared_prefix := self.process_raw_log_field_prefix(field=f, source_mapping=source_mapping):
+                    list_of_prefix.extend(prepared_prefix)
+            return list_of_prefix
+        if raw_log_field_type := source_mapping.raw_log_fields.get(field):
+            return [self.process_raw_log_field(field=field, field_type=raw_log_field_type)]
+
     def generate_raw_log_fields(self, fields: list[Field], source_mapping: SourceMapping) -> str:
+        if self.raw_log_field_pattern_map is None:
+            return ""
         defined_raw_log_fields = []
         for field in fields:
             mapped_field = source_mapping.fields_mapping.get_platform_field_name(generic_field_name=field.source_name)
@@ -340,10 +367,8 @@ def generate_raw_log_fields(self, fields: list[Field], source_mapping: SourceMap
                 )
             if not mapped_field and self.is_strict_mapping:
                 raise StrictPlatformException(field_name=field.source_name, platform_name=self.details.name)
-            if mapped_field not in source_mapping.raw_log_fields:
-                continue
-            field_prefix = self.raw_log_field_pattern.format(field=mapped_field)
-            defined_raw_log_fields.append(field_prefix)
+            if field_prefix := self.process_raw_log_field_prefix(field=mapped_field, source_mapping=source_mapping):
+                defined_raw_log_fields.extend(field_prefix)
         return "\n".join(set(defined_raw_log_fields))
 
     def _generate_from_tokenized_query_container(self, query_container: TokenizedQueryContainer) -> str:
@@ -374,6 +399,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]

uncoder-core/app/translator/core/render_cti.py

@@ -17,7 +17,6 @@
 -----------------------------------------------------------------
 """
 
-
 from app.translator.core.models.iocs import IocsChunkValue
 from app.translator.core.models.platform_details import PlatformDetails