Fix count_oval_object get wrong count

utils/count_oval_objects.py

@@ -28,14 +28,31 @@ def load_xml(file_name):
     try:
         it = ET.iterparse(file_name)
         for _, el in it:
-            el.tag = el.tag.split('}', 1)[1]  # strip all namespaces
+            if '}'in el.tag:
+                el.tag = el.tag.split('}', 1)[1]  # strip all namespaces
         root = it.root
         return root
     except:
         sys.stderr.write("Error while loading file " + file_name + ".\n")
         exit(-1)
 
 
+def get_ext_def_tests(oval_root, def_refs):
+    t = []
+    for d in oval_root.findall(".//definition"):
+        if d.attrib.get('id') == def_refs:
+            definition = d
+            break
+    if definition is not None:
+        for criterion in definition.findall(".//criterion"):
+            test_ref = criterion.attrib["test_ref"]
+            t.append(test_ref)
+        for extend_def in definition.findall(".//extend_definition"):
+            extend_ref = extend_def.attrib["definition_ref"]
+            t = t + (get_ext_def_tests(oval_root, extend_ref))
+    return t
+
+
 def find_oval_objects(oval_refs):
     ''' Finds OVAL objects according to definitions ID '''
     tests = []
@@ -54,29 +71,38 @@ def find_oval_objects(oval_refs):
                 definition = d
                 break
         if definition is not None:
+            extend_refs = []
             for criterion in definition.findall(".//criterion"):
                 test_ref = criterion.attrib["test_ref"]
                 tests.append(test_ref)
+            for extend_def in definition.findall(".//extend_definition"):
+                extend_ref = extend_def.attrib["definition_ref"]
+                t = get_ext_def_tests(oval_root, extend_ref)
+                tests += t
 
     # find references to objects in tests
-    for test in tests:
-        test_element = None
-        for t in oval_root.findall("tests/*"):
-            if t.attrib.get('id') == test:
-                test_element = t
-                break
-        if test_element is not None:
-            for object_element in test_element.findall(".//*"):
-                if 'object_ref' in object_element.attrib:
-                    object_ref = object_element.attrib['object_ref']
-                    object_refs.append(object_ref)
+    for key in oval_files:
+        oval_root = oval_files[key]
+        for test in tests:
+            test_element = None
+            for t in oval_root.findall("tests/*"):
+                if t.attrib.get('id') == test:
+                    test_element = t
+                    break
+            if test_element is not None:
+                for object_element in test_element.findall(".//*"):
+                    if 'object_ref' in object_element.attrib:
+                        object_ref = object_element.attrib['object_ref']
+                        object_refs.append(object_ref)
 
     # find objects
-    for r in object_refs:
-        for obj in oval_root.findall("objects/*"):
-            if obj.attrib.get('id') == r:
-                objects.append(obj.tag)
-                break
+    for key in oval_files:
+        oval_root = oval_files[key]
+        for r in object_refs:
+            for obj in oval_root.findall("objects/*"):
+                if obj.attrib.get('id') == r:
+                    objects.append(obj.tag)
+                    break
 
     return set(objects)