# Copyright (c) 2017 VMware, Inc.  All rights reserved.
# -- VMware Confidential

""" XML editing made simple.

This simple XML editing wrapper is made only with Python 3 in mind.

The Reverse Proxy B2B script needs to switch to TLS1.2 in a 6.5 to 6.6 upgrade.
We have two ways of achieving this:
- By explicitly setting vmacore/ssl/protocols and possibly others to tls1.2.
- By unsetting the config tag vmacore/ssl/protocols and possibly others.

Both XML operations are supported by the current XML wrapper.

@see Bug 1915931 for more information.

@todo Consider using the XML wrapper of the TLS Reconfigurator when it becomes
      standalone and stop depending on the 'vmware' package:
      bora/vim/apps/vSphereTlsReconfigurator/VcTlsReconfigurator/vmware/xml.py

@note In Python3 ET.XMLParser is overwritten in xml.etree.ElementTree.py
      when the _elementtree module is loaded which we think is incorrect.
      Therefore we create an empty _elementtree module in order to cheat Python
      that it's already loaded.
"""

import imp
import sys
sys.modules['_elementtree'] = imp.new_module('_elementtree')

import xml.etree.ElementTree as ET


class _XmlParser(ET.XMLParser):
    """Keeps the comments in the edited XML files.

    Unfortunately, this doesn't work on VCSA because there is no self.parser.
    """

    def __init__(self):
        ET.XMLParser.__init__(self)
        # assumes ElementTree 1.2.X
        self.parser.CommentHandler = self.HandleComment

    def HandleComment(self, data):
        self._target.start(ET.Comment, {})
        self._target.data(data)
        self._target.end(ET.Comment)


def Parse(string):
    "Creates an XML tree from a given string."
    return ET.XML(string, parser=_XmlParser())


def Serialize(root):
    "Serializes an XML tree to string."
    return ET.tostring(root, encoding="unicode")


def _AppendIndent(element, depth, unitIndent=" " * 3):
    """Adds the necessary identation to a given element.

    Looks like the element.tail is the only place to apply indentation.
    """

    children = list(element)

    if children:
        last = children[-1]
        indent = last.tail
        last.tail = element.text
        return indent

    element.text = '\n' + (depth + 1) * unitIndent
    return '\n' + depth * unitIndent


def SetText(root, path, value):
    """Adds an element with a given value to the tree.

    The path to this element is a '/'-separated string of node names.
    For example, "vmacore/ssl/protocols".
    """

    element = root
    for i, tag in enumerate(path.split('/')):
        match = element.find(tag)
        if match is None:
            subelement = ET.Element(tag)
            subelement.tail = _AppendIndent(element, i)
            element.append(subelement)
            element = subelement
        else:
            element = match

    element.text = value


def UnsetPath(root, path):
    """Removes all occurances of an element from an XML tree.

    Returns a list of the the removed values.
    """

    text = []
    cpm = {c: p for p in root.iter() for c in p}

    for match in root.findall(path):
        text.append(match.text)

        parent = cpm[match]
        children = list(parent)

        if match == children[-1]:  # last
            if len(children) > 1:  # has siblings
                children[-2].tail = match.tail
            else:
                parent.text = match.tail

        parent.remove(match)

    return text


import unittest


class TestEditXML(unittest.TestCase):
    "Unit tests the few cases necessary for RHTTPProxy B2B."

    def test_AddNewTag(self):
        "Add a non-existing element."

        src = """
<root>
   <boo>BOO</boo>
</root>
"""

        dst = """
<root>
   <boo>BOO</boo>
   <moo>
      <doo>DOO</doo>
   </moo>
</root>
""".strip()

        root = Parse(src)
        SetText(root, "moo/doo", "DOO")

        self.assertEqual(dst, Serialize(root))

    def test_ChangeExistingTag(self):
        "Modify an existing element."

        src = """
<root>
   <boo>BOO</boo>
</root>
"""

        dst = """
<root>
   <boo>DOO</boo>
</root>
""".strip()

        root = Parse(src)
        SetText(root, "boo", "DOO")

        self.assertEqual(dst, Serialize(root))

    def test_RemoveExistingTag(self):
        "Modify an existing element."

        src = """
<root>
   <boo>BOO</boo>
   <moo>
      <doo>DOO</doo>
      <foo>FOO</foo>
   </moo>
</root>
"""

        dst = """
<root>
   <boo>BOO</boo>
   <moo>
      <foo>FOO</foo>
   </moo>
</root>
""".strip()

        root = Parse(src)
        values = UnsetPath(root, "moo/doo")

        self.assertEqual(values, ["DOO"])
        self.assertEqual(dst, Serialize(root))


if __name__ == "__main__":
    unittest.main()
