From a5190d272f7a64732b9fa5e283ab04b1e80626bb Mon Sep 17 00:00:00 2001
From: ben-githubs <38414634+ben-githubs@users.noreply.github.com>
Date: Wed, 25 Sep 2024 12:38:45 -0500
Subject: [PATCH] Add Format Checker for MITRE ATT&CK Matrix Report Mappings
 (#1360)

---
 .github/workflows/check-mitre.yml | 36 +++++++++++++++++++
 .scripts/mitre_mapping_check.py   | 57 +++++++++++++++++++++++++++++++
 Makefile                          |  3 ++
 3 files changed, 96 insertions(+)
 create mode 100644 .github/workflows/check-mitre.yml
 create mode 100644 .scripts/mitre_mapping_check.py

diff --git a/.github/workflows/check-mitre.yml b/.github/workflows/check-mitre.yml
new file mode 100644
index 000000000..0da059229
--- /dev/null
+++ b/.github/workflows/check-mitre.yml
@@ -0,0 +1,36 @@
+on:
+    pull_request:
+  
+permissions:
+  contents: read
+
+jobs:
+  lint:
+    name: Check MITRE Mappings
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: step-security/harden-runner@91182cccc01eb5e619899d80e4e971d6181294a7 # v2.10.1
+        with:
+          disable-sudo: true
+          egress-policy: block
+          allowed-endpoints: >
+            files.pythonhosted.org:443
+            github.com:443
+            pypi.org:443
+      - name: Checkout panther-analysis
+        uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 #v4.1.7
+
+      - name: Set python version
+        uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3 #v5.2.0
+        with:
+          python-version: "3.11"
+
+      - name: Install pipenv
+        run: pip install pipenv
+
+      - name: Setup venv
+        run: make venv
+
+      - name: make lint-mitre
+        run: make lint-mitre
diff --git a/.scripts/mitre_mapping_check.py b/.scripts/mitre_mapping_check.py
new file mode 100644
index 000000000..d856ec7a7
--- /dev/null
+++ b/.scripts/mitre_mapping_check.py
@@ -0,0 +1,57 @@
+""" This script checks all the MITRE Mappings in the Reports section of each analysis item to
+ensure they follow the formal TAXXXX:TXXXX. If MITRE mappings aren't in this format, they don't
+display properly in Panther's UI. """
+
+import re
+import sys
+from pathlib import Path
+
+from panther_analysis_tool.analysis_utils import load_analysis_specs
+
+# All MITRE Tags must match this regex pattern
+MITRE_PATTERN = re.compile("^TA\d+\:T\d+(\.\d+)?$")
+
+def main(path: Path) -> bool:
+    # Load Repo
+    analysis_items = load_analysis_specs([path], ignore_files=[])
+
+    items_with_invalid_mappings = [] # Record all items with bad tags
+    for analysis_item in analysis_items:
+        rel_path = analysis_item[0] # Relative path to YAML file
+        spec = analysis_item[2] # YAML spec as a dict
+
+        bad_tags = [] # Record the invalid tags for this analysis item
+        if reports := spec.get("Reports"):
+            if mitre := reports.get("MITRE ATT&CK"):
+                for mapping in mitre:
+                    if not MITRE_PATTERN.match(mapping):
+                        bad_tags.append(mapping)
+        
+        if bad_tags:
+            items_with_invalid_mappings.append({
+                "rel_path": rel_path,
+                "bad_tags": bad_tags
+            })
+    
+    if items_with_invalid_mappings:
+        print("❌ Some items had invalid MITRE mapping formats:")
+        print()
+        for invalid_item in items_with_invalid_mappings:
+            print(invalid_item.get("rel_path", "<UNKNOWN PATH>"))
+            for bad_tag in invalid_item.get("bad_tags", []):
+                print("\t" + bad_tag)
+            print()
+
+        print(("To ensure that your MITRE mappings are correctly displayed in the Panther "
+               "console, make sure your MITRE mappings are formatted like 'TA0000:T0000'."))
+    else:
+        print("✅ No invalid MITRE mappings found! You're in the clear! 👍")
+    
+    return bool(items_with_invalid_mappings)
+
+if __name__ == "__main__":
+    path = Path.cwd() # Default to current directory
+    if len(sys.argv) > 1:
+        path = Path(sys.argv[1])
+    if main(path):
+        exit(1) # Exit with error if issues were found
\ No newline at end of file
diff --git a/Makefile b/Makefile
index 8cf5815a1..08377a946 100644
--- a/Makefile
+++ b/Makefile
@@ -43,6 +43,9 @@ lint-fmt:
 	@echo Checking python file formatting with the black code style checker
 	pipenv run black --line-length=100 --check $(dirs)
 
+lint-mitre:
+	pipenv run python3 ./.scripts/mitre_mapping_check.py
+
 venv:
 	pipenv sync --dev