UncoderIO · tarnopolskyi · Sep 11, 2024 · Aug 21, 2024 · Aug 21, 2024 · Aug 21, 2024
diff --git a/uncoder-core/app/translator/core/exceptions/core.py b/uncoder-core/app/translator/core/exceptions/core.py
@@ -1,10 +1,6 @@
 from typing import Optional
 
 
-class NotImplementedException(BaseException):
-    ...
-
-
 class BasePlatformException(BaseException):
     ...
 

diff --git a/uncoder-core/app/translator/core/exceptions/render.py b/uncoder-core/app/translator/core/exceptions/render.py
@@ -14,5 +14,5 @@ class FunctionRenderException(BaseRenderException):
 
 class UnsupportedRenderMethod(BaseRenderException):
     def __init__(self, platform_name: str, method: str):
-        message = f"Cannot translate. {platform_name} backend does not support {method}."
+        message = f'Cannot translate. {platform_name} backend does not support "{method}".'
         super().__init__(message)
diff --git a/uncoder-core/app/translator/core/mitre.py b/uncoder-core/app/translator/core/mitre.py
@@ -3,21 +3,90 @@
 import ssl
 import urllib.request
 from json import JSONDecodeError
-from typing import Optional
+from typing import Optional, Union
 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
 
 
+class TrieNode:
+    def __init__(self):
+        self.children = {}
+        self.is_end_of_word = False
+        self.result = None
+
+
+class Trie:
+    """
+    Trie (prefix tree) data structure for storing and searching Mitre ATT&CK Techniques and Tactics strings.
+
+    This class handles the insertion and searching of strings related to Mitre ATT&CK Techniques and Tactics, even when
+    the strings have variations in spacing, case, or underscores. By normalizing the text—converting it to lowercase and
+    removing spaces and underscores—different variations of the same logical string are treated as equivalent.
+
+    It means strings 'CredentialAccess', 'credential Access', and 'credential_access' will be processed identically,
+    leading to the same result.
+    """
+
+    def __init__(self):
+        self.root = TrieNode()
+
+    def normalize_text(self, text: str) -> str:
+        return text.replace(" ", "").lower().replace("_", "").lower()
+
+    def insert(self, text: str, result: Union[MitreTacticContainer, MitreTechniqueContainer]) -> None:
+        node = self.root
+        normalized_text = self.normalize_text(text)
+
+        for char in normalized_text:
+            if char not in node.children:
+                node.children[char] = TrieNode()
+            node = node.children[char]
+
+        node.is_end_of_word = True
+        node.result = result
+
+
+class TacticsTrie(Trie):
+    def __init__(self):
+        self.root = TrieNode()
+
+    def search(self, text: str) -> Optional[MitreTacticContainer]:
+        node: TrieNode = self.root
+        normalized_text = self.normalize_text(text)
+
+        for char in normalized_text:
+            if char not in node.children:
+                return
+            node = node.children[char]
+
+        if node.is_end_of_word:
+            return node.result
+
+
+class TechniquesTrie(Trie):
+    def search(self, text: str) -> Optional[MitreTechniqueContainer]:
+        node: TrieNode = self.root
+        normalized_text = self.normalize_text(text)
+
+        for char in normalized_text:
+            if char not in node.children:
+                return
+            node = node.children[char]
+
+        if node.is_end_of_word:
+            return node.result
+
+
 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",)
 
     def __init__(self, server: bool = False):
-        self.tactics = {}
-        self.techniques = {}
+        self.tactics: TacticsTrie = TacticsTrie()
+        self.techniques: TechniquesTrie = TechniquesTrie()
         if not server:
             self.__load_mitre_configs_from_files()
 
@@ -44,7 +113,6 @@ def update_mitre_config(self) -> None:  # noqa: PLR0912
             return
 
         tactic_map = {}
-        technique_map = {}
 
         # Map the tactics
         for entry in mitre_json["objects"]:
@@ -53,11 +121,12 @@ def update_mitre_config(self) -> None:  # noqa: PLR0912
             for ref in entry["external_references"]:
                 if ref["source_name"] == "mitre-attack":
                     tactic_map[entry["x_mitre_shortname"]] = entry["name"]
-                    self.tactics[entry["name"].replace(" ", "_").lower()] = {
-                        "external_id": ref["external_id"],
-                        "url": ref["url"],
-                        "tactic": entry["name"],
-                    }
+
+                    tactic_data = MitreTacticContainer(
+                        external_id=ref["external_id"], url=ref["url"], name=entry["name"]
+                    )
+                    self.tactics.insert(entry["name"], tactic_data)
+
                     break
 
         # Map the techniques
@@ -68,19 +137,15 @@ def update_mitre_config(self) -> None:  # noqa: PLR0912
                 continue
             for ref in entry["external_references"]:
                 if ref["source_name"] in self.mitre_source_types:
-                    technique_map[ref["external_id"]] = entry["name"]
                     sub_tactics = []
-                    # Get Mitre Tactics (Kill-Chains)
                     for tactic in entry["kill_chain_phases"]:
                         if tactic["kill_chain_name"] in self.mitre_source_types:
-                            # Map the short phase_name to tactic name
                             sub_tactics.append(tactic_map[tactic["phase_name"]])
