updating to latest mitre attack and atomic red team

.github/workflows/book.yml

@@ -4,34 +4,33 @@ name: deploy-book
 on:
   push:
     branches:
-    - master
+      - master
 
 # This job installs dependencies, build the book, and pushes it to `gh-pages`
 jobs:
   deploy-book:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+      - uses: actions/checkout@v2
 
-    # Install dependencies
-    - name: Set up Python 3.10
-      uses: actions/setup-python@v2
-      with:
-        python-version: '3.10'
+      # Install dependencies
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.10"
 
-    - name: Install dependencies
-      run: |
-        pip install -r requirements.txt
+      - name: Install dependencies
+        run: |
+          pip install -r requirements.txt
 
+      # Build the book
+      - name: Build the book
+        run: |
+          jupyter-book build playbook
 
-    # Build the book
-    - name: Build the book
-      run: |
-        jupyter-book build playbook
-
-    # Push the book's HTML to github-pages
-    - name: GitHub Pages action
-      uses: peaceiris/[email protected]
-      with:
-        github_token: ${{ secrets.GITHUB_TOKEN }}
-        publish_dir: ./playbook/_build/html
+      # Push the book's HTML to github-pages
+      - name: GitHub Pages action
+        uses: peaceiris/[email protected]
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          publish_dir: ./playbook/_build/html

playbook-generator/generator.py

@@ -1,27 +1,28 @@
 import json
 import os
+from pathlib import Path
 
 import nbformat.v4 as nbf
 import yaml
-from attackcti import attack_client
+from mitreattack.stix20 import MitreAttackData
 from models import AggregatedAtomic, AttackTactic, folder_paths
 
 
 def generate_notebook(file_name, cells):
-    file = f'{file_name}.ipynb'
+    file = f"{file_name}.ipynb"
     metadata = {
         "kernelspec": {
             "display_name": ".NET (PowerShell)",
-            "language": "PowerShell",
-            "name": ".net-powershell"
+            "language": "pwsh",
+            "name": ".net-powershell",
         },
         "language_info": {
             "file_extension": ".ps1",
             "mimetype": "text/x-powershell",
-            "name": "PowerShell",
+            "name": "pwsh",
             "pygments_lexer": "powershell",
-            "version": "7.0"
-        }
+            "version": "7.0",
+        },
     }
     nb = nbf.new_notebook(metadata=metadata)
     nb["cells"] = cells
@@ -31,64 +32,102 @@ def generate_notebook(file_name, cells):
 
 class PlaybookGenerator:
     def __init__(self):
-        self.attack_client = attack_client()
-        # We use only Enterprise ATT&CK matrix.
-        self.attack_client.COMPOSITE_DS = self.attack_client.TC_ENTERPRISE_SOURCE
+        self.attack_client = MitreAttackData(
+            f"{Path(folder_paths.ATOMICS_FOLDER).parent.absolute()}/atomic_red_team/enterprise-attack.json"
+        )
 
     def generate_toc(self):
-        # TODO: Generate TOC automatically
-        pass
+        arr = [
+            "initial-access",
+            "execution",
+            "persistence",
+            "privilege-escalation",
+            "defense-evasion",
+            "credential-access",
+            "discovery",
+            "lateral-movement",
+            "collection",
+            "command-and-control",
+            "exfiltration",
+            "impact",
+        ]
+        sections = []
+        for j in arr:
+            path = os.path.join(os.getcwd(), "playbook", "tactics", j)
+            sections.append(
+                {
+                    "file": f"tactics/{j}",
+                    "sections": [
+                        {"file": f"tactics/{j}/{i}"} for i in sorted(os.listdir(path))
+                    ],
+                }
+            )
+
+        with open(os.path.join(os.getcwd(), "playbook", "_toc.yaml"), "w") as f:
+            toc = {
+                "format": "jb-article",
+                "root": "intro",
+                "sections": [{"file": "tactics", "sections": sections}],
+            }
+            f.write(yaml.dump(toc, indent=2, sort_keys=False))
 
     def get_tactics(self):
