Merge pull request #152 from UncoderIO/gis-7984

Gis 7984

Author: alexvolha

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"""
+"""Set to True to return only first query if rendered multiple options"""

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,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/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

@@ -16,9 +16,10 @@
 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
@@ -165,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,19 +200,19 @@ class PlatformQueryRender(QueryRender):
 
     field_value_map = BaseQueryFieldValue(or_token=or_token)
 
-    query_pattern = "{table} {query} {functions}"
-    raw_log_field_pattern_map: dict = 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} ",
             LogicalOperatorType.NOT: f" {self.not_token} ",
         }
 
-    def generate_prefix(self, log_source_signature: LogSourceSignature, functions_prefix: str = "") -> str:  # noqa: ARG002
-        if str(log_source_signature):
-            return f"{log_source_signature!s} {self.and_token}"
+    def generate_prefix(self, log_source_signature: Optional[LogSourceSignature], functions_prefix: str = "") -> str:  # noqa: ARG002
+        if log_source_signature and str(log_source_signature):
+            return f"{log_source_signature} {self.and_token}"
         return ""
 
     def generate_functions(self, functions: list[Function], source_mapping: SourceMapping) -> RenderedFunctions:
@@ -272,6 +280,10 @@ def wrap_query_with_meta_info(self, meta_info: MetaInfoContainer, query: str) ->
             query = f"{query}\n\n{query_meta_info}"
         return query
 
+    @staticmethod
+    def _finalize_search_query(query: str) -> str:
+        return query
+
     def finalize_query(
         self,
         prefix: str,
@@ -283,8 +295,8 @@ def finalize_query(
         *args,  # noqa: ARG002
         **kwargs,  # noqa: ARG002
     ) -> str:
-        query = self.query_pattern.format(prefix=prefix, query=query, functions=functions).strip()
-
+        parts = filter(lambda s: bool(s), map(str.strip, [prefix, self._finalize_search_query(query), functions]))
+        query = " ".join(parts)
         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)
@@ -327,15 +339,15 @@ def _generate_from_raw_query_container(self, query_container: RawQueryContainer)
 
     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)
+            return raw_log_field_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 = []
+            prefix_list = []
             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 _prefix_list := self.process_raw_log_field_prefix(field=f, source_mapping=source_mapping):
+                    prefix_list.extend(_prefix_list)
+            return prefix_list
         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)]
 
@@ -352,9 +364,11 @@ 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 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))
+            if prefix_list := self.process_raw_log_field_prefix(field=mapped_field, source_mapping=source_mapping):
+                for prefix in prefix_list:
+                    if prefix not in defined_raw_log_fields:
+                        defined_raw_log_fields.append(prefix)
+        return "\n".join(defined_raw_log_fields)
 
     def _generate_from_tokenized_query_container(self, query_container: TokenizedQueryContainer) -> str:
         queries_map = {}

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/core/str_value_manager.py

@@ -16,6 +16,7 @@
 limitations under the License.
 -----------------------------------------------------------------
 """
+
 from typing import ClassVar, Optional, TypeVar, Union
 
 from app.translator.core.custom_types.values import ValueType