contentctl 4.3.4__py3-none-any.whl → 4.3.5__py3-none-any.whl

contentctl/enrichments/attack_enrichment.py

@@ -1,14 +1,13 @@
 
 from __future__ import annotations
-import csv
-import os
 import sys
 from attackcti import attack_client
 import logging
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 from dataclasses import field
-from typing import Annotated,Any
-from contentctl.objects.mitre_attack_enrichment import MitreAttackEnrichment
+from typing import Any
+from pathlib import Path
+from contentctl.objects.mitre_attack_enrichment import MitreAttackEnrichment, MitreTactics
 from contentctl.objects.config import validate
 from contentctl.objects.annotated_types import MITRE_ATTACK_ID_TYPE
 logging.getLogger('taxii2client').setLevel(logging.CRITICAL)
@@ -21,12 +20,12 @@ class AttackEnrichment(BaseModel):
     @staticmethod
     def getAttackEnrichment(config:validate)->AttackEnrichment:
         enrichment = AttackEnrichment(use_enrichment=config.enrichments)
-        _ = enrichment.get_attack_lookup(str(config.path))
+        _ = enrichment.get_attack_lookup(config.mitre_cti_repo_path, config.enrichments)
         return enrichment
     
     def getEnrichmentByMitreID(self, mitre_id:MITRE_ATTACK_ID_TYPE)->MitreAttackEnrichment:
         if not self.use_enrichment:
-            raise Exception(f"Error, trying to add Mitre Enrichment, but use_enrichment was set to False")
+            raise Exception("Error, trying to add Mitre Enrichment, but use_enrichment was set to False")
         
         enrichment = self.data.get(mitre_id, None)
         if enrichment is not None:
@@ -34,71 +33,69 @@ class AttackEnrichment(BaseModel):
         else:
             raise Exception(f"Error, Unable to find Mitre Enrichment for MitreID {mitre_id}")
         
-    def addMitreIDViaGroupNames(self, technique:dict, tactics:list[str], groupNames:list[str])->None:
+    def addMitreIDViaGroupNames(self, technique:dict[str,Any], tactics:list[str], groupNames:list[str])->None:
         technique_id = technique['technique_id']
         technique_obj = technique['technique']
         tactics.sort()
         
         if technique_id in self.data:
             raise Exception(f"Error, trying to redefine MITRE ID '{technique_id}'")
-        self.data[technique_id] = MitreAttackEnrichment(mitre_attack_id=technique_id, 
-                                                        mitre_attack_technique=technique_obj, 
-                                                        mitre_attack_tactics=tactics, 
-                                                        mitre_attack_groups=groupNames,
-                                                        mitre_attack_group_objects=[]) 
+        self.data[technique_id] = MitreAttackEnrichment.model_validate({'mitre_attack_id':technique_id, 
+                                                                        'mitre_attack_technique':technique_obj, 
+                                                                        'mitre_attack_tactics':tactics, 
+                                                                        'mitre_attack_groups':groupNames,
+                                                                        'mitre_attack_group_objects':[]}) 
 
-    def addMitreIDViaGroupObjects(self, technique:dict, tactics:list[str],  groupObjects:list[dict[str,Any]])->None:
+    def addMitreIDViaGroupObjects(self, technique:dict[str,Any], tactics:list[MitreTactics],  groupDicts:list[dict[str,Any]])->None:
         technique_id = technique['technique_id']
         technique_obj = technique['technique']
         tactics.sort()
         
-        groupNames:list[str] = sorted([group['group'] for group in groupObjects])
+        groupNames:list[str] = sorted([group['group'] for group in groupDicts])
         
         if technique_id in self.data:
             raise Exception(f"Error, trying to redefine MITRE ID '{technique_id}'")
