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

import stix
import stix.utils
from stix.common import StructuredText, VocabString
import stix.bindings.ttp as ttp_binding
from stix.common.vocabs import MalwareType

class MalwareInstance(stix.Entity):
    _binding = ttp_binding
    _binding_class = _binding.MalwareInstanceType
    _namespace = "http://stix.mitre.org/TTP-1"

    def __init__(self, id_=None, title=None, description=None, short_description=None):
        self.id_ = id_
        self.title = title
        self.description = description
        self.short_description = short_description
        self.names = None
        self.types = None

    @property
    def title(self):
        return self._title

    @title.setter
    def title(self, value):
        self._title = value

    @property
    def description(self):
        return self._description

    @description.setter
    def description(self, value):
        if value:
            if isinstance(value, StructuredText):
                self._description = value
            else:
                self._description = StructuredText(value=value)
        else:
            self._description = None

    @property
    def short_description(self):
        return self._short_description

    @short_description.setter
    def short_description(self, value):
        if value:
            if isinstance(value, StructuredText):
                self._short_description = value
            else:
                self._short_description = StructuredText(value=value)
        else:
            self._short_description = None

    @property
    def names(self):
        return self._names

    @names.setter
    def names(self, value):
        self._names = MalwareNames()
        if not value:
            return
        elif isinstance(value, list):
            for v in value:
                self.add_name(v)
        else:
            self.add_name(value)

    def add_name(self, name):
        if not name:
            return
        elif isinstance(name, VocabString):
            self._names.append(name)
        else:
            self._names.append(VocabString(value=name))

    @property
    def types(self):
        return self._types

    @types.setter
    def types(self, value):
        self._types = MalwareTypes()
        if not value:
            return
        elif isinstance(value, list):
            for v in value:
                self.add_type(v)
        else:
            self.add_type(value)

    def add_type(self, type_):
        if not type_:
            return
        elif isinstance(type_, VocabString):
            self._types.append(type_)
        else:
            self._types.append(MalwareType(value=type_))

    @staticmethod
    def lookup_class(xsi_type):
        if not xsi_type:
            raise ValueError("xsi:type is required")
        for (k, v) in _EXTENSION_MAP.iteritems():
            # TODO: for now we ignore the prefix and just check for
            # a partial match
            if xsi_type in k:
                return v

        raise ValueError("Unregistered xsi:type %s" % xsi_type)
    

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

        if not return_obj:
            return_obj = self._binding_class()

        return_obj.id = self.id_
        return_obj.Title = self.title

        if self.description:
            return_obj.Description = self.description.to_obj(ns_info=ns_info)
        if self.short_description:
            return_obj.Short_Description = self.short_description.to_obj(ns_info=ns_info)
        if self.names:
            return_obj.Name = [x.to_obj(ns_info=ns_info) for x in self.names]
        if self.types:
            return_obj.Type = [x.to_obj(ns_info=ns_info) for x in self.types]

        return return_obj

    @classmethod
    def from_obj(cls, obj, return_obj=None):
        from stix.extensions.malware import maec_4_1_malware # register the extension
        
        if not obj:
            return None
        
        if not return_obj:
            try:
                klass = MalwareInstance.lookup_class(obj.xml_type)
                return_obj = klass.from_obj(obj)
            except AttributeError:
                return_obj = MalwareInstance.from_obj(obj, cls())
        else:
            return_obj.id_ = obj.id
            return_obj.title = obj.Title
            return_obj.description = StructuredText.from_obj(obj.Description)
            return_obj.short_description = StructuredText.from_obj(obj.Short_Description)
            return_obj.names = [VocabString.from_obj(x) for x in obj.Name]
            return_obj.types = [VocabString.from_obj(x) for x in obj.Type]

        return return_obj

    def to_dict(self):
        d = {}
        if self.id_:
            d['id'] = self.id_
        if self.title:
            d['title'] = self.title
        if self.description:
            d['description'] = self.description.to_dict()
        if self.short_description:
            d['short_description'] = self.short_description.to_dict()
        if self.names:
            d['names'] = [x.to_dict() for x in self.names]
        if self.types:
            d['types'] = [x.to_dict() for x in self.types]

        try:
            d['xsi:type'] = self._XSI_TYPE
        except AttributeError:
            pass

        return d

    @classmethod
    def from_dict(cls, dict_repr, return_obj=None):
        if not dict_repr:
            return None
        
        if not return_obj:
            xsi_type = dict_repr.get('xsi:type')
            if xsi_type:
                klass = MalwareInstance.lookup_class(dict_repr.get('xsi:type'))
                return_obj = klass.from_dict(dict_repr)
            else:
                return_obj = MalwareInstance.from_dict(dict_repr, cls())
        else:
            return_obj.id_ = dict_repr.get('id')
            return_obj.title = dict_repr.get('title')
            return_obj.description = StructuredText.from_dict(dict_repr.get('description'))
            return_obj.short_description = StructuredText.from_dict(dict_repr.get('short_description'))
            return_obj.names = [VocabString.from_dict(x) for x in dict_repr.get('names', [])]
            return_obj.types = [VocabString.from_dict(x) for x in dict_repr.get('types', [])]

        return return_obj


class MalwareNames(stix.EntityList):
    _namespace = "http://stix.mitre.org/TTP-1"
    _contained_type = VocabString


class MalwareTypes(stix.EntityList):
    _namespace = "http://stix.mitre.org/TTP-1"
    _contained_type = VocabString

    def _fix_value(self, value):
        return MalwareType(value)


#: Mapping of malware instance extension types to classes
_EXTENSION_MAP = {}
def add_extension(cls):
    _EXTENSION_MAP[cls._XSI_TYPE] = cls

