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

import stix
import stix.utils
from stix.utils import dates

import stix.bindings.exploit_target as exploit_target_binding
from stix.common.related import (GenericRelationshipList, RelatedCOA)
from stix.common import StructuredText, VocabString, InformationSource, Statement
from stix.coa import CourseOfAction
from stix.data_marking import Marking
from .vulnerability import Vulnerability

class ExploitTarget(stix.Entity):
    _binding = exploit_target_binding
    _binding_class = _binding.ExploitTargetType
    _namespace = "http://stix.mitre.org/ExploitTarget-1"
    _version = "1.1"

    def __init__(self, id_=None, idref=None, timestamp=None, title=None, description=None, short_description=None):
        self.id_ = id_ or stix.utils.create_id("et")
        self.idref = idref
        self.timestamp = timestamp
        self.version = self._version
        self.title = title
        self.description = description
        self.short_description = short_description
        self.information_source = None
        self.handling = None
        self.potential_coas = PotentialCOAs()
        self.vulnerabilities = None

    @property
    def timestamp(self):
        return self._timestamp

    @timestamp.setter
    def timestamp(self, value):
        self._timestamp = dates.parse_value(value)

    @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 information_source(self):
        return self._information_source
    
    @information_source.setter
    def information_source(self, value):
        if not value:
            self._information_source = None
        elif isinstance(value, InformationSource):
            self._information_source = value
        else:
            raise ValueError('value must be instance of InformationSource')

    @property
    def handling(self):
        return self._handling

    @handling.setter
    def handling(self, value):
        if value and not isinstance(value, Marking):
            raise ValueError('value must be instance of Marking')

        self._handling = value
 
    @property
    def vulnerabilities(self):
        return self._vulnerabilities
    
    @vulnerabilities.setter
    def vulnerabilities(self, value):
        self._vulnerabilities = []
        if not value:
            return
        elif isinstance(value, list):
            for v in value:
                self.add_vulnerability(v)
        else:
            self.add_vulnerability(value)
            
    def add_vulnerability(self, v):
        if not v:
            return
        elif isinstance(v, Vulnerability):
            self.vulnerabilities.append(v)
        else:
            raise ValueError('Cannot add type %s to vulnerabilities list' % type(v))
 
    def to_obj(self, return_obj=None):
        if not return_obj:
            return_obj = self._binding_class()

        return_obj.set_id(self.id_)
        return_obj.set_idref(self.idref)
        return_obj.set_timestamp(dates.serialize_value(self.timestamp))
        return_obj.set_version(self.version)
        return_obj.set_Title(self.title)

        if self.description:
            return_obj.set_Description(self.description.to_obj())
        if self.short_description:
            return_obj.set_Short_Description(self.short_description.to_obj())
        if self.information_source:
            return_obj.set_Information_Source(self.information_source.to_obj())
        if self.handling:
            return_obj.set_Handling(self.handling.to_obj())
        if self.potential_coas:
            return_obj.set_Potential_COAs(self.potential_coas.to_obj())
        if self.vulnerabilities:
            return_obj.set_Vulnerability([x.to_obj() for x in self.vulnerabilities])
            
        return return_obj

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

        return_obj.id_ = obj.get_id()
        return_obj.idref = obj.get_idref()
        return_obj.timestamp = obj.get_timestamp() # not yet implemented

        if isinstance(obj, cls._binding_class): # TTPType properties
            return_obj.version = obj.get_version() or cls._version
            return_obj.title = obj.get_Title()
            return_obj.description = StructuredText.from_obj(obj.get_Description())
            return_obj.short_description = StructuredText.from_obj(obj.get_Short_Description())
            return_obj.information_source = InformationSource.from_obj(obj.get_Information_Source())
            return_obj.handling = Marking.from_obj(obj.get_Handling())
            return_obj.potential_coas = PotentialCOAs.from_obj(obj.get_Potential_COAs())
            return_obj.vulnerabilities = [Vulnerability.from_obj(x) for x in obj.get_Vulnerability()]

        return return_obj

    def to_dict(self):
        d = {}
        if self.id_:
            d['id'] = self.id_
        if self.idref:
            d['idref'] = self.idref
        if self.timestamp:
            d['timestamp'] = dates.serialize_value(self.timestamp)
        if self.version:
            d['version'] = self.version or self._version
        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.information_source:
            d['information_source'] = self.information_source.to_dict()
        if self.handling:
            d['handling'] = self.handling.to_dict()
        if self.potential_coas:
            d['potential_coas'] = self.potential_coas.to_dict()
        if self.vulnerabilities:
            d['vulnerabilities'] = [x.to_dict() for x in self.vulnerabilities]

        return d

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

        return_obj.id_ = dict_repr.get('id')
        return_obj.idref = dict_repr.get('idref')
        return_obj.timestamp = dict_repr.get('timestamp')
        return_obj.version = dict_repr.get('version', cls._version)
        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.information_source = InformationSource.from_dict(dict_repr.get('information_source'))
        return_obj.handling = Marking.from_dict(dict_repr.get('handling'))
        return_obj.potential_coas = PotentialCOAs.from_dict(dict_repr.get('potential_coas'))
        return_obj.vulnerabilities = [Vulnerability.from_dict(x) for x in dict_repr.get('vulnerabilities', [])]
        
        return return_obj

class PotentialCOAs(GenericRelationshipList):
    _namespace = "http://stix.mitre.org/ExploitTarget-1"
    _binding = exploit_target_binding
    _binding_class = exploit_target_binding.PotentialCOAsType
    _binding_var = "Potential_COA"
    _contained_type = RelatedCOA
    _inner_name = "coas"

    def __init__(self, coas=None, scope=None):
        if coas is None:
            coas = []
        super(PotentialCOAs, self).__init__(*coas, scope=scope)