-                    self.techniques[ref["external_id"].lower()] = {
-                        "technique_id": ref["external_id"],
-                        "technique": entry["name"],
-                        "url": ref["url"],
-                        "tactic": sub_tactics,
-                    }
+
+                    technique_data = MitreTechniqueContainer(
+                        technique_id=ref["external_id"], name=entry["name"], url=ref["url"], tactic=sub_tactics
+                    )
+                    self.techniques.insert(ref["external_id"], technique_data)
                     break
 
         # Map the sub-techniques
@@ -92,58 +157,60 @@ def update_mitre_config(self) -> None:  # noqa: PLR0912
                     if ref["source_name"] in self.mitre_source_types:
                         sub_technique_id = ref["external_id"]
                         sub_technique_name = entry["name"]
-                        parent_technique_name = technique_map[sub_technique_id.split(".")[0]]
-                        parent_tactics = self.techniques.get(sub_technique_id.split(".")[0].lower(), {}).get(
-                            "tactic", []
-                        )
-                        sub_technique_name = f"{parent_technique_name} : {sub_technique_name}"
-                        self.techniques[ref["external_id"].lower()] = {
-                            "technique_id": ref["external_id"],
-                            "technique": sub_technique_name,
-                            "url": ref["url"],
-                            "tactic": parent_tactics,
-                        }
+                        if parent_technique := self.techniques.search(sub_technique_id.split(".")[0]):
+                            sub_technique_name = f"{parent_technique.name} : {sub_technique_name}"
+                            sub_technique_data = MitreTechniqueContainer(
+                                technique_id=ref["external_id"],
+                                name=sub_technique_name,
+                                url=ref["url"],
+                                tactic=parent_technique.tactic,
+                            )
+                            self.techniques.insert(sub_technique_id, sub_technique_data)
                         break
 
     def __load_mitre_configs_from_files(self) -> None:
         try:
             with open(os.path.join(ROOT_PROJECT_PATH, "app/dictionaries/tactics.json")) as file:
-                self.tactics = json.load(file)
+                loaded = json.load(file)
+
+                for tactic_name, tactic_data in loaded.items():
+                    tactic = MitreTacticContainer(
+                        external_id=tactic_data["external_id"], url=tactic_data["url"], name=tactic_data["tactic"]
+                    )
+                    self.tactics.insert(tactic_name, tactic)
         except JSONDecodeError:
-            self.tactics = {}
+            print("Unable to load MITRE Tactics")
 
         try:
             with open(os.path.join(ROOT_PROJECT_PATH, "app/dictionaries/techniques.json")) as file:
-                self.techniques = json.load(file)
+                loaded = json.load(file)
+                for technique_id, technique_data in loaded.items():
+                    technique = MitreTechniqueContainer(
+                        technique_id=technique_data["technique_id"],
+                        name=technique_data["technique"],
+                        url=technique_data["url"],
+                        tactic=technique_data["tactic"],
+                    )
+                    self.techniques.insert(technique_id, technique)
         except JSONDecodeError:
-            self.techniques = {}
+            print("Unable to load MITRE Techniques")
 
     def get_tactic(self, tactic: str) -> Optional[MitreTacticContainer]:
-        tactic = tactic.replace(".", "_")
-        if tactic_found := self.tactics.get(tactic):
-            return MitreTacticContainer(
-                external_id=tactic_found["external_id"], url=tactic_found["url"], name=tactic_found["tactic"]
-            )
+        return self.tactics.search(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"],
-            )
+        return self.techniques.search(technique_id)
 
     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()):
+            if tactic_found := self.tactics.search(tactic):
                 tactics_list.append(tactic_found)
         for technique in techniques or []:
-            if technique_found := self.get_technique(technique_id=technique.lower()):
+            if technique_found := self.techniques.search(technique):
                 techniques_list.append(technique_found)
         return MitreInfoContainer(
             tactics=sorted(tactics_list, key=lambda x: x.name),

diff --git a/uncoder-core/app/translator/core/models/query_container.py b/uncoder-core/app/translator/core/models/query_container.py
@@ -31,6 +31,21 @@ class MitreInfoContainer:
     techniques: list[MitreTechniqueContainer] = field(default_factory=list)
 
 
+class RawMetaInfoContainer:
+    def __init__(
+        self,
+        *,
+        trigger_operator: Optional[str] = None,
+        trigger_threshold: Optional[str] = None,
+        query_frequency: Optional[str] = None,
+        query_period: Optional[str] = None,
+    ) -> None:
+        self.trigger_operator = trigger_operator
+        self.trigger_threshold = trigger_threshold
+        self.query_frequency = query_frequency
+        self.query_period = query_period
+
+
 class MetaInfoContainer:
     def __init__(
         self,
@@ -52,7 +67,9 @@ def __init__(
         source_mapping_ids: Optional[list[str]] = None,
         parsed_logsources: Optional[dict] = None,
         timeframe: Optional[timedelta] = None,
+        query_period: Optional[timedelta] = None,
         mitre_attack: MitreInfoContainer = MitreInfoContainer(),
+        raw_metainfo_container: Optional[RawMetaInfoContainer] = None,
     ) -> None:
         self.id = id_ or str(uuid.uuid4())
         self.title = title or ""
@@ -72,6 +89,8 @@ def __init__(
         self._source_mapping_ids = source_mapping_ids or [DEFAULT_MAPPING_NAME]
         self.parsed_logsources = parsed_logsources or {}
         self.timeframe = timeframe
+        self.query_period = query_period
+        self.raw_metainfo_container = raw_metainfo_container
 
     @property
     def author_str(self) -> str: