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

# internal
import stix
from stix.common import vocabs
from stix.common import StructuredText, VocabString

# bindings
import stix.bindings.ttp as ttp_binding

from mixbox import fields, entities


class MalwareInstance(stix.Entity):
    _binding = ttp_binding
    _binding_class = _binding.MalwareInstanceType
    _namespace = "http://stix.mitre.org/TTP-1"
    _XSI_TYPE = None  # defined by subclasses

    id_ = fields.IdField("id")
    idref = fields.IdrefField("idref")
    title = fields.TypedField("Title")
    description = fields.TypedField("Description", type_="stix.common.StructuredText")
    short_description = fields.TypedField("Short_Description", type_="stix.common.StructuredText")
    names = vocabs.VocabField("Name", type_=VocabString, multiple=True, key_name="names")
    types = vocabs.VocabField("Type", type_=vocabs.MalwareType, multiple=True, key_name="types")

    def __init__(self, id_=None, idref=None, title=None, description=None, short_description=None):
        super(MalwareInstance, self).__init__()
        self.id_ = id_
        self.idref = idref
        self.title = title
        self.description = description
        self.short_description = short_description

    def add_name(self, name):
        self.names.append(name)

    def add_type(self, type_):
        self.types.append(type_)

    @staticmethod
    def lookup_class(xsi_type):
        if not xsi_type:
            raise ValueError("xsi:type is required")

        return stix.lookup_extension(xsi_type)

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

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

        return d


class MalwareInstanceFactory(entities.EntityFactory):
    @classmethod
    def entity_class(cls, key):
        from stix.extensions.malware.maec_4_1_malware import MAECInstance  # noqa
        return stix.lookup_extension(key, default=MalwareInstance)


# Backwards compatibility
add_extension = stix.add_extension