-        self.data[technique_id] = MitreAttackEnrichment(mitre_attack_id=technique_id, 
-                                                        mitre_attack_technique=technique_obj, 
-                                                        mitre_attack_tactics=tactics, 
-                                                        mitre_attack_groups=groupNames,
-                                                        mitre_attack_group_objects=groupObjects)
+        
+        self.data[technique_id] = MitreAttackEnrichment.model_validate({'mitre_attack_id': technique_id, 
+                                                                        'mitre_attack_technique': technique_obj, 
+                                                                        'mitre_attack_tactics': tactics, 
+                                                                        'mitre_attack_groups': groupNames,
+                                                                        'mitre_attack_group_objects': groupDicts})
 
     
-    def get_attack_lookup(self, input_path: str, store_csv: bool = False, force_cached_or_offline: bool = False, skip_enrichment:bool = False) -> dict:
-        if not self.use_enrichment:
-            return {}
-        print("Getting MITRE Attack Enrichment Data. This may take some time...")
-        attack_lookup = dict()
-        file_path = os.path.join(input_path, "app_template", "lookups", "mitre_enrichment.csv")
-
-        if skip_enrichment is True:
-            print("Skipping enrichment")
+    def get_attack_lookup(self, input_path: Path, enrichments:bool = False) -> dict[str,MitreAttackEnrichment]:            
+        attack_lookup:dict[str,MitreAttackEnrichment] = {}
+        if not enrichments:
             return attack_lookup
+        
         try:
-
-            if force_cached_or_offline is True:
-                raise(Exception("WARNING - Using cached MITRE Attack Enrichment.  Attack Enrichment may be out of date. Only use this setting for offline environments and development purposes."))
-            print(f"\r{'Client'.rjust(23)}: [{0:3.0f}%]...", end="", flush=True)
-            lift = attack_client()
-            print(f"\r{'Client'.rjust(23)}: [{100:3.0f}%]...Done!", end="\n", flush=True)
+            print(f"Performing MITRE Enrichment using the repository at {input_path}...",end="", flush=True)
+            # The existence of the input_path is validated during cli argument validation, but it is 
+            # possible that the repo is in the wrong format. If the following directories do not
+            # exist, then attack_client will fall back to resolving via REST API. We do not 
+            # want this as it is slow and error prone, so we will force an exception to 
+            # be generated.
+            enterprise_path = input_path/"enterprise-attack"
+            mobile_path = input_path/"ics-attack"
+            ics_path = input_path/"mobile-attack"
+            if not (enterprise_path.is_dir() and mobile_path.is_dir() and ics_path.is_dir()):
+                raise FileNotFoundError("One or more of the following paths does not exist: "
+                                        f"{[str(enterprise_path),str(mobile_path),str(ics_path)]}. "
+                                        f"Please ensure that the {input_path} directory "
+                                        "has been git cloned correctly.") 
+            lift = attack_client(
+                local_paths= {
+                    "enterprise":str(enterprise_path),
+                    "mobile":str(mobile_path),
+                    "ics":str(ics_path)
+                }
+            )
             
-            print(f"\r{'Techniques'.rjust(23)}: [{0.0:3.0f}%]...", end="", flush=True)
             all_enterprise_techniques = lift.get_enterprise_techniques(stix_format=False)
-            
-            print(f"\r{'Techniques'.rjust(23)}: [{100:3.0f}%]...Done!", end="\n", flush=True)
-            
-            print(f"\r{'Relationships'.rjust(23)}: [{0.0:3.0f}%]...", end="", flush=True)
             enterprise_relationships = lift.get_enterprise_relationships(stix_format=False)
-            print(f"\r{'Relationships'.rjust(23)}: [{100:3.0f}%]...Done!", end="\n", flush=True)
-            
-            print(f"\r{'Groups'.rjust(23)}: [{0:3.0f}%]...", end="", flush=True)
             enterprise_groups = lift.get_enterprise_groups(stix_format=False)
-            print(f"\r{'Groups'.rjust(23)}: [{100:3.0f}%]...Done!", end="\n", flush=True)
-            
             
-            for index, technique in enumerate(all_enterprise_techniques):
-                progress_percent = ((index+1)/len(all_enterprise_techniques)) * 100
-                if (sys.stdout.isatty() and sys.stdin.isatty() and sys.stderr.isatty()):
-                    print(f"\r\t{'MITRE Technique Progress'.rjust(23)}: [{progress_percent:3.0f}%]...", end="", flush=True)
+            for technique in all_enterprise_techniques:
                 apt_groups:list[dict[str,Any]] = []
                 for relationship in enterprise_relationships:
                     if (relationship['target_object'] == technique['id']) and relationship['source_object'].startswith('intrusion-set'):
