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

import unittest

from stix.test import EntityTestCase, TypedListTestCase, assert_warnings
from stix.test import data_marking_test
from stix.test.common import information_source_test, related_test

from stix.core import STIXPackage
import stix.exploit_target as et
from stix.exploit_target import weakness, vulnerability, configuration

class CVSSVectorTests(EntityTestCase, unittest.TestCase):
    klass = vulnerability.CVSSVector

    _full_dict = {
        'overall_score': "9.3",
        'base_score': "8.3",
        'base_vector': "(AV:N/AC:L/Au:N/C:N/I:N/A:P)",
        'temporal_score': "7.7",
        'temporal_vector': "(E:U/RL:T/RC:C)",
        'environmental_score': "2.0",
        'environmental_vector': "(CPD:ND/TD:ND)",
    }


class AffectedSoftwareTests(EntityTestCase, unittest.TestCase):
    klass = vulnerability.AffectedSoftware

    _full_dict = {
        'scope': 'inclusive',
        'affected_software': [
            related_test.RelatedObservableTests._full_dict
        ]
    }


class VulnerabilityTests(EntityTestCase, unittest.TestCase):
    klass = vulnerability.Vulnerability

    _full_dict = {
        'is_known': False,
        'is_publicly_acknowledged': True,
        'title': "CVE-2012-0158",
        'description': "Vulnerability Description",
        'short_description': "MSCOMCTL.OCX Memory Corruption",
        'cvss_score': CVSSVectorTests._full_dict,
        'discovered_datetime': {
            'value': '2010-02-21T00:00:00',
            'precision': 'day',
        },
        'published_datetime': {
            'value': '2010-03-01T00:00:00',
            'precision': 'month',
        },
        'affected_software': AffectedSoftwareTests._full_dict,
        'references': ['foo','bar']
    }


class VulnerabilitiesTests(TypedListTestCase, unittest.TestCase):
    klass = vulnerability._Vulnerabilities

    _full_dict = [
        VulnerabilityTests._full_dict
    ]


class PotentialCOAsTests(EntityTestCase, unittest.TestCase):
    klass = et.PotentialCOAs

    _full_dict = {
        'scope': 'inclusive',
        'coas': [
            related_test.RelatedCOATests._full_dict
        ]
    }


class RelatedExploitTargetsTests(EntityTestCase, unittest.TestCase):
    klass = et.RelatedExploitTargets

    _full_dict = {
        'scope': 'inclusive',
        'related_exploit_targets': [
            related_test.RelatedExploitTargetTests._full_dict
        ]
    }

class WeaknessTests(EntityTestCase, unittest.TestCase):
    klass = weakness.Weakness

    _full_dict = {
        'description': "Deadlock",
        'cwe_id': "CWE-833",
    }


class WeaknessesTests(TypedListTestCase, unittest.TestCase):
    klass = weakness._Weaknesses

    _full_dict = [
        WeaknessTests._full_dict
    ]



class ConfigurationTests(EntityTestCase, unittest.TestCase):
    klass = configuration.Configuration

    _full_dict =  {
        'description': "The 'Games' features should be configured"
                       "correctly.",
        'short_description': "Games feature",
        'cce_id': "CCE-18880-5",
    }


class ExploitTargetTests(EntityTestCase, unittest.TestCase):
    klass = et.ExploitTarget
    _full_dict = {
        'id': 'example:test-1',
        #idref omitted since it should not have both an ID and an IDREF.
        'timestamp': "2014-04-01T03:17:45",
        'version': '1.1',
        'title': "ExploitTarget 1",
        'description': "This is a long description about an ExploitTarget",
        'short_description': "an ExploitTarget",
        'vulnerabilities': VulnerabilitiesTests._full_dict,
        'weaknesses': WeaknessesTests._full_dict,
        'configuration': [ConfigurationTests._full_dict],
        'potential_coas': PotentialCOAsTests._full_dict,
        'information_source': information_source_test.InformationSourceTests._full_dict,
        'handling': data_marking_test.MarkingTests._full_dict,
        'related_exploit_targets': RelatedExploitTargetsTests._full_dict,
        'related_packages': related_test.RelatedPackageRefsTests._full_dict
    }

    def test_add_description(self):
        o1 = self.klass()
        o2 = self.klass()

        o1.add_description("Test")
        o2.descriptions.add("Test")

        self.assertEqual(
            o1.descriptions.to_dict(),
            o2.descriptions.to_dict()
        )

    def test_add_short_description(self):
        o1 = self.klass()
        o2 = self.klass()

        o1.add_short_description("Test")
        o2.short_descriptions.add("Test")

        self.assertEqual(
            o1.short_descriptions.to_dict(),
            o2.short_descriptions.to_dict()
        )

    @assert_warnings
    def test_deprecated_related_packages(self):
        e = et.ExploitTarget()
        e.related_packages.append(STIXPackage())
        self.assertEqual(len(e.related_packages), 1)

if __name__ == "__main__":
    unittest.main()
