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

import stix
import stix.utils as utils
import stix.bindings.exploit_target as exploit_target_binding
import stix.bindings.stix_common as stix_common_binding
from stix.common import DateTimeWithPrecision, StructuredTextList
from stix.common.related import GenericRelationshipList, RelatedObservable


class Vulnerability(stix.Entity):
    """Implementation of STIX ``Vulnerability``.

    Args:
        title (optional): A string title.
        description (optional): A string description.
        short_description (optional): A string short description.

    """
    _binding = exploit_target_binding
    _binding_class = _binding.VulnerabilityType
    _namespace = "http://stix.mitre.org/ExploitTarget-1"

    def __init__(self, title=None, description=None, short_description=None):
        self.is_known = None
        self.is_publicly_acknowledged = None
        self.title = title
        self.description = description
        self.short_description = short_description
        self.cve_id = None
        self.osvdb_id = None
        self.source = None
        self.cvss_score = None
        self.discovered_datetime = None
        self.published_datetime = None
        self.affected_software = AffectedSoftware()
        self.references = []

    @property
    def title(self):
        """
        String representing the Vulnerability Title
        """
        return self._title

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

    @property
    def description(self):
        """A single description about the contents or purpose of this object.

        Default Value: ``None``

        Note:
            If this object has more than one description set, this will return
            the description with the lowest ordinality value.

        Returns:
            An instance of
            :class:`.StructuredText`

        """
        return next(iter(self.descriptions), None)

    @description.setter
    def description(self, value):
        self.descriptions = value

    @property
    def descriptions(self):
        """A :class:`.StructuredTextList` object, containing descriptions about
        the purpose or intent of this object.

        This is typically used for the purpose of providing multiple
        descriptions with different classificaton markings.

        Iterating over this object will yield its contents sorted by their
        ``ordinality`` value.

        Default Value: Empty :class:`.StructuredTextList` object.

        Note:
            IF this is set to a value that is not an instance of
            :class:`.StructuredText`, an effort will ne made to convert it.
            If this is set to an iterable, any values contained that are not
            an instance of :class:`.StructuredText` will be be converted.

        Returns:
            An instance of
            :class:`.StructuredTextList`

        """
        return self._description

    @descriptions.setter
    def descriptions(self, value):
        self._description = StructuredTextList(value)

    def add_description(self, description):
        """Adds a description to the ``descriptions`` collection.

        This is the same as calling "foo.descriptions.add(bar)".

        """
        self.descriptions.add(description)

    @property
    def short_description(self):
        """A single short description about the contents or purpose of this
        object.

        Default Value: ``None``

        Note:
            If this object has more than one short description set, this will
            return the description with the lowest ordinality value.

        Returns:
            An instance of :class:`.StructuredText`

        """
        return next(iter(self.short_descriptions), None)

    @short_description.setter
    def short_description(self, value):
        self.short_descriptions = value

    @property
    def short_descriptions(self):
        """A :class:`.StructuredTextList` object, containing short descriptions
        about the purpose or intent of this object.

        This is typically used for the purpose of providing multiple
        short descriptions with different classificaton markings.

        Iterating over this object will yield its contents sorted by their
        ``ordinality`` value.

        Default Value: Empty :class:`.StructuredTextList` object.

        Note:
            IF this is set to a value that is not an instance of
            :class:`.StructuredText`, an effort will ne made to convert it.
            If this is set to an iterable, any values contained that are not
            an instance of :class:`.StructuredText` will be be converted.

        Returns:
            An instance of :class:`.StructuredTextList`

        """
        return self._short_description

    @short_descriptions.setter
    def short_descriptions(self, value):
        self._short_description = StructuredTextList(value)

    def add_short_description(self, description):
        """Adds a description to the ``short_descriptions`` collection.

        This is the same as calling "foo.short_descriptions.add(bar)".

        """
        self.short_descriptions.add(description)

    @property
    def discovered_datetime(self):
        """
        Returns: 
            The time this vulnerability was discovered, represented as 
            class:`DateTimeWithPrecision` 
        """
        return self._discovered_datetime

    @discovered_datetime.setter
    def discovered_datetime(self, value):
        """
        Sets the time this vulnerability was discovered, represented as 
        class:`DateTimeWithPrecision` 

        Default Value: ``None``

        Returns: 
            None

        """
        self._set_var(DateTimeWithPrecision, discovered_datetime=value)

    @property
    def references(self):
        return self._references

    @references.setter
    def references(self, value):
        self._references = []

        if not value:
            return
        elif utils.is_sequence(value):
            self._references.extend(x for x in value if x)
        else:
            self._references.append(value)

    def add_reference(self, reference):
        if not reference:
            return

        self.references.append(reference)

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

        if not return_obj:
            return_obj = self._binding_class()

        return_obj.Title = self.title
        return_obj.CVE_ID = self.cve_id
        return_obj.OSVDB_ID = self.osvdb_id
        return_obj.Source = self.source
        return_obj.is_known = utils.xml_bool(self.is_known)
        return_obj.is_publicly_acknowledged = utils.xml_bool(self.is_publicly_acknowledged)

        if self.descriptions:
            return_obj.Description = self.descriptions.to_obj(ns_info=ns_info)
        if self.short_descriptions:
            return_obj.Short_Description = self.short_descriptions.to_obj(ns_info=ns_info)
        if self.cvss_score:
            return_obj.CVSS_Score = self.cvss_score.to_obj(ns_info=ns_info)
        if self.discovered_datetime:
            return_obj.Discovered_DateTime = self.discovered_datetime.to_obj(ns_info=ns_info)
        if self.published_datetime:
            return_obj.Published_DateTime = self.published_datetime.to_obj(ns_info=ns_info)
        if self.affected_software:
            return_obj.Affected_Software = self.affected_software.to_obj(ns_info=ns_info)
        if self.references:
            return_obj.References = stix_common_binding.ReferencesType(Reference=self.references)

        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.is_known = utils.xml_bool(obj.is_known)
        return_obj.is_publicly_acknowledged = utils.xml_bool(obj.is_publicly_acknowledged)
        return_obj.title = obj.Title
        return_obj.descriptions = StructuredTextList.from_obj(obj.Description)
        return_obj.short_descriptions = StructuredTextList.from_obj(obj.Short_Description)
        return_obj.cve_id = obj.CVE_ID
        return_obj.osvdb_id = obj.OSVDB_ID
        return_obj.source = obj.Source
        return_obj.cvss_score = CVSSVector.from_obj(obj.CVSS_Score)
        return_obj.discovered_datetime = DateTimeWithPrecision.from_obj(obj.Discovered_DateTime)
        return_obj.published_datetime = DateTimeWithPrecision.from_obj(obj.Published_DateTime)
        return_obj.affected_software = AffectedSoftware.from_obj(obj.Affected_Software)

        if obj.References:
            return_obj.references = obj.References.Reference

        return return_obj

    def to_dict(self):
        return super(Vulnerability, self).to_dict()

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

        get = dict_repr.get
        return_obj.is_known = utils.xml_bool(get('is_known'))
        return_obj.is_publicly_acknowledged = utils.xml_bool(get('is_publicly_acknowledged'))
        return_obj.title = get('title')
        return_obj.descriptions = StructuredTextList.from_dict(get('description'))
        return_obj.short_descriptions = StructuredTextList.from_dict(get('short_description'))
        return_obj.cve_id = get('cve_id')
        return_obj.osvdb_id = get('osvdb_id')
        return_obj.source = get('source')
        return_obj.cvss_score = CVSSVector.from_dict(get('cvss_score'))
        return_obj.discovered_datetime = DateTimeWithPrecision.from_dict(get('discovered_datetime'))
        return_obj.published_datetime = DateTimeWithPrecision.from_dict(get('published_datetime'))
        return_obj.affected_software = AffectedSoftware.from_dict(get('affected_software'))
        return_obj.references = get('references')

        return return_obj


