Add ability to test local STIX files

mitre-attack · Oct 31, 2023 · c43174d · c43174d
1 parent 3a11239
commit c43174d
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -29,4 +29,5 @@ isort
 pyinstrument
 pytest
 pytest-cov
+pytest-dotenv
 ruff
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,3 +1,4 @@
+import os
 import shutil
 
 import pytest
@@ -11,10 +12,14 @@
 
 from .resources.testing_data import example_layer_v3_all, example_layer_v43_dict
 
+STIX_LOCATION_ENTERPRISE = os.getenv("STIX_LOCATION_ENTERPRISE")
+STIX_LOCATION_MOBILE = os.getenv("STIX_LOCATION_MOBILE")
+STIX_LOCATION_ICS = os.getenv("STIX_LOCATION_ICS")
+
 
 @pytest.fixture(autouse=True, scope="session")
 def attack_stix_dir():
-    logger.debug("Downloading the ATT&CK STIX!!!")
+    logger.debug("Downloading the ATT&CK STIX 2.0!!!")
     download_dir = "attack-releases"
     download_domains(
         domains=["enterprise", "mobile", "ics"],
@@ -30,16 +35,22 @@ def attack_stix_dir():
 
 @pytest.fixture(scope="session")
 def stix_file_enterprise_latest(attack_stix_dir):
+    if STIX_LOCATION_ENTERPRISE:
+        return STIX_LOCATION_ENTERPRISE
     return f"{attack_stix_dir}/v{LATEST_VERSION}/enterprise-attack.json"
 
 
 @pytest.fixture(scope="session")
 def stix_file_mobile_latest(attack_stix_dir):
+    if STIX_LOCATION_MOBILE:
+        return STIX_LOCATION_MOBILE
     return f"{attack_stix_dir}/v{LATEST_VERSION}/mobile-attack.json"
 
 
 @pytest.fixture(scope="session")
 def stix_file_ics_latest(attack_stix_dir):
+    if STIX_LOCATION_ICS:
+        return STIX_LOCATION_ICS
     return f"{attack_stix_dir}/v{LATEST_VERSION}/ics-attack.json"
 
 
@@ -69,21 +80,21 @@ def memstore_ics_latest(stix_file_ics_latest):
 
 @pytest.fixture(scope="session")
 def mitre_attack_data_enterprise(memstore_enterprise_latest):
-    logger.debug("Loading STIX memstore for Enterprise ATT&CK")
+    logger.debug("Creating MitreAttackData() for Enterprise ATT&CK")
     mitre_attack_data = MitreAttackData(src=memstore_enterprise_latest)
     return mitre_attack_data
 
 
 @pytest.fixture(scope="session")
 def mitre_attack_data_mobile(memstore_mobile_latest):
-    logger.debug("Loading STIX memstore for Mobile ATT&CK")
+    logger.debug("Creating MitreAttackData() for Mobile ATT&CK")
     mitre_attack_data = MitreAttackData(src=memstore_mobile_latest)
     return mitre_attack_data
 
 
 @pytest.fixture(scope="session")
 def mitre_attack_data_ics(memstore_ics_latest):
-    logger.debug("Loading STIX memstore for ICS ATT&CK")
+    logger.debug("Creating MitreAttackData() for ICS ATT&CK")
     mitre_attack_data = MitreAttackData(src=memstore_ics_latest)
     return mitre_attack_data