153 lines
		
	
	
		
			5.8 KiB
		
	
	
	
		
			Python
		
	
	
			
		
		
	
	
			153 lines
		
	
	
		
			5.8 KiB
		
	
	
	
		
			Python
		
	
	
| 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 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 == ["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)
 |