fix: handle missing/deprecated data sources, separate data component parsing in v18+

Author: elucchesileon

mitreattack/attackToExcel/attackToExcel.py

@@ -2,6 +2,7 @@
 
 import argparse
 import os
+import re
 from typing import Dict, List, Optional
 
 import pandas as pd
@@ -80,10 +81,10 @@ def get_stix_data(
     return mem_store
 
 
-def build_dataframes(src: MemoryStore, domain: str) -> Dict:
+def build_dataframes_pre_v18(src: MemoryStore, domain: str) -> Dict:
     """Build pandas dataframes for each attack type, and return a dictionary lookup for each type to the relevant dataframe.
 
-    :returns:
+    This version of the function is used for ATT&CK versions prior to v18, to account for changes to data components/data sources.
 
     Parameters
     ----------
@@ -114,6 +115,38 @@ def build_dataframes(src: MemoryStore, domain: str) -> Dict:
     return df
 
 
+def build_dataframes(src: MemoryStore, domain: str) -> Dict:
+    """Build pandas dataframes for each attack type, and return a dictionary lookup for each type to the relevant dataframe.
+
+    Parameters
+    ----------
+    src : MemoryStore
+        MemoryStore or other stix2 DataSource object
+    domain : str
+        domain of ATT&CK src corresponds to, e.g "enterprise-attack"
+
+    Returns
+    -------
+    dict
+        A dict lookup of each ATT&CK type to dataframes for the given type to be ingested by write_excel
+    """
+    df = {
+        "techniques": stixToDf.techniquesToDf(src, domain),
+        "tactics": stixToDf.tacticsToDf(src),
+        "software": stixToDf.softwareToDf(src),
+        "groups": stixToDf.groupsToDf(src),
+        "campaigns": stixToDf.campaignsToDf(src),
+        "assets": stixToDf.assetsToDf(src),
+        "mitigations": stixToDf.mitigationsToDf(src),
+        "matrices": stixToDf.matricesToDf(src, domain),
+        "relationships": stixToDf.relationshipsToDf(src),
+        "datacomponents": stixToDf.datacomponentsToDf(src),
+        "analytics": stixToDf.analyticsToDf(src),
+        "detectionstrategies": stixToDf.detectionstrategiesToDf(src),
+    }
+    return df
+
+
 def write_excel(dataframes: Dict, domain: str, version: Optional[str] = None, output_dir: str = ".") -> List:
     """Given a set of dataframes from build_dataframes, write the ATT&CK dataset to output directory.
 
@@ -148,7 +181,7 @@ def write_excel(dataframes: Dict, domain: str, version: Optional[str] = None, ou
         os.makedirs(output_directory)
     # master dataset file
     master_fp = os.path.join(output_directory, f"{domain_version_string}.xlsx")
-    with pd.ExcelWriter(master_fp, engine="xlsxwriter") as master_writer:
+    with pd.ExcelWriter(path=master_fp, engine="xlsxwriter") as master_writer:
         # master list of citations
         citations = pd.DataFrame()
 
@@ -324,6 +357,15 @@ def export(
     logger.info(f"************ Exporting {domain} to Excel ************")
 
     # build dataframes
+    if version:
+        version_pattern = r"v(\d+)\.(\d+)$"
+        match = re.search(version_pattern, version)
+        if match:
+            major_version = int(match.group(1))
+            if major_version < 18:
+                dataframes = build_dataframes_pre_v18(src=mem_store, domain=domain)
+                write_excel(dataframes=dataframes, domain=domain, version=version, output_dir=output_dir)
+
     dataframes = build_dataframes(src=mem_store, domain=domain)
     write_excel(dataframes=dataframes, domain=domain, version=version, output_dir=output_dir)
 

mitreattack/attackToExcel/stixToDf.py

@@ -259,6 +259,8 @@ def tacticsToDf(src):
 def datasourcesToDf(src):
     """Parse STIX Data Sources and their Data components from the given data and return corresponding pandas dataframes.
 
+    This is only used in versions of ATT&CK before v18.
+
     :param src: MemoryStore or other stix2 DataSource object holding the domain data
     :returns: a lookup of labels (descriptors/names) to dataframes
     """
@@ -290,7 +292,7 @@ def datasourcesToDf(src):
             if "x_mitre_aliases" in data_object:
                 row["aliases"] = ", ".join(sorted(data_object["x_mitre_aliases"][1:]))
             if data_object["type"] == "x-mitre-data-component":
-                if "x_mitre_data_source_ref" in data_object:
+                if "x_mitre_data_source_ref" in data_object and data_object["x_mitre_data_source_ref"] in source_lookup:
                     row["name"] = f"{source_lookup[data_object['x_mitre_data_source_ref']]}: {data_object['name']}"
                 row["type"] = "datacomponent"
             else:
@@ -333,6 +335,29 @@ def datasourcesToDf(src):
     return dataframes
 
 
+def datacomponentsToDf(src):
+    """Parse STIX Data components from the given data and return corresponding pandas dataframes.
+
+    :param src: MemoryStore or other stix2 DataSource object holding the domain data
+    :returns: a lookup of labels (descriptors/names) to dataframes
+    """
+    data_components = src.query([Filter("type", "=", "x-mitre-data-component")])
+    data_components = remove_revoked_deprecated(data_components)
+
+    data_component_rows = []
+    for data_component in tqdm(data_components, desc="parsing data components"):
+        data_component_rows.append(parseBaseStix(data_component))
+
+    citations = get_citations(data_components)
+    dataframes = {
+        "datacomponents": pd.DataFrame(data_component_rows).sort_values("name"),
+    }
+    if not citations.empty:
+        dataframes["citations"] = citations.sort_values("reference")
+
+    return dataframes
+
+
 def analyticsToDf(src):
     """Parse STIX Analytics from the given data and return corresponding pandas dataframes.
 
@@ -380,22 +405,9 @@ def detectionstrategiesToDf(src):
         dataframes = {
             "detectionstrategies": pd.DataFrame(detection_strategy_rows).sort_values("name"),
         }
-
-        # add relationships
-        codex = relationshipsToDf(src, relatedType="detectionstrategy")
-        dataframes.update(codex)
-        # add relationship references
-        dataframes["detectionstrategies"]["relationship citations"] = _get_relationship_citations(
-            dataframes["detectionstrategies"], codex
-        )
-        # add/merge citations
         if not citations.empty:
             if "citations" in dataframes:  # append to existing citations from references
-                dataframes["citations"] = pd.concat([dataframes["citations"], citations])
-            else:  # add citations
-                dataframes["citations"] = citations
-
-            dataframes["citations"].sort_values("reference")
+                dataframes["citations"] = citations.sort_values("reference")
 
     else:
         logger.warning("No detection strategies found - nothing to parse")