class CVSSVector(stix.Entity):
    """
    Common Vulnerabilit Scoring System object, representing its component measures
    
    """
    _binding = exploit_target_binding
    _binding_class = exploit_target_binding.CVSSVectorType
    _namespace = "http://stix.mitre.org/ExploitTarget-1"

    def __init__(self):
        self.overall_score = None
        self.base_score = None
        self.base_vector = None
        self.temporal_score = None
        self.temporal_vector = None
        self.environmental_score = None
        self.environmental_vector = None

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

        if not return_obj:
            return_obj = self._binding_class()

        return_obj.Overall_Score = self.overall_score
        return_obj.Base_Score = self.base_score
        return_obj.Base_Vector = self.base_vector
        return_obj.Temporal_Score = self.temporal_score
        return_obj.Temporal_Vector = self.temporal_vector
        return_obj.Environmental_Score = self.environmental_score
        return_obj.Environmental_Vector = self.environmental_vector

        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.overall_score = obj.Overall_Score
        return_obj.base_score = obj.Base_Score
        return_obj.base_vector = obj.Base_Vector
        return_obj.temporal_score = obj.Temporal_Score
        return_obj.temporal_vector = obj.Temporal_Vector
        return_obj.environmental_score = obj.Environmental_Score
        return_obj.environmental_vector = obj.Environmental_Vector

        return return_obj

    def to_dict(self):
        d = {}

        if self.overall_score:
            d['overall_score'] = self.overall_score
        if self.base_score:
            d['base_score'] = self.base_score
        if self.base_vector:
            d['base_vector'] = self.base_vector
        if self.temporal_score:
            d['temporal_score'] = self.temporal_score
        if self.temporal_vector:
            d['temporal_vector'] = self.temporal_vector
        if self.environmental_score:
            d['environmental_score'] = self.environmental_score
        if self.environmental_vector:
            d['environmental_vector'] = self.environmental_vector

        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.overall_score = dict_repr.get('overall_score')
        return_obj.base_score = dict_repr.get('base_score')
        return_obj.base_vector = dict_repr.get('base_vector')
        return_obj.temporal_score = dict_repr.get('temporal_score')
        return_obj.temporal_vector = dict_repr.get('temporal_vector')
        return_obj.environmental_score = dict_repr.get('environmental_score')
        return_obj.environmental_vector = dict_repr.get('environmental_vector')

        return return_obj


class AffectedSoftware(GenericRelationshipList):
    _binding = exploit_target_binding
    _binding_class = exploit_target_binding.AffectedSoftwareType
    _namespace = "http://stix.mitre.org/ExploitTarget-1"
    _binding_var = "Affected_Software"
    _contained_type = RelatedObservable
    _inner_name = "affected_software"


# NOT AN ACTUAL STIX TYPE!
class _Vulnerabilities(stix.TypedList):
    _contained_type = Vulnerability
