github_service.py

import csv
import glob
import logging
import os
import pathlib
import subprocess
import sys
from typing import Union
from docker import types
import datetime
import git
import yaml
from git.objects import base
from modules import testing_service
import pathlib

# Logger
logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
LOGGER = logging.getLogger(__name__)

SECURITY_CONTENT_URL = "https://github.com/splunk/security_content"


DETECTION_ROOT_PATH = "security_content/detections"
TEST_ROOT_PATH = "security_content/tests"
DETECTION_FILE_EXTENSION = ".yml"
TEST_FILE_EXTENSION = ".test.yml"
SSA_PREFIX = "ssa___"


class GithubService:
    def __init__(
        self,
        security_content_branch: str,
        commit_hash: Union[str, None],
        PR_number: Union[int, None] = None,
        persist_security_content: bool = False,
    ):
        self.security_content_branch = security_content_branch
        if persist_security_content:
            print("Getting handle on existing security_content repo!")
            self.security_content_repo_obj = git.Repo("security_content")
        else:
            print("Checking out security_content repo!")
            self.security_content_repo_obj = self.clone_project(
                SECURITY_CONTENT_URL, f"security_content", f"develop"
            )

        # Ensure that the branch name is valid
        # Get all the branch names, prefixed with "origin/"
        branch_names = [
            branch.name for branch in self.security_content_repo_obj.remote().refs
        ]

        if "origin/%s" % (security_content_branch) not in branch_names:
            raise (
                Exception(
                    "Branch name [%s] not found in valid branches.  Try running \n"
                    "'git branch -a' to examine [%d] branches"
                    % (security_content_branch, len(branch_names))
                )
            )

        if commit_hash is not None and PR_number is not None:
            print(
                f"\n************\nWARNING - both the PR_number {PR_number} and the commit_hash {commit_hash} were provided. "
                f"You should only pass neither or one of these.  We will ASSUME you want to use the PR_number, not the commit_hash. "
                f"Removing the commit_hash...\n************\n"
            )
            commit_hash = None

        if PR_number:
            ret = subprocess.run(
                [
                    "git",
                    "-C",
                    "security_content/",
                    "fetch",
                    "origin",
                    "refs/pull/%d/head:%s" % (PR_number, security_content_branch),
                ],
                capture_output=True,
            )
            # ret = subprocess.call(["git", "-C", "security_content/", "fetch", "origin",
            #                "refs/pull/%d/head:%s" % (PR_number, security_content_branch)])

            if ret.returncode != 0:
                raise (
                    Exception(
                        "Error checking out repository: [%s]"
                        % (
                            ret.stdout.decode("utf-8")
                            + "\n"
                            + ret.stderr.decode("utf-8")
                        )
                    )
                )

        # No checking to see if the hash is to a commit inside of the branch - the user
        # has to do that by hand.

        # -- ensures that we check out the appropriate branch or commit hash.
        # Without --, there can be ambiguity if a file/folder exists with the
        # same name as the branch, causing the checkout to fail with error
        if commit_hash is not None:
            print("Checking out commit hash: [%s]" % (commit_hash))
            self.security_content_repo_obj.git.checkout(commit_hash, "--")
        else:
            # Even if we have fetched a PR, we still MUST check out the branch to
            # be able to do anything with it. Otherwise we won't have the files
            print("Checking out branch: [%s]..." % (security_content_branch), end="")
            sys.stdout.flush()
            self.security_content_repo_obj.git.checkout(security_content_branch, "--")
            commit_hash = self.security_content_repo_obj.head.object.hexsha
            print("commit_hash %s" % (commit_hash))

        self.commit_hash = commit_hash

    def update_and_commit_passed_tests(self, results: list[dict]) -> bool:
        changed_file_paths = []
        for result in results:
            detection_obj_path = os.path.join(
                "security_content", "detections", result["detection_file"]
            )

            test_obj_path = detection_obj_path.replace("detections", "tests", 1)
            test_obj_path = test_obj_path.replace(".yml", ".test.yml")

            detection_obj = testing_service.load_file(detection_obj_path)
            test_obj = testing_service.load_file(test_obj_path)
            detection_obj["tags"]["automated_detection_testing"] = "passed"
            # detection_obj['tags']['automated_detection_testing_date'] = datetime.datetime.today().strftime('%Y-%m-%d-%H:%M:%S')

            for o in test_obj["tests"]:
                if "attack_data" in o:
                    datasets = []
                    for dataset in o["attack_data"]:
                        datasets.append(dataset["data"])
                    detection_obj["tags"]["dataset"] = datasets
            with open(detection_obj_path, "w") as f:
                yaml.dump(detection_obj, f, sort_keys=False, allow_unicode=True)

            changed_file_paths.append(detection_obj_path)

        relpaths = [
            pathlib.Path(*pathlib.Path(p).parts[1:]).as_posix()
            for p in changed_file_paths
        ]
        newpath = relpaths[0] + ".wow"
        relpaths.append(newpath)
        with open("security_content/" + newpath, "w") as d:
            d.write("fake file")
        print("status results:")
        print(
            self.security_content_repo_obj.index.diff(
                self.security_content_repo_obj.head.commit
            )
        )

        if len(relpaths) > 0:
            print("there is at least one changed file")
            print(relpaths)
            self.security_content_repo_obj.index.add(relpaths)
            print("status results after add:")
            print(
                self.security_content_repo_obj.index.diff(
                    self.security_content_repo_obj.head.commit
                )
            )

            commit_message = (
                "The following detections passed detection testing.  Their YAMLs have been updated and their datasets linked:\n - %s"
                % ("\n - ".join(relpaths))
            )
            self.security_content_repo_obj.index.commit(commit_message)
            return True
        else:
            return False

        return True

    def clone_project(self, url, project, branch):
        LOGGER.info(f"Clone Security Content Project")
        repo_obj = git.Repo.clone_from(url, project, branch=branch)
        return repo_obj

    def prune_detections(
        self,
        detection_files: list[str],
        types_to_test: list[str],
        exclude_ssa: bool = True,
    ) -> list[str]:
        pruned_tests = []

        for detection in detection_files:
            if os.path.basename(detection).startswith(SSA_PREFIX) and exclude_ssa:
                continue
            with open(detection, "r") as d:
                description: dict = yaml.safe_load(d)

                detection_filepath_without_security_content = str(
                    pathlib.Path(*pathlib.Path(detection).parts[1:])
                )
                # If no   types are provided, then we will get everything
                if (
                    description.get("type", None) in types_to_test
                    and "production" in description.get("status", "")
                    and not (
                        description.get("tags", False)
                        and description["tags"].get("manual_test", False)
                    )
                ):
                    if len(description.get("tests", [])) == 0:
                        print(
                            Exception(
                                f"Detection {detection_filepath_without_security_content} has no tests/test section defined. Detection must include at least one test."
                            )
                        )
                        continue
                        raise (
                            Exception(
                                f"Detection {detection_filepath_without_security_content} has no tests/test section defined. Detection must include at least one test."
                            )
                        )

                    pruned_tests.append(detection_filepath_without_security_content)

                else:
                    if description.get("tags", False):
                        manual_test = description["tags"].get("manual_test", False)
                    else:
                        manual_test = False

                    print(
                        f"Ignore {detection}:\n - [status:'{description.get('status', None)}'] [type:'{description.get('type', None)}'] [manual_test:'{manual_test}']"
                    )
                    # Don't do anything with these files
                    pass

        # if not self.ensure_paired_detection_and_test_files([], [os.path.join("security_content", p) for p in pruned_tests], exclude_ssa):
        #    raise(Exception("Missing one or more test/detection files. Please see the output above."))

        return pruned_tests

    def ensure_paired_detection_and_test_files(
        self,
        detection_files: list[str],
        test_files: list[str],
        exclude_ssa: bool = True,
    ) -> bool:
        """
        The security_content repo contains two folders: detections and test.
        For EVERY detection in the detections folder, there must be a test.
        for EVERY test in the tests folder, there MUST be a detection.

        If this requirement is not met, then throw an error
        """

        MISSING_TEMPLATE = (
            "Missing {type} file:" "\n\tEXISTS  - {exists}" "\n\tMISSING - {missing}"
        )

        no_missing_files = True
        # Check that all detection files have a test file
        for detection_file in detection_files:
            test_file = self.convert_detection_filename_into_test_filename(
                detection_file
            )
            if not os.path.exists(test_file):
                if (
                    os.path.basename(detection_file).startswith(SSA_PREFIX)
                    and exclude_ssa is True
                ):
                    print(
                        MISSING_TEMPLATE.format(
                            type="test", exists=detection_file, missing=test_file
                        )
                    )
                    print(
                        "\tSince exclude_ssa is TRUE, this is not an error, just a warning"
                    )
                else:
                    print(
                        MISSING_TEMPLATE.format(
                            type="test", exists=detection_file, missing=test_file
                        )
                    )
                    no_missing_files = False

        # Check that all test files have a detection file
        for test_file in test_files:
            detection_file = self.convert_test_filename_into_detection_filename(
                test_file
            )
            if not os.path.exists(detection_file):
                if (
                    os.path.basename(test_file).startswith(SSA_PREFIX)
                    and exclude_ssa is True
                ):
                    print(
                        MISSING_TEMPLATE.format(
                            type="detection", exists=test_file, missing=detection_file
                        )
                    )
                    print(
                        "\tSince exclude_ssa is TRUE, this is not an error, just a warning"
                    )
                else:
                    print(
                        MISSING_TEMPLATE.format(
                            type="detection", exists=test_file, missing=detection_file
                        )
                    )
                    no_missing_files = False

        return no_missing_files

    def convert_detection_filename_into_test_filename(
        self, detection_filename: str
    ) -> str:
        head, tail = os.path.split(detection_filename)

        assert head.startswith(
            DETECTION_ROOT_PATH
        ), f"Error - Expected detection filename to start with [{DETECTION_ROOT_PATH}] but instead got {detection_filename}"

        updated_head = head.replace(DETECTION_ROOT_PATH, TEST_ROOT_PATH, 1)

        assert tail.endswith(
            DETECTION_FILE_EXTENSION
        ), f"Error - Expected detection filename to end with [{DETECTION_FILE_EXTENSION}] but instead got [{detection_filename}]"
        updated_tail = TEST_FILE_EXTENSION.join(tail.rsplit(DETECTION_FILE_EXTENSION))

        return os.path.join(updated_head, updated_tail)

    def convert_test_filename_into_detection_filename(self, test_filename: str) -> str:
        head, tail = os.path.split(test_filename)

        assert head.startswith(
            TEST_ROOT_PATH
        ), f"Error - Expected test filename to start with [{TEST_ROOT_PATH}] but instead got {test_filename}"

        updated_head = head.replace(TEST_ROOT_PATH, DETECTION_ROOT_PATH, 1)

        assert tail.endswith(
            TEST_FILE_EXTENSION
        ), f"Error - Expected test filename to end with [{TEST_FILE_EXTENSION}] but instead got [{test_filename}]"
        updated_tail = DETECTION_FILE_EXTENSION.join(tail.rsplit(TEST_FILE_EXTENSION))

        return os.path.join(updated_head, updated_tail)

    def get_detection_files(
        self,
        mode: str,
        folders: list[str],
        types: list[str],
        detections_list: Union[list[str], None],
    ) -> list[str]:
        if mode == "changes":
            tests = self.get_changed_detection_files(folders, types)
        elif mode == "selected":
            if detections_list is None:
                # It's actually valid to supply an EMPTY list of files and the test should pass.
                # This can occur when we try to test, for example, 1 detection but start 2 containers.
                # We still want this to pass testing, so we shouldn't fail there!
                print(
                    "Trying to test a list of files, but None were provided",
                    file=sys.stderr,
                )
                sys.exit(1)

            elif detections_list is not None:
                tests = self.get_selected_test_files(detections_list, types)
            else:
                # impossible to get here
                print(
                    "Impossible to get here.  Just kept to make the if/elif more self describing",
                    file=sys.stderr,
                )
                sys.exit(1)

        elif mode == "all":
            tests = self.get_all_tests_and_detections(folders, types)
        elif mode == "smoketest":
            tests = self.get_everything_including_experimental_and_deprecated(types)
        else:
            print(
                "Error, unsupported mode [%s].  Mode must be one of %s", file=sys.stderr
            )
            sys.exit(1)

        return tests

    def get_selected_test_files(
        self,
        detection_file_list: list[str],
        types_to_test: list[str] = ["Anomaly", "Hunting", "TTP"],
    ) -> list[str]:
        return self.prune_detections(detection_file_list, types_to_test)

    def get_everything_including_experimental_and_deprecated(
        self,
        types_to_test: list[str] = ["Anomaly", "Hunting", "TTP"],
        exclude_ssa: bool = True,
    ) -> list[str]:
        all_detections = glob.glob(
            os.path.join(DETECTION_ROOT_PATH, "**", f"*{DETECTION_FILE_EXTENSION}"),
            recursive=True,
        )
        only_selected_detection_types: list[str] = []
        for detection in all_detections:
            if os.path.basename(detection).startswith(SSA_PREFIX) and exclude_ssa:
                continue
            with open(detection, "r") as d:
                description: dict = yaml.safe_load(d)

                detection_filepath_without_security_content = str(
                    pathlib.Path(*pathlib.Path(detection).parts[1:])
                )

                if description.get("type", None) not in types_to_test:
                    continue
                only_selected_detection_types.append(
                    detection_filepath_without_security_content
                )

        print(
            f"Number of non-ssa tests including experimental and deprecated: {len(only_selected_detection_types)}"
        )

        return only_selected_detection_types

    def get_all_tests_and_detections(
        self,
        folders: list[str] = ["endpoint", "cloud", "network"],
        types_to_test: list[str] = ["Anomaly", "Hunting", "TTP"],
    ) -> list[str]:
        detections = []
        for folder in folders:
            detections.extend(
                self.get_all_files_in_folder(
                    os.path.join(DETECTION_ROOT_PATH, folder), "*"
                )
            )

        # Prune this down to only the subset of detections we can test
        return self.prune_detections(detections, types_to_test)

    def get_all_files_in_folder(self, foldername: str, extension: str) -> list[str]:
        filenames = glob.glob(os.path.join(foldername, extension))
        return filenames

    def get_changed_detection_files(
        self,
        folders=["endpoint", "cloud", "network"],
        types_to_test=["Anomaly", "Hunting", "TTP"],
    ) -> list[str]:
        branch1 = self.security_content_branch
        branch2 = "develop"
        g = git.Git("security_content")
        all_changed_test_files = []

        all_changed_detection_files = []
        if branch1 != "develop":
            if self.commit_hash is None:
                differ = g.diff("--name-status", branch2 + "..." + branch1)
            else:
                differ = g.diff("--name-status", branch2 + "..." + self.commit_hash)

            changed_files = differ.splitlines()

            for file_path in changed_files:
                # added or changed test files
                if file_path.startswith("A") or file_path.startswith("M"):
                    # changed detections
                    if "detections" in file_path and os.path.basename(
                        file_path
                    ).endswith(".yml"):
                        all_changed_detection_files.append(file_path)
        else:
            print(
                "Looking for changed detections by diffing [%s] against [%s].  They are the same branch, so none were returned."
                % (branch1, branch2),
                file=sys.stderr,
            )
            return []

        all_changed_detection_files = [
            os.path.join("security_content", name.split("\t")[1])
            for name in all_changed_detection_files
            if len(name.split("\t")) == 2
        ]

        # Trim out any of the tests/detection  that are not in the selected folders, but at least print a notice
        # to the user.

        changed_detection_files = [
            x
            for x in all_changed_detection_files
            if (len(pathlib.Path(x).parts) > 3 and pathlib.Path(x).parts[2] in folders)
        ]

        for missing in set(changed_detection_files).symmetric_difference(
            all_changed_detection_files
        ):
            print(
                "Ignoring modified detecton [%s] not in set of selected folders: %s"
                % (missing, folders)
            )

        # Convert the test files to the detection file equivalent.
        # Note that some of these tests may be baselines and their associated
        # detection could be in experimental or not in the experimental folder
        converted_test_files = []
        # for test_filepath in changed_test_files:
        #    detection_filename = str(pathlib.Path(
        #        *pathlib.Path(test_filepath).parts[-2:])).replace("tests", "detections", 1)
        #    converted_test_files.append(detection_filename)

        return self.prune_detections(changed_detection_files, types_to_test)

        # detections_to_test,_,_ = self.filter_test_types(changed_detection_files)
        # for f in detections_to_test:
        #    file_path_base = os.path.splitext(f)[0].replace('detections', 'tests') + '.test'
        #    file_path_new = file_path_base + '.yml'
        #    if file_path_new not in changed_test_files:
        #        changed_test_files.append(file_path_new)

        # print("Total things to test (test files and detection files changed): [%d]"%(len(changed_test_files)))
        # for l in changed_test_files:
        #    print(l)
        # print(len(changed_test_files))
        # import time
        # time.sleep(5)

    def filter_test_types(self, test_files, test_types=["Anomaly", "Hunting", "TTP"]):
        files_to_test = []
        files_not_to_test = []
        error_files = []
        for filename in test_files:
            try:
                with open(os.path.join("security_content", filename), "r") as fileData:
                    yaml_dict = list(yaml.safe_load_all(fileData))[0]
                    if "type" not in yaml_dict.keys():
                        print(
                            "Failed to find 'type' in the yaml for: [%s]" % (filename)
                        )
                        error_files.append(filename)
                    if yaml_dict["type"] in test_types:
                        files_to_test.append(filename)
                    else:
                        files_not_to_test.append(filename)
            except Exception as e:
                print("Error on trying to scan [%s]: [%s]" % (filename, str(e)))
                error_files.append(filename)
        print(
            "***Detection Information***\n"
            "\tTotal Files       : %d"
            "\tFiles to test     : %d"
            "\tFiles not to test : %d"
            "\tError files       : %d"
            % (
                len(test_files),
                len(files_to_test),
                len(files_not_to_test),
                len(error_files),
            )
        )
        import time

        time.sleep(5)
        return files_to_test, files_not_to_test, error_files