merge main

Author: dtsprm

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/functions.py

@@ -9,16 +9,20 @@ class FunctionType(CustomEnum):
     min = "min"
     sum = "sum"
 
-    divide = "divide"
+    values = "values"
 
     earliest = "earliest"
     latest = "latest"
 
+    divide = "divide"
+
     lower = "lower"
+    split = "split"
     upper = "upper"
 
-    compare_boolean = "compare_boolean"
-
+    array_length = "array_length"
+    compare = "compare"
+    extract_time = "extract_time"
     ipv4_is_in_range = "ipv4_is_in_range"
 
     bin = "bin"
@@ -30,4 +34,6 @@ class FunctionType(CustomEnum):
     stats = "stats"
     table = "table"
     timeframe = "timeframe"
-    values = "values"
+    union = "union"
+
+    aggregation_data_function = "aggregation_data_function"

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

@@ -0,0 +1,23 @@
+from app.translator.tools.custom_enum import CustomEnum
+
+
+class TimeFrameType(CustomEnum):
+    years = "years"
+    months = "months"
+    days = "days"
+    hours = "hours"
+    minutes = "minutes"
+
+
+class TimePartType(CustomEnum):
+    day = "day"
+    day_of_week = "day_of_week"
+    day_of_year = "day_of_year"
+    hour = "hour"
+    microsecond = "microsecond"
+    millisecond = "millisecond"
+    minute = "minute"
+    month = "month"
+    quarter = "quarter"
+    second = "second"
+    year = "year"

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

@@ -1,3 +1,6 @@
+from typing import Optional
+
+
 class NotImplementedException(BaseException):
     ...
 
@@ -7,8 +10,19 @@ class BasePlatformException(BaseException):
 
 
 class StrictPlatformException(BasePlatformException):
-    def __init__(self, platform_name: str, field_name: str):
-        message = f"Platform {platform_name} has strict mapping. Source field {field_name} has no mapping."
+    field_name: str = None
+
+    def __init__(
+        self, platform_name: str, field_name: str, mapping: Optional[str] = None, detected_fields: Optional[list] = None
+    ):
+        message = (
+            f"Platform {platform_name} has strict mapping. "
+            f"Source fields: {', '.join(detected_fields) if detected_fields else field_name} has no mapping."
+            f" Mapping file: {mapping}."
+            if mapping
+            else ""
+        )
+        self.field_name = field_name
         super().__init__(message)
 
 

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

@@ -21,13 +21,13 @@
 import re
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from functools import cached_property
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
 
-from app.translator.core.exceptions.functions import InvalidFunctionSignature, NotSupportedFunctionException
+from app.translator.core.exceptions.functions import NotSupportedFunctionException
 from app.translator.core.mapping import SourceMapping
-from app.translator.core.models.field import Field
+from app.translator.core.models.field import Alias, Field
 from app.translator.core.models.functions.base import Function, ParsedFunctions, RenderedFunctions
+from app.translator.tools.utils import execute_module
 from settings import INIT_FUNCTIONS
 
 if TYPE_CHECKING:
@@ -41,6 +41,14 @@ class FunctionMatchContainer:
 
 
 class BaseFunctionParser(ABC):
+    function_names_map: ClassVar[dict[str, str]] = {}
+    functions_group_name: str = None
+    manager: PlatformFunctionsManager = None
+
+    def set_functions_manager(self, manager: PlatformFunctionsManager) -> BaseFunctionParser:
+        self.manager = manager
+        return self
+
     @abstractmethod
     def parse(self, *args, **kwargs) -> Function:
         raise NotImplementedError
@@ -73,21 +81,25 @@ def parse(self, func_body: str, raw: str) -> Function:
 
 
 class FunctionRender(ABC):
+    function_names_map: ClassVar[dict[str, str]] = {}
+    order_to_render: int = 0
+    in_query_render: bool = False
+    render_to_prefix: bool = False
+    manager: PlatformFunctionsManager = None
+
+    def set_functions_manager(self, manager: PlatformFunctionsManager) -> FunctionRender:
+        self.manager = manager
+        return self
+
     @abstractmethod
     def render(self, function: Function, source_mapping: SourceMapping) -> str:
         raise NotImplementedError
 
     @staticmethod
-    def concat_kwargs(kwargs: dict[str, str]) -> str:
-        result = ""
-        for key, value in kwargs.items():
-            if value:
-                result = f"{result}, {key}={value}" if result else f"{key}={value}"
+    def map_field(field: Union[Alias, Field], source_mapping: SourceMapping) -> str:
+        if isinstance(field, Alias):
+            return field.name
 
