More cleanup and error handling

to how enrichment is done.
It now works as expected
and the proper errors are
thrown for AtomicEnrichment
depending on the
--enrichments cli setting.

Author: pyth0n1c

contentctl/actions/validate.py

@@ -5,20 +5,16 @@
 from contentctl.objects.config import validate
 from contentctl.enrichments.attack_enrichment import AttackEnrichment
 from contentctl.enrichments.cve_enrichment import CveEnrichment
-from contentctl.objects.atomic import AtomicTest
+from contentctl.objects.atomic import AtomicEnrichment
 from contentctl.helper.utils import Utils
 from contentctl.objects.data_source import DataSource
 from contentctl.helper.splunk_app import SplunkApp
 
 
 class Validate:
-    def execute(self, input_dto: validate) -> DirectorOutputDto:
-
+    def execute(self, input_dto: validate) -> DirectorOutputDto: 
         director_output_dto = DirectorOutputDto(
-            AtomicTest.getAtomicTestsFromArtRepo(
-                repo_path=input_dto.atomic_red_team_repo_path,
-                enabled=input_dto.enrichments,
-            ),
+            AtomicEnrichment.getAtomicEnrichment(input_dto),
             AttackEnrichment.getAttackEnrichment(input_dto),
             CveEnrichment.getCveEnrichment(input_dto),
             [],

contentctl/enrichments/attack_enrichment.py

@@ -25,7 +25,7 @@ def getAttackEnrichment(config:validate)->AttackEnrichment:
 
     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:

contentctl/input/director.py

@@ -18,7 +18,7 @@
 from contentctl.objects.deployment import Deployment
 from contentctl.objects.macro import Macro
 from contentctl.objects.lookup import Lookup
-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
@@ -39,7 +39,7 @@
 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]
@@ -145,7 +145,7 @@ def createSecurityContent(self, contentType: SecurityContentType) -> None:
                 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/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,47 +87,23 @@ 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",
                           auto_generated_guid=auto_generated_guid,
                           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,45 +112,36 @@ def parseArtRepo(cls, repo_path:pathlib.Path)->List[AtomicFile]:
                 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"
                    "Please raise an issue so that they can be fixed at https://github.com/redcanaryco/atomic-red-team/issues.\n"
                    "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!")
 
 
 

contentctl/objects/detection_tags.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import uuid
-from typing import TYPE_CHECKING, List, Optional, Annotated, Union
+from typing import TYPE_CHECKING, List, Optional, Union
 from pydantic import (
     BaseModel,
     Field,
@@ -32,7 +32,7 @@
     RiskLevel,
     SecurityContentProductName
 )
-from contentctl.objects.atomic import AtomicTest
+from contentctl.objects.atomic import AtomicEnrichment, AtomicTest
 from contentctl.objects.annotated_types import MITRE_ATTACK_ID_TYPE, CVE_TYPE
 
 # TODO (#266): disable the use_enum_values configuration
@@ -240,7 +240,7 @@ def mapAtomicGuidsToAtomicTests(cls, v: List[UUID4], info: ValidationInfo) -> Li
         if output_dto is None:
             raise ValueError("Context not provided to detection.detection_tags.atomic_guid validator")
 
-        all_tests: None | List[AtomicTest] = output_dto.atomic_tests
+        atomic_enrichment: AtomicEnrichment = output_dto.atomic_enrichment
 
         matched_tests: List[AtomicTest] = []
         missing_tests: List[UUID4] = []
@@ -254,7 +254,7 @@ def mapAtomicGuidsToAtomicTests(cls, v: List[UUID4], info: ValidationInfo) -> Li
                 badly_formatted_guids.append(str(atomic_guid_str))
                 continue
             try:
-                matched_tests.append(AtomicTest.getAtomicByAtomicGuid(atomic_guid, all_tests))
+                matched_tests.append(atomic_enrichment.getAtomic(atomic_guid))
             except Exception:
                 missing_tests.append(atomic_guid)
 
@@ -265,7 +265,7 @@ def mapAtomicGuidsToAtomicTests(cls, v: List[UUID4], info: ValidationInfo) -> Li
                 f"\n\tPlease review the output above for potential exception(s) when parsing the "
                 "Atomic Red Team Repo."
                 "\n\tVerify that these auto_generated_guid exist and try updating/pulling the "
-                f"repo again.: {[str(guid) for guid in missing_tests]}"
+                f"repo again: {[str(guid) for guid in missing_tests]}"
             )
         else:
             missing_tests_string = ""
@@ -278,6 +278,6 @@ def mapAtomicGuidsToAtomicTests(cls, v: List[UUID4], info: ValidationInfo) -> Li
             raise ValueError(f"{bad_guids_string}{missing_tests_string}")
 
         elif len(missing_tests) > 0:
-            print(missing_tests_string)
+            raise ValueError(missing_tests_string)
 
         return matched_tests + [AtomicTest.AtomicTestWhenTestIsMissing(test) for test in missing_tests]