From ce132fb4045b021387272b3e49fcd68f4e3dc86e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20=C4=8Cern=C3=BD?= <jcerny@redhat.com>
Date: Thu, 15 Feb 2024 15:13:33 +0100
Subject: [PATCH] Add stigrefs after references from controls

Move the step of adding a stigref to a later stage in the process,
specifically from the moment rule.yml files are loaded to the time when
references from controls are already added to rules.

This change will allow us in future to use control file as the source of
stigids. Up until now we can't do it because adding stigrefs depend on
existence of stigid key in references in the rule object. If we want to
add stigids from control files we need to add stigrefs after stigids are
added, which means after controls are processed instead of during
original rule.yml files.
---
 build-scripts/compile_all.py | 12 +++++++++++-
 ssg/build_yaml.py            |  9 +--------
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/build-scripts/compile_all.py b/build-scripts/compile_all.py
index 4dbb8b3de56..7c8020ed82f 100644
--- a/build-scripts/compile_all.py
+++ b/build-scripts/compile_all.py
@@ -121,6 +121,14 @@ def find_existing_rules(project_root):
     return rules
 
 
+def add_stig_references(stig_reference_path, all_rules):
+    if not stig_reference_path:
+        return
+    stig_references = ssg.build_stig.map_versions_to_rule_ids(stig_reference_path)
+    for rule in all_rules:
+        rule.add_stig_references(stig_references)
+
+
 def main():
     parser = create_parser()
     args = parser.parse_args()
@@ -144,7 +152,7 @@ def main():
     product_cpes.load_content_cpes(env_yaml)
 
     loader = ssg.build_yaml.BuildLoader(
-        None, env_yaml, product_cpes, args.sce_metadata, args.stig_references)
+        None, env_yaml, product_cpes, args.sce_metadata)
     loader.load_components()
     load_benchmark_source_data_from_directory_tree(loader, env_yaml, product_yaml)
 
@@ -158,6 +166,8 @@ def main():
     controls_manager.remove_selections_not_known(loader.all_rules)
     controls_manager.add_references(loader.all_rules)
 
+    add_stig_references(args.stig_references, loader.all_rules.values())
+
     profiles_by_id = get_all_resolved_profiles_by_id(
         env_yaml, product_yaml, loader, product_cpes, controls_manager, controls_dir)
 
diff --git a/ssg/build_yaml.py b/ssg/build_yaml.py
index 312f81e5958..804e4f9ff71 100644
--- a/ssg/build_yaml.py
+++ b/ssg/build_yaml.py
@@ -1366,15 +1366,12 @@ def save_entities(self, entities, destdir):
 class BuildLoader(DirectoryLoader):
     def __init__(
             self, profiles_dir, env_yaml, product_cpes,
-            sce_metadata_path=None, stig_reference_path=None):
+            sce_metadata_path=None):
         super(BuildLoader, self).__init__(profiles_dir, env_yaml, product_cpes)
 
         self.sce_metadata = None
         if sce_metadata_path and os.path.getsize(sce_metadata_path):
             self.sce_metadata = json.load(open(sce_metadata_path, 'r'))
-        self.stig_references = None
-        if stig_reference_path:
-            self.stig_references = ssg.build_stig.map_versions_to_rule_ids(stig_reference_path)
         self.components_dir = None
         self.rule_to_components = None
 
@@ -1405,8 +1402,6 @@ def _process_rule(self, rule):
         self.loaded_group.add_rule(
             rule, env_yaml=self.env_yaml, product_cpes=self.product_cpes)
         rule.normalize(self.env_yaml["product"])
-        if self.stig_references:
-            rule.add_stig_references(self.stig_references)
         if self.rule_to_components is not None:
             rule.components = self.rule_to_components[rule.id_]
         return True
@@ -1427,8 +1422,6 @@ def _get_new_loader(self):
             self.profiles_dir, self.env_yaml, self.product_cpes)
         # Do it this way so we only have to parse the SCE metadata once.
         loader.sce_metadata = self.sce_metadata
-        # Do it this way so we only have to parse the STIG references once.
-        loader.stig_references = self.stig_references
         # Do it this way so we only have to parse the component metadata once.
         loader.rule_to_components = self.rule_to_components
         return loader