feat: add version to excel file generation (#234)

Co-authored-by: Jared Ondricek <jondricek@mitre.org>

Author: adpare

examples/generate_excel_files.py

@@ -1,19 +1,99 @@
-from mitreattack.attackToExcel import attackToExcel
+"""Generate ATT&CK Excel exports from local STIX bundles."""
+
+import argparse
+from os import environ
+from pathlib import Path
+
 from stix2 import MemoryStore
-import os
 
-def main():
+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."""
+    output_dir = Path(output_dir)
+    versioned_dir = output_dir / f"{domain}-{version}"
+    domain_dir = output_dir / domain
+
+    if not versioned_dir.is_dir():
+        return
+
+    domain_dir.mkdir(parents=True, exist_ok=True)
+
+    for source_path in versioned_dir.iterdir():
+        if not source_path.is_file():
+            continue
+
+        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
+
+
+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):
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        prog="generate_excel_files.py",
+        description="Generate ATT&CK Excel exports from local STIX bundles.",
+    )
+    parser.add_argument(
+        "-a",
+        "--attack-version",
+        default=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)
+
+
+def main(argv=None):
+    """Generate excel files for specific versions of ATT&CK."""
+    args = parse_args(argv)
+    attack_version = args.attack_version
+
     # List of domains and version to process
     domains = ["enterprise-attack", "mobile-attack", "ics-attack"]
-    output_dir = "output/"
+    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", "attack-releases/stix-2.0/v18.0")
+    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]
@@ -26,9 +106,14 @@ def main():
         # Export to Excel
         attackToExcel.export(
             domain=domain,
+            version=attack_version,
             output_dir=output_dir,
             mem_store=mem_store,
         )
 
+        if attack_version and not VERSIONED_OUTPUT_DIR:
+            move_versioned_exports_to_domain_dir(output_dir=output_dir, domain=domain, version=attack_version)
+
+
 if __name__ == "__main__":
     main()