@@ -115,39 +112,10 @@ class AttackEnrichment(BaseModel):
                 self.addMitreIDViaGroupObjects(technique, tactics, apt_groups)
                 attack_lookup[technique['technique_id']] = {'technique': technique['technique'], 'tactics': tactics, 'groups': apt_groups}
 
-            if store_csv:
-                f = open(file_path, 'w')
-                writer = csv.writer(f)
-                writer.writerow(['mitre_id', 'technique', 'tactics' ,'groups'])
-                for key in attack_lookup.keys():
-                    if len(attack_lookup[key]['groups']) == 0:
-                        groups = 'no'
-                    else:
-                        groups = '|'.join(attack_lookup[key]['groups'])
-                    
-                    writer.writerow([
-                        key,
-                        attack_lookup[key]['technique'],
-                        '|'.join(attack_lookup[key]['tactics']),
-                        groups
-                    ])
-                
-                f.close()
 
+        
         except Exception as err:
-            print(f'\nError: {str(err)}')
-            print('Use local copy app_template/lookups/mitre_enrichment.csv')
-            with open(file_path, mode='r') as inp:
-                reader = csv.reader(inp)
-                attack_lookup = {rows[0]:{'technique': rows[1], 'tactics': rows[2].split('|'), 'groups': rows[3].split('|')} for rows in reader}
-            attack_lookup.pop('mitre_id')
-            for key in attack_lookup.keys():
-                technique_input = {'technique_id': key , 'technique': attack_lookup[key]['technique'] }
-                tactics_input = attack_lookup[key]['tactics']
-                groups_input = attack_lookup[key]['groups']
-                self.addMitreIDViaGroupNames(technique=technique_input, tactics=tactics_input, groups=groups_input)
-            
-                
-
+            raise Exception(f"Error getting MITRE Enrichment: {str(err)}")
+        
         print("Done!")
         return attack_lookup

contentctl/helper/splunk_app.py

@@ -1,20 +1,20 @@
-import os
-import time
 import json
+from typing import Optional, Collection
+from pathlib import Path
 import xml.etree.ElementTree as ET
-from typing import List, Tuple, Optional
 from urllib.parse import urlencode
 
 import requests
 import urllib3
 import xmltodict
 from requests.adapters import HTTPAdapter
-from requests.packages.urllib3.util.retry import Retry
+from urllib3.util.retry import Retry
 
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
 MAX_RETRY = 3
 
+
 class APIEndPoint:
     """
     Class which contains Static Endpoint
@@ -27,6 +27,7 @@ class APIEndPoint:
     SPLUNK_BASE_GET_UID_REDIRECT = "https://apps.splunk.com/apps/id/{app_name_id}"
     SPLUNK_BASE_APP_INFO = "https://splunkbase.splunk.com/api/v1/app/{app_uid}"
 
+
 class RetryConstant:
     """
     Class which contains Retry Constant
