From 24dbb32521b603a2351f0ad10458e1eaaa26e955 Mon Sep 17 00:00:00 2001 From: adpare Date: Mon, 4 May 2026 17:21:40 -0400 Subject: [PATCH 1/2] feat: add version to excle file generation --- examples/generate_excel_files.py | 74 +++++++++++++++++++++++++++++--- 1 file changed, 69 insertions(+), 5 deletions(-) diff --git a/examples/generate_excel_files.py b/examples/generate_excel_files.py index ab166111..db7a3406 100644 --- a/examples/generate_excel_files.py +++ b/examples/generate_excel_files.py @@ -1,14 +1,73 @@ -from mitreattack.attackToExcel import attackToExcel -from stix2 import MemoryStore +"""Generate ATT&CK Excel exports from local STIX bundles.""" + +import argparse import os -def main(): +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" +# 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) + + if not os.path.isdir(versioned_dir): + return + + os.makedirs(domain_dir, 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): + continue + + if os.path.exists(target_path): + os.remove(target_path) + + os.replace(source_path, target_path) + + os.rmdir(versioned_dir) + + +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=( + "ATT&CK version to export, such as v19.0. " + f"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 = f"{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 = os.environ.get("STIX_BASE_DIR", f"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"), @@ -26,9 +85,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() From f4afb1c07ef2352da7be83429ae1a0975ad1b40c Mon Sep 17 00:00:00 2001 From: Jared Ondricek Date: Tue, 5 May 2026 08:05:42 -0500 Subject: [PATCH 2/2] fix: improve file handling and path management in Excel export script --- examples/generate_excel_files.py | 67 +++++++++++++++++++++----------- 1 file changed, 44 insertions(+), 23 deletions(-) diff --git a/examples/generate_excel_files.py b/examples/generate_excel_files.py index db7a3406..bf995a1f 100644 --- a/examples/generate_excel_files.py +++ b/examples/generate_excel_files.py @@ -1,7 +1,8 @@ """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 @@ -9,6 +10,8 @@ # 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 @@ -16,27 +19,47 @@ 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]