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", "infrastructure_types": [ "command-and-control" ] }""" 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", infrastructure_types=["command-and-control"], ) assert str(infra) == 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 == ["infrastructure_types", "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", "infrastructure_types": ["command-and-control"], "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.infrastructure_types == ['command-and-control'] 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)