@@ -53,11 +54,11 @@ class SplunkApp:
 
     @staticmethod
     def requests_retry_session(
-        retries=RetryConstant.RETRY_COUNT,
-        backoff_factor=1,
-        status_forcelist=(500, 502, 503, 504),
-        session=None,
-    ):
+        retries: int = RetryConstant.RETRY_COUNT,
+        backoff_factor: int = 1,
+        status_forcelist: Collection[int] = (500, 502, 503, 504),
+        session: requests.Session | None = None,
+    ) -> requests.Session:
         session = session or requests.Session()
         retry = Retry(
             total=retries,
@@ -260,4 +261,134 @@ class SplunkApp:
 
         # parse out the version number and fetch the download URL
         self.latest_version = info_url.split("/")[-1]
-        self.latest_version_download_url = self.__fetch_url_latest_version_download(info_url)
+        self.latest_version_download_url = self.__fetch_url_latest_version_download(info_url)
+
+    def __get_splunk_base_session_token(self, username: str, password: str) -> str:
+        """
+        This method will generate Splunk base session token
+
+        :param username: Splunkbase username
+        :type username: str
+        :param password: Splunkbase password
+        :type password: str
+
+        :return: Splunk base session token
+        :rtype: str
+        """
+        # Data payload for fetch splunk base session token
+        payload = urlencode(
+            {
+                "username": username,
+                "password": password,
+            }
+        )
+
+        headers = {
+            "content-type": "application/x-www-form-urlencoded",
+            "cache-control": "no-cache",
+        }
+
+        response = requests.request(
+            "POST",
+            APIEndPoint.SPLUNK_BASE_AUTH_URL,
+            data=payload,
+            headers=headers,
+        )
+
+        token_value = ""
+
+        if response.status_code != 200:
+            msg = (
+                f"Error occurred while executing the rest call for splunk base authentication api,"
+                f"{response.content}"
+            )
+            raise Exception(msg)
+        else:
+            root = ET.fromstring(response.content)
+            token_value = root.find("{http://www.w3.org/2005/Atom}id").text.strip()
+        return token_value
+
+    def download(
+            self,
+            out: Path,
+            username: str,
+            password: str,
+            is_dir: bool = False,
+            overwrite: bool = False
+    ) -> Path:
+        """
+        Given an output path, download the app to the specified location
+
+        :param out: the Path to download the app to
+        :type out: :class:`pathlib.Path`
+        :param username: Splunkbase username
+        :type username: str
+        :param password: Splunkbase password
+        :type password: str
+        :param is_dir: a flag indicating whether out is directory, otherwise a file (default: False)
+        :type is_dir: bool
+        :param overwrite: a flag indicating whether we can overwrite the file at out or not
+        :type overwrite: bool
+
+        :returns path: the Path the download was written to (needed when is_dir is True)
+        :rtype: :class:`pathlib.Path`
+        """
+        # Get the Splunkbase session token
+        token = self.__get_splunk_base_session_token(username, password)
+        response = requests.request(
+            "GET",
+            self.latest_version_download_url,
+            cookies={
+                "sessionid": token
+            }
+        )
+
+        # If the provided output path was a directory we need to try and pull the filename from the
+        # response headers
+        if is_dir:
+            try:
+                # Pull 'Content-Disposition' from the headers
+                content_disposition: str = response.headers['Content-Disposition']
+
+                # Attempt to parse the filename as a KV
+                key, value = content_disposition.strip().split("=")
+                if key != "attachment;filename":
+                    raise ValueError(f"Unexpected key in 'Content-Disposition' KV pair: {key}")
+
+                # Validate the filename is the expected .tgz file
+                filename = Path(value.strip().strip('"'))
+                if filename.suffixes != [".tgz"]:
+                    raise ValueError(f"Filename has unexpected extension(s): {filename.suffixes}")
+                out = Path(out, filename)
+            except KeyError as e:
+                raise KeyError(
+                    f"Unable to properly extract 'Content-Disposition' from response headers: {e}"
+                ) from e
+            except ValueError as e:
+                raise ValueError(
+                    f"Unable to parse filename from 'Content-Disposition' header: {e}"
+                ) from e
+
+        # Ensure the output path is not already occupied
+        if out.exists() and not overwrite:
+            msg = (
+                f"File already exists at {out}, cannot download the app."
+            )
+            raise Exception(msg)
+
+        # Make any parent directories as needed
+        out.parent.mkdir(parents=True, exist_ok=True)
+
+        # Check for HTTP errors
+        if response.status_code != 200:
+            msg = (
+                f"Error occurred while executing the rest call for splunk base authentication api,"
+                f"{response.content}"
+            )
+            raise Exception(msg)
+
+        # Write the app to disk
+        with open(out, "wb") as file:
+            file.write(response.content)
+
+        return out

contentctl/input/director.py

@@ -18,8 +18,7 @@ from contentctl.objects.playbook import Playbook
 from contentctl.objects.deployment import Deployment
 from contentctl.objects.macro import Macro
 from contentctl.objects.lookup import Lookup
-from contentctl.objects.ssa_detection import SSADetection
-from contentctl.objects.atomic import AtomicTest
+from contentctl.objects.atomic import AtomicEnrichment
 from contentctl.objects.security_content_object import SecurityContentObject
 from contentctl.objects.data_source import DataSource
 from contentctl.objects.event_source import EventSource
@@ -33,17 +32,14 @@ from contentctl.objects.enums import SecurityContentType
 from contentctl.objects.enums import DetectionStatus
 from contentctl.helper.utils import Utils
 
-from contentctl.objects.enums import SecurityContentType
 
-from contentctl.objects.enums import DetectionStatus 
-from contentctl.helper.utils import Utils
 
 
 @dataclass
 class DirectorOutputDto:
     # Atomic Tests are first because parsing them 
     # is far quicker than attack_enrichment
-    atomic_tests: None | list[AtomicTest]
+    atomic_enrichment: AtomicEnrichment
     attack_enrichment: AttackEnrichment
     cve_enrichment: CveEnrichment
     detections: list[Detection]
@@ -60,10 +56,7 @@ class DirectorOutputDto:
 
     def addContentToDictMappings(self, content: SecurityContentObject):
         content_name = content.name
-        if isinstance(content, SSADetection):
-            # Since SSA detections may have the same name as ESCU detection,
-            # for this function we prepend 'SSA ' to the name.
-            content_name = f"SSA {content_name}"
+        
         
         if content_name in self.name_to_content_map:
             raise ValueError(
@@ -149,10 +142,10 @@ class Director():
                 os.path.join(self.input_dto.path, str(contentType.name))
             )
             security_content_files = [
-                f for f in files if not f.name.startswith("ssa___")
+                f for f in files 
             ]
         else:
-            raise (Exception(f"Cannot createSecurityContent for unknown product."))
+            raise (Exception(f"Cannot createSecurityContent for unknown product {contentType}."))
 
         validation_errors = []
 

contentctl/objects/abstract_security_content_objects/detection_abstract.py

@@ -83,15 +83,13 @@ class Detection_Abstract(SecurityContentObject):
 
 
         Args:
-            value (Union[str, dict[str,Any]]): The search. It can either be a string (and should be
-                SPL or a dict, in which case it is Sigma-formatted.
+            value (str): The SPL search. It must be an SPL-formatted string.
             info (ValidationInfo): The validation info can contain a number of different objects.
                 Today it only contains the director.
 
         Returns:
-            Union[str, dict[str,Any]]: The search, either in sigma or SPL format.
-        """        
-        
+            str: The search, as an SPL formatted string.
+        """
         
         # Otherwise, the search is SPL.
 
@@ -390,7 +388,11 @@ class Detection_Abstract(SecurityContentObject):
         # NOTE: we ignore the type error around self.status because we are using Pydantic's
         # use_enum_values configuration
         # https://docs.pydantic.dev/latest/api/config/#pydantic.config.ConfigDict.populate_by_name
-        
+
+        # NOTE: The `inspect` action is HIGHLY sensitive to the structure of the metadata line in
+        # the detection stanza in savedsearches.conf. Additive operations (e.g. a new field in the
+        # dict below) should not have any impact, but renaming or removing any of these fields will
+        # break the `inspect` action.
         return {
             'detection_id': str(self.id),
             'deprecated': '1' if self.status == DetectionStatus.deprecated.value else '0',          # type: ignore

contentctl/objects/atomic.py

@@ -1,12 +1,15 @@
 from __future__ import annotations
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from contentctl.objects.config import validate
+
 from contentctl.input.yml_reader import YmlReader
 from pydantic import BaseModel, model_validator, ConfigDict, FilePath, UUID4
+import dataclasses
 from typing import List, Optional, Dict, Union, Self
 import pathlib
-
-
 from enum import StrEnum, auto
-
+import uuid
 
 class SupportedPlatform(StrEnum):        
     windows = auto()
@@ -84,15 +87,6 @@ class AtomicTest(BaseModel):
     dependencies: Optional[List[AtomicDependency]] = None
     dependency_executor_name: Optional[DependencyExecutorType] = None
 
-    @staticmethod
-    def AtomicTestWhenEnrichmentIsDisabled(auto_generated_guid: UUID4) -> AtomicTest:
-        return AtomicTest(name="Placeholder Atomic Test (enrichment disabled)",
-                          auto_generated_guid=auto_generated_guid,
-                          description="This is a placeholder AtomicTest. Because enrichments were not enabled, it has not been validated against the real Atomic Red Team Repo.",
-                          supported_platforms=[],
-                          executor=AtomicExecutor(name="Placeholder Executor (enrichment disabled)", 
-                                                  command="Placeholder command (enrichment disabled)"))
-    
     @staticmethod
     def AtomicTestWhenTestIsMissing(auto_generated_guid: UUID4) -> AtomicTest:
         return AtomicTest(name="Missing Atomic",
@@ -100,31 +94,16 @@ class AtomicTest(BaseModel):
                           description="This is a placeholder AtomicTest. Either the auto_generated_guid is incorrect or it there was an exception while parsing its AtomicFile.",
                           supported_platforms=[],
                           executor=AtomicExecutor(name="Placeholder Executor (failed to find auto_generated_guid)", 
-                                                  command="Placeholder command (failed to find auto_generated_guid)"))
-
-
-    @classmethod
-    def getAtomicByAtomicGuid(cls, guid: UUID4, all_atomics:list[AtomicTest] | None)->AtomicTest:
-        if all_atomics is None:
-            return AtomicTest.AtomicTestWhenEnrichmentIsDisabled(guid)
-        matching_atomics = [atomic for atomic in all_atomics if atomic.auto_generated_guid == guid]
-        if len(matching_atomics) == 0:
-            raise ValueError(f"Unable to find atomic_guid {guid} in {len(all_atomics)} atomic_tests from ART Repo")
-        elif len(matching_atomics) > 1:
-            raise ValueError(f"Found {len(matching_atomics)} matching tests for atomic_guid {guid} in {len(all_atomics)} atomic_tests from ART Repo")
-        
-        return matching_atomics[0]
+                                                  command="Placeholder command (failed to find auto_generated_guid)"))    
     
     @classmethod
-    def parseArtRepo(cls, repo_path:pathlib.Path)->List[AtomicFile]:
-        if not repo_path.is_dir():
-            print(f"WARNING: Atomic Red Team repo does NOT exist at {repo_path.absolute()}. You can check it out with:\n * git clone --single-branch https://github.com/redcanaryco/atomic-red-team. This will ONLY throw a validation error if you reference atomid_guids in your detection(s).")
-            return []
+    def parseArtRepo(cls, repo_path:pathlib.Path)->dict[uuid.UUID, AtomicTest]:
+        test_mapping: dict[uuid.UUID, AtomicTest] = {}
         atomics_path = repo_path/"atomics"
         if not atomics_path.is_dir():
-            print(f"WARNING: Atomic Red Team repo exists at {repo_path.absolute}, but atomics directory does NOT exist at {atomics_path.absolute()}. Was it deleted or renamed? This will ONLY throw a validation error if you reference atomid_guids in your detection(s).")
-            return []
-        
+            raise FileNotFoundError(f"WARNING: Atomic Red Team repo exists at {repo_path}, "
+                                    f"but atomics directory does NOT exist at {atomics_path}. "
+                                    "Was it deleted or renamed?")
 
         atomic_files:List[AtomicFile] = []
         error_messages:List[str] = []
@@ -133,6 +112,7 @@ class AtomicTest(BaseModel):
                 atomic_files.append(cls.constructAtomicFile(obj_path))
             except Exception as e:
                 error_messages.append(f"File [{obj_path}]\n{str(e)}")
+
         if len(error_messages) > 0:
             exceptions_string = '\n\n'.join(error_messages)
             print(f"WARNING: The following [{len(error_messages)}] ERRORS were generated when parsing the Atomic Red Team Repo.\n"
@@ -140,38 +120,28 @@ class AtomicTest(BaseModel):
                    "Note that this is only a warning and contentctl will ignore Atomics contained in these files.\n"
                   f"However, if you have written a detection that references them, 'contentctl build --enrichments' will fail:\n\n{exceptions_string}")
         
-        return atomic_files
+        # Now iterate over all the files, collect all the tests, and return the dict mapping
+        redefined_guids:set[uuid.UUID] = set()
+        for atomic_file in atomic_files:
+            for atomic_test in atomic_file.atomic_tests:
+                if atomic_test.auto_generated_guid in test_mapping:
+                    redefined_guids.add(atomic_test.auto_generated_guid)
+                else:
+                    test_mapping[atomic_test.auto_generated_guid] = atomic_test
+        if len(redefined_guids) > 0:
+            guids_string = '\n\t'.join([str(guid) for guid in redefined_guids])
+            raise Exception(f"The following [{len(redefined_guids)}] Atomic Test"
+                            " auto_generated_guid(s) were defined more than once. "
+                            f"auto_generated_guids MUST be unique:\n\t{guids_string}")
+
+        print(f"Successfully parsed [{len(test_mapping)}] Atomic Red Team Tests!")
+        return test_mapping
     
     @classmethod
     def constructAtomicFile(cls, file_path:pathlib.Path)->AtomicFile:
         yml_dict = YmlReader.load_file(file_path) 
         atomic_file = AtomicFile.model_validate(yml_dict)
         return atomic_file
-    
-    @classmethod
-    def getAtomicTestsFromArtRepo(cls, repo_path:pathlib.Path, enabled:bool=True)->list[AtomicTest] | None:
-        # Get all the atomic files.  Note that if the ART repo is not found, we will not throw an error,
-        # but will not have any atomics. This means that if atomic_guids are referenced during validation,
-        # validation for those detections will fail
-        if not enabled:
-            return None
-        
-        atomic_files = cls.getAtomicFilesFromArtRepo(repo_path)
-            
-        atomic_tests:List[AtomicTest] = []
-        for atomic_file in atomic_files:
-            atomic_tests.extend(atomic_file.atomic_tests)
-        print(f"Found [{len(atomic_tests)}] Atomic Simulations in the Atomic Red Team Repo!")
-        return atomic_tests
-
-    
-    @classmethod
-    def getAtomicFilesFromArtRepo(cls, repo_path:pathlib.Path)->List[AtomicFile]:
-        return cls.parseArtRepo(repo_path)
-
-    
-    
-
 
 
 class AtomicFile(BaseModel):
@@ -182,27 +152,31 @@ class AtomicFile(BaseModel):
     atomic_tests: List[AtomicTest]
 
 
+class AtomicEnrichment(BaseModel):
+    data: dict[uuid.UUID,AtomicTest] = dataclasses.field(default_factory = dict)
+    use_enrichment: bool = False
 
+    @classmethod
+    def getAtomicEnrichment(cls, config:validate)->AtomicEnrichment:
+        enrichment = AtomicEnrichment(use_enrichment=config.enrichments)
+        if config.enrichments:
+            enrichment.data = AtomicTest.parseArtRepo(config.atomic_red_team_repo_path)
+
+        return enrichment
+
+    def getAtomic(self, atomic_guid: uuid.UUID)->AtomicTest:
+        if self.use_enrichment:
+            if atomic_guid in self.data:
+                return self.data[atomic_guid]
+            else:
+                raise Exception(f"Atomic with GUID {atomic_guid} not found.")
+        else:
+            # If enrichment is not enabled, for the sake of compatability
+            # return a stub test with no useful or meaningful information.
+            return AtomicTest.AtomicTestWhenTestIsMissing(atomic_guid)
 
-# ATOMICS_PATH = pathlib.Path("./atomics")
-# atomic_objects = []
-# atomic_simulations = []
-# for obj_path in ATOMICS_PATH.glob("**/T*.yaml"):
-#     try:
-#         with open(obj_path, 'r', encoding="utf-8") as obj_handle:
-#             obj_data = yaml.load(obj_handle, Loader=yaml.CSafeLoader)
-#             atomic_obj = AtomicFile.model_validate(obj_data)
-#     except Exception as e:
-#         print(f"Error parsing object at path {obj_path}: {str(e)}")
-#         print(f"We have successfully parsed {len(atomic_objects)}, however!")
-#         sys.exit(1)
-
-#     print(f"Successfully parsed {obj_path}!")
-#     atomic_objects.append(atomic_obj)
-#     atomic_simulations += atomic_obj.atomic_tests
+    
 
-# print(f"Successfully parsed all {len(atomic_objects)} files!")
-# print(f"Successfully parsed all {len(atomic_simulations)} simulations!")