-        return result
-
-    @staticmethod
-    def map_field(field: Field, source_mapping: SourceMapping) -> str:
         generic_field_name = field.get_generic_field_name(source_mapping.source_id)
         mapped_field = source_mapping.fields_mapping.get_platform_field_name(generic_field_name=generic_field_name)
         if isinstance(mapped_field, list):
@@ -97,23 +109,48 @@ def map_field(field: Field, source_mapping: SourceMapping) -> str:
 
 
 class PlatformFunctionsManager:
+    platform_functions: PlatformFunctions = None
+
     def __init__(self):
-        self._parsers_map: dict[str, HigherOrderFunctionParser] = {}
-        self._renders_map: dict[str, FunctionRender] = {}
-        self._in_query_renders_map: dict[str, FunctionRender] = {}
-        self._names_map: dict[str, str] = {}
-        self._order_to_render: dict[str, int] = {}
-        self._render_to_prefix_functions: list[str] = []
-
-    def post_init_configure(self, platform_render: PlatformQueryRender) -> None:
-        raise NotImplementedError
+        # {platform_func_name: HigherOrderFunctionParser}
+        self._hof_parsers_map: dict[str, HigherOrderFunctionParser] = {}
+        self._parsers_map: dict[str, FunctionParser] = {}  # {platform_func_name: FunctionParser}
+
+        self._renders_map: dict[str, FunctionRender] = {}  # {generic_func_name: FunctionRender}
+        self._in_query_renders_map: dict[str, FunctionRender] = {}  # {generic_func_name: FunctionRender}
+        self._order_to_render: dict[str, int] = {}  # {generic_func_name: int}
+
+    def register_render(self, render_class: type[FunctionRender]) -> type[FunctionRender]:
+        render = render_class()
+        render.manager = self
+        for generic_function_name in render.function_names_map:
+            self._renders_map[generic_function_name] = render
+            self._order_to_render[generic_function_name] = render.order_to_render
+            if render.in_query_render:
+                self._in_query_renders_map[generic_function_name] = render
+
+        return render_class
+
+    def register_parser(self, parser_class: type[BaseFunctionParser]) -> type[BaseFunctionParser]:
+        parser = parser_class()
+        parser.manager = self
+        parsers_map = self._hof_parsers_map if isinstance(parser, HigherOrderFunctionParser) else self._parsers_map
+        for platform_function_name in parser.function_names_map:
+            parsers_map[platform_function_name] = parser
+
+        if parser.functions_group_name:
+            parsers_map[parser.functions_group_name] = parser
+
+        return parser_class
+
+    def get_hof_parser(self, platform_func_name: str) -> HigherOrderFunctionParser:
+        if INIT_FUNCTIONS and (parser := self._hof_parsers_map.get(platform_func_name)):
+            return parser
 
-    @cached_property
-    def _inverted_names_map(self) -> dict[str, str]:
-        return {value: key for key, value in self._names_map.items()}
+        raise NotSupportedFunctionException
 
-    def get_parser(self, generic_func_name: str) -> HigherOrderFunctionParser:
-        if INIT_FUNCTIONS and (parser := self._parsers_map.get(generic_func_name)):
+    def get_parser(self, platform_func_name: str) -> FunctionParser:
+        if INIT_FUNCTIONS and (parser := self._parsers_map.get(platform_func_name)):
             return parser
 
         raise NotSupportedFunctionException
@@ -130,56 +167,29 @@ def get_in_query_render(self, generic_func_name: str) -> FunctionRender:
 
         raise NotSupportedFunctionException
 
-    def get_generic_func_name(self, platform_func_name: str) -> Optional[str]:
-        if INIT_FUNCTIONS and (generic_func_name := self._names_map.get(platform_func_name)):
-            return generic_func_name
-
-        raise NotSupportedFunctionException
-
-    def get_platform_func_name(self, generic_func_name: str) -> Optional[str]:
-        if INIT_FUNCTIONS:
-            return self._inverted_names_map.get(generic_func_name)
-
     @property
     def order_to_render(self) -> dict[str, int]:
         if INIT_FUNCTIONS:
             return self._order_to_render
 
         return {}
 
-    @property
-    def render_to_prefix_functions(self) -> list[str]:
-        if INIT_FUNCTIONS:
-            return self._render_to_prefix_functions
-
-        return []
-
 
 class PlatformFunctions:
+    dir_path: str = None
+    platform_query_render: PlatformQueryRender = None
     manager: PlatformFunctionsManager = PlatformFunctionsManager()
+
     function_delimiter = "|"
 
