import stix
import stix.ttp.malware_instance
from stix.ttp.malware_instance import MalwareInstance
import stix.bindings.extensions.malware.maec_4_1 as ext_binding
from lxml import etree
from StringIO import StringIO

try:
    from maec.package.package import Package as maecPackage
    _MAEC_INSTALLED = True
except ImportError:
    _MAEC_INSTALLED = False

class MAECInstance(MalwareInstance):
    _binding = ext_binding
    _binding_class = _binding.MAEC4_1InstanceType
    _namespace = 'http://stix.mitre.org/extensions/Malware#MAEC4.1-1'
    _xml_ns_prefix = "stix-maec"
    _XSI_TYPE = "stix-maec:MAEC4.1InstanceType"

    def __init__(self, maec=None):
        super(MAECInstance, self).__init__()
        self.maec = maec
        
    @property
    def maec(self):
        return self._maec
    
    @maec.setter
    def maec(self, value):
        if value is None:
            self._maec = None
        elif _MAEC_INSTALLED and isinstance(value, maecPackage):
            self._parse_maec_package(value)
            self._maec = value
        elif isinstance(value, (etree._ElementTree, etree._Element)):
            root = stix.utils.parser.get_etree_root(value)
            self._parse_etree(root)
            self._maec = etree.ElementTree(root)
        else:
            raise ValueError(
                'maec must be instance of lxml.etree._Element or '
                'lxml.etree._ElementTree or maec.package.package.Package. '
                'Received: %s' % (type(value))
            )

    def _parse_etree(self, root):
        ns_maec = "http://maec.mitre.org/XMLSchema/maec-package-2"
        node_ns = etree.QName(root).namespace
        expected_node_tag = "{%s}MAEC" % (self._namespace)
        node_tag = root.tag

        if node_tag != expected_node_tag:
            if node_ns != ns_maec:
                raise ValueError(
                    "Cannot set maec property. Expected tag %s but found %s" %
                    (expected_node_tag, node_tag)
                )

            # attempt to cast
            etree.register_namespace(self._xml_ns_prefix, self._namespace)
            root.tag = expected_node_tag

        self.__input_namespaces__ = dict(
            (ns, alias) for alias, ns in root.nsmap.iteritems()
        )

        try:
            schemaloc = stix.utils.parser.get_schemaloc_pairs(root)
            self.__input_schemalocations__ = dict(schemaloc)
        except KeyError:
            self.__input_schemalocations__ = {}


    def _parse_maec_package(self, package):
        # Handle the namespace/schemaLocation processing for the MAEC content
        nsset = package._get_namespaces()

        ns_dict = dict(
            (ns.name, ns.prefix) for ns in nsset if ns
        )
        schemaloc_dict = dict(
            (ns.name, ns.schema_location) for ns in nsset if ns and ns.schema_location
        )

        # Add the MAEC Vocabs ns/schemalocation by default
        vocab_ns = 'http://maec.mitre.org/default_vocabularies-1'
        vocab_schemaloc = 'http://maec.mitre.org/language/version4.1/maec_default_vocabularies.xsd'

        ns_dict[vocab_ns] = 'maecVocabs'
        schemaloc_dict[vocab_ns] = vocab_schemaloc

        self.__input_namespaces__ = ns_dict
        self.__input_schemalocations__ = schemaloc_dict

    @classmethod
    def from_obj(cls, obj, return_obj=None):
        if not obj:
            return None
        if not return_obj:
            return_obj = cls()
        
        super(MAECInstance, cls).from_obj(obj, return_obj)

        if _MAEC_INSTALLED:
            return_obj.maec = maecPackage.from_obj(obj.MAEC)
        else:
            return_obj.maec = obj.MAEC

        return return_obj
    
    def to_obj(self, return_obj=None, ns_info=None):
        if not return_obj:
            return_obj = self._binding_class()
            
        super(MAECInstance, self).to_obj(return_obj=return_obj, ns_info=ns_info)

        if _MAEC_INSTALLED and isinstance(self.maec, maecPackage):
            return_obj.MAEC = self.maec.to_obj(ns_info=ns_info)
        else:
            return_obj.MAEC = self.maec

        return return_obj
    
    @classmethod
    def from_dict(cls, d, return_obj=None):
        if not d:
            return None
        if not return_obj:
            return_obj = cls()
            
        super(MAECInstance, cls).from_dict(d, return_obj)

        if 'maec' in d:
            maec = d['maec']

            if isinstance(maec, dict):
                if not _MAEC_INSTALLED:
                    raise ValueError(
                        "Unable to parse 'maec' value in dictionary. Please "
                        "install python-maec to parse dictionary value."
                    )

                return_obj.maec = maecPackage.from_dict(maec)
            else:
                parser = stix.utils.parser.get_xml_parser()
                return_obj.maec = etree.parse(StringIO(maec), parser=parser)
        
        return return_obj


    def to_dict(self):
        d = super(MAECInstance, self).to_dict()

        if self.maec:
            if _MAEC_INSTALLED and isinstance(self.maec, maecPackage):
                d['maec'] = self.maec.to_dict()
            else:
                d['maec'] = etree.tostring(self.maec)

        return d

stix.ttp.malware_instance.add_extension(MAECInstance)