fix

Author: tarnopolskyi

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

@@ -2,36 +2,15 @@
 import os
 import ssl
 import urllib.request
-from dataclasses import dataclass
 from json import JSONDecodeError
-from typing import Optional, Union
+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
 
 
-@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",)
@@ -157,16 +136,13 @@ def get_technique(self, technique_id: str) -> Optional[MitreTechniqueContainer]:
 
     def get_mitre_info(
         self, tactics: Optional[list[str]] = None, techniques: Optional[list[str]] = None
-    ) -> Optional[MitreInfoContainer]:
+    ) -> 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)
+        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/models/query_container.py

@@ -6,11 +6,31 @@
 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
 
 
+@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,

uncoder-core/app/translator/tools/utils.py

@@ -85,7 +85,10 @@ def parse_rule_description_str(description: str) -> dict:
     pattern = r"___name___:\s*(?P<value>.+)\."
     for key, name in keys_map.items():
         if search := re.search(pattern.replace("___name___", name), description):
-            parsed[key] = search.group("value")
+            if key == "author":
+                parsed[key] = search.group("value").split(",")
+            else:
+                parsed[key] = search.group("value")
             description = description[: search.start()]
 
     parsed["description"] = description.strip()