-    def parse(self, query: str) -> ParsedFunctions:
-        parsed = []
-        not_supported = []
-        invalid = []
-        functions = query.split(self.function_delimiter)
-        for func in functions:
-            split_func = func.strip().split(" ")
-            func_name, func_body = split_func[0], " ".join(split_func[1:])
-            try:
-                func_parser = self.manager.get_parser(self.manager.get_generic_func_name(func_name))
-                parsed.append(func_parser.parse(func_body, func))
-            except NotSupportedFunctionException:
-                not_supported.append(func)
-            except InvalidFunctionSignature:
-                invalid.append(func)
+    def __init__(self):
+        self.manager.platform_functions = self
+        if self.dir_path:
+            execute_module(f"{self.dir_path}/parsers/__init__.py")
+            execute_module(f"{self.dir_path}/renders/__init__.py")
 
-        return ParsedFunctions(
-            functions=parsed,
-            not_supported=[self.wrap_function_with_delimiter(func) for func in not_supported],
-            invalid=invalid,
-        )
+    def parse(self, query: str) -> ParsedFunctions:  # noqa: ARG002
+        return ParsedFunctions()
 
     def _sort_functions_to_render(self, functions: list[Function]) -> list[Function]:
         return sorted(functions, key=lambda func: self.manager.order_to_render.get(func.name, 0))
@@ -193,7 +203,7 @@ def render(self, functions: list[Function], source_mapping: SourceMapping) -> Re
             try:
                 func_render = self.manager.get_render(func.name)
                 _rendered = func_render.render(func, source_mapping)
-                if func.name in self.manager.render_to_prefix_functions:
+                if func_render.render_to_prefix:
                     rendered_prefix += _rendered
                 else:
                     rendered += self.wrap_function_with_delimiter(_rendered)

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

@@ -197,6 +197,7 @@ class PlatformQueryRender(QueryRender):
     not_token = "not"
 
     group_token = "(%s)"
+    query_parts_delimiter = " "
 
     field_value_map = BaseQueryFieldValue(or_token=or_token)
 
@@ -262,8 +263,14 @@ def apply_token(self, token: Union[FieldValue, Keyword, Identifier], source_mapp
 
     def generate_query(self, tokens: list[TOKEN_TYPE], source_mapping: SourceMapping) -> str:
         result_values = []
+        unmapped_fields = set()
         for token in tokens:
-            result_values.append(self.apply_token(token=token, source_mapping=source_mapping))
+            try:
+                result_values.append(self.apply_token(token=token, source_mapping=source_mapping))
+            except StrictPlatformException as err:
+                unmapped_fields.add(err.field_name)
+        if unmapped_fields:
+            raise StrictPlatformException(self.details.name, "", source_mapping.source_id, sorted(unmapped_fields))
         return "".join(result_values)
 
     def wrap_query_with_meta_info(self, meta_info: MetaInfoContainer, query: str) -> str:
@@ -284,6 +291,10 @@ def wrap_query_with_meta_info(self, meta_info: MetaInfoContainer, query: str) ->
     def _finalize_search_query(query: str) -> str:
         return query
 
+    def _join_query_parts(self, prefix: str, query: str, functions: str) -> str:
+        parts = filter(lambda s: bool(s), map(str.strip, [prefix, self._finalize_search_query(query), functions]))
+        return self.query_parts_delimiter.join(parts)
+
     def finalize_query(
         self,
         prefix: str,
@@ -295,8 +306,7 @@ def finalize_query(
         *args,  # noqa: ARG002
         **kwargs,  # noqa: ARG002
     ) -> str:
-        parts = filter(lambda s: bool(s), map(str.strip, [prefix, self._finalize_search_query(query), functions]))
-        query = " ".join(parts)
+        query = self._join_query_parts(prefix, query, functions)
         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)
@@ -383,7 +393,7 @@ def _generate_from_tokenized_query_container(self, query_container: TokenizedQue
                     defined_raw_log_fields = self.generate_raw_log_fields(
                         fields=query_container.meta_info.query_fields, source_mapping=source_mapping
                     )
-                    prefix += f"\n{defined_raw_log_fields}\n"
+                    prefix += f"\n{defined_raw_log_fields}"
                 result = self.generate_query(tokens=query_container.tokens, source_mapping=source_mapping)
             except StrictPlatformException as err:
                 errors.append(err)

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
 

uncoder-core/app/translator/mappings/platforms/palo_alto_cortex/default.yml

@@ -125,3 +125,4 @@ field_mapping:
   SourceOS: xdm.source.host.os
   DestinationOS: xdm.target.host.os
   url_category: xdm.network.http.url_category
+  EventSeverity: xdm.alert.severity