resolve conflicts

Author: alexvolha

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

@@ -3,8 +3,10 @@
 import ssl
 import urllib.request
 from json import JSONDecodeError
+from typing import Optional
 from urllib.error import HTTPError
 
+from app.translator.core.models.query_container import MitreInfoContainer, MitreTacticContainer, MitreTechniqueContainer
 from app.translator.tools.singleton_meta import SingletonMeta
 from const import ROOT_PROJECT_PATH
 
@@ -116,9 +118,31 @@ def __load_mitre_configs_from_files(self) -> None:
         except JSONDecodeError:
             self.techniques = {}
 
-    def get_tactic(self, tactic: str) -> dict:
+    def get_tactic(self, tactic: str) -> Optional[MitreTacticContainer]:
         tactic = tactic.replace(".", "_")
-        return self.tactics.get(tactic, {})
-
-    def get_technique(self, technique_id: str) -> dict:
-        return self.techniques.get(technique_id, {})
+        if tactic_found := self.tactics.get(tactic):
+            return MitreTacticContainer(
+                external_id=tactic_found["external_id"], url=tactic_found["url"], name=tactic_found["tactic"]
+            )
+
+    def get_technique(self, technique_id: str) -> Optional[MitreTechniqueContainer]:
+        if technique_found := self.techniques.get(technique_id):
+            return MitreTechniqueContainer(
+                technique_id=technique_found["technique_id"],
+                name=technique_found["technique"],
+                url=technique_found["url"],
+                tactic=technique_found["tactic"],
+            )
+
+    def get_mitre_info(
+        self, tactics: Optional[list[str]] = None, techniques: Optional[list[str]] = None
+    ) -> MitreInfoContainer:
+        tactics_list = []
+        techniques_list = []
+        for tactic in tactics or []:
+            if tactic_found := self.get_tactic(tactic=tactic.lower()):
+                tactics_list.append(tactic_found)
+        for technique in techniques or []:
+            if technique_found := self.get_technique(technique_id=technique.lower()):
+                techniques_list.append(technique_found)
+        return MitreInfoContainer(tactics=tactics_list, techniques=techniques_list)

uncoder-core/app/translator/core/mixins/rule.py

@@ -5,10 +5,12 @@
 import yaml
 
 from app.translator.core.exceptions.core import InvalidJSONStructure, InvalidXMLStructure, InvalidYamlStructure
-from app.translator.core.mitre import MitreConfig
+from app.translator.core.mitre import MitreConfig, MitreInfoContainer
 
 
 class JsonRuleMixin:
+    mitre_config: MitreConfig = MitreConfig()
+
     @staticmethod
     def load_rule(text: str) -> dict:
         try:
@@ -27,18 +29,19 @@ def load_rule(text: str) -> dict:
         except yaml.YAMLError as err:
             raise InvalidYamlStructure(error=str(err)) from err
 
-    def parse_mitre_attack(self, tags: list[str]) -> dict[str, list]:
-        result = {"tactics": [], "techniques": []}
+    def parse_mitre_attack(self, tags: list[str]) -> MitreInfoContainer:
+        parsed_techniques = []
+        parsed_tactics = []
         for tag in set(tags):
             tag = tag.lower()
             if tag.startswith("attack."):
                 tag = tag[7::]
             if tag.startswith("t"):
                 if technique := self.mitre_config.get_technique(tag):
-                    result["techniques"].append(technique)
+                    parsed_techniques.append(technique)
             elif tactic := self.mitre_config.get_tactic(tag):
-                result["tactics"].append(tactic)
-        return result
+                parsed_tactics.append(tactic)
+        return MitreInfoContainer(tactics=parsed_tactics, techniques=parsed_techniques)
 
 
 class XMLRuleMixin:

uncoder-core/app/translator/core/models/query_container.py

@@ -1,6 +1,6 @@
 import uuid
 from dataclasses import dataclass, field
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Optional
 
 from app.translator.core.const import QUERY_TOKEN_TYPE
