From 927aeafe9f6f129c993191ecf140c88d5b0b0796 Mon Sep 17 00:00:00 2001
From: Watson Sato <wsato@redhat.com>
Date: Tue, 16 Jul 2024 15:59:43 +0200
Subject: [PATCH] Allow multiple control files to add same ref type

The build system is currently limited to a single ref type per control
file. Two cotnrol files cannot add the same reference type.

With this commit, the references added by the control file are tracked
separately from the references loaded from the rule.yml.
This allows us to differentiate references coming from the rule, and
references coming from the control file.
---
 ssg/build_yaml.py | 18 +++++++++++++-----
 ssg/controls.py   | 11 ++++++++++-
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/ssg/build_yaml.py b/ssg/build_yaml.py
index 5f0d3a41d77..7e55341838b 100644
--- a/ssg/build_yaml.py
+++ b/ssg/build_yaml.py
@@ -893,6 +893,7 @@ class Rule(XCCDFEntity, Templatable):
         rationale=lambda: "",
         severity=lambda: "",
         references=lambda: dict(),
+        control_references=lambda: dict(),
         components=lambda: list(),
         identifiers=lambda: dict(),
         ocil_clause=lambda: None,
@@ -1212,16 +1213,23 @@ def _add_ident_elements(self, rule):
             ident.set("system", SSG_IDENT_URIS[ident_type])
             ident.text = ident_val
 
-    def add_extra_reference(self, ref_type, ref_value):
-        if ref_type in self.references:
-            if ref_value in self.references[ref_type]:
+    def add_control_reference(self, ref_type, ref_value):
+        if ref_type in self.control_references:
+            if ref_value in self.control_references[ref_type]:
                 msg = (
                     "Rule %s already contains a '%s' reference with value '%s'." % (
                         self.id_, ref_type, ref_value))
                 raise ValueError(msg)
-            self.references[ref_type].append(ref_value)
+            self.control_references[ref_type].append(ref_value)
         else:
-            self.references[ref_type] = [ref_value]
+            self.control_references[ref_type] = [ref_value]
+
+    def merge_control_references(self):
+        for ref_type in self.control_references:
+            if ref_type in self.references:
+                self.references[ref_type].append(self.control_references[ref_type])
+            else:
+                self.references[ref_type] = self.control_references[ref_type]
 
     def to_xml_element(self, env_yaml=None):
         rule = ET.Element('{%s}Rule' % XCCDF12_NS)
diff --git a/ssg/controls.py b/ssg/controls.py
index 68268dbb7ed..80f8f583b5d 100644
--- a/ssg/controls.py
+++ b/ssg/controls.py
@@ -184,7 +184,7 @@ def add_references(self, reference_type, rules):
             if not rule:
                 continue
             try:
-                rule.add_extra_reference(reference_type, self.id)
+                rule.add_control_reference(reference_type, self.id)
             except ValueError as exc:
                 msg = (
                     "Please remove any duplicate listing of rule '%s' in "
@@ -527,5 +527,14 @@ def save_everything(self, output_dir):
             policy.dump_yaml(filename)
 
     def add_references(self, rules):
+        # First we add all control references into a separate attribute
         for policy in self.policies.values():
             policy.add_references(rules)
+        # Then we unify them under references attribute
+        # This allows multiple control files to add references of the same type, and still track
+        # what references already existed in the rule.
+        self._merge_references(rules)
+
+    def _merge_references(self, rules):
+        for rule in rules.values():
+            rule.merge_control_references()