fix(pain001): updated against XML vulnerabilities

pain001/xml/create_root_element.py

@@ -13,31 +13,45 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import xml.etree.ElementTree as ET
+# pylint: disable=C0301
+"""
+Module for creating secure XML payment initiation message documents.
 
-# Create the root element and set its attributes (XML tags and CSV
-# columns mapping)
+Uses the defusedxml library to prevent XML vulnerabilities.
+"""
 
+import defusedxml.ElementTree as et
 
-def create_root_element(payment_initiation_message_type):
-    # Create the namespace for the payment initiation message type.
-    namespace = (
-        "urn:iso:std:iso:20022:tech:xsd:"
-        + payment_initiation_message_type
-    )
+NAMESPACE = "urn:iso:std:iso:20022:tech:xsd:"
+
+XSI_NAMESPACE = "http://www.w3.org/2001/XMLSchema-instance"
+
+def create_payment_initiation_root(message_type: str) -> et.XML:
+    """
+    Create the root Element for a payment initiation XML document.
+
+    Args:
+        message_type: The message type string, e.g. "pain.001.001.09".
 
-    # Create the root element.
-    root = ET.Element("Document")
-    root.set("xmlns", namespace)
-    root.set("xmlns:xsi", "http://www.w3.org/2001/XMLSchema-instance")
+    Returns:
+        The root Element node for the XML document.
+    """
 
-    # Set the schema location.
-    schema_location = (
-        namespace + " " + payment_initiation_message_type + ".xsd"
+    parser = et.DefusedXMLParser()
+
+    processing_instruction = "version=\"1.0\" encoding=\"UTF-8\""
+    root = parser.ProcessingInstruction("xml", processing_instruction)
+
+    root = parser.Entity("Document",
+        xmlns=NAMESPACE + message_type,
+        xmlns_xsi=XSI_NAMESPACE
     )
+
+    schema_location = f"{NAMESPACE}{message_type} {message_type}.xsd"
     root.set("xsi:schemaLocation", schema_location)
 
+    # Remove namespaces from child elements
     for elem in root.iter():
-        elem.tag = elem.tag.split("}", 1)[-1]
+        elem.tag = elem.tag.split('}', 1)[-1]
 
-    return root
+    return root