cti-python-stix2/stix2/test/v21/test_infrastructure.py

import datetime as dt

import pytest
import pytz

import stix2

from .constants import FAKE_TIME, INFRASTRUCTURE_ID, INFRASTRUCTURE_KWARGS

EXPECTED_INFRASTRUCTURE = """{
    "type": "infrastructure",
    "spec_version": "2.1",
    "id": "infrastructure--3000ae1b-784c-f03d-8abc-0a625b2ff018",
    "created": "2017-01-01T12:34:56.000Z",
    "modified": "2017-01-01T12:34:56.000Z",
    "name": "Poison Ivy C2"
}"""


def test_infrastructure_with_all_required_properties():
    now = dt.datetime(2017, 1, 1, 12, 34, 56, tzinfo=pytz.utc)

    infra = stix2.v21.Infrastructure(
        type="infrastructure",
        id=INFRASTRUCTURE_ID,
        created=now,
        modified=now,
        name="Poison Ivy C2",
    )

    assert infra.serialize(pretty=True) == EXPECTED_INFRASTRUCTURE


def test_infrastructure_autogenerated_properties(infrastructure):
    assert infrastructure.type == 'infrastructure'
    assert infrastructure.id == 'infrastructure--00000000-0000-4000-8000-000000000001'
    assert infrastructure.created == FAKE_TIME
    assert infrastructure.modified == FAKE_TIME
    assert infrastructure.infrastructure_types == ['command-and-control']
    assert infrastructure.name == "Poison Ivy C2"

    assert infrastructure['type'] == 'infrastructure'
    assert infrastructure['id'] == 'infrastructure--00000000-0000-4000-8000-000000000001'
    assert infrastructure['created'] == FAKE_TIME
    assert infrastructure['modified'] == FAKE_TIME
    assert infrastructure['infrastructure_types'] == ['command-and-control']
    assert infrastructure['name'] == "Poison Ivy C2"


def test_infrastructure_type_must_be_infrastructure():
    with pytest.raises(stix2.exceptions.InvalidValueError) as excinfo:
        stix2.v21.Infrastructure(type='xxx', **INFRASTRUCTURE_KWARGS)

    assert excinfo.value.cls == stix2.v21.Infrastructure
    assert excinfo.value.prop_name == "type"
    assert excinfo.value.reason == "must equal 'infrastructure'."
    assert str(excinfo.value) == "Invalid value for Infrastructure 'type': must equal 'infrastructure'."


def test_infrastructure_id_must_start_with_infrastructure():
    with pytest.raises(stix2.exceptions.InvalidValueError) as excinfo:
        stix2.v21.Infrastructure(id='my-prefix--', **INFRASTRUCTURE_KWARGS)

    assert excinfo.value.cls == stix2.v21.Infrastructure
    assert excinfo.value.prop_name == "id"
    assert excinfo.value.reason == "must start with 'infrastructure--'."
    assert str(excinfo.value) == "Invalid value for Infrastructure 'id': must start with 'infrastructure--'."


def test_infrastructure_required_properties():
    with pytest.raises(stix2.exceptions.MissingPropertiesError) as excinfo:
        stix2.v21.Infrastructure()

    assert excinfo.value.cls == stix2.v21.Infrastructure
    assert excinfo.value.properties == ["name"]


def test_infrastructure_required_property_name():
    with pytest.raises(stix2.exceptions.MissingPropertiesError) as excinfo:
        stix2.v21.Infrastructure(infrastructure_types=['command-and-control'])

    assert excinfo.value.cls == stix2.v21.Infrastructure
    assert excinfo.value.properties == ["name"]


def test_invalid_kwarg_to_infrastructure():
    with pytest.raises(stix2.exceptions.ExtraPropertiesError) as excinfo:
        stix2.v21.Infrastructure(my_custom_property="foo", **INFRASTRUCTURE_KWARGS)

    assert excinfo.value.cls == stix2.v21.Infrastructure
    assert excinfo.value.properties == ['my_custom_property']
    assert str(excinfo.value) == "Unexpected properties for Infrastructure: (my_custom_property)."


@pytest.mark.parametrize(
    "data", [
        EXPECTED_INFRASTRUCTURE,
        {
            "type": "infrastructure",
            "spec_version": "2.1",
            "id": INFRASTRUCTURE_ID,
            "created": "2017-01-01T12:34:56.000Z",
            "modified": "2017-01-01T12:34:56.000Z",
            "name": "Poison Ivy C2",
        },
    ],
)
def test_parse_infrastructure(data):
    infra = stix2.parse(data)

    assert infra.type == 'infrastructure'
    assert infra.spec_version == '2.1'
    assert infra.id == INFRASTRUCTURE_ID
    assert infra.created == dt.datetime(2017, 1, 1, 12, 34, 56, tzinfo=pytz.utc)
    assert infra.modified == dt.datetime(2017, 1, 1, 12, 34, 56, tzinfo=pytz.utc)
    assert infra.name == 'Poison Ivy C2'


def test_parse_infrastructure_kill_chain_phases():
    kill_chain = """
      "kill_chain_phases": [
        {
          "kill_chain_name": "lockheed-martin-cyber-kill-chain",
          "phase_name": "reconnaissance"
        }
      ]"""
    data = EXPECTED_INFRASTRUCTURE.replace('infrastructure"', 'infrastructure",%s' % kill_chain)
    infra = stix2.parse(data, version="2.1")
    assert infra.kill_chain_phases[0].kill_chain_name == "lockheed-martin-cyber-kill-chain"
    assert infra.kill_chain_phases[0].phase_name == "reconnaissance"
    assert infra['kill_chain_phases'][0]['kill_chain_name'] == "lockheed-martin-cyber-kill-chain"
    assert infra['kill_chain_phases'][0]['phase_name'] == "reconnaissance"


def test_parse_infrastructure_clean_kill_chain_phases():
    kill_chain = """
      "kill_chain_phases": [
        {
          "kill_chain_name": "lockheed-martin-cyber-kill-chain",
          "phase_name": 1
        }
      ]"""
    data = EXPECTED_INFRASTRUCTURE.replace('2.1"', '2.1",%s' % kill_chain)
    infra = stix2.parse(data, version="2.1")
    assert infra['kill_chain_phases'][0]['phase_name'] == "1"


def test_infrastructure_invalid_last_before_first():
    with pytest.raises(ValueError) as excinfo:
        stix2.v21.Infrastructure(first_seen="2017-01-01T12:34:56.000Z", last_seen="2017-01-01T12:33:56.000Z", **INFRASTRUCTURE_KWARGS)

    assert "'last_seen' must be greater than or equal to 'first_seen'" in str(excinfo.value)