@@ -10,43 +10,72 @@
 from app.translator.core.models.query_tokens.field import Field
 
 
+@dataclass
+class MitreTechniqueContainer:
+    technique_id: str
+    name: str
+    url: str
+    tactic: list[str]
+
+
+@dataclass
+class MitreTacticContainer:
+    external_id: str
+    url: str
+    name: str
+
+
+@dataclass
+class MitreInfoContainer:
+    tactics: list[MitreTacticContainer] = field(default_factory=list)
+    techniques: list[MitreTechniqueContainer] = field(default_factory=list)
+
+
 class MetaInfoContainer:
     def __init__(
         self,
         *,
         id_: Optional[str] = None,
         title: Optional[str] = None,
         description: Optional[str] = None,
-        author: Optional[str] = None,
+        author: Optional[list[str]] = None,
         date: Optional[str] = None,
         output_table_fields: Optional[list[Field]] = None,
         query_fields: Optional[list[Field]] = None,
         license_: Optional[str] = None,
         severity: Optional[str] = None,
         references: Optional[list[str]] = None,
         tags: Optional[list[str]] = None,
-        mitre_attack: Optional[dict[str, list]] = None,
+        raw_mitre_attack: Optional[list[str]] = None,
         status: Optional[str] = None,
         false_positives: Optional[list[str]] = None,
         source_mapping_ids: Optional[list[str]] = None,
         parsed_logsources: Optional[dict] = None,
+        timeframe: Optional[timedelta] = None,
+        mitre_attack: MitreInfoContainer = MitreInfoContainer(),
     ) -> None:
         self.id = id_ or str(uuid.uuid4())
         self.title = title or ""
         self.description = description or ""
-        self.author = author or ""
+        self.author = [v.strip() for v in author] if author else []
         self.date = date or datetime.now().date().strftime("%Y-%m-%d")
         self.output_table_fields = output_table_fields or []
         self.query_fields = query_fields or []
         self.license = license_ or "DRL 1.1"
         self.severity = severity or SeverityType.low
         self.references = references or []
         self.tags = tags or []
-        self.mitre_attack = mitre_attack or {}
+        self.mitre_attack = mitre_attack or None
+        self.raw_mitre_attack = raw_mitre_attack or []
         self.status = status or "stable"
         self.false_positives = false_positives or []
-        self.source_mapping_ids = source_mapping_ids or [DEFAULT_MAPPING_NAME]
+        self.source_mapping_ids = sorted(source_mapping_ids) if source_mapping_ids else [DEFAULT_MAPPING_NAME]
         self.parsed_logsources = parsed_logsources or {}
+        self.timeframe = timeframe
+
+    @property
+    def author_str(self) -> str:
+        return ", ".join(self.author)
 
 
 @dataclass

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

