From b19b18dc3e4f064104e484fa4f460b26f1a17ef2 Mon Sep 17 00:00:00 2001 From: Greg Elin Date: Mon, 26 Oct 2015 16:06:46 -0400 Subject: [PATCH] Handle no control_enhancement case --- compliancelib/seccontrol.py | 9 ++++++--- compliancelib/tests/test_seccontrol.py | 5 +++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/compliancelib/seccontrol.py b/compliancelib/seccontrol.py index 02372e0..88de4d6 100644 --- a/compliancelib/seccontrol.py +++ b/compliancelib/seccontrol.py @@ -75,9 +75,12 @@ def _load_control_from_xml(self): self.description = ''.join(sc.find('{http://scap.nist.gov/schema/sp800-53/2.0}statement').itertext()) self.description = re.sub(r'[ ]{2,}','',re.sub(r'^[ ]', '',re.sub(r'\n','',re.sub(r'[ ]{2,}',' ',self.description)))) self.description = self.description.replace(self.id, '\n').strip() - self.control_enhancements = ''.join(sc.find('{http://scap.nist.gov/schema/sp800-53/2.0}control-enhancements').itertext()) - self.control_enhancements = re.sub(r'[ ]{2,}','',re.sub(r'^[ ]', '',re.sub(r'[\n ]{2,}','\n',re.sub(r'[ ]{2,}',' ',self.control_enhancements)))) - # self.control_enhancements = self.control_enhancements.replace(self.id, '\n') + if (sc.find('{http://scap.nist.gov/schema/sp800-53/2.0}control-enhancements')) is not None: + self.control_enhancements = ''.join(sc.find('{http://scap.nist.gov/schema/sp800-53/2.0}control-enhancements').itertext()) + self.control_enhancements = re.sub(r'[ ]{2,}','',re.sub(r'^[ ]', '',re.sub(r'[\n ]{2,}','\n',re.sub(r'[ ]{2,}',' ',self.control_enhancements)))) + # self.control_enhancements = self.control_enhancements.replace(self.id, '\n') + else: + self.control_enhancements = None self.sg = sc.find('{http://scap.nist.gov/schema/sp800-53/2.0}supplemental-guidance') self.supplemental_guidance = self.sg.find('{http://scap.nist.gov/schema/sp800-53/2.0}description').text.strip() related_controls = [] diff --git a/compliancelib/tests/test_seccontrol.py b/compliancelib/tests/test_seccontrol.py index a24ca60..4fc1fda 100644 --- a/compliancelib/tests/test_seccontrol.py +++ b/compliancelib/tests/test_seccontrol.py @@ -45,6 +45,11 @@ def test_details_nonexistent_control(self): c = SecControl(id) self.assertTrue(c.title == None) + def test_nonexistent_control_enhancements(self): + id = "AC-1" + c = SecControl(id) + self.assertTrue(c.control_enhancements == None) + def test_supplemental_guidance(self): id = "AC-16" c = SecControl(id)