fix: improve file handling and path management in Excel export script

Author: jondricek

examples/generate_excel_files.py

@@ -1,42 +1,65 @@
 """Generate ATT&CK Excel exports from local STIX bundles."""
 
 import argparse
-import os
+from os import environ
+from pathlib import Path
 
 from stix2 import MemoryStore
 
 from mitreattack.attackToExcel import attackToExcel
 
 # Pass attack version via the command line or update the variable below
 DEFAULT_ATTACK_VERSION = "v19.0"
+# Parent directory where ATT&CK version export folders are written.
+OUTPUT_DIR = Path("output")
 # Set to true if you want the parent subfolder of the excel files to have a version.
 # Example - If you want the folder to be named enterprise-attack-v19.0 instead of enterprise-attack, set to True
 VERSIONED_OUTPUT_DIR = False
 
 
 def move_versioned_exports_to_domain_dir(output_dir, domain, version):
     """Move versioned Excel exports into the unversioned domain folder."""
-    versioned_dir = os.path.join(output_dir, f"{domain}-{version}")
-    domain_dir = os.path.join(output_dir, domain)
+    output_dir = Path(output_dir)
+    versioned_dir = output_dir / f"{domain}-{version}"
+    domain_dir = output_dir / domain
 
-    if not os.path.isdir(versioned_dir):
+    if not versioned_dir.is_dir():
         return
 
-    os.makedirs(domain_dir, exist_ok=True)
+    domain_dir.mkdir(parents=True, exist_ok=True)
 
-    for filename in os.listdir(versioned_dir):
-        source_path = os.path.join(versioned_dir, filename)
-        target_path = os.path.join(domain_dir, filename)
-
-        if not os.path.isfile(source_path):
+    for source_path in versioned_dir.iterdir():
+        if not source_path.is_file():
             continue
 
-        if os.path.exists(target_path):
-            os.remove(target_path)
+        target_path = domain_dir / source_path.name
+        if target_path.exists():
+            target_path.unlink()
+
+        source_path.replace(target_path)
+
+    versioned_dir.rmdir()
+
+
+def format_missing_stix_bundle_error(stix_file, attack_version):
+    """Format a concise missing STIX bundle error."""
+    message = (
+        f"STIX bundle not found: {stix_file}\n"
+        "Download the STIX bundles before running this script, or set STIX_BASE_DIR to the directory containing "
+        "enterprise-attack.json, mobile-attack.json, and ics-attack.json."
+    )
+
+    if attack_version and not attack_version.startswith("v"):
+        message = f"{message}\nDid you mean -a v{attack_version}?"
+
+    return message
 
-        os.replace(source_path, target_path)
 
-    os.rmdir(versioned_dir)
+def validate_stix_files(stix_files, attack_version):
+    """Exit with a clean error if any expected STIX bundle is missing."""
+    for stix_file in stix_files.values():
+        if not stix_file.is_file():
+            raise SystemExit(format_missing_stix_bundle_error(stix_file, attack_version))
 
 
 def parse_args(argv=None):
@@ -49,10 +72,7 @@ def parse_args(argv=None):
         "-a",
         "--attack-version",
         default=DEFAULT_ATTACK_VERSION,
-        help=(
-            "ATT&CK version to export, such as v19.0. "
-            f"Defaults to {DEFAULT_ATTACK_VERSION}."
-        ),
+        help=(f"ATT&CK version to export, such as v19.0. Defaults to {DEFAULT_ATTACK_VERSION}."),
     )
     return parser.parse_args(args=argv)
 
@@ -64,15 +84,16 @@ def main(argv=None):
 
     # List of domains and version to process
     domains = ["enterprise-attack", "mobile-attack", "ics-attack"]
-    output_dir = f"{attack_version}/"
+    output_dir = OUTPUT_DIR / attack_version
 
     # Path to the STIX bundles for each domain (assumes STIX files are downloaded)
-    stix_base_dir = os.environ.get("STIX_BASE_DIR", f"attack-releases/stix-2.0/{attack_version}")
+    stix_base_dir = Path(environ.get("STIX_BASE_DIR", Path("attack-releases") / "stix-2.0" / attack_version))
     stix_files = {
-        "enterprise-attack": os.path.join(stix_base_dir, "enterprise-attack.json"),
-        "mobile-attack": os.path.join(stix_base_dir, "mobile-attack.json"),
-        "ics-attack": os.path.join(stix_base_dir, "ics-attack.json"),
+        "enterprise-attack": stix_base_dir / "enterprise-attack.json",
+        "mobile-attack": stix_base_dir / "mobile-attack.json",
+        "ics-attack": stix_base_dir / "ics-attack.json",
     }
+    validate_stix_files(stix_files, attack_version)
 
     for domain in domains:
         stix_file = stix_files[domain]