-        if not os.path.exists(f'{os.getcwd()}/playbook/tactics'):
-            os.mkdir(f'{os.getcwd()}/playbook/tactics')
+        if not os.path.exists(f"{os.getcwd()}/playbook/tactics"):
+            os.mkdir(f"{os.getcwd()}/playbook/tactics")
 
-        stix_tactics = self.attack_client.get_enterprise_tactics()
-        stix_tactics.sort(
-            key=lambda x: x["external_references"][0]["external_id"])
+        stix_tactics = self.attack_client.get_tactics()
+        stix_tactics.sort(key=lambda x: x["external_references"][0]["external_id"])
         # TODO: Swap Exfiltration and C2C
         markdown = [
             "| ID      | Name | Description |",
-            "| -------- | --------- | --------- |"
+            "| -------- | --------- | --------- |",
         ]
-        stix_tactics = filter(lambda x: x["external_references"][0]["external_id"] not in ["TA0042", "TA0043"],
-                              stix_tactics)
+        stix_tactics = filter(
+            lambda x: x["external_references"][0]["external_id"]
+            not in ["TA0042", "TA0043"],
+            stix_tactics,
+        )
         for i in stix_tactics:
             techniques = self.get_techniques_by_tactic(i["x_mitre_shortname"])
             tactic = AttackTactic(stix=i, techniques=techniques)
             generate_notebook(
-                file_name=f'{os.getcwd()}/playbook/tactics/{tactic.stix.short_name}',
-                cells=tactic.__repr__())
+                file_name=f"{os.getcwd()}/playbook/tactics/{tactic.stix.short_name}",
+                cells=tactic.__repr__(),
+            )
             desc = tactic.stix.description.split("\n")[0]
-            markdown.append(f'| {tactic.stix.id} | {tactic.stix.name} | {desc}|')
+            markdown.append(f"| {tactic.stix.id} | {tactic.stix.name} | {desc}|")
         cells = [nbf.new_markdown_cell("\n".join(markdown))]
-        generate_notebook(
-            file_name=f'{os.getcwd()}/playbook/tactics',
-            cells=cells)
+        generate_notebook(file_name=f"{os.getcwd()}/playbook/tactics", cells=cells)
 
     def get_techniques_by_tactic(self, tactic):
-        folder_path = f'{os.getcwd()}/playbook/tactics/{tactic}'
+        folder_path = f"{os.getcwd()}/playbook/tactics/{tactic}"
         if not os.path.exists(folder_path):
             os.makedirs(folder_path)
-        techniques = self.attack_client.get_techniques_by_tactic(tactic)
+        techniques = self.attack_client.get_techniques_by_tactic(
+            tactic, domain="enterprise-attack", remove_revoked_deprecated=True
+        )
         json_techniques = []
         for stix in techniques:
             technique_id = stix["external_references"][0]["external_id"]
-            file_path = f'{folder_paths.ATOMICS_FOLDER}/{technique_id}/{technique_id}.yaml'
+            file_path = (
+                f"{folder_paths.ATOMICS_FOLDER}/{technique_id}/{technique_id}.yaml"
+            )
             if os.path.exists(file_path):
                 with open(file_path, "r") as f:
                     atomic = yaml.load(f.read(), Loader=yaml.SafeLoader)
                     playbook_technique = AggregatedAtomic(stix=stix, atomic=atomic)
             else:
                 playbook_technique = AggregatedAtomic(stix=stix)
             generate_notebook(
-                file_name=f'{folder_path}/{technique_id}',
-                cells=playbook_technique.__repr__())
+                file_name=f"{folder_path}/{technique_id}",
+                cells=playbook_technique.__repr__(),
+            )
             json_techniques.append(playbook_technique)
         return json_techniques
 
     def start(self):
         self.get_tactics()
