update in render

tarnopolskyi · tarnopolskyi · commit 4abc013fae9d · 2024-06-11T09:31:05.000+02:00
diff --git a/uncoder-core/app/translator/platforms/sumo_logic/__init__.py b/uncoder-core/app/translator/platforms/sumo_logic/__init__.py
@@ -1,2 +1,3 @@
 from app.translator.platforms.sumo_logic.renders.sumologic_cti import SumologicCTI  # noqa: F401
-from app.translator.platforms.sumo_logic.renders.sumologic import SumologicSearchQueryRender
+from app.translator.platforms.sumo_logic.renders.sumologic import SumologicSearchQueryRender
+from app.translator.platforms.sumo_logic.renders.sumologic_cse import SumologicCSERender, SumologicCSERuleRender
diff --git a/uncoder-core/app/translator/platforms/sumo_logic/const.py b/uncoder-core/app/translator/platforms/sumo_logic/const.py
@@ -26,17 +26,43 @@
 SUMOLOGIC_CSE_QUERY_DETAILS = {
     "platform_id": "sumologic-cse-query",
     "name": "Sumo Logic CSE",
-    "platform_name": "Query",
+    "platform_name": "CSE Query",
     **PLATFORM_DETAILS
 }
 
 SUMOLOGIC_CSE_RULE_DETAILS = {
     "platform_id": "sumologic-cse-rule",
     "name": "Sumo Logic CSE Rule",
-    "platform_name": "Query",
+    "platform_name": "CSE Rule",
     **PLATFORM_DETAILS
 }
 
+SEVERITY_MAP = {
+    'informational': 1,
+    'low': 2,
+    'medium': 3,
+    'high': 4,
+    'critical': 5
+}
+
+ALLOWED_CATEGORIES = [
+    "Threat Intelligence",
+    "Initial Access",
+    "Execution",
+    "Persistence",
+    "Privilege Escalation",
+    "Defense Evasion",
+    "Credential Access",
+    "Discovery",
+    "Lateral Movement",
+    "Collection",
+    "Command and Control",
+    "Exfiltration",
+    "Impact"
+]
+
+
+
 sumologic_search_query_details = PlatformDetails(**SUMOLOGIC_SEARCH_QUERY_DETAILS)
 sumologic_cse_query_details = PlatformDetails(**SUMOLOGIC_CSE_QUERY_DETAILS)
 sumologic_cse_rule_details = PlatformDetails(**SUMOLOGIC_CSE_RULE_DETAILS)
diff --git a/uncoder-core/app/translator/platforms/sumo_logic/renders/sumologic_cse.py b/uncoder-core/app/translator/platforms/sumo_logic/renders/sumologic_cse.py
@@ -16,19 +16,29 @@
 limitations under the License.
 -----------------------------------------------------------------
 """
+import copy
+import json
+from typing import Optional
+from app.translator.core.mapping import SourceMapping
 from app.translator.core.models.platform_details import PlatformDetails
+from app.translator.core.models.query_container import MetaInfoContainer
 from app.translator.managers import render_manager
 from app.translator.core.render import PlatformQueryRender
 
-from app.translator.platforms.sumo_logic.const import sumologic_search_query_details
+from app.translator.platforms.sumo_logic.const import (
+    SEVERITY_MAP,
+    sumologic_cse_query_details,
+    sumologic_cse_rule_details,
+    DEFAULT_SUMOLOGIC_CSE_RULE
+)
 from app.translator.platforms.sumo_logic.mapping import SumoLogicMappings, sumologic_mappings
 from app.translator.platforms.sumo_logic.renders.sumologic import SumologicFieldValue
 
 
 
 @render_manager.register
 class SumologicCSERender(PlatformQueryRender):
-    details: PlatformDetails = sumologic_search_query_details
+    details: PlatformDetails = sumologic_cse_query_details
     mappings: SumoLogicMappings = sumologic_mappings
 
     or_token = "OR"
@@ -41,12 +51,36 @@ class SumologicCSERender(PlatformQueryRender):
 
 @render_manager.register
 class SumologicCSERuleRender(PlatformQueryRender):
-    details: PlatformDetails = sumologic_search_query_details
+    details: PlatformDetails = sumologic_cse_rule_details
     mappings: SumoLogicMappings = sumologic_mappings
 
     or_token = "OR"
     and_token = "AND"
     not_token = "NOT"
 
     field_value_map = SumologicFieldValue(or_token=or_token)
-    query_pattern = "{prefix} {query} {functions}"
+    query_pattern = "{prefix} {query} {functions}"
+
+
+    def finalize_query(
+        self,
+        prefix: str,
+        query: str,
+        functions: str,
+        meta_info: Optional[MetaInfoContainer] = None,
+        source_mapping: Optional[SourceMapping] = None,  # noqa: ARG002
+        not_supported_functions: Optional[list] = None,
+        *args,  # noqa: ARG002
+        **kwargs,  # noqa: ARG002
+    ) -> str:
+        query = super().finalize_query(prefix=prefix, query=query, functions=functions)
+        rule = copy.deepcopy(DEFAULT_SUMOLOGIC_CSE_RULE)
+        rule["name"] = meta_info.title or rule["name"] if meta_info else rule['name']
+        rule["description"] = meta_info.description or rule["description"] if meta_info else rule['description']
+        rule["score"] = SEVERITY_MAP.get(meta_info.severity) if meta_info else SEVERITY_MAP['medium']
+        rule["expression"] = query
+        rule_str = json.dumps(rule, indent=4, sort_keys=False, ensure_ascii=False)
+        if not_supported_functions:
+            rendered_not_supported = self.render_not_supported_functions(not_supported_functions)
+            return rule_str + rendered_not_supported
+        return rule_str
