linter: add the possibility to enable or disable mbc and att&ck linting

scripts/lint.py

@@ -235,10 +235,20 @@ class InvalidAttckOrMbcTechnique(Lint):
     def __init__(self):
         super(InvalidAttckOrMbcTechnique, self).__init__()
 
-        # This regex match the format defined in the recommandation attribute
+        try:
+            with open("scripts/linter-data.json", "r") as fd:
+                self.data = json.load(fd)
+            self.enabled_frameworks = self.data.keys()
+        except BaseException:
+            # If linter-data.json is not present, or if an error happen
+            # we log an error and lint nothing.
+            logger.warning(
+                "Could not load 'scripts/linter-data.json'. The att&ck and mbc information will not be linted."
+            )
+            self.enabled_frameworks = []
+
+        # This regex match the format defined in the recommendation attribute
         self.reg = re.compile("^([a-zA-Z| ]+)::(.*) \[([A-Za-z0-9.]+)\]$")
-        with open("scripts/linter-data.json", "r") as fd:
-            self.data = json.load(fd)
 
     def _entry_check(self, framework, category, entry, eid):
         if category not in self.data[framework].keys():
@@ -255,7 +265,7 @@ class InvalidAttckOrMbcTechnique(Lint):
         return False
 
     def check_rule(self, ctx: Context, rule: Rule):
-        for framework in ["mbc"]:
+        for framework in self.enabled_frameworks:
             if framework in rule.meta.keys():
                 for r in rule.meta[framework]:
                     m = self.reg.match(r)

scripts/setup-linter-dependencies.py

@@ -1,10 +1,16 @@
+import argparse
 import json
+import logging
 from os.path import dirname
+from sys import argv
 
 import requests
 from stix2 import Filter, MemoryStore, AttackPattern
 
 
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+
+
 class MitreExtractor:
     url = ""
     kill_chain_name = ""
@@ -16,6 +22,7 @@ class MitreExtractor:
         if self.url == "":
             raise ValueError(f"URL not specified in class {self.__class__.__name__}")
 
+        logging.info(f"Downloading STIX data at: {self.url}")
         stix_json = requests.get(self.url).json()
         self._memory_store = MemoryStore(stix_data=stix_json["objects"])
 
@@ -67,6 +74,7 @@ class MitreExtractor:
         return parent_technique
 
     def run(self) -> dict[str, dict[str, str]]:
+        logging.info("Starting extraction...")
         data: dict[str, dict[str, str]] = {}
         for tactic in self._get_tactics():
             data[tactic["name"]] = {}
@@ -98,12 +106,33 @@ class MbcExtractor(MitreExtractor):
         return tactics
 
 
-def main():
-    data = {"att&ck": AttckExtractor().run(), "mbc": MbcExtractor().run()}
+def main(args: argparse.Namespace) -> None:
+    data = {}
+    if args.extractor == "att&ck" or args.extractor == "both":
+        logging.info("Extracting Mitre Att&ck techniques...")
+        data["att&ck"] = AttckExtractor().run()
+    if args.extractor == "mbc" or args.extractor == "both":
+        logging.info("Extracting MBC behaviors...")
+        data["mbc"] = MbcExtractor().run()
 
-    with open(f"{dirname(__file__)}/linter-data.json", "w") as jf:
-        json.dump(data, jf, indent=2)
+    logging.info(f"Writing results to {args.output}")
+    try:
+        with open(args.output, "w") as jf:
+            json.dump(data, jf, indent=2)
+    except BaseException as e:
+        logging.error(f"Exception encountered when writing results: {e}")
 
 
 if __name__ == "__main__":
-    main()
+    parser = argparse.ArgumentParser(description="Setup linter dependencies.")
+    parser.add_argument(
+        "--extractor", type=str, choices=["both", "mbc", "att&ck"], default="both", help="Extractor that will be run"
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        default=f"{dirname(__file__)}/linter-data.json",
+        help="Path to output file (lint.py will be looking for linter-data.json)",
+    )
+    main(parser.parse_args(args=argv[1:]))