Python conversion of maintenance commands (#2739)

* updating atomics count in README.md [ci skip] * converting python * rename * fix path * minor refactor --------- Co-authored-by: publish bot <[email protected]>
redcanaryco · Apr 7, 2024 · 3bf390b · 3bf390b
1 parent 3bcc943
commit 3bf390b
Show file tree

Hide file tree

Showing 17 changed files with 900 additions and 751 deletions.
diff --git a/.github/workflows/generate-counter.yml b/.github/workflows/generate-counter.yml
@@ -2,7 +2,7 @@ name: generate-svg-counter
 
 on:
   push:
-    branches: ["master"]
+    branches: [ "master" ]
 
 jobs:
   generate-counter:
@@ -15,26 +15,28 @@ jobs:
         run: pipx install poetry
       - uses: actions/setup-python@v5
         with:
-          python-version: '3.11.2' 
+          python-version: '3.11.2'
           cache: 'poetry'
       - name: Install dependencies
         run: poetry install --no-interaction --no-root
       - name: Generate shields.io URL
-        run: poetry run python bin/generate_counter.py -f atomics/
+        run: poetry run python atomic_red_team/runner.py generate-counter
         id: counter
       - name: Update README
         run: |
           echo ${{ steps.counter.outputs.result }}
           sed -i "s|https://img.shields.io/badge/Atomics-.*-flat.svg|${{ steps.counter.outputs.result }}|" README.md
         shell: bash
+      - name: Generate and commit unique GUIDs for each atomic test
+        run: poetry run python atomic_red_team/runner.py generate-guids
       - name: update github with new site
         run: |
           # configure git to prep for commit
           git config user.email "[email protected]"
           git config user.name "publish bot"
           git config --global push.default simple
           git add README.md
-          git commit --allow-empty -m "updating atomics count in README.md [ci skip]"
+          git commit --allow-empty -m "updating atomics count and guids [ci skip]"
           # push quietly to prevent showing the token in log
           # no need to provide any credentials
           git push --force
diff --git a/.github/workflows/generate-docs.yml b/.github/workflows/generate-docs.yml
@@ -1,7 +1,7 @@
 name: generate-docs
 on:
   push:
-    branches: ["master"]
+    branches: [ "master" ]
 
 jobs:
   generate-docs:
@@ -15,29 +15,7 @@ jobs:
         uses: ruby/setup-ruby@v1
         with:
           ruby-version: 2.7
-          bundler-cache: true 
-
-      - name: Generate and commit unique GUIDs for each atomic test
-        run: |
-            bin/generate-guids.rb
-
-            echo ""
-            echo ""
-            git status
-            echo ""
-            echo ""
-            git diff-index HEAD --
-
-            if git diff-index --quiet HEAD -- ; then
-              echo "Not committing GUID changes because there are no changes"
-            else
-              git config credential.helper 'cache --timeout=120'
-              git config user.email "[email protected]"
-              git config user.name "Atomic Red Team GUID generator"
-              git add atomics
-              git commit -am "Generate GUIDs from job=$GITHUB_JOB branch=$GITHUB_REF_NAME [skip ci]"
-              git push origin $GITHUB_REF_NAME -f 
-            fi
+          bundler-cache: true
 
       - name: generate markdown docs for atomics
         run: |

diff --git a/.github/workflows/validate-atomics.yml b/.github/workflows/validate-atomics.yml
@@ -24,8 +24,7 @@ jobs:
         run: poetry install --no-interaction --no-root
 
       - name: validate the format of atomics tests against the spec
-        run: |
-          poetry run python bin/validate/validate.py
+        run: poetry run python atomic_red_team/runner.py validate
 
   upload:
     runs-on: ubuntu-latest
@@ -62,7 +61,7 @@ jobs:
         run: poetry install --no-interaction --no-root
       - name: save labels and reviewers into a file.
         run: |
-          poetry run python bin/generate_labels.py  -t  ${{ secrets.GITHUB_TOKEN }} -pr '${{steps.get_pr_number.outputs.result}}'
+          poetry run python atomic_red_team/runner.py generate-labels --pr '${{steps.get_pr_number.outputs.result}}' --token ${{ secrets.GITHUB_TOKEN }}
       - uses: actions/upload-artifact@v4
         with:
           name: labels.json

diff --git a/atomic_red_team/__init__.py b/atomic_red_team/__init__.py
diff --git a/atomic_red_team/common.py b/atomic_red_team/common.py
@@ -0,0 +1,5 @@
+from os.path import dirname, realpath
+
+base_path = dirname(dirname(realpath(__file__)))
+atomics_path = f"{base_path}/atomics"
+used_guids_file = f"{atomics_path}/used_guids.txt"
diff --git a/atomic_red_team/guid.py b/atomic_red_team/guid.py
@@ -0,0 +1,45 @@
+import re
+import uuid
+from typing import List
+
+from ruamel.yaml import YAML
+
+from common import used_guids_file
+
+yaml = YAML(typ="safe")
+
+
+def get_unique_guid(guids: List[str]):
+    # This function should return a unique GUID that's not in the used_guids_file.
+    guid = str(uuid.uuid4())
+    if guid not in guids:
+        with open(used_guids_file, "a") as f:  # append mode
+            f.write(guid)
+        return guid
+    else:
+        return get_unique_guid(guids)
+
+
+def generate_guids_for_yaml(path, get_guid):
+    with open(path, "r") as file:
+        og_text = file.read()
+
+    # Add the "auto_generated_guid:" element after the "- name:" element if it isn't already there
+    text = re.sub(
+        r"(?i)(^([ \t]*-[ \t]*)name:.*$(?!\s*auto_generated_guid))",
+        lambda m: f"{m.group(1)}\n{m.group(2).replace('-', ' ')}auto_generated_guid:",
+        og_text,
+        flags=re.MULTILINE,
+    )
+
+    # Fill the "auto_generated_guid:" element in if it doesn't contain a guid
+    text = re.sub(
+        r"(?i)^([ \t]*auto_generated_guid:)(?!([ \t]*[a-f0-9]{8}-[a-f0-9]{4}-4[a-f0-9]{3}-[89aAbB][a-f0-9]{3}-[a-f0-9]{12})).*$",
+        lambda m: f"{m.group(1)} {get_guid()}",
+        text,
+        flags=re.MULTILINE,
+    )
+    if text != og_text:
+        with open(path, "wb") as file:
+            # using wb mode instead of w. If not, the end of line characters are auto-converted to OS specific ones.
+            file.write(text.encode())
diff --git a/bin/generate_labels.py → atomic_red_team/labels.py b/bin/generate_labels.py → atomic_red_team/labels.py
@@ -1,4 +1,3 @@
-import argparse
 import fnmatch
 import json
 import os
@@ -18,6 +17,7 @@ def get_technique_from_filename(filename):
 @dataclass
 class ChangedAtomic:
     """Returns atomic technique with test number which can be later used to run atomics in CI/CD pipelines."""
+
     technique: str
     test_number: int
     data: dict
@@ -28,7 +28,7 @@ def construct_mapping(self, node, deep=False):
         """Add line number to each block of the atomic test."""
         mapping = super(SafeLineLoader, self).construct_mapping(node, deep=deep)
         # Add 1 so line numbering starts at 1
-        mapping['__line__'] = node.start_mark.line + 1
+        mapping["__line__"] = node.start_mark.line + 1
         return mapping
 
 
@@ -43,7 +43,7 @@ class GithubAPI:
         "iaas:aws": "cloud",
         "iaas:azure": "cloud",
         "office-365": "cloud",
-        "google-workspace":"cloud"
+        "google-workspace": "cloud",
     }
 
     maintainers = {
@@ -56,7 +56,7 @@ class GithubAPI:
         "iaas:azure": ["patel-bhavin"],
         "azure-ad": ["patel-bhavin"],
         "google-workspace": ["patel-bhavin"],
-        "office-365": ["patel-bhavin"]
+        "office-365": ["patel-bhavin"],
     }
 
     def __init__(self, token):
