feat: enhance export_release function to handle explicit local STIX base directory and improve error handling

Author: jondricek

mitreattack/attackToExcel/attackToExcel.py

@@ -143,13 +143,19 @@ def export_release(
         expected_stix_versions = ", ".join(VALID_STIX_VERSIONS)
         raise ValueError(f"Invalid STIX version: {stix_version}. Expected one of: {expected_stix_versions}")
 
-    attack_version = normalize_attack_version(version or release_info.LATEST_VERSION)
+    has_explicit_local_stix_base_dir = stix_base_dir is not None or os.environ.get("STIX_BASE_DIR") is not None
+    attack_version = normalize_attack_version(version) if version else None
+    release_version = attack_version or normalize_attack_version(release_info.LATEST_VERSION)
     release_domains = _validate_release_domains(domains)
     local_release_dir = Path(
-        stix_base_dir or os.environ.get("STIX_BASE_DIR") or _default_release_dir(attack_version, stix_version)
+        stix_base_dir or os.environ.get("STIX_BASE_DIR") or _default_release_dir(release_version, stix_version)
     )
     local_release_dir = local_release_dir.resolve()
-    release_output_dir = Path(output_dir) / attack_version
+    release_output_dir = (
+        Path(output_dir)
+        if has_explicit_local_stix_base_dir and attack_version is None
+        else Path(output_dir) / release_version
+    )
 
     local_stix_files = {domain: _release_stix_file(local_release_dir, domain) for domain in release_domains}
     missing_domains = [domain for domain, stix_file in local_stix_files.items() if not stix_file.is_file()]
@@ -163,6 +169,13 @@ def export_release(
         )
         return
 
+    if attack_version is None:
+        missing_domains_text = ", ".join(missing_domains)
+        raise FileNotFoundError(
+            f"Missing local STIX file(s) for domain(s): {missing_domains_text}. "
+            "Pass --version to download missing ATT&CK release bundles."
+        )
+
     with tempfile.TemporaryDirectory() as temporary_directory:
         temporary_release_dir = _download_missing_release_domains(
             missing_domains=missing_domains,
@@ -186,7 +199,7 @@ def export_release(
 
 def _export_release_domains(
     *,
-    version: str,
+    version: Optional[str],
     output_dir: Path,
     stix_files: Dict[str, Path],
     versioned_output_dir: bool,

tests/test_to_excel.py

@@ -167,6 +167,33 @@ def fake_export(**kwargs):
     assert calls["exports"][0]["output_dir"] == str(tmp_path / "output" / "v19.0")
 
 
+def test_export_release_with_explicit_local_stix_base_dir_without_version_is_unversioned(tmp_path: Path, monkeypatch):
+    """Explicit local STIX bundle directories should not be labelled as ATT&CK releases unless a version is given."""
+    stix_base_dir = tmp_path / "attack-releases" / "stix-2.0" / "attackwb"
+    stix_base_dir.mkdir(parents=True)
+    for domain in ["enterprise-attack", "mobile-attack"]:
+        (stix_base_dir / f"{domain}.json").write_text("{}", encoding="utf-8")
+
+    calls = {}
+
+    def fake_export(**kwargs):
+        calls.setdefault("exports", []).append(kwargs)
+
+    monkeypatch.setattr(attackToExcel, "export", fake_export)
+
+    attackToExcel.export_release(
+        stix_base_dir=str(stix_base_dir),
+        output_dir=str(tmp_path / "output" / "attackwb"),
+        domains=["enterprise-attack", "mobile-attack"],
+    )
+
+    assert [call["domain"] for call in calls["exports"]] == ["enterprise-attack", "mobile-attack"]
+    assert calls["exports"][0]["version"] is None
+    assert calls["exports"][0]["output_dir"] == str(tmp_path / "output" / "attackwb")
+    assert calls["exports"][1]["version"] is None
+    assert calls["exports"][1]["output_dir"] == str(tmp_path / "output" / "attackwb")
+
+
 def test_export_release_downloads_only_missing_domains_to_temporary_directory(tmp_path: Path, monkeypatch):
     """Missing release STIX files should be downloaded per missing domain into a temporary tree."""
     stix_base_dir = tmp_path / "attack-releases" / "stix-2.0" / "v19.0"