# Copyright (c) 2017, The MITRE Corporation. All rights reserved.
# See LICENSE.txt for complete terms.

# stdlib
from distutils.version import LooseVersion

# external
from lxml import etree
import mixbox.xml
from mixbox.vendor.six import BytesIO, iteritems, binary_type

# internal
import stix
import stix.utils as utils
import stix.ttp.malware_instance
from stix.ttp.malware_instance import MalwareInstance
import stix.bindings.extensions.malware.maec_4_1 as ext_binding
from mixbox import fields
from stix.bindings.extensions.malware.maec_4_1 import maec_installed
from lxml.etree import _ElementTree

_MIN_PYTHON_MAEC_VERSION = '4.1.0.12'


class UnsupportedVersion(Exception):
    def __init__(self, message, expected, found):
        super(UnsupportedVersion, self).__init__(message)
        self.expected = expected
        self.found = found


def _check_maec_version():
    """Checks that the installed python-maec has a version greater than or
    equal to the minimum supported version.

    Note:
        We do this rather than having a python-maec dependency requirement
        listed in setup.py because MAEC is used as an extension to STIX and
        not a core component to STIX (like CybOX).

    Raises:
        ImportError: If python-maec is not installed.
        UnsupportedVersion: If the python-maec installation does not satisfy
            the version requirements.

    """
    import maec

    found = maec.__version__
    expected = _MIN_PYTHON_MAEC_VERSION

    if LooseVersion(found) >= LooseVersion(expected):
        return

    fmt = ("Unsupported python-maec version installed: '%s'. Minimum version "
           "is '%s'.")
    error = fmt % (found, expected)
    raise UnsupportedVersion(error, expected=expected, found=found)


try:
    # Check that the correct version of python-maec is installed.
    _check_maec_version()

    # Import maecPackage into global space
    from maec.package.package import Package as maecPackage

    _MAEC_INSTALLED = True
except ImportError:
    maecPackage, Package = None, None
    _MAEC_INSTALLED = False


def is_maec(obj):
    """Checks if the input object is python-maec object.

    Returns:
        True if python-maec is ins

    """
    if not _MAEC_INSTALLED:
        return False

    return isinstance(obj, maecPackage)

def validate_maec_input(instance, value):
    if value is None:
        return
    elif _MAEC_INSTALLED and is_maec(value):
        return
    elif mixbox.xml.is_element(value) or mixbox.xml.is_etree(value):
        return
    else:
        error = (
            "Cannot set maec to '{0}'. Expected 'lxml.etree._Element' or "
            "'maec.package.package.Package'."
        )
        error = error.format(type(value))
        raise ValueError(error)

@stix.register_extension
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"
    _TAG_MAEC = "{%s}MAEC" % _namespace

    maec = fields.TypedField("MAEC", preset_hook=validate_maec_input)

    def __init__(self, maec=None):
        super(MAECInstance, self).__init__()
        self.__input_namespaces__ = {}
        self.__input_schemalocations__ = {}
        self.maec = maec

    def _parse_etree(self, root):
        node_tag = root.tag

        if node_tag != self._TAG_MAEC:
            self._cast_maec(root)

        self._collect_namespaces(root)
        self._collect_schemalocs(root)

    def _cast_maec(self, node):
        ns_maec = "http://maec.mitre.org/XMLSchema/maec-package-2"
        node_ns = etree.QName(node).namespace

        if node_ns == ns_maec:
            etree.register_namespace(self._xml_ns_prefix, self._namespace)
            node.tag = self._TAG_MAEC
        else:
            error = "Cannot set maec. Expected tag '{0}' found '{1}'."
            error = error.format(self._TAG_MAEC, node.tag)
            raise ValueError(error)

    def _collect_schemalocs(self, node):
        try:
            schemaloc = mixbox.xml.get_schemaloc_pairs(node)
            self.__input_schemalocations__ = dict(schemaloc)
        except KeyError:
            self.__input_schemalocations__ = {}

    def _collect_namespaces(self, node):
        self.__input_namespaces__ = dict(iteritems(node.nsmap))

    @classmethod
    def from_obj(cls, obj):
        if not obj:
            return None

        return_obj = cls()

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

        return_obj = super(MAECInstance, cls).from_obj(obj)

        return return_obj

    def to_obj(self, ns_info=None):
        return_obj = super(MAECInstance, self).to_obj(ns_info=ns_info)

        if mixbox.xml.is_element(self.maec) or mixbox.xml.is_etree(self.maec):
            tree = mixbox.xml.get_etree(self.maec)
            root = mixbox.xml.get_etree_root(tree)
            self._parse_etree(root)
            self.maec = root

        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 _maec_from_dict(cls, d):
        if _MAEC_INSTALLED:
            return maecPackage.from_dict(d)

        raise ValueError(
            "Unable to parse 'maec' value in dictionary. Please "
            "install python-maec to parse dictionary value."
        )

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

        d = d.copy()

        maec = d.get('maec')

        if maec is None:
            pass
        elif isinstance(maec, dict):
            d['maec'] = cls._maec_from_dict(maec)
        elif isinstance(maec, binary_type):
            d['maec'] = mixbox.xml.get_etree_root(BytesIO(maec))
        else:
            raise TypeError("Unknown type for 'maec' entry.")

        return_obj = super(MAECInstance, cls).from_dict(d)

        return return_obj

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

        if self.maec is not None:
            if mixbox.xml.is_element(self.maec) or mixbox.xml.is_etree(self.maec):
                tree = mixbox.xml.get_etree(self.maec)
                root = mixbox.xml.get_etree_root(tree)
                self._parse_etree(root)
                self.maec = root

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

            if self._XSI_TYPE:
                d['xsi:type'] = self._XSI_TYPE

        return d
