fix forti_siem

tarnopolskyi · tarnopolskyi · commit 4c14e94bc321 · 2024-03-22T10:18:48.000+01:00
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
@@ -18,6 +18,7 @@
 from typing import Optional, Union
 
 from app.translator.const import DEFAULT_VALUE_TYPE
+from app.translator.core.context_vars import with_meta_info_annotation_ctx_var
 from app.translator.core.custom_types.meta_info import SeverityType
 from app.translator.core.custom_types.tokens import OperatorType
 from app.translator.core.custom_types.values import ValueType
@@ -192,8 +193,9 @@ def render_query(
         functions: str,
         meta_info: Optional[MetaInfoContainer] = None,
         source_mapping: Optional[SourceMapping] = None,  # noqa: ARG002
+        fields: Optional[set[str]] = None,
         *args,  # noqa: ARG002
-        **kwargs,
+        **kwargs,  # noqa: ARG002
     ) -> str:
         query = super().render_query(prefix=prefix, query=query, functions=functions)
         rule = FORTI_SIEM_RULE.replace("<header_placeholder>", self.generate_rule_header(meta_info))
@@ -203,11 +205,39 @@ def render_query(
         description = meta_info.description.replace("\n", " ") or _AUTOGENERATED_TEMPLATE
         rule = rule.replace("<description_placeholder>", description)
         rule = rule.replace("<incident_def_placeholder>", self.generate_event_type(title, meta_info.severity))
-        args_list = self.get_args_list(kwargs.get("fields", set()).copy())
+        args_list = self.get_args_list(fields.copy())
         rule = rule.replace("<args_list_placeholder>", self.get_args_str(args_list))
         rule = rule.replace("<query_placeholder>", query)
         rule = rule.replace("<group_by_attr_placeholder>", ", ".join(args_list))
-        return rule.replace("<attr_list_placeholder>", self.get_attr_str(kwargs.get("fields", set()).copy()))
+        return rule.replace("<attr_list_placeholder>", self.get_attr_str(fields.copy()))
+
+    def finalize_query(
+        self,
+        prefix: str,
+        query: str,
+        functions: str,
+        meta_info: Optional[MetaInfoContainer] = None,
+        source_mapping: Optional[SourceMapping] = None,
+        not_supported_functions: Optional[list] = None,
+        fields: Optional[set[str]] = None,
+        *args,  # noqa: ARG002
+        **kwargs,  # noqa: ARG002
+    ) -> str:
+        query = self.render_query(
+            prefix=prefix,
+            query=query,
+            functions=functions,
+            meta_info=meta_info,
+            source_mapping=source_mapping,
+            fields=fields,
+        )
+        if with_meta_info_annotation_ctx_var.get() is False:
+            return query
+        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)
+            return query + rendered_not_supported
+        return query
 
     @staticmethod
     def get_attr_str(fields: set[str]) -> str:
