refactor: adjust ODV generation and behavior

src/mscp/classes/baseline.py

@@ -142,7 +142,6 @@ class Baseline(BaseModelWithAccessors):
                         section_data.get("name", "").strip(),
                         baseline_tag,
                         language,
-                        custom,
                     ),
                 )
             )
@@ -366,6 +365,7 @@ class Baseline(BaseModelWithAccessors):
 
         create_yaml(output_path, ordered_data)
         logger.success("Created baseline yaml: {}", output_path)
+        print(f"Generated new baseline file: {output_path}")
 
     @classmethod
     def load_all_from_folder(

src/mscp/classes/macsecurityrule.py

@@ -2,7 +2,7 @@
 
 # Standard python modules
 import base64
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
 from enum import StrEnum
 from pathlib import Path
 from typing import Any
@@ -241,10 +241,9 @@ class Macsecurityrule(BaseModelWithAccessors):
         os_version: float,
         parent_values: str,
         section: str,
+        tailoring: bool | None = False,
         baseline_tag: str | None = None,
         language: str = "en",
-        custom: bool = False,
-        generate_baseline: bool = False,
     ) -> list["Macsecurityrule"]:
         """
         Load Macsecurityrule objects from YAML files for the given rule IDs.
@@ -257,7 +256,6 @@ class Macsecurityrule(BaseModelWithAccessors):
             section (str): Section name for the rules.
             language (str): Language used for rule text.
             custom (bool): Whether to include custom rules.
-            generate_baseline (bool): Whether to generate a baseline.
 
         Returns:
             list[Macsecurityrule]: A list of loaded Macsecurityrule objects.
@@ -279,6 +277,12 @@ class Macsecurityrule(BaseModelWithAccessors):
             Path(config["defaults"]["rules_dir"]),
         ]
 
+        # collect custom rules if they exist
+        if tailoring:
+            custom_rule_dict = {}
+        else:
+            custom_rule_dict = collect_overrides(Path(config["custom"]["rules_dir"]))
+
         for rule_id in rule_ids:
             logger.debug("Transforming rule: {}", rule_id)
 
@@ -330,7 +334,6 @@ class Macsecurityrule(BaseModelWithAccessors):
 
             # process any customized rules
             customized_fields = []
-            custom_rule_dict = collect_overrides(Path(config["custom"]["rules_dir"]))
 
             if rule_yaml["rule_id"] in custom_rule_dict:
                 logger.info(f"Found customization for {rule_yaml['rule_id']}")
@@ -345,7 +348,8 @@ class Macsecurityrule(BaseModelWithAccessors):
                         rule_yaml[custom_rule_key].update(custom_rule_value)
                         continue
                     if custom_rule_key == "tags":
-                        rule_yaml[custom_rule_key] += custom_rule_value
+                        if custom_rule_value not in rule_yaml[custom_rule_key]:
+                            rule_yaml[custom_rule_key] += custom_rule_value
                         continue
                     rule_yaml[custom_rule_key] = custom_rule_value
 
@@ -549,11 +553,7 @@ class Macsecurityrule(BaseModelWithAccessors):
             #     rule._format_mobileconfig_fix()
             #     logger.success("Formatted mobileconfig_info for rule: {}", rule.rule_id)
 
-            if (
-                rule.odv is not None
-                and not generate_baseline
-                and parent_values is not None
-            ):
+            if rule.odv is not None and parent_values is not None:
                 rule._fill_in_odv(parent_values)
 
             logger.success("Transformed rule: {}", rule_id)
@@ -570,7 +570,7 @@ class Macsecurityrule(BaseModelWithAccessors):
         cls,
         os_type: str,
         os_version: int,
-        generate_baseline: bool = True,
+        tailoring: bool = False,
         parent_values: str = "default",
     ) -> list["Macsecurityrule"]:
         """
@@ -589,6 +589,7 @@ class Macsecurityrule(BaseModelWithAccessors):
         logger.info("=== LOADING ALL RULES ===")
 
         rules: list[Macsecurityrule] = []
+        rules_to_collect = defaultdict(list)
 
         section_dirs: list[Path] = [
             Path(config["custom"]["sections_dir"]),
@@ -629,10 +630,6 @@ class Macsecurityrule(BaseModelWithAccessors):
                             Sectionmap[folder_name.upper()], ""
                         )
                         logger.debug("Section Name: {}", section_name)
-                        custom: bool = False
-
-                        if "custom" in str(folder).lower():
-                            custom = True
 
                         if not section_name:
                             logger.warning(
@@ -640,21 +637,22 @@ class Macsecurityrule(BaseModelWithAccessors):
                             )
                             continue
 
-                        rules += cls.load_rules(
-                            rule_ids=[rule_name],
-                            os_type=os_type,
-                            os_version=os_version,
-                            parent_values=parent_values,
-                            section=section_name,
-                            custom=custom,
-                            generate_baseline=generate_baseline,
-                        )
+                        rules_to_collect[section_name].append(rule_name)
 
                     except Exception as e:
                         logger.error(
                             "Failed to load rule from file {}: {}", rule_file, e
                         )
 
+        for section, collected_rules in rules_to_collect.items():
+            rules += cls.load_rules(
+                rule_ids=collected_rules,
+                os_type=os_type,
+                os_version=os_version,
+                parent_values=parent_values,
+                section=section,
+                tailoring=tailoring,
+            )
         logger.info("=== ALL RULES LOADED ===")
 
         return rules
@@ -967,7 +965,7 @@ class Macsecurityrule(BaseModelWithAccessors):
 
     def write_odv_custom_rule(self, odv: Any) -> None:
         """
-        Writes a custom ODV (Object Data Value) rule to a YAML file.
+        Writes a custom ODV (Organization Defined Value) rule to a YAML file.
 
         Args:
             odv (Any): The custom ODV data to be written.
@@ -985,11 +983,11 @@ class Macsecurityrule(BaseModelWithAccessors):
 
         self.customized = []
 
-        self.to_yaml(rule_file_path)
+        self.to_yaml(rule_file_path, odv=True)
 
     def remove_odv_custom_rule(self) -> None:
         """
-        Removes the custom rule from the ODV (Object Data Value) and updates the corresponding YAML file.
+        Removes the custom rule from the ODV (Organization Defined Value) and updates the corresponding YAML file.
 
         This method performs the following steps:
         1. Constructs the file path for the custom rule YAML file based on the configuration and rule ID.
@@ -1004,17 +1002,11 @@ class Macsecurityrule(BaseModelWithAccessors):
             f"{config['custom']['rules_dir']}", f"{self.rule_id}.yaml"
         )
 
-        if not rule_file_path.exists():
+        try:
+            rule_file_path.unlink()
+            logger.info("Custome rule file deleted: {}", rule_file_path)
+        except FileNotFoundError:
             logger.warning("Rule file not found: {}", rule_file_path)
-            return
-
-        if self.odv is not None and "custom" in self.odv:
-            self.odv.pop("custom", 0)
-
-        # if "references" in self.customized:
-        #     self.references.pop("custom_refs")
-
-        self.to_yaml(rule_file_path)
 
     @classmethod
     def odv_query(
@@ -1058,7 +1050,6 @@ class Macsecurityrule(BaseModelWithAccessors):
                     include = "y"
                     get_odv = True
                     queried_rule_ids.append(rule.rule_id)
-                    rule.remove_odv_custom_rule()
             else:
                 if rule.rule_id not in queried_rule_ids:
                     include = sanitize_input(
@@ -1077,7 +1068,6 @@ class Macsecurityrule(BaseModelWithAccessors):
                         )
                     queried_rule_ids.append(rule.rule_id)
                     get_odv = True
-                    rule.remove_odv_custom_rule()
                     if include.lower() == "all":
                         include_all = True
                         include = "y"
@@ -1087,6 +1077,8 @@ class Macsecurityrule(BaseModelWithAccessors):
                 if rule.odv == "missing":
                     continue
                 elif get_odv and rule.odv:
+                    # remove custom odv if there
+                    rule.remove_odv_custom_rule()
                     odv_hint = rule.odv.get("hint", "")
                     odv_recommended = rule.odv.get("recommended")
                     odv_benchmark = rule.odv.get(benchmark)
@@ -1111,7 +1103,7 @@ class Macsecurityrule(BaseModelWithAccessors):
 
         return included_rules
 
-    def to_yaml(self, output_path: Path) -> None:
+    def to_yaml(self, output_path: Path, odv: bool = False) -> None:
         key_order: list[str] = [
             "id",
             "title",
@@ -1161,11 +1153,21 @@ class Macsecurityrule(BaseModelWithAccessors):
             if key in serialized_data:
                 ordered_data[key] = serialized_data[key]
 
-        clean_dict: dict = {
-            key: value
-            for key, value in ordered_data.items()
-            if value or key in required_keys
-        }
+        if odv:
+            odv_fields = ["hint", "custom"]
+            clean_dict: dict = {
+                key: value for key, value in ordered_data.items() if key == "odv"
+            }
+            for key in list(clean_dict["odv"].keys()):
+                if key not in odv_fields:
+                    del clean_dict["odv"][key]
+
+        else:
+            clean_dict: dict = {
+                key: value
+                for key, value in ordered_data.items()
+                if value or key in required_keys
+            }
 
         create_yaml(rule_file_path, clean_dict)
 

src/mscp/generate/baseline.py

@@ -143,7 +143,7 @@ def generate_baseline(args: argparse.Namespace) -> None:
         make_dir(build_path)
 
     all_rules: list[Macsecurityrule] = Macsecurityrule.collect_all_rules(
-        args.os_name, args.os_version, parent_values="Default"
+        args.os_name, args.os_version, args.tailor, parent_values="Default"
     )
 
     all_tags, benchmark_map = collect_tags_and_benchmarks(all_rules)
@@ -234,7 +234,9 @@ def generate_baseline(args: argparse.Namespace) -> None:
         )
         custom_author_name: str = sanitize_input("Enter your name: ")
         custom_author_org: str = sanitize_input("Enter your organization: ")
-        baseline_output_file = build_path / f"{tailored_filename}.yaml"
+        baseline_output_file: Path = (
+            build_path / f"{tailored_filename}_{args.os_name}_{args.os_version}.yaml"
+        )
         authors.append(Author(name=custom_author_name, organization=custom_author_org))
 
         if tailored_filename == args.keyword: