import unittest
from stix.core import STIXPackage
from mixbox.vendor.six import StringIO, BytesIO, text_type

from lxml import etree
import mixbox.xml

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'}
                         }
                    }
                ]
        }
    }

    def test_add_name_type(self):
        maec_malware_instance = MAECInstance()
        maec_malware_instance.add_name("Poison Ivy Variant v4392-acc")
        maec_malware_instance.add_type("Remote Access Trojan")
        maec_xml = text_type(maec_malware_instance.to_xml())
        self.assertTrue("Poison Ivy Variant v4392-acc" in maec_xml)
        self.assertTrue("Remote Access Trojan" in maec_xml)


class PythonMAECEtreeTests(unittest.TestCase):
    XML = (
        """
        <stix-maec:MAEC
            xmlns:cyboxCommon="http://docs.oasis-open.org/cti/ns/cybox/common-2"
            xmlns:cybox="http://docs.oasis-open.org/cti/ns/cybox/core-2"
            xmlns:cyboxVocabs="http://docs.oasis-open.org/cti/ns/cybox/vocabularies-2"
            xmlns:FileObj="http://docs.oasis-open.org/cti/ns/cybox/objects/file-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://docs.oasis-open.org/cti/ns/stix/ttp-1"
            xmlns:stixCommon="http://docs.oasis-open.org/cti/ns/stix/common-1"
            xmlns:stixVocabs="http://docs.oasis-open.org/cti/ns/stix/vocabularies-1"
            xmlns:stix-maec="http://docs.oasis-open.org/cti/ns/stix/extensions/malware/maec-4.1-1"
            xmlns:stix="http://docs.oasis-open.org/cti/ns/stix/core-1"
            xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
            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 = mixbox.xml.get_xml_parser()
        tree = etree.parse(BytesIO(xml), parser=parser)
        root = tree.getroot()

        xpath = "//cyboxCommon:Type"
        nodes = root.xpath(xpath, namespaces={'cyboxCommon': 'http://docs.oasis-open.org/cti/ns/cybox/common-2'})

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


    def test_etree(self):
        parser = mixbox.xml.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 = mixbox.xml.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)


class PythonMAECInPackageTests(unittest.TestCase):
    XML = StringIO(
        """
        <stix:STIX_Package 
            xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
            xmlns:stix="http://docs.oasis-open.org/cti/ns/stix/core-1"
            xmlns:stixCommon="http://docs.oasis-open.org/cti/ns/stix/common-1"
            xmlns:stixVocabs="http://docs.oasis-open.org/cti/ns/stix/vocabularies-1"
            xmlns:ttp="http://docs.oasis-open.org/cti/ns/stix/ttp-1"
            xmlns:example="http://example.com"
            xsi:schemaLocation="
            http://docs.oasis-open.org/cti/ns/stix/ttp-1 http://docs.oasis-open.org/cti/stix/v1.2.1/csd01/xml-schemas/ttp.xsd
            http://docs.oasis-open.org/cti/ns/stix/common-1 http://docs.oasis-open.org/cti/stix/v1.2.1/csd01/xml-schemas/common.xsd
            http://docs.oasis-open.org/cti/ns/stix/vocabularies-1 http://docs.oasis-open.org/cti/stix/v1.2.1/csd01/xml-schemas/vocabularies.xsd
            http://docs.oasis-open.org/cti/ns/stix/core-1 http://docs.oasis-open.org/cti/stix/v1.2.1/csd01/xml-schemas/core.xsd"
            id="example:Package-2b8fb66f-b6b3-4d40-865a-33e4a5ee1246" 
            version="1.2.1"
            >
          <stix:TTPs>
            <stix:TTP xsi:type="ttp:TTPType" id="example:ttp-7d9fe1f7-429d-077e-db51-92c70b8da45a">
              <ttp:Title>Poison Ivy Variant v4392-acc</ttp:Title>
              <ttp:Behavior>
                <ttp:Malware>
                  <ttp:Malware_Instance>
                    <ttp:Type xsi:type="stixVocabs:MalwareTypeVocab-1.0">Remote Access Trojan</ttp:Type>
                    <ttp:Name>Poison Ivy Variant v4392-acc</ttp:Name>
                  </ttp:Malware_Instance>
                </ttp:Malware>
              </ttp:Behavior>
            </stix:TTP>
          </stix:TTPs>
        </stix:STIX_Package>
        """
    )
    XML_MAEC = StringIO(
        """
        <stix:STIX_Package 
            xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
            xmlns:stix="http://docs.oasis-open.org/cti/ns/stix/core-1"
            xmlns:stixCommon="http://docs.oasis-open.org/cti/ns/stix/common-1"
            xmlns:stixVocabs="http://docs.oasis-open.org/cti/ns/stix/vocabularies-1"
            xmlns:ttp="http://docs.oasis-open.org/cti/ns/stix/ttp-1"
            xmlns:stix-maec="http://docs.oasis-open.org/cti/ns/stix/extensions/malware/maec-4.1-1"
            xmlns:maecPackage="http://maec.mitre.org/XMLSchema/maec-package-2"
            xmlns:example="http://example.com"
            xsi:schemaLocation="
            http://docs.oasis-open.org/cti/ns/stix/ttp-1 http://docs.oasis-open.org/cti/stix/v1.2.1/csd01/xml-schemas/ttp.xsd
            http://docs.oasis-open.org/cti/ns/stix/common-1 http://docs.oasis-open.org/cti/stix/v1.2.1/csd01/xml-schemas/common.xsd
            http://docs.oasis-open.org/cti/ns/stix/vocabularies-1 http://docs.oasis-open.org/cti/stix/v1.2.1/csd01/xml-schemas/vocabularies.xsd
            http://docs.oasis-open.org/cti/ns/stix/core-1 http://docs.oasis-open.org/cti/stix/v1.2.1/csd01/xml-schemas/core.xsd
            http://docs.oasis-open.org/cti/ns/stix/extensions/malware/maec-4.1-1 http://docs.oasis-open.org/cti/stix/v1.2.1/csd01/xml-schemas/extensions/malware/maec-4.1-malware.xsd
            http://maec.mitre.org/XMLSchema/maec-package-2 http://maec.mitre.org/language/version4.1/maec_package_schema.xsd"
            id="example:Package-2b8fb66f-b6b3-4d40-865a-33e4a5ee1246" 
            version="1.2.1"
            >
          <stix:TTPs>
            <stix:TTP xsi:type="ttp:TTPType" id="example:ttp-7d9fe1f7-429d-077e-db51-92c70b8da45a">
              <ttp:Title>Poison Ivy Variant v4392-acc</ttp:Title>
              <ttp:Behavior>
                <ttp:Malware>
                  <ttp:Malware_Instance xsi:type="stix-maec:MAEC4.1InstanceType">
                    <ttp:Type xsi:type="stixVocabs:MalwareTypeVocab-1.0">Remote Access Trojan</ttp:Type>
                    <ttp:Name>Poison Ivy Variant v4392-acc</ttp:Name>
                    <stix-maec:MAEC id="example:package-2fb96bef-1b11-436e-af4a-15588ac3198b" schema_version="2.1">
                      <maecPackage:Malware_Subjects>
                        <maecPackage:Malware_Subject id="example:Subject-57cd4839-436e-1b11-af4a-15588ac3198b">
                          <maecPackage:Malware_Instance_Object_Attributes>
                          </maecPackage:Malware_Instance_Object_Attributes>
                        </maecPackage:Malware_Subject>
                      </maecPackage:Malware_Subjects>
                    </stix-maec:MAEC>
                  </ttp:Malware_Instance>
                </ttp:Malware>
              </ttp:Behavior>
            </stix:TTP>
          </stix:TTPs>
        </stix:STIX_Package>
        """
    )

    def test_parse_malware(self):
        """Test parsing a normal MalwareInstance from XML
        """
        stix_pkg = STIXPackage.from_xml(self.XML)
        mw = stix_pkg.ttps[0].behavior.malware_instances[0].to_dict()
        self.assertTrue('names' in mw)

    def test_parse_malware_maec(self):
        """Test parsing a MaecInstance from XML
        """
        stix_pkg = STIXPackage.from_xml(self.XML_MAEC)
        mw = stix_pkg.ttps[0].behavior.malware_instances[0].to_dict()
        self.assertTrue('names' in mw)


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