From 435ad2a81f4d56da42a2b24cad9e0e04d0523fd4 Mon Sep 17 00:00:00 2001
From: Vojtech Polasek <vpolasek@redhat.com>
Date: Tue, 9 Jun 2020 18:55:42 +0200
Subject: [PATCH] fix for older python versions lacking .iter method

Signed-off-by: HuaxinLu <luhuaxin1@huawei.com>
---
 build-scripts/verify_references.py            |  8 +++++-
 .../pcidss/transform_benchmark_to_pcidss.py   |  8 +++++-
 ssg/build_cpe.py                              | 28 +++++++++++++++++--
 ssg/id_translate.py                           |  9 +++++-
 4 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/build-scripts/verify_references.py b/build-scripts/verify_references.py
index 95d387f46f..003aa4c48c 100755
--- a/build-scripts/verify_references.py
+++ b/build-scripts/verify_references.py
@@ -179,7 +179,13 @@ def main():
     check_content_refs = xccdftree.findall(".//{%s}check-content-ref"
                                            % xccdf_ns)
 
-    xccdf_parent_map = dict((c, p) for p in xccdftree.iter() for c in p)
+    # decide on usage of .iter or .getiterator method of elementtree class.
+    # getiterator is deprecated in Python 3.9, but iter is not available in
+    # older versions
+    if getattr(xccdftree, 'iter', None) == None:
+        xccdf_parent_map = dict((c, p) for p in xccdftree.getiterator() for c in p)
+    else:
+        xccdf_parent_map = dict((c, p) for p in xccdftree.iter() for c in p)
     # now we can actually do the verification work here
     if options.rules_with_invalid_checks or options.all_checks:
         for check_content_ref in check_content_refs:
diff --git a/shared/transforms/pcidss/transform_benchmark_to_pcidss.py b/shared/transforms/pcidss/transform_benchmark_to_pcidss.py
index c94b12c45c..05df5a9889 100755
--- a/shared/transforms/pcidss/transform_benchmark_to_pcidss.py
+++ b/shared/transforms/pcidss/transform_benchmark_to_pcidss.py
@@ -111,7 +111,13 @@ def main():
             benchmark.findall(".//{%s}Value" % (XCCDF_NAMESPACE)):
         values.append(value)
 
-    parent_map = dict((c, p) for p in benchmark.iter() for c in p)
+    # decide on usage of .iter or .getiterator method of elementtree class.
+    # getiterator is deprecated in Python 3.9, but iter is not available in
+    # older versions
+    if getattr(benchmark, "iter", None) == None:
+        parent_map = dict((c, p) for p in benchmark.getiterator() for c in p)
+    else:
+        parent_map = dict((c, p) for p in benchmark.iter() for c in p)
     for rule in \
             benchmark.findall(".//{%s}Rule" % (XCCDF_NAMESPACE)):
         parent_map[rule].remove(rule)
diff --git a/ssg/build_cpe.py b/ssg/build_cpe.py
index 8c046777a5..26f182395e 100644
--- a/ssg/build_cpe.py
+++ b/ssg/build_cpe.py
@@ -17,7 +17,14 @@ def extract_subelement(objects, sub_elem_type):
     """
 
     for obj in objects:
-        for subelement in obj.iter():
+        # decide on usage of .iter or .getiterator method of elementtree class.
+        # getiterator is deprecated in Python 3.9, but iter is not available in
+        # older versions
+        if getattr(obj, "iter", None) == None:
+            obj_iterator = obj.getiterator()
+        else:
+            obj_iterator = obj.iter()
+        for subelement in obj_iterator:
             if subelement.get(sub_elem_type):
                 sub_element = subelement.get(sub_elem_type)
                 return sub_element
@@ -44,12 +51,27 @@ def extract_referred_nodes(tree_with_refs, tree_with_ids, attrname):
     reflist = []
     elementlist = []
 
-    for element in tree_with_refs.iter():
+
+    # decide on usage of .iter or .getiterator method of elementtree class.
+    # getiterator is deprecated in Python 3.9, but iter is not available in
+    # older versions
+    if getattr(tree_with_refs, "iter", None) == None:
+        tree_with_refs_iterator = tree_with_refs.getiterator()
+    else:
+        tree_with_refs_iterator = tree_with_refs.iter()
+    for element in tree_with_refs_iterator:
         value = element.get(attrname)
         if value is not None:
             reflist.append(value)
 
-    for element in tree_with_ids.iter():
+    # decide on usage of .iter or .getiterator method of elementtree class.
+    # getiterator is deprecated in Python 3.9, but iter is not available in
+    # older versions
+    if getattr(tree_with_ids, "iter", None) == None:
+        tree_with_ids_iterator = tree_with_ids.getiterator()
+    else:
+        tree_with_ids_iterator = tree_with_ids.iter()
+    for element in tree_with_ids_iterator:
         if element.get("id") in reflist:
             elementlist.append(element)
 
diff --git a/ssg/id_translate.py b/ssg/id_translate.py
index ba92259049..520c47fbfd 100644
--- a/ssg/id_translate.py
+++ b/ssg/id_translate.py
@@ -64,7 +64,14 @@ class IDTranslator(object):
         )
 
     def translate(self, tree, store_defname=False):
-        for element in tree.iter():
+        # decide on usage of .iter or .getiterator method of elementtree class.
+        # getiterator is deprecated in Python 3.9, but iter is not available in
+        # older versions
+        if getattr(tree, "iter", None) == None:
+            tree_iterator = tree.getiterator()
+        else:
+            tree_iterator = tree.iter()
+        for element in tree_iterator:
             idname = element.get("id")
             if idname:
                 # store the old name if requested (for OVAL definitions)
-- 
2.30.0