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

from stix.test import EntityTestCase
from stix.extensions.malware.maec_4_1_malware import MAECInstance

try:
    import maec
    maec_present = True
except ImportError:
    maec_present = False


@unittest.skipIf(condition=maec_present is False, reason="These tests require the 'maec' library.")
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)


@unittest.skipIf(condition=maec_present is False, reason="These tests require the 'maec' library.")
class PythonMAECInPackageTests(unittest.TestCase):
    XML = StringIO(
        """
        <stix:STIX_Package
            xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
            xmlns:stix="http://stix.mitre.org/stix-1"
            xmlns:stixCommon="http://stix.mitre.org/common-1"
            xmlns:stixVocabs="http://stix.mitre.org/default_vocabularies-1"
            xmlns:ttp="http://stix.mitre.org/TTP-1"
            xmlns:example="http://example.com"
            xsi:schemaLocation="
            http://stix.mitre.org/TTP-1 http://stix.mitre.org/XMLSchema/ttp/1.2/ttp.xsd
            http://stix.mitre.org/common-1 http://stix.mitre.org/XMLSchema/common/1.2/stix_common.xsd
            http://stix.mitre.org/default_vocabularies-1 http://stix.mitre.org/XMLSchema/default_vocabularies/1.2.0/stix_default_vocabularies.xsd
            http://stix.mitre.org/stix-1 http://stix.mitre.org/XMLSchema/core/1.2/stix_core.xsd"
            id="example:Package-2b8fb66f-b6b3-4d40-865a-33e4a5ee1246"
            version="1.2"
            >
          <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://stix.mitre.org/stix-1"
            xmlns:stixCommon="http://stix.mitre.org/common-1"
            xmlns:stixVocabs="http://stix.mitre.org/default_vocabularies-1"
            xmlns:ttp="http://stix.mitre.org/TTP-1"
            xmlns:stix-maec="http://stix.mitre.org/extensions/Malware#MAEC4.1-1"
            xmlns:maecPackage="http://maec.mitre.org/XMLSchema/maec-package-2"
            xmlns:example="http://example.com"
            xsi:schemaLocation="
            http://stix.mitre.org/TTP-1 http://stix.mitre.org/XMLSchema/ttp/1.2/ttp.xsd
            http://stix.mitre.org/common-1 http://stix.mitre.org/XMLSchema/common/1.2/stix_common.xsd
            http://stix.mitre.org/default_vocabularies-1 http://stix.mitre.org/XMLSchema/default_vocabularies/1.2.0/stix_default_vocabularies.xsd
            http://stix.mitre.org/stix-1 http://stix.mitre.org/XMLSchema/core/1.2/stix_core.xsd
            http://stix.mitre.org/extensions/Malware#MAEC4.1-1 http://stix.mitre.org/XMLSchema/extensions/malware/maec_4.1/1.0/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"
            >
          <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.ttp[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.ttp[0].behavior.malware_instances[0].to_dict()
        self.assertTrue('names' in mw)


@unittest.skipIf(condition=maec_present is True, reason="These tests require the 'maec' library to be missing.")
class PythonMAECNotInstalledTest(unittest.TestCase):

    def test_parsing_maec_fails(self):
        try:
            STIXPackage.from_xml(PythonMAECInPackageTests.XML_MAEC)
        except ImportError as e:
            self.assertTrue(all(x in str(e) for x in ("No module named", "maec")))

    def test_handling_maec_object_fails(self):
        try:
            MAECInstance().from_dict(PythonMAECTests._full_dict)
        except ImportError as e:
            self.assertTrue(all(x in str(e) for x in ("No module named", "maec")))

    def test_setting_maec_property_fails(self):
        try:
            m = MAECInstance()
            m.maec = "foo"
        except ImportError as e:
            self.assertTrue(all(x in str(e) for x in ("No module named", "maec")))


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