upd mitre process

Author: tarnopolskyi

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

@@ -2,13 +2,36 @@
 import os
 import ssl
 import urllib.request
+from dataclasses import dataclass
 from json import JSONDecodeError
+from typing import Optional, Union
 from urllib.error import HTTPError
 
 from app.translator.tools.singleton_meta import SingletonMeta
 from const import ROOT_PROJECT_PATH
 
 
+@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: Union[list[MitreTacticContainer], list]
+    techniques: Union[list[MitreTechniqueContainer], list]
+
+
 class MitreConfig(metaclass=SingletonMeta):
     config_url: str = "https://raw.githubusercontent.com/mitre/cti/master/enterprise-attack/enterprise-attack.json"
     mitre_source_types: tuple = ("mitre-attack",)
@@ -116,9 +139,34 @@ 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
+    ) -> Optional[MitreInfoContainer]:
+        tactics_list = []
+        techniques_list = []
+        if tactics:
+            for tactic in tactics:
+                if tactic_found := self.get_tactic(tactic=tactic.lower()):
+                    tactics_list.append(tactic_found)
+        if techniques:
+            for technique in techniques:
+                if technique_found := self.get_technique(technique_id=technique.lower()):
+                    techniques_list.append(technique_found)
+        if tactics_list or techniques_list:
+            return MitreInfoContainer(tactics=tactics_list, techniques=techniques_list)

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

@@ -1,11 +1,11 @@
 import json
-from typing import Union
+from typing import Optional, Union
 
 import xmltodict
 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:
@@ -29,18 +29,20 @@ 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]) -> Optional[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)
+        if parsed_techniques or parsed_tactics:
+            return MitreInfoContainer(tactics=parsed_tactics, techniques=parsed_techniques)
 
 
 class XMLRuleMixin:

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

@@ -6,6 +6,7 @@
 from app.translator.core.const import QUERY_TOKEN_TYPE
 from app.translator.core.custom_types.meta_info import SeverityType
 from app.translator.core.mapping import DEFAULT_MAPPING_NAME
+from app.translator.core.mitre import MitreInfoContainer
 from app.translator.core.models.functions.base import ParsedFunctions
 from app.translator.core.models.query_tokens.field import Field
 
@@ -17,15 +18,15 @@ def __init__(
         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,
+        mitre_attack: Optional[MitreInfoContainer] = None,
         raw_mitre_attack: Optional[list[str]] = None,
         status: Optional[str] = None,
         false_positives: Optional[list[str]] = None,
@@ -44,7 +45,7 @@ def __init__(
         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 []

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

@@ -31,14 +31,16 @@ class ElasticSearchRuleParser(ElasticSearchQueryParser, JsonRuleMixin):
 
     def parse_raw_query(self, text: str, language: str) -> RawQueryContainer:
         rule = self.load_rule(text=text)
-        mitre_attack = {"tactics": [], "techniques": []}
         parsed_description = parse_rule_description_str(rule.get("description", ""))
+
+        mitre_attack = None
         if rule_mitre_attack_data := rule.get("threat"):
-            for threat_data in rule_mitre_attack_data:
-                if technique := self.mitre_config.get_technique(threat_data["technique"][0]["id"].lower()):
-                    mitre_attack["techniques"].append(technique)
-                if tactic := self.mitre_config.get_tactic(threat_data["tactic"]["name"].replace(" ", "_").lower()):
-                    mitre_attack["tactics"].append(tactic)
+            mitre_attack = self.mitre_config.get_mitre_info(
+                tactics=[
+                    threat_data["tactic"]["name"].replace(" ", "_").lower() for threat_data in rule_mitre_attack_data
+                ],
+                techniques=[threat_data["technique"][0]["id"].lower() for threat_data in rule_mitre_attack_data],
+            )
 
         return RawQueryContainer(
             query=rule["query"],
@@ -52,6 +54,6 @@ def parse_raw_query(self, text: str, language: str) -> RawQueryContainer:
                 severity=rule.get("severity"),
                 license_=parsed_description.get("license"),
                 tags=rule.get("tags"),
-                mitre_attack=mitre_attack if mitre_attack["tactics"] or mitre_attack["techniques"] else None,
+                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
@@ -53,25 +53,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:
@@ -100,7 +100,7 @@ def finalize_query(
                 "description": meta_info.description or rule["description"] or _AUTOGENERATED_TEMPLATE,
                 "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)

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

@@ -62,6 +62,10 @@ def finalize_query(
     ) -> str:
         query = super().finalize_query(prefix=prefix, query=query, functions=functions)
         rule = copy.deepcopy(XPACK_WATCHER_RULE)
+        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["metadata"].update(
             {
                 "query": query,
@@ -70,9 +74,9 @@ def finalize_query(
                     description=meta_info.description or _AUTOGENERATED_TEMPLATE,
                     author=meta_info.author,
                     license_=meta_info.license,
-                    mitre_attack=meta_info.mitre_attack,
+                    mitre_attack=mitre_attack,
                 ),
-                "tags": meta_info.mitre_attack,
+                "tags": meta_info.tags,
             }
         )
         rule["input"]["search"]["request"]["body"]["query"]["bool"]["must"][0]["query_string"]["query"] = query

uncoder-core/app/translator/platforms/forti_siem/renders/forti_siem_rule.py

@@ -24,6 +24,7 @@
 from app.translator.core.custom_types.values import ValueType
 from app.translator.core.exceptions.render import UnsupportedRenderMethod
 from app.translator.core.mapping import SourceMapping
+from app.translator.core.mitre import MitreInfoContainer
 from app.translator.core.models.platform_details import PlatformDetails
 from app.translator.core.models.query_container import MetaInfoContainer, TokenizedQueryContainer
 from app.translator.core.models.query_tokens.field_value import FieldValue
@@ -363,16 +364,18 @@ def generate_rule_name(title: str) -> str:
         return re.sub(r'[\'"()+,]*', "", rule_name)
 
     @staticmethod
-    def get_mitre_info(mitre_attack: dict) -> tuple[str, str]:
+    def get_mitre_info(mitre_attack: Union[MitreInfoContainer, None]) -> tuple[str, str]:
+        if not mitre_attack:
+            return "", ""
         tactics = set()
         techniques = set()
-        for tactic in mitre_attack.get("tactics", []):
-            if tactic_name := tactic.get("tactic"):
+        for tactic in mitre_attack.tactics:
+            if tactic_name := tactic.name:
                 tactics.add(tactic_name)
 
-        for tech in mitre_attack.get("techniques", []):
-            techniques.add(tech["technique_id"])
-            tactics = tactics.union(set(tech.get("tactic", [])))
+        for tech in mitre_attack.techniques:
+            techniques.add(tech.technique_id)
+            tactics = tactics.union(set(tech.tactic))
 
         return ", ".join(sorted(tactics)), ", ".join(sorted(techniques))