import os
import unittest

import io
from xml.etree import ElementTree

import pyaxml

test_dir = os.path.dirname(os.path.abspath(__file__))


def is_valid_manifest(tree):
    # We can not really check much more...
    if tree.tag == "manifest" and "package" in tree.attrib:
        return True
    return False


# text_compare and xml_compare are modified from
# https://bitbucket.org/ianb/formencode/src/tip/formencode/doctest_xml_compare.py
def text_compare(t1, t2):
    if not t1 and not t2:
        return True
    if t1 == "*" or t2 == "*":
        return True
    return (t1 or "").strip() == (t2 or "").strip()


def xml_compare(x1, x2, reporter=None):
    """
    Compare two XML files
    x1 must be the plain version
    x2 must be the version generated by aapt
    """
    if x1.tag != x2.tag:
        if reporter:
            reporter("Tags do not match: {} and {}".format(x1.tag, x2.tag))
        return False
    for name, value in x1.attrib.items():
        if value[0] == "@" and x2.attrib.get(name)[0] == "@":
            # Can not be sure...
            pass
        elif x2.attrib.get(name) != value:
            if reporter:
                reporter(
                    "Attributes do not match: %s=%r, %s=%r"
                    % (name, value, name, x2.attrib.get(name))
                )
            return False
    for name in x2.attrib.keys():
        if name not in x1.attrib:
            if (
                x2.tag == "application"
                and name == "{http://schemas.android.com/apk/res/android}debuggable"
            ):
                # Debug attribute might be added by aapt
                pass
            else:
                if reporter:
                    reporter("x2 has an attribute x1 is missing: %s" % name)
                return False
    if not text_compare(x1.text, x2.text):
        if reporter:
            reporter("text: {!r} != {!r}".format(x1.text, x2.text))
        return False
    if not text_compare(x1.tail, x2.tail):
        if reporter:
            reporter("tail: {!r} != {!r}".format(x1.tail, x2.tail))
        return False
    cl1 = x1.getchildren()
    cl2 = x2.getchildren()
    if len(cl1) != len(cl2):
        if reporter:
            reporter("children length differs, %i != %i" % (len(cl1), len(cl2)))
        return False
    i = 0
    for c1, c2 in zip(cl1, cl2):
        i += 1
        if not xml_compare(c1, c2, reporter=reporter):
            if reporter:
                reporter("children %i do not match: %s" % (i, c1.tag))
            return False
    return True


