Merge branch 'master' into feature/assets

mitreattack/attackToExcel/stixToDf.py

@@ -112,7 +112,7 @@ def parseBaseStix(sdo):
         if sdo["external_references"][0]["source_name"] in MITRE_ATTACK_ID_SOURCE_NAMES:
             row["ID"] = sdo["external_references"][0]["external_id"]
             url = sdo["external_references"][0]["url"]
-    if "id" in sdo: # required for workbench collection import
+    if "id" in sdo:  # required for workbench collection import
         row["STIX ID"] = sdo["id"]
     if "name" in sdo:
         row["name"] = sdo["name"]
@@ -124,7 +124,7 @@ def parseBaseStix(sdo):
         row["created"] = format_date(sdo["created"])
     if "modified" in sdo:
         row["last modified"] = format_date(sdo["modified"])
-    if "x_mitre_domains" in sdo: # required for workbench collection import
+    if "x_mitre_domains" in sdo:  # required for workbench collection import
         row["domain"] = ",".join(sdo["x_mitre_domains"])
     if "x_mitre_version" in sdo:
         row["version"] = sdo["x_mitre_version"]
@@ -171,7 +171,7 @@ def techniquesToDf(src, domain):
 
         # sub-technique properties
         if "kill_chain_phases" not in technique:
-            attack_id = technique['external_references'][0]['external_id']
+            attack_id = technique["external_references"][0]["external_id"]
             logger.error(f"Skipping {attack_id} [{technique['id']}] because it does't have kill chain phases")
             continue
         tactic_shortnames = []
@@ -929,7 +929,7 @@ def add_side(label, sdo):
                 row[f"{label} name"] = sdo["name"]
             if "id" in sdo:
                 # "source ref" or "target ref"
-                row[f"{label} ref"] = sdo ["id"]
+                row[f"{label} ref"] = sdo["id"]
             # "source type" or "target type"
             row[f"{label} type"] = stixToAttackTerm[sdo["type"]]
 
@@ -948,10 +948,20 @@ def add_side(label, sdo):
 
     citations = get_citations(relationships)
     relationships = pd.DataFrame(relationship_rows).sort_values(
-        ["mapping type", "source type", "target type", "source name", "target name", "source ref", "target ref", "created", "last modified"]
+        [
+            "mapping type",
+            "source type",
+            "target type",
+            "source name",
+            "target name",
+            "source ref",
+            "target ref",
+            "created",
+            "last modified",
+        ]
     )
 
