XPath Control Origination fix

ssp/control.py

@@ -28,11 +28,16 @@ def __iter__(self):
 
     def get_parts(self):
         self.parts = []
+        if bool(re.match(r'\w+-\d',self.implementation_table.cell(1,0).text)):
+            self.parts.append('Main')
         for row in self.implementation_table.rows:
             if "Part" in row.cells[0].text and len(row.cells[0].text) < 8:
                 self.parts.append(row.cells[0].text.replace('Part ', '').strip())
             elif 'Req.' in row.cells[0].text and len(row.cells[0].text) < 8:
                 self.parts.append(row.cells[0].text.replace('Req. ', '').strip())
+            elif 'Ext.' in row.cells[0].text and len(row.cells[0].text) < 8:
+                self.parts.append(row.cells[0].text.strip())
+
 
         if len(self.parts) < 1 and len(self.implementation_table.rows) < 3:
             self.parts.append(None)
@@ -46,6 +51,11 @@ def part(self, part_id):
                 return self.implementation_table.cell(1,0)
             except IndexError:
                 raise IndexError('Control table %s does not have enough rows to contain a control response.' % (self.number))
+        elif part_id.upper() == 'Main'.upper():
+            try:
+                return self.implementation_table.cell(1,1)
+            except IndexError:
+                raise IndexError('Control table %s does not have enough rows to contain a control response.' % (self.number))
 
         if part_id in self.parts:
             split_alpha_part = re.compile(r'\w\d')
@@ -56,6 +66,12 @@ def part(self, part_id):
                     return self.implementation_table.cell(self.LETTERS[letter]+row_difference, 1)
                 except:
                     ValueError('Control does not have part ' + part_id)
+            elif bool(re.match(r'\w{3,}', part_id)):
+                try:
+                    row = [cell.text.strip() for cell in self.implementation_table.columns[0].cells].index(part_id)
+                    return self.implementation_table.cell(row, 1)
+                except ValueError:
+                    raise ValueError('Control does not have part ' + part_id)
             else:
                 try:
                     return self.implementation_table.cell(self.LETTERS[part_id], 1)
@@ -109,6 +125,18 @@ def get_implementation_status(self, cell):
                     self.implementation_status.append('Alternative Implementation')
                 elif 'Not' in implementation_status:
                     self.implementation_status.append('Not Applicable')
+            elif '☒' in paragraph.text:
+                implementation_status = paragraph.text.strip()
+                if 'Partially' in implementation_status:
+                    self.implementation_status.append('Partially Implemented')
+                elif 'Implemented' in implementation_status:
+                    self.implementation_status.append('Implemented')
+                elif 'Planned' in implementation_status:
+                    self.implementation_status.append('Planned')
+                elif 'Alternative' in implementation_status:
+                    self.implementation_status.append('Alternative Implementation')
+                elif 'Not' in implementation_status:
+                    self.implementation_status.append('Not Applicable')
 
     def get_control_origination(self, cell):
         """
@@ -119,9 +147,9 @@ def get_control_origination(self, cell):
             p = paragraph._element
             if 'w14:checked w14:val="1"' in p.xml:
                 xpath_elements = p.xpath('.//w:t')
-                control_origination = xpath_elements[len(xpath_elements)-1].text.strip()
-                if not control_origination:
-                    control_origination = xpath_elements[len(xpath_elements)-2].text.strip() #TODO: this is really ugly, but had to be done because inherited checkboxes werent being captured.
+                for text in [element.text.strip() for element in xpath_elements]:
+                    if any(origination in text for origination in ['Service Provider', 'Inherited', 'Not', 'Customer']):
+                        control_origination = xpath_elements[1].text.strip()
                 if "Service Provider Corporate" in control_origination:
                     self.control_origination.append("Service Provider Corporate")
                 elif "Service Provider System Specific" in control_origination:
@@ -156,3 +184,26 @@ def get_control_origination(self, cell):
                     self.control_origination.append("Inherited")
                 elif "Not" in control_origination:
                     self.control_origination.append("Not Applicable")
+            elif '☒' in paragraph.text:
+                control_origination = paragraph.text.strip()
+                if "Service Provider Corporate" in control_origination:
+                    self.control_origination.append("Service Provider Corporate")
+                elif "Service Provider System Specific" in control_origination:
+                    self.control_origination.append("Service Provider System Specific")
+                elif "Hybrid" in control_origination:
+                    self.control_origination.append("Service Provider Hybrid")
+                elif "Configured" in control_origination:
+                    self.control_origination.append("Configured by Customer")
+                elif "Provided" in control_origination:
+                    self.control_origination.append("Provided by Customer")
+                elif "Shared" in control_origination:
+                    self.control_origination.append("Shared")
+                elif "Inherited" in control_origination:
+                    self.control_origination.append("Inherited")
+                elif "Not" in control_origination:
+                    self.control_origination.append("Not Applicable")
+
+#class Part():
+    #def __init__(self, id, implementation):
+        #self.id = id
+        #self.implementation = implementation