@@ -67,7 +67,7 @@ def headers(self):
         return {
             "Authorization": f"Bearer {self.token}",
             "X-GitHub-Api-Version": "2022-11-28",
-            "Accept": "application/vnd.github+json"
+            "Accept": "application/vnd.github+json",
         }
 
     def get_atomic_with_lines(self, file_url: str):
@@ -78,13 +78,18 @@ def get_atomic_with_lines(self, file_url: str):
 
     def get_files_for_pr(self, pr):
         """Get new and modified files in the `atomics` directory changed in a PR."""
-        response = requests.get(f"https://api.github.com/repos/{os.getenv('GITHUB_REPOSITORY')}/pulls/{pr}/files",
-                                headers=self.headers, timeout=15)
+        response = requests.get(
+            f"https://api.github.com/repos/{os.getenv('GITHUB_REPOSITORY')}/pulls/{pr}/files",
+            headers=self.headers,
+            timeout=15,
+        )
         assert response.status_code == 200
         files = response.json()
         return filter(
-            lambda x: x["status"] in ["added", "modified"] and fnmatch.fnmatch(x["filename"], "atomics/T*/T*.yaml"),
-            files)
+            lambda x: x["status"] in ["added", "modified"]
+            and fnmatch.fnmatch(x["filename"], "atomics/T*/T*.yaml"),
+            files,
+        )
 
     def get_tests_changed(self, pr: str):
         """Get all the tests changed in a PR"""
