From df5dd5e58f0a8deb63cb0a11b49e7f2adcfcee6f Mon Sep 17 00:00:00 2001
From: Milan Lysonek <mlysonek@redhat.com>
Date: Wed, 31 Jan 2024 12:25:01 +0100
Subject: [PATCH] Add helper functions for work with controls

---
 ctf/DiffStruct.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/ctf/DiffStruct.py b/ctf/DiffStruct.py
index 4ab6d60..0a5cd8e 100644
--- a/ctf/DiffStruct.py
+++ b/ctf/DiffStruct.py
@@ -61,6 +61,42 @@ def find_rule_profiles(self, rule):
                         if find_rule.search(line):
                             yield profile_file
 
+    def find_rule_controls(self, rule):
+        controls = []
+        find_rule = re.compile(r"^\s*-\s*" + rule + r"\s*$", re.MULTILINE)
+        control_folder = git_wrapper.repo_path + "/" + "controls/"
+        # Check all yaml files in controls/
+        for control in os.listdir(control_folder):
+            if not control.endswith(".yml"):
+                continue
+            control_path = control_folder + control
+            with open(control_path) as f:
+                control_content = f.read()
+                # If controls in separate directory, merge them to one string
+                controls_dir = re.search(r"controls_dir:\s*(\w+)", control_content)
+                if controls_dir:
+                    controls_dir = controls_dir.group(1)
+                    for c in os.listdir(control_folder + controls_dir):
+                        with open(control_folder + controls_dir + "/" + c) as cf:
+                            control_content += cf.read()
+            # Search for rule in control content
+            if find_rule.search(control_content):
+                yield control.rstrip(".yml")
+
+    def find_control_products(self, control):
+        products_folder = git_wrapper.repo_path + "/" + "products"
+        find_control = re.compile(r"^\s*-\s*" + control + r":", re.MULTILINE)
+        # Find dirs with profile files
+        for dir_path, _, files in os.walk(products_folder):
+            for file in files:
+                if not file.endswith(".profile"):
+                    continue
+                # Search if desired control is used and if so, return product
+                with open(dir_path + "/" + file) as f:
+                    for line in f:
+                        if find_control.search(line):
+                            yield re.match(r".*/products/([^/]+)", dir_path).group(1)
+
     def get_rule_ruleyml(self, rule):
         # Find a directory with a rule name and check if it has rule.yml file
         for root, dirs, files in os.walk(git_wrapper.repo_path):