@@ -208,15 +208,17 @@ def wrap_with_not_supported_functions(self, query: str, not_supported_functions:
         return query
 
     def wrap_with_unmapped_fields(self, query: str, fields: Optional[list[str]]) -> str:
-        if fields:
+        if wrap_query_with_meta_info_ctx_var.get() and fields:
             return query + "\n\n" + self.wrap_with_comment(f"{self.unmapped_fields_text}{', '.join(fields)}")
         return query
 
     def wrap_with_comment(self, value: str) -> str:
         return f"{self.comment_symbol} {value}"
 
     @abstractmethod
-    def generate(self, query_container: Union[RawQueryContainer, TokenizedQueryContainer]) -> str:
+    def generate(
+        self, raw_query_container: RawQueryContainer, tokenized_query_container: Optional[TokenizedQueryContainer]
+    ) -> str:
         raise NotImplementedError("Abstract method")
 
 
@@ -318,7 +320,7 @@ def wrap_with_meta_info(self, query: str, meta_info: Optional[MetaInfoContainer]
             meta_info_dict = {
                 "name: ": meta_info.title,
                 "uuid: ": meta_info.id,
-                "author: ": meta_info.author if meta_info.author else "not defined in query/rule",
+                "author: ": meta_info.author_str or "not defined in query/rule",
                 "licence: ": meta_info.license,
             }
             query_meta_info = "\n".join(
@@ -370,7 +372,7 @@ def finalize(self, queries_map: dict[str, str]) -> str:
 
         return result
 
-    def _get_source_mappings(self, source_mapping_ids: list[str]) -> list[SourceMapping]:
+    def _get_source_mappings(self, source_mapping_ids: list[str]) -> Optional[list[SourceMapping]]:
         source_mappings = []
         for source_mapping_id in source_mapping_ids:
             if source_mapping := self.mappings.get_source_mapping(source_mapping_id):
@@ -468,8 +470,9 @@ def generate_from_tokenized_query_container(self, query_container: TokenizedQuer
             raise errors[0]
         return self.finalize(queries_map)
 
-    def generate(self, query_container: Union[RawQueryContainer, TokenizedQueryContainer]) -> str:
-        if isinstance(query_container, RawQueryContainer):
-            return self.generate_from_raw_query_container(query_container)
-
-        return self.generate_from_tokenized_query_container(query_container)
+    def generate(
+        self, raw_query_container: RawQueryContainer, tokenized_query_container: Optional[TokenizedQueryContainer]
+    ) -> str:
+        if tokenized_query_container:
+            return self.generate_from_tokenized_query_container(tokenized_query_container)
+        return self.generate_from_raw_query_container(raw_query_container)

uncoder-core/app/translator/platforms/chronicle/parsers/chronicle_rule.py

@@ -31,10 +31,10 @@
 @parser_manager.register
 class ChronicleRuleParser(ChronicleQueryParser):
     details: PlatformDetails = chronicle_rule_details
-    rule_name_pattern = "rule\s(?P<rule_name>[a-z0-9_]+)\s{"
-    meta_info_pattern = "meta:\n(?P<meta_info>[a-zA-Z0-9_\\\.*,>–<—~#$’`:;%+^\|?!@\s\"/=\-&'\(\)\[\]]+)\n\s+events:"  # noqa: RUF001
-    rule_pattern = "events:\n\s*(?P<query>[a-zA-Z\w0-9_%{}\|\.,!#^><:~\s\"\/=+?\-–&;$()`\*@\[\]'\\\]+)\n\s+condition:"  # noqa: RUF001
-    event_name_pattern = "condition:\n\s*(?P<event_name>\$[a-zA-Z_0-9]+)\n"
+    rule_name_pattern = r"rule\s+(?P<rule_name>[a-zA-Z0-9_]+)\s+{"
+    meta_info_pattern = r"meta:\n(?P<meta_info>[a-zA-Z0-9_\\\.*,>–<—~#$’`:;%+^\|?!@\s\"/=\-&'\(\)\[\]]+)\n\s+events:"  # noqa: RUF001
+    rule_pattern = r"events:\n\s*(?P<query>[a-zA-Z\w0-9_%{}\|\.,!#^><:~\s\"\/=+?\-–&;$()`\*@\[\]'\\]+)\n\s+condition:"  # noqa: RUF001
+    event_name_pattern = r"condition:\n\s*(?P<event_name>\$[a-zA-Z_0-9]+)\n"
     mappings: ChronicleMappings = chronicle_rule_mappings
     tokenizer = ChronicleRuleTokenizer()
 

uncoder-core/app/translator/platforms/chronicle/renders/chronicle_rule.py

@@ -119,7 +119,7 @@ def finalize_query(
         rule = DEFAULT_CHRONICLE_SECURITY_RULE.replace("<query_placeholder>", query)
         rule = rule.replace("<title_place_holder>", self.prepare_title(meta_info.title) or _AUTOGENERATED_TEMPLATE)
         description = meta_info.description or _AUTOGENERATED_TEMPLATE
-        rule = rule.replace("<author_place_holder>", meta_info.author)
+        rule = rule.replace("<author_place_holder>", ", ".join(meta_info.author))
         rule = rule.replace("<description_place_holder>", description)
         rule = rule.replace("<licence_place_holder>", meta_info.license)
         rule = rule.replace("<rule_id_place_holder>", meta_info.id)

uncoder-core/app/translator/platforms/elasticsearch/parsers/detection_rule.py

@@ -22,6 +22,7 @@
 from app.translator.managers import parser_manager
 from app.translator.platforms.elasticsearch.const import elasticsearch_rule_details
 from app.translator.platforms.elasticsearch.parsers.elasticsearch import ElasticSearchQueryParser
+from app.translator.tools.utils import parse_rule_description_str
 
 
 @parser_manager.register
@@ -30,8 +31,25 @@ class ElasticSearchRuleParser(ElasticSearchQueryParser, JsonRuleMixin):
 
     def parse_raw_query(self, text: str, language: str) -> RawQueryContainer:
         rule = self.load_rule(text=text)
+        parsed_description = parse_rule_description_str(rule.get("description", ""))
+
+        mitre_attack = self.mitre_config.get_mitre_info(
+            tactics=[threat_data["tactic"]["name"].replace(" ", "_").lower() for threat_data in rule.get("threat", [])],
+            techniques=[threat_data["technique"][0]["id"].lower() for threat_data in rule.get("threat", [])],
+        )
+
         return RawQueryContainer(
             query=rule["query"],
             language=language,
-            meta_info=MetaInfoContainer(title=rule["name"], description=rule["description"]),
+            meta_info=MetaInfoContainer(
+                id_=rule.get("rule_id"),
+                title=rule.get("name"),
+                description=parsed_description.get("description") or rule.get("description"),
+                references=rule.get("references", []),
+                author=parsed_description.get("author") or rule.get("author"),
+                severity=rule.get("severity"),
+                license_=parsed_description.get("license"),
+                tags=rule.get("tags"),
+                mitre_attack=mitre_attack,
+            ),
         )

uncoder-core/app/translator/platforms/elasticsearch/renders/detection_rule.py

@@ -22,7 +22,7 @@
 from typing import Optional, Union
 
 from app.translator.core.mapping import SourceMapping
-from app.translator.core.mitre import MitreConfig
+from app.translator.core.mitre import MitreConfig, MitreInfoContainer
 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
@@ -33,6 +33,7 @@
     ElasticSearchFieldValue,
     ElasticSearchQueryRender,
 )
+from app.translator.tools.utils import get_rule_description_str
 
 _AUTOGENERATED_TEMPLATE = "Autogenerated Elastic Rule"
 
@@ -53,25 +54,25 @@ class ElasticSearchRuleRender(ElasticSearchQueryRender):
 
     field_value_render = ElasticSearchRuleFieldValue(or_token=or_token)
 
-    def __create_mitre_threat(self, mitre_attack: dict) -> Union[list, list[dict]]:
-        if not mitre_attack.get("techniques"):
+    def __create_mitre_threat(self, mitre_attack: MitreInfoContainer) -> Union[list, list[dict]]:
+        if not mitre_attack.techniques:
             return []
         threat = []
 
-        for tactic in mitre_attack["tactics"]:
-            tactic_render = {"id": tactic["external_id"], "name": tactic["tactic"], "reference": tactic["url"]}
+        for tactic in mitre_attack.tactics:
+            tactic_render = {"id": tactic.external_id, "name": tactic.name, "reference": tactic.url}
             sub_threat = {"tactic": tactic_render, "framework": "MITRE ATT&CK", "technique": []}
-            for technique in mitre_attack["techniques"]:
-                technique_id = technique["technique_id"].lower()
+            for technique in mitre_attack.techniques:
+                technique_id = technique.technique_id.lower()
                 if "." in technique_id:
-                    technique_id = technique_id[: technique["technique_id"].index(".")]
+                    technique_id = technique_id[: technique.technique_id.index(".")]
                 main_technique = self.mitre.get_technique(technique_id)
-                if tactic["tactic"] in main_technique["tactic"]:
+                if tactic.name in main_technique.tactic:
                     sub_threat["technique"].append(
                         {
-                            "id": main_technique["technique_id"],
-                            "name": main_technique["technique"],
-                            "reference": main_technique["url"],
+                            "id": main_technique.technique_id,
+                            "name": main_technique.name,
+                            "reference": main_technique.url,
                         }
                     )
             if len(sub_threat["technique"]) > 0:
@@ -94,13 +95,17 @@ def finalize_query(
         query = super().finalize_query(prefix=prefix, query=query, functions=functions)
         rule = copy.deepcopy(ELASTICSEARCH_DETECTION_RULE)
         index = source_mapping.log_source_signature.default_source.get("index") if source_mapping else None
+        description_str = get_rule_description_str(
+            description=meta_info.description or rule["description"] or _AUTOGENERATED_TEMPLATE,
+            license_=meta_info.license,
+        )
         rule.update(
             {
                 "query": query,
-                "description": meta_info.description or rule["description"] or _AUTOGENERATED_TEMPLATE,
+                "description": description_str,
                 "name": meta_info.title or _AUTOGENERATED_TEMPLATE,
                 "rule_id": meta_info.id,
-                "author": [meta_info.author],
+                "author": meta_info.author,
                 "severity": meta_info.severity,
                 "references": meta_info.references,
                 "license": meta_info.license,

uncoder-core/app/translator/platforms/elasticsearch/renders/elast_alert.py

@@ -66,13 +66,18 @@ def finalize_query(
     ) -> str:
         query = super().finalize_query(prefix=prefix, query=query, functions=functions)
         rule = ELASTICSEARCH_ALERT.replace("<query_placeholder>", query)
+        mitre_attack = []
+        if meta_info and meta_info.mitre_attack:
+            mitre_attack.extend([technique.technique_id for technique in meta_info.mitre_attack.techniques])
+            mitre_attack.extend([tactic.name for tactic in meta_info.mitre_attack.tactics])
         rule = rule.replace(
             "<description_place_holder>",
             get_rule_description_str(
                 author=meta_info.author,
                 description=meta_info.description or _AUTOGENERATED_TEMPLATE,
                 license_=meta_info.license,
                 rule_id=meta_info.id,
+                mitre_attack=mitre_attack,
             ),
         )
         rule = rule.replace("<title_place_holder>", meta_info.title or _AUTOGENERATED_TEMPLATE)

uncoder-core/app/translator/platforms/elasticsearch/renders/kibana.py

@@ -68,12 +68,17 @@ def finalize_query(
         rule["_source"]["kibanaSavedObjectMeta"]["searchSourceJSON"] = dumped_rule
         rule["_source"]["title"] = meta_info.title or _AUTOGENERATED_TEMPLATE
         descr = meta_info.description or rule["_source"]["description"] or _AUTOGENERATED_TEMPLATE
+        mitre_attack = []
+        if meta_info and meta_info.mitre_attack:
+            mitre_attack.extend([technique.technique_id for technique in meta_info.mitre_attack.techniques])
+            mitre_attack.extend([tactic.name for tactic in meta_info.mitre_attack.tactics])
         rule["_source"]["description"] = get_rule_description_str(
             description=descr,
             author=meta_info.author,
             rule_id=meta_info.id,
             license_=meta_info.license,
             references=meta_info.references,
+            mitre_attack=mitre_attack,
         )
         rule_str = json.dumps(rule, indent=4, sort_keys=False)
         rule_str = self.wrap_with_unmapped_fields(rule_str, unmapped_fields)