@@ -96,8 +101,10 @@ def get_tests_changed(self, pr: str):
             technique = get_technique_from_filename(file["filename"])
             if file["status"] == "added":
                 # New file; run the entire technique; Invoke-AtomicTest Txxxx
-                tests += [ChangedAtomic(technique=technique, test_number=index + 1, data=t)
-                          for index, t in enumerate(data["atomic_tests"])]
+                tests += [
+                    ChangedAtomic(technique=technique, test_number=index + 1, data=t)
+                    for index, t in enumerate(data["atomic_tests"])
+                ]
             else:
                 changed_lines = []
                 count = 0
@@ -114,14 +121,21 @@ def get_tests_changed(self, pr: str):
                 atomics = data["atomic_tests"]
                 for index, t in enumerate(atomics):
                     curr_atomic_start = atomics[index]["__line__"]
-                    if index+1<len(atomics):
-                        curr_atomic_end = atomics[index+1]["__line__"]
+                    if index + 1 < len(atomics):
+                        curr_atomic_end = atomics[index + 1]["__line__"]
                     else:
-                        curr_atomic_end = start+60
-                    changes_in_current_atomic = [i for i in changed_lines if i > curr_atomic_start and i < curr_atomic_end]
+                        curr_atomic_end = start + 60
+                    changes_in_current_atomic = [
+                        i
+                        for i in changed_lines
+                        if i > curr_atomic_start and i < curr_atomic_end
+                    ]
                     if len(changes_in_current_atomic) > 0:
-                        tests.append(ChangedAtomic(technique=technique, test_number=index + 1,
-                                                   data=t))
+                        tests.append(
+                            ChangedAtomic(
+                                technique=technique, test_number=index + 1, data=t
+                            )
+                        )
 
         return tests
 
@@ -145,20 +159,5 @@ def save_labels_and_maintainers(self, pr):
             f.write(json.dumps(x))
 
         with open("pr/labels.json", "w") as f:
-            j = {
-                "pr": pr,
-                "labels": labels,
-                "maintainers": maintainers
-            }
+            j = {"pr": pr, "labels": labels, "maintainers": maintainers}
             f.write(json.dumps(j))
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description='Generate an SVG counter for a folder with a list of YAML files.')
-    parser.add_argument('-pr', '--pull-request', type=str,
-                        help="Current pull request number")
-    parser.add_argument('-t', '--token', type=str, help="Github Token to be used")
-    args = parser.parse_args()
-
-    api = GithubAPI(args.token)
-    api.save_labels_and_maintainers(args.pull_request)