-    # return all relationships and citations 
+    # return all relationships and citations
     if not relatedType:
         dataframes = {
             "relationships": relationships,

mitreattack/stix20/MitreAttackData.py

@@ -66,6 +66,8 @@ def __init__(self, stix_filepath: str):
         if not isinstance(stix_filepath, str):
             raise TypeError(f"Argument stix_filepath must be of type str, not {type(stix_filepath)}")
 
+        self.stix_filepath = stix_filepath
+
         self.src = MemoryStore()
         self.src.load_from_file(stix_filepath)
 
@@ -321,6 +323,9 @@ def get_objects_by_type(self, stix_type: str, remove_revoked_deprecated=False) -
         if remove_revoked_deprecated:
             objects = self.remove_revoked_deprecated(objects)
 
+        if not objects:
+            return []
+
         # since ATT&CK has custom objects, we need to reconstruct the query results
         return [StixObjectFactory(o) for o in objects]
 
@@ -539,6 +544,10 @@ def get_object_by_stix_id(self, stix_id: str) -> object:
             the STIX Domain Object specified by the STIX ID
         """
         object = self.src.get(stix_id)
+
+        if not object:
+            raise ValueError(f"{stix_id} not found in {self.stix_filepath}")
+
         return StixObjectFactory(object)
 
     def get_object_by_attack_id(self, attack_id: str, stix_type: str) -> object:
@@ -605,6 +614,9 @@ def get_objects_by_name(self, name: str, stix_type: str) -> list:
         filter = [Filter("type", "=", stix_type), Filter("name", "=", name)]
         objects = self.src.query(filter)
 
+        if not objects:
+            return []
+
         # since ATT&CK has custom objects, we need to reconstruct the query results
         return [StixObjectFactory(o) for o in objects]
 

mitreattack/stix20/custom_attack_objects.py

@@ -47,8 +47,10 @@ def StixObjectFactory(data: dict) -> object:
         "x-mitre-data-component": DataComponent,
     }
 
-    if "type" in data and data["type"] in stix_type_to_custom_class:
-        return stix_type_to_custom_class[data["type"]](**data, allow_custom=True)
+    stix_type = data.get("type")
+
+    if data and stix_type in stix_type_to_custom_class:
+        return stix_type_to_custom_class[stix_type](**data, allow_custom=True)
     return data
 
 

pyproject.toml

@@ -8,6 +8,7 @@ profile = "black"
 line-length = 120
 ignore = [
     "E501",  # line-too-long
+    "D100",  # Missing docstring in public module
     "D105",  # undocumented-magic-method
     ]
 # On top of the defaults (E, F), enable pydocstyle (D) and isort (I).

requirements-dev.txt

@@ -1,7 +1,7 @@
 black>=22.3.0
 check-wheel-contents
 colour>=0.1.5
-deepdiff==6.3.0
+deepdiff==6.3.1
 drawsvg==2.1.1
 flake8>=5.0.4
 flake8-docstrings>=1.6.0
@@ -18,7 +18,7 @@ pytest-cov>=3.0.0
 python-dateutil
 Pillow>=7.1.2
 requests>=2.21.0
-rich==13.4.1
+rich==13.5.2
 ruff>=0.0.275
 stix2>=3.0.1
 stix2-elevator>=4.0.1

tests/conftest.py

@@ -2,10 +2,9 @@
 from pathlib import Path
 
 import pytest
-from stix2 import MemoryStore
 from loguru import logger
-
 from resources.testing_data import example_layer_v3_all, example_layer_v43_dict
+from stix2 import MemoryStore
 
 from mitreattack.download_stix import download_domains
 from mitreattack.navlayers import Layer

tests/test_collections.py

@@ -1,4 +1,5 @@
 import json
+import os
 
 from resources.testing_data import collection, index
 
@@ -21,10 +22,14 @@ def test_collection_to_index():
 
 def test_stix_to_collection():
     """Test converting stix bundle file to a collection"""
-    with open("resources/ics-bundle.json", "r") as fio:
+    dir = os.path.dirname(__file__)
+    ics_bundle_collection = os.path.join(dir, "resources", "ics-bundle.json")
+    enterprise_bundle_collection = os.path.join(dir, "resources", "enterprise-bundle.json")
+
+    with open(ics_bundle_collection, "r") as fio:
         v20 = json.load(fio)
 
-    with open("resources/enterprise-bundle.json", "r") as fio:
+    with open(enterprise_bundle_collection, "r") as fio:
         v21 = json.load(fio)
 
     STIXToCollection.stix_to_collection(v20, name="v20_test", version="9.0", description="testing")

tests/test_layers.py

@@ -1,10 +1,10 @@
 import json
+import os
 from pathlib import Path
 
 import pytest
-from stix2 import MemoryStore
-
 from resources import testing_data
+from stix2 import MemoryStore
 
 from mitreattack.navlayers import (
     Layer,
@@ -37,7 +37,8 @@ def test_depreciated_tactics_export(tmp_path: Path, memstore_enterprise_latest:
 def test_colormap_export(tmp_path: Path, memstore_enterprise_latest: MemoryStore):
     """Test exporting a layer with a gradiant of scores"""
     lay = Layer()
-    lay.from_file("resources/heatmap_example.json")
+    dir = os.path.dirname(__file__)
+    lay.from_file(os.path.join(dir, "resources", "heatmap_example.json"))
     xlsx_output = tmp_path / "layer.xlsx"
     svg_output = tmp_path / "layer.svg"