+        self.generate_toc()
 
 
 if __name__ == "__main__":

playbook-generator/models.py

@@ -6,12 +6,11 @@
 from shield.api import shield
 
 
-def replace_command_with_input_args(
-        command: str, input_arguments: dict):
-    """ Replace the input arguments in the commands."""
+def replace_command_with_input_args(command: str, input_arguments: dict):
+    """Replace the input arguments in the commands."""
     if command:
         for k, v in input_arguments.items():
-            command = command.replace(f'#{{{k}}}', str(v["default"]))
+            command = command.replace(f"#{{{k}}}", str(v["default"]))
         return command
     else:
         return None
@@ -24,6 +23,7 @@ class FolderPaths(BaseSettings):
 try:
     folder_paths = FolderPaths()
 except ValidationError:
+    # local development
     folder_paths = FolderPaths(ATOMICS_FOLDER="~/Pro/atomic-red-team/atomics")
     folder_paths.ATOMICS_FOLDER = os.path.expanduser(folder_paths.ATOMICS_FOLDER)
 
@@ -68,22 +68,32 @@ class AtomicTest(BaseModel):
     def __init__(self, **data):
         super().__init__(**data)
         if args := self.input_arguments:
-            self.executor.command = replace_command_with_input_args(self.executor.command, args)
-            self.executor.cleanup = replace_command_with_input_args(self.executor.cleanup, args)
+            self.executor.command = replace_command_with_input_args(
+                self.executor.command, args
+            )
+            self.executor.cleanup = replace_command_with_input_args(
+                self.executor.cleanup, args
+            )
             if deps := self.dependencies:
                 for i in deps:
                     i.get_prereq_command = replace_command_with_input_args(
-                        i.get_prereq_command, args)
+                        i.get_prereq_command, args
+                    )
                     i.prereq_command = replace_command_with_input_args(
-                        i.prereq_command, args)
+                        i.prereq_command, args
+                    )
 
     def __repr__(self):
         cells = []
         markdown = []
-        markdown += [f"### Atomic Test #{self.test_number} - {self.name}", self.description]
+        markdown += [
+            f"### Atomic Test #{self.test_number} - {self.name}",
+            self.description,
+        ]
         if self.platforms:
             markdown.append(
-                "**Supported Platforms:** {0}".format(", ".join(self.platforms)))
+                "**Supported Platforms:** {0}".format(", ".join(self.platforms))
+            )
         if self.executor.elevation_required:
             markdown.append("\nElevation Required (e.g. root or admin)")
         if self.executor.name == "manual":
@@ -93,7 +103,8 @@ def __repr__(self):
         else:
             if self.dependencies:
                 markdown.append(
-                    f"#### Dependencies:  Run with `{self.dependency_executor_name or self.executor.name}`!")
+                    f"#### Dependencies:  Run with `{self.dependency_executor_name or self.executor.name}`!"
+                )
                 for dep in self.dependencies:
                     executor_name = self.executor.name
                     if executor_name == "command_prompt":
@@ -103,26 +114,37 @@ def __repr__(self):
                         "##### Check Prereq Commands:",
                         f"```{executor_name}\n{dep.prereq_command}\n```",
                         "##### Get Prereq Commands:",
-                        f"```{executor_name}\n{dep.get_prereq_command}\n```"
+                        f"```{executor_name}\n{dep.get_prereq_command}\n```",
                     ]
                 cells.append(nbf.new_markdown_cell("\n".join(markdown)))
-                cells.append(nbf.new_code_cell(
-                    f'Invoke-AtomicTest {self.technique_id} -TestNumbers {self.test_number} -GetPreReqs'))
+                cells.append(
+                    nbf.new_code_cell(
+                        f"Invoke-AtomicTest {self.technique_id} -TestNumbers {self.test_number} -GetPreReqs"
+                    )
+                )
                 markdown = []
             markdown.append(f"#### Attack Commands: Run with `{self.executor.name}`\n")
             markdown.append(f"```{self.executor.name}\n{self.executor.command}```")
