diff --git a/scripts/lint.py b/scripts/lint.py index fb5ab266..efb5a023 100644 --- a/scripts/lint.py +++ b/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) diff --git a/scripts/setup-linter-dependencies.py b/scripts/setup-linter-dependencies.py index 46d5bb57..eb299a8b 100644 --- a/scripts/setup-linter-dependencies.py +++ b/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:]))