cti-python-stix2/stix2/test/test_vulnerability.py

import datetime as dt

import pytest
import pytz
import stix2

from .constants import VULNERABILITY_ID

EXPECTED = """{
    "created": "2016-05-12T08:17:27Z",
    "external_references": [
        {
            "external_id": "CVE-2016-1234",
            "source_name": "cve"
        }
    ],
    "id": "vulnerability--0c7b5b88-8ff7-4a4d-aa9d-feb398cd0061",
    "modified": "2016-05-12T08:17:27Z",
    "name": "CVE-2016-1234",
    "type": "vulnerability"
}"""


def test_vulnerability_example():
    vulnerability = stix2.Vulnerability(
        id="vulnerability--0c7b5b88-8ff7-4a4d-aa9d-feb398cd0061",
        created="2016-05-12T08:17:27Z",
        modified="2016-05-12T08:17:27Z",
        name="CVE-2016-1234",
        external_references=[
            stix2.ExternalReference(source_name='cve',
                                    external_id="CVE-2016-1234"),
        ],
    )

    assert str(vulnerability) == EXPECTED


@pytest.mark.parametrize("data", [
    EXPECTED,
    {
        "created": "2016-05-12T08:17:27Z",
        "external_references": [
            {
                "external_id": "CVE-2016-1234",
                "source_name": "cve"
            }
        ],
        "id": "vulnerability--0c7b5b88-8ff7-4a4d-aa9d-feb398cd0061",
        "modified": "2016-05-12T08:17:27Z",
        "name": "CVE-2016-1234",
        "type": "vulnerability"
    },
])
def test_parse_vulnerability(data):
    vuln = stix2.parse(data)

    assert vuln.type == 'vulnerability'
    assert vuln.id == VULNERABILITY_ID
    assert vuln.created == dt.datetime(2016, 5, 12, 8, 17, 27, tzinfo=pytz.utc)
    assert vuln.modified == dt.datetime(2016, 5, 12, 8, 17, 27, tzinfo=pytz.utc)
    assert vuln.name == "CVE-2016-1234"
    assert vuln.external_references[0].external_id == "CVE-2016-1234"
    assert vuln.external_references[0].source_name == "cve"

# TODO: Add other examples