add support for the entire mitre

group object that is returned.

Author: pyth0n1c

contentctl/enrichments/attack_enrichment.py

@@ -7,7 +7,7 @@
 import logging
 from pydantic import BaseModel, Field
 from dataclasses import field
-from typing import Annotated
+from typing import Annotated,Any
 from contentctl.objects.mitre_attack_enrichment import MitreAttackEnrichment
 from contentctl.objects.config import validate
 logging.getLogger('taxii2client').setLevel(logging.CRITICAL)
@@ -34,20 +34,21 @@ def getEnrichmentByMitreID(self, mitre_id:Annotated[str, Field(pattern=r"^T\d{4}
             raise Exception(f"Error, Unable to find Mitre Enrichment for MitreID {mitre_id}")
 
 
-    def addMitreID(self, technique:dict, tactics:list[str], groups:list[str])->None:
+    def addMitreID(self, technique:dict, tactics:list[str], groups:list[dict[str,Any]])->None:
 
         technique_id = technique['technique_id']
         technique_obj = technique['technique']
         tactics.sort()
-        groups.sort()
+        group_names_only:list[str] = sorted([group['group'] for group in groups])
+        
 
         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=groups)
+                                                        mitre_attack_groups=group_names_only,
+                                                        mitre_attack_group_objects=groups)
 
 
     def get_attack_lookup(self, input_path: str, store_csv: bool = False, force_cached_or_offline: bool = False, skip_enrichment:bool = False) -> dict:
@@ -86,12 +87,13 @@ def get_attack_lookup(self, input_path: str, store_csv: bool = False, force_cach
                 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)
-                apt_groups = []
+                apt_groups:list[dict[str,Any]] = []
                 for relationship in enterprise_relationships:
                     if (relationship['target_object'] == technique['id']) and relationship['source_object'].startswith('intrusion-set'):
                         for group in enterprise_groups:
                             if relationship['source_object'] == group['id']:
-                                apt_groups.append(group['group'])
+                                apt_groups.append(group)
+                                #apt_groups.append(group['group'])
 
                 tactics = []
                 if ('tactic' in technique):

contentctl/objects/mitre_attack_enrichment.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
-from pydantic import BaseModel, Field, ConfigDict
+from pydantic import BaseModel, Field, ConfigDict, HttpUrl
 from typing import List, Annotated
 from enum import StrEnum
-
+import datetime
 
 class MitreTactics(StrEnum):
     RECONNAISSANCE = "Reconnaissance"
@@ -20,13 +20,55 @@ class MitreTactics(StrEnum):
     EXFILTRATION = "Exfiltration"
     IMPACT = "Impact"
 
+from enum import StrEnum
+class AttackGroupMatrix(StrEnum):
+    mitre_attack = "mitre-attack"
+
+
+class AttackGroupType(StrEnum):
+    intrusion_set = "intrusion-set"
+
+class MitreDomain(StrEnum):
+    intrusion_set = "enterprise-attack"
+    mobile_attack = "mobile-attack"
+    ics_attack = "ics-attack"
+
+class MitreExternalReference(BaseModel):
+    model_config = ConfigDict(extra='forbid')
+    source_name: str
+    external_id: None | str = None 
+    url: None | HttpUrl = None
+    description: None | str = None
+
+
+class MitreAttackGroup(BaseModel):
+    model_config = ConfigDict(extra='forbid')
+    created: datetime.datetime
+    created_by_ref: str
+    external_references: list[MitreExternalReference]
+    group: str
+    group_aliases: list[str]
+    group_description: str
+    id: str
+    matrix: AttackGroupMatrix
+    modified: datetime.datetime
+    object_marking_refs: list[str]
+    type: AttackGroupType
+    url: HttpUrl
+    x_mitre_attack_spec_version: None | str = None
+    x_mitre_deprecated: None | bool = None
+    x_mitre_domains: list[MitreDomain]
+    x_mitre_modified_by_ref: str
+    x_mitre_version: str
+    contributors: list[str] = []
+
 
 class MitreAttackEnrichment(BaseModel):
     ConfigDict(use_enum_values=True)
     mitre_attack_id: Annotated[str, Field(pattern=r"^T\d{4}(.\d{3})?$")] = Field(...)
     mitre_attack_technique: str = Field(...)
     mitre_attack_tactics: List[MitreTactics] = Field(...)
     mitre_attack_groups: List[str] = Field(...)
-
+    mitre_attack_group_objects: list[MitreAttackGroup] = Field(...)
     def __hash__(self) -> int:
         return id(self)