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

import stix
import stix.bindings.exploit_target as exploit_target_binding
from stix.common import DateTimeWithPrecision, StructuredText
from stix.common.related import GenericRelationshipList, RelatedObservable

# TODO document these
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()
        # 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):
        """The ``description`` property for this :class:`Vulnerability`.

        Default Value: ``None``

        Note:
            If set to a value that is not an instance of
            :class:`stix.common.structured_text.StructuredText`, an attempt to
            will be made to convert the value into an instance of
            :class:`stix.common.structured_text.StructuredText`.

        Returns:
            An instance of
            :class:`stix.common.structured_text.StructuredText`

        """

        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):
        """The ``short_description`` property for this :class:`Vulnerability`.

        Default Value: ``None``

        Note:
            If set to a value that is not an instance of
            :class:`stix.common.structured_text.StructuredText`, an attempt to
            will be made to convert the value into an instance of
            :class:`stix.common.structured_text.StructuredText`.

        Returns:
            An instance of
            :class:`stix.common.structured_text.StructuredText`

        """
        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 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

        """
        if value:
            if isinstance(value, DateTimeWithPrecision):
                self._discovered_datetime = value
            else:
                self._discovered_datetime = DateTimeWithPrecision(value=value)
        else:
            self._discovered_datetime = None

    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()

        if self.is_known is not None:
            return_obj.is_known = self.is_known
        if self.is_publicly_acknowledged is not None:
            return_obj.is_publicly_acknowledged = self.is_publicly_acknowledged
        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)
        return_obj.CVE_ID = self.cve_id
        return_obj.OSVDB_ID = self.osvdb_id
        return_obj.Source = self.source
        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)

        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 = obj.is_known
        return_obj.is_publicly_acknowledged = obj.is_publicly_acknowledged
        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.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)

        return return_obj

    def to_dict(self):
        d = {}

        if self.is_known is not None:
            d['is_known'] = self.is_known
        if self.is_publicly_acknowledged is not None:
            d['is_publicly_acknowledged'] = self.is_publicly_acknowledged
        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.cve_id:
            d['cve_id'] = self.cve_id
        if self.osvdb_id:
            d['osvdb_id'] = self.osvdb_id
        if self.source:
            d['source'] = self.source
        if self.cvss_score:
            d['cvss_score'] = self.cvss_score.to_dict()
        if self.discovered_datetime:
            d['discovered_datetime'] = self.discovered_datetime.to_dict()
        if self.published_datetime:
            d['published_datetime'] = self.published_datetime.to_dict()
        if self.affected_software:
            d['affected_software'] = self.affected_software.to_dict()

        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.is_known = dict_repr.get('is_known')
        return_obj.is_publicly_acknowledged = dict_repr.get('is_publicly_acknowledged')
        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.cve_id = dict_repr.get('cve_id')
        return_obj.osvdb_id = dict_repr.get('osvdb_id')
        return_obj.source = dict_repr.get('source')
        return_obj.cvss_score = CVSSVector.from_dict(dict_repr.get('cvss_score'))
        return_obj.discovered_datetime = DateTimeWithPrecision.from_dict(dict_repr.get('discovered_datetime'))
        return_obj.published_datetime = DateTimeWithPrecision.from_dict(dict_repr.get('published_datetime'))
        return_obj.affected_software = AffectedSoftware.from_dict(dict_repr.get('affected_software'))

        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"
