From 4bfc0f1d9cfe21ec672fc806f5421272f1c0b41f Mon Sep 17 00:00:00 2001
From: Wesley Ceraso Prudencio <wcerasop@redhat.com>
Date: Wed, 1 Nov 2017 14:17:24 +0100
Subject: [PATCH] Enables the STIG Rule ID to be output

Signed-off-by: Wesley Ceraso Prudencio <wcerasop@redhat.com>
---
 cmake/SSGCommon.cmake               |  5 ++++
 shared/utils/add_stig_references.py | 57 +++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)
 create mode 100755 shared/utils/add_stig_references.py

diff --git a/cmake/SSGCommon.cmake b/cmake/SSGCommon.cmake
index 8ac826ef6..786e07532 100644
--- a/cmake/SSGCommon.cmake
+++ b/cmake/SSGCommon.cmake
@@ -130,10 +130,15 @@ macro(ssg_build_shorthand_xml PRODUCT)
 endmacro()
 
 macro(ssg_build_xccdf_unlinked PRODUCT)
+    file(GLOB STIG_REFERENCE_FILE_LIST "${SSG_SHARED_REFS}/disa-stig-${PRODUCT}-*-xccdf-manual.xml")
+    list(APPEND STIG_REFERENCE_FILE_LIST "not-found")
+    list(GET STIG_REFERENCE_FILE_LIST 0 STIG_REFERENCE_FILE)
+    
     add_custom_command(
         OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/xccdf-unlinked-resolved.xml"
         COMMAND "${XSLTPROC_EXECUTABLE}" --stringparam ssg_version "${SSG_VERSION}" --output "${CMAKE_CURRENT_BINARY_DIR}/xccdf-unlinked-resolved.xml" "${CMAKE_CURRENT_SOURCE_DIR}/transforms/shorthand2xccdf.xslt" "${CMAKE_CURRENT_BINARY_DIR}/shorthand.xml"
         COMMAND "${OPENSCAP_OSCAP_EXECUTABLE}" xccdf resolve -o "${CMAKE_CURRENT_BINARY_DIR}/xccdf-unlinked-resolved.xml" "${CMAKE_CURRENT_BINARY_DIR}/xccdf-unlinked-resolved.xml"
+        COMMAND "${SSG_SHARED_UTILS}/add_stig_references.py" --disa-stig "${STIG_REFERENCE_FILE}" --unlinked-xccdf "${CMAKE_CURRENT_BINARY_DIR}/xccdf-unlinked-resolved.xml"
         DEPENDS generate-internal-${PRODUCT}-shorthand.xml
         DEPENDS "${CMAKE_CURRENT_BINARY_DIR}/shorthand.xml"
         DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/transforms/shorthand2xccdf.xslt"
diff --git a/shared/utils/add_stig_references.py b/shared/utils/add_stig_references.py
new file mode 100755
index 000000000..0ab208793
--- /dev/null
+++ b/shared/utils/add_stig_references.py
@@ -0,0 +1,57 @@
+#!/usr/bin/env python2
+
+try:
+    from xml.etree import cElementTree as etree
+except ImportError:
+    import cElementTree as etree
+
+import re
+import sys
+import argparse
+
+parser = argparse.ArgumentParser(
+    description='Add STIG references to XCCDF files.')
+parser.add_argument(
+    "--disa-stig", help="DISA STIG Reference XCCDF file",dest="reference")
+parser.add_argument(
+    "--unlinked-xccdf", help="unlinked SSG XCCDF file", dest="destination")
+args = parser.parse_args()
+
+reference = args.reference
+destination = args.destination
+
+xccdf_namespace = "http://checklists.nist.gov/xccdf/1.1"
+stig_href = 'http://iase.disa.mil/stigs/Pages/stig-viewing-guidance.aspx'
+stig_references_beginning = 'http://iase.disa.mil/stigs/'
+
+try:
+    reference_root = etree.parse(reference)
+except IOError as exception:
+    print 'INFO: DISA STIG Reference file not found for this platform'
+    sys.exit(0)
+
+reference_rules = reference_root.findall('.//{%s}Rule' % xccdf_namespace)
+
+dictionary = {}
+
+for rule in reference_rules:
+    version = rule.find('.//{%s}version' % xccdf_namespace)
+    if version is not None and version.text:
+        dictionary[version.text] = rule.get('id')
+
+target_root = etree.parse(destination)
+target_rules = target_root.findall('.//{%s}Rule' % xccdf_namespace)
+
+for rule in target_rules:
+    refs = rule.findall('.//{%s}reference' % xccdf_namespace)
+    for ref in refs:
+        if (ref.get('href').startswith(stig_references_beginning) and
+                ref.text in dictionary):
+            index = rule.getchildren().index(ref)
+            new_ref = etree.Element(
+                '{%s}reference' % xccdf_namespace, {'href': stig_href})
+            new_ref.text = dictionary[ref.text]
+            new_ref.tail = ref.tail
+            rule.insert(index + 1, new_ref)
+
+target_root.write(destination)