+            cells.append(nbf.new_markdown_cell(markdown))
             cells.append(
-                nbf.new_markdown_cell(markdown))
-            cells.append(
-                nbf.new_code_cell(f'Invoke-AtomicTest {self.technique_id} -TestNumbers {self.test_number}'))
+                nbf.new_code_cell(
+                    f"Invoke-AtomicTest {self.technique_id} -TestNumbers {self.test_number}"
+                )
+            )
             if cleanup := self.executor.cleanup:
                 executor_name = self.executor.name
                 if executor_name == "command_prompt":
                     executor_name = "cmd"
-                cells.append(nbf.new_markdown_cell(
-                    f"#### Cleanup: \n```{executor_name}\n{cleanup}```"))
-                cells.append(nbf.new_code_cell(
-                    f'Invoke-AtomicTest {self.technique_id} -TestNumbers {self.test_number} -Cleanup'))
+                cells.append(
+                    nbf.new_markdown_cell(
+                        f"#### Cleanup: \n```{executor_name}\n{cleanup}```"
+                    )
+                )
+                cells.append(
+                    nbf.new_code_cell(
+                        f"Invoke-AtomicTest {self.technique_id} -TestNumbers {self.test_number} -Cleanup"
+                    )
+                )
         return cells
 
 
@@ -154,15 +176,20 @@ class AggregatedAtomic(BaseModel):
 
     def __repr__(self):
         cells = [
-            nbf.new_markdown_cell(f'# {self.stix.technique_id} - {self.stix.name}\n{self.stix.description}')
+            nbf.new_markdown_cell(
+                f"# {self.stix.technique_id} - {self.stix.name}\n{self.stix.description}"
+            )
         ]
         if atomic := self.atomic:
             cells.append(nbf.new_markdown_cell("## Atomic Tests"))
             for i in atomic.tests:
                 cells += i.__repr__()
         else:
             cells.append(
-                nbf.new_markdown_cell("## Atomic Tests:\nCurrently, no tests are available for this technique."))
+                nbf.new_markdown_cell(
+                    "## Atomic Tests:\nCurrently, no tests are available for this technique."
+                )
+            )
         if detection := self.stix.detection:
             cells.append(nbf.new_markdown_cell(f"## Detection\n{detection}"))
         if shield_obj := shield.get_shield_obj(self.stix.technique_id):
@@ -179,16 +206,21 @@ class AttackTactic(BaseModel):
 
     def __repr__(self):
         markdown = [
-            f'# {self.stix.name}',
+            f"# {self.stix.name}",
             self.stix.description,
             "## Techniques",
             "| ID      | Name | Description |",
-            "| :--------: | :---------: | :---------: |"
+            "| :--------: | :---------: | :---------: |",
         ]
         for i in self.techniques:
-            markdown.append(f'{i.stix.technique_id} | {i.stix.name} | {i.stix.description}')
-        cells = [nbf.new_markdown_cell("\n".join(markdown)),
-                 nbf.new_code_cell(
-                     f"#Invoke-AtomicTest-By can be downloaded from "
-                     f"https://github.com/cyberbuff/ART-Utils/\nInvoke-AtomicTest-By -Tactic {self.stix.short_name}")]
+            markdown.append(
+                f"{i.stix.technique_id} | {i.stix.name} | {i.stix.description}"
+            )
+        cells = [
+            nbf.new_markdown_cell("\n".join(markdown)),
+            nbf.new_code_cell(
+                f"#Invoke-AtomicTest-By can be downloaded from "
+                f"https://github.com/cyberbuff/ART-Utils/\nInvoke-AtomicTest-By -Tactic {self.stix.short_name}"
+            ),
+        ]
         return cells