class AXMLTest(unittest.TestCase):

    def testAndroidManifest(self):
        filenames = [
            "data/AXML/AndroidManifest.xml",
            "data/AXML/AndroidManifest-Chinese.xml",
            "data/AXML/AndroidManifestDoubleNamespace.xml",
            "data/AXML/AndroidManifestExtraNamespace.xml",
            "data/AXML/AndroidManifest_InvalidCharsInAttribute.xml",
            "data/AXML/AndroidManifestLiapp.xml",
            "data/AXML/AndroidManifestMaskingNamespace.xml",
            "data/AXML/AndroidManifest_NamespaceInAttributeName.xml",
            "data/AXML/AndroidManifest_NamespaceInAttributeName2.xml",
            "data/AXML/AndroidManifestNonZeroStyle.xml",
            "data/AXML/AndroidManifestNullbytes.xml",
            "data/AXML/AndroidManifestTextChunksXML.xml",
            # "data/AXML/AndroidManifestUTF8Strings.xml",
            "data/AXML/AndroidManifestWithComment.xml",
            "data/AXML/AndroidManifest_WrongChunkStart.xml",
            "data/AXML/AndroidManifest-xmlns.xml",
        ]

        for filename in filenames:
            with open(os.path.join(test_dir, filename), "rb") as fd:
                ap = pyaxml.AXML.from_axml(fd.read())[0]
            self.assertIsNotNone(ap)

    def testAXMLnotAligned(self):
        filenames = [
            "data/AXML/StringBlocksNotAligned.xml",
        ]

        for filename in filenames:
            with open(os.path.join(test_dir, filename), "rb") as fp:
                ap = pyaxml.AXML.from_axml(fp.read())[0]
                ap.compute()
                sb = pyaxml.StringBlocks(proto=ap.proto.stringblocks)
                sb_s = (
                    b"".join(
                        pyaxml.StringBlock(
                            proto=elt,
                            utf8=sb.proto.hnd.flag & pyaxml.axml_pb2.UTF8_FLAG
                            == pyaxml.axml_pb2.UTF8_FLAG,
                        ).pack()
                        for elt in sb.proto.stringblocks
                    )
                    + sb.proto.pad_stringblocks
                )
            self.assertEqual(len(sb_s) % 4, 0)

    def testUnpackingRepackingXMLFormCompute(self):
        filenames = [
            "data/AXML/AndroidManifest.xml",
            "data/AXML/AndroidManifest-Chinese.xml",
            "data/AXML/AndroidManifestDoubleNamespace.xml",
            "data/AXML/AndroidManifestExtraNamespace.xml",
            "data/AXML/AndroidManifestLiapp.xml",
            "data/AXML/AndroidManifestMaskingNamespace.xml",
            "data/AXML/AndroidManifest_NamespaceInAttributeName.xml",
            "data/AXML/AndroidManifest_NamespaceInAttributeName2.xml",
            "data/AXML/AndroidManifestNonZeroStyle.xml",
            "data/AXML/AndroidManifestTextChunksXML.xml",
            # "data/AXML/AndroidManifestUTF8Strings.xml", # TODO: check if attribute exist or not
            "data/AXML/AndroidManifestWithComment.xml",
            "data/AXML/AndroidManifest_WrongChunkStart.xml",
            "data/AXML/AndroidManifest-xmlns.xml",
            "data/AXML/StringBlocksNotAligned.xml",
        ]

        for filename in filenames:
            content = ""
            with open(os.path.join(test_dir, filename), "rb") as fd:
                ap = pyaxml.AXML.from_axml(fd.read())[0]
                content = ElementTree.tostring(ap.to_xml(), encoding="utf8")
                ap.compute()
                new = ap.pack()
                ap2 = pyaxml.AXML.from_axml(new)[0]
                content2 = ElementTree.tostring(ap2.to_xml(), encoding="utf8")
                self.assertEqual(content, content2)

    def testUnpackingRepackingXMLForm(self):
        filenames = [
            "data/AXML/AndroidManifest.xml",
            "data/AXML/AndroidManifest-Chinese.xml",
            "data/AXML/AndroidManifestDoubleNamespace.xml",
            "data/AXML/AndroidManifestExtraNamespace.xml",
            "data/AXML/AndroidManifestLiapp.xml",
            "data/AXML/AndroidManifestMaskingNamespace.xml",
            "data/AXML/AndroidManifest_NamespaceInAttributeName.xml",
            "data/AXML/AndroidManifest_NamespaceInAttributeName2.xml",
            "data/AXML/AndroidManifestNonZeroStyle.xml",
            "data/AXML/AndroidManifestTextChunksXML.xml",
            # "data/AXML/AndroidManifestUTF8Strings.xml", # TODO: check if attribute exist or not
            "data/AXML/AndroidManifestWithComment.xml",
            "data/AXML/AndroidManifest_WrongChunkStart.xml",
            "data/AXML/AndroidManifest-xmlns.xml",
        ]

        for filename in filenames:
            content = ""
            with open(os.path.join(test_dir, filename), "rb") as fd:
                ap = pyaxml.AXML.from_axml(fd.read())[0]
                content = ElementTree.tostring(ap.to_xml(), encoding="utf8")
                new = ap.pack()
                ap2 = pyaxml.AXML.from_axml(new)[0]
                content2 = ElementTree.tostring(ap2.to_xml(), encoding="utf8")
                self.assertEqual(content, content2)

    def testUnpackingRepacking(self):
        filenames = [
            "data/AXML/AndroidManifest.xml",
            "data/AXML/AndroidManifest-Chinese.xml",
            "data/AXML/AndroidManifestDoubleNamespace.xml",
            "data/AXML/AndroidManifestExtraNamespace.xml",
            "data/AXML/AndroidManifest_InvalidCharsInAttribute.xml",
            "data/AXML/AndroidManifestLiapp.xml",
            "data/AXML/AndroidManifestMaskingNamespace.xml",
            "data/AXML/AndroidManifest_NamespaceInAttributeName.xml",
            "data/AXML/AndroidManifest_NamespaceInAttributeName2.xml",
            "data/AXML/AndroidManifestNonZeroStyle.xml",
            "data/AXML/AndroidManifestNullbytes.xml",
            "data/AXML/AndroidManifestTextChunksXML.xml",
            "data/AXML/AndroidManifestUTF8Strings.xml",
            "data/AXML/AndroidManifestWithComment.xml",
            "data/AXML/AndroidManifest_WrongChunkStart.xml",
            "data/AXML/AndroidManifest-xmlns.xml",
            "data/AXML/StringBlocksNotAligned.xml",
        ]

        for filename in filenames:
            with open(os.path.join(test_dir, filename), "rb") as fd:
                old = fd.read()
                ap = pyaxml.AXML.from_axml(old)[0]
                new = ap.pack()
                self.assertEqual(old, new)

    def testNonManifest(self):
        filenames = [
            "data/AXML/test.xml",
            "data/AXML/test1.xml",
            "data/AXML/test2.xml",
            "data/AXML/test3.xml",
        ]

        for filename in filenames:
            with open(os.path.join(test_dir, filename), "rb") as fp:
                ap = pyaxml.AXML.from_axml(fp.read())[0]

            self.assertEqual(ap.to_xml().tag, "LinearLayout")

    def testNonZeroStyleOffset(self):
        """
        Test if a nonzero style offset in the string section causes problems
        if the counter is 0
        """
        filename = "data/AXML/AndroidManifestNonZeroStyle.xml"

        with open(os.path.join(test_dir, filename), "rb") as f:
            ap = pyaxml.AXML.from_axml(f.read())[0].to_xml()

    def testNonTerminatedString(self):
        """
        Test if non-null terminated strings are detected.
        This sample even segfaults aapt...
        """
        filename = "data/AXML/AndroidManifest_StringNotTerminated.xml"

        # TODO
        # with open(os.path.join(test_dir, filename), "rb") as f:
        #    ap = pyaxml.AXML.from_axml(f.read())[0].to_xml()

    def testExtraNamespace(self):
        """
        Test if extra namespaces cause problems
        """
        filename = "data/AXML/AndroidManifestExtraNamespace.xml"

        with open(os.path.join(test_dir, filename), "rb") as f:
            ap = pyaxml.AXML.from_axml(f.read())[0].to_xml()

    def testTextChunksWithXML(self):
        """
        Test for Text chunks containing XML
        """
        filename = "data/AXML/AndroidManifestTextChunksXML.xml"

        # TODO
        # with open(os.path.join(test_dir, filename), "rb") as f:
        #    ap = pyaxml.AXML.from_axml(f.read())[0].to_xml()

    def testWrongFilesize(self):
        """
        Assert that files with a broken filesize are not parsed
        """
        filename = "data/AXML/AndroidManifestWrongFilesize.xml"

        try:
            with open(os.path.join(test_dir, filename), "rb") as f:
                ap = pyaxml.AXML.from_axml(f.read())[0].to_xml()
        except IndexError:
            return
        self.assertTrue(False, "should raise an error")

    def testNullbytes(self):
        """
        Assert that Strings with nullbytes are handled correctly
        """
        filename = "data/AXML/AndroidManifestNullbytes.xml"

        try:
            with open(os.path.join(test_dir, filename), "rb") as f:
                ap = pyaxml.AXML.from_axml(f.read())[0].to_xml()
        except ValueError as e:
            self.assertEqual(
                str(e),
                "All strings must be XML compatible: Unicode or ASCII, no NULL bytes or control characters",
            )
            return
        self.assertTrue(False, "should raise an error")

    def testMaskingNamespace(self):
        """
        Assert that Namespaces which are used in a tag and the tag is closed
        are actually correctly parsed.
        """
        filename = "data/AXML/AndroidManifestMaskingNamespace.xml"

        with open(os.path.join(test_dir, filename), "rb") as f:
            ap = pyaxml.AXML.from_axml(f.read())[0].to_xml()

    def testDoubleNamespace(self):
        """
        Test if weird namespace constelations cause problems
        """
        filename = "data/AXML/AndroidManifestDoubleNamespace.xml"

        with open(os.path.join(test_dir, filename), "rb") as f:
            ap = pyaxml.AXML.from_axml(f.read())[0].to_xml()

    def testPackers(self):
        """
        Assert that Packed files are read
        """
        filename = "data/AXML/AndroidManifestLiapp.xml"

        with open(os.path.join(test_dir, filename), "rb") as f:
            ap = pyaxml.AXML.from_axml(f.read())[0].to_xml()


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