import unittest
from StringIO import StringIO
from lxml import etree
import stix.bindings as bindings
import stix.utils.parser
from stix.test import EntityTestCase
from stix.extensions.malware.maec_4_1_malware import MAECInstance

class PythonMAECTests(EntityTestCase, unittest.TestCase):
    klass = MAECInstance

    _full_dict = {
        'xsi:type': 'stix-maec:MAEC4.1InstanceType',
        'maec': {
            'malware_subjects':
                [
                    {'malware_instance_object_attributes':
                         {'id': 'maec-tst-obj-1',
                          'properties': {
                              'hashes':
                                  [
                                      {
                                          'simple_hash_value': '9d7006e30fdf15e9c8e03e62534b3a3e',
                                          'type': 'MD5'
                                      }
                                  ],
                              'xsi:type': 'FileObjectType'}
                         }
                    }
                ]
        }
    }

class PythonMAECEtreeTests(unittest.TestCase):
    XML = (
        """
        <stix-maec:MAEC
            xmlns:cyboxCommon="http://cybox.mitre.org/common-2"
            xmlns:cybox="http://cybox.mitre.org/cybox-2"
            xmlns:cyboxVocabs="http://cybox.mitre.org/default_vocabularies-2"
            xmlns:FileObj="http://cybox.mitre.org/objects#FileObject-2"
            xmlns:example="http://example.com"
            xmlns:maecPackage="http://maec.mitre.org/XMLSchema/maec-package-2"
            xmlns:maecVocabs="http://maec.mitre.org/default_vocabularies-1"
            xmlns:ttp="http://stix.mitre.org/TTP-1"
            xmlns:stixCommon="http://stix.mitre.org/common-1"
            xmlns:stixVocabs="http://stix.mitre.org/default_vocabularies-1"
            xmlns:stix-maec="http://stix.mitre.org/extensions/Malware#MAEC4.1-1"
            xmlns:stix="http://stix.mitre.org/stix-1"
            xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
            xsi:schemaLocation="
            http://cybox.mitre.org/common-2 http://cybox.mitre.org/XMLSchema/common/2.1/cybox_common.xsd
            http://cybox.mitre.org/cybox-2 http://cybox.mitre.org/XMLSchema/core/2.1/cybox_core.xsd
            http://cybox.mitre.org/default_vocabularies-2 http://cybox.mitre.org/XMLSchema/default_vocabularies/2.1/cybox_default_vocabularies.xsd
            http://cybox.mitre.org/objects#FileObject-2 http://cybox.mitre.org/XMLSchema/objects/File/2.1/File_Object.xsd
            http://maec.mitre.org/XMLSchema/maec-package-2 http://maec.mitre.org/language/version4.1/maec_package_schema.xsd
            http://maec.mitre.org/default_vocabularies-1 http://maec.mitre.org/language/version4.1/maec_default_vocabularies.xsd
            http://stix.mitre.org/TTP-1 http://stix.mitre.org/XMLSchema/ttp/1.1.1/ttp.xsd
            http://stix.mitre.org/common-1 http://stix.mitre.org/XMLSchema/common/1.1.1/stix_common.xsd
            http://stix.mitre.org/default_vocabularies-1 http://stix.mitre.org/XMLSchema/default_vocabularies/1.1.1/stix_default_vocabularies.xsd
            http://stix.mitre.org/extensions/Malware#MAEC4.1-1 http://stix.mitre.org/XMLSchema/extensions/malware/maec_4.1/1.0.1/maec_4.1_malware.xsd
            http://stix.mitre.org/stix-1 http://stix.mitre.org/XMLSchema/core/1.1.1/stix_core.xsd" xsi:type='stix-maec:MAEC4.1InstanceType'
            id="example:package-2fb96bef-1b11-436e-af4a-15588ac3198b" schema_version="2.1">
            <maecPackage:Malware_Subjects>
                <maecPackage:Malware_Subject>
                    <maecPackage:Malware_Instance_Object_Attributes id="maec-tst-obj-1">
                        <cybox:Properties xsi:type="FileObj:FileObjectType">
                            <FileObj:Hashes>
                                <cyboxCommon:Hash>
                                    <cyboxCommon:Type xsi:type="cyboxVocabs:HashNameVocab-1.0">MD5</cyboxCommon:Type>
                                    <cyboxCommon:Simple_Hash_Value>9d7006e30fdf15e9c8e03e62534b3a3e</cyboxCommon:Simple_Hash_Value>
                                </cyboxCommon:Hash>
                            </FileObj:Hashes>
                        </cybox:Properties>
                    </maecPackage:Malware_Instance_Object_Attributes>
                </maecPackage:Malware_Subject>
            </maecPackage:Malware_Subjects>
        </stix-maec:MAEC>
        """
    )

    def _test_xml(self, obj):
        xml = obj.to_xml()
        parser = stix.utils.parser.get_xml_parser()
        tree = etree.parse(StringIO(xml), parser=parser)
        root = tree.getroot()

        xpath = "//cyboxCommon:Type"
        nodes = root.xpath(xpath, namespaces={'cyboxCommon': 'http://cybox.mitre.org/common-2'})

        self.assertTrue(nodes != None)
        self.assertEqual(len(nodes), 1)
        self.assertEqual(nodes[0].text, "MD5")


    def test_etree(self):
        parser = stix.utils.parser.get_xml_parser()
        tree = etree.parse(StringIO(self.XML), parser=parser)

        ext = MAECInstance()
        ext.maec = tree
        self._test_xml(ext)


    def test_etree_dict(self):
        parser = stix.utils.parser.get_xml_parser()
        tree = etree.parse(StringIO(self.XML), parser=parser)
        ext = MAECInstance()
        ext.maec = tree

        d = ext.to_dict()
        ext2 = MAECInstance.from_dict(d)
        self._test_xml(ext2)


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