diff --git a/stix2/test/constants.py b/stix2/test/constants.py index d8b66fe..3db39d6 100644 --- a/stix2/test/constants.py +++ b/stix2/test/constants.py @@ -28,6 +28,11 @@ MARKING_IDS = [ "marking-definition--68520ae2-fefe-43a9-84ee-2c2a934d2c7d", "marking-definition--2802dfb1-1019-40a8-8848-68d0ec0e417f", ] +RELATIONSHIP_IDS = [ + 'relationship--06520621-5352-4e6a-b976-e8fa3d437ffd', + 'relationship--181c9c09-43e6-45dd-9374-3bec192f05ef', + 'relationship--a0cbb21c-8daf-4a7f-96aa-7155a4ef8f70' +] # All required args for a Campaign instance CAMPAIGN_KWARGS = dict( diff --git a/stix2/test/test_environment.py b/stix2/test/test_environment.py index 086e446..e595eb9 100644 --- a/stix2/test/test_environment.py +++ b/stix2/test/test_environment.py @@ -4,11 +4,7 @@ import stix2 from .constants import (CAMPAIGN_ID, CAMPAIGN_KWARGS, FAKE_TIME, IDENTITY_ID, IDENTITY_KWARGS, INDICATOR_ID, INDICATOR_KWARGS, - MALWARE_ID, MALWARE_KWARGS) - -RELATIONSHIP_ID1 = 'relationship--06520621-5352-4e6a-b976-e8fa3d437ffd' -RELATIONSHIP_ID2 = 'relationship--181c9c09-43e6-45dd-9374-3bec192f05ef' -RELATIONSHIP_ID3 = 'relationship--a0cbb21c-8daf-4a7f-96aa-7155a4ef8f70' + MALWARE_ID, MALWARE_KWARGS, RELATIONSHIP_IDS) @pytest.fixture @@ -17,9 +13,9 @@ def ds(): idy = stix2.Identity(id=IDENTITY_ID, **IDENTITY_KWARGS) ind = stix2.Indicator(id=INDICATOR_ID, **INDICATOR_KWARGS) mal = stix2.Malware(id=MALWARE_ID, **MALWARE_KWARGS) - rel1 = stix2.Relationship(ind, 'indicates', mal, id=RELATIONSHIP_ID1) - rel2 = stix2.Relationship(mal, 'targets', idy, id=RELATIONSHIP_ID2) - rel3 = stix2.Relationship(cam, 'uses', mal, id=RELATIONSHIP_ID3) + rel1 = stix2.Relationship(ind, 'indicates', mal, id=RELATIONSHIP_IDS[0]) + rel2 = stix2.Relationship(mal, 'targets', idy, id=RELATIONSHIP_IDS[1]) + rel3 = stix2.Relationship(cam, 'uses', mal, id=RELATIONSHIP_IDS[2]) stix_objs = [cam, idy, ind, mal, rel1, rel2, rel3] yield stix2.MemoryStore(stix_objs) @@ -242,9 +238,9 @@ def test_relationships(ds): resp = env.relationships(mal) assert len(resp) == 3 - assert any(x['id'] == RELATIONSHIP_ID1 for x in resp) - assert any(x['id'] == RELATIONSHIP_ID2 for x in resp) - assert any(x['id'] == RELATIONSHIP_ID3 for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[0] for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[1] for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[2] for x in resp) def test_relationships_by_type(ds): @@ -253,7 +249,7 @@ def test_relationships_by_type(ds): resp = env.relationships(mal, relationship_type='indicates') assert len(resp) == 1 - assert resp[0]['id'] == RELATIONSHIP_ID1 + assert resp[0]['id'] == RELATIONSHIP_IDS[0] def test_relationships_by_source(ds): @@ -261,7 +257,7 @@ def test_relationships_by_source(ds): resp = env.relationships(MALWARE_ID, source_only=True) assert len(resp) == 1 - assert resp[0]['id'] == RELATIONSHIP_ID2 + assert resp[0]['id'] == RELATIONSHIP_IDS[1] def test_relationships_by_target(ds): @@ -269,8 +265,8 @@ def test_relationships_by_target(ds): resp = env.relationships(MALWARE_ID, target_only=True) assert len(resp) == 2 - assert any(x['id'] == RELATIONSHIP_ID1 for x in resp) - assert any(x['id'] == RELATIONSHIP_ID3 for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[0] for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[2] for x in resp) def test_relationships_by_target_and_type(ds): @@ -278,7 +274,7 @@ def test_relationships_by_target_and_type(ds): resp = env.relationships(MALWARE_ID, relationship_type='uses', target_only=True) assert len(resp) == 1 - assert any(x['id'] == RELATIONSHIP_ID3 for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[2] for x in resp) def test_relationships_by_target_and_source(ds): diff --git a/stix2/test/test_filesystem.py b/stix2/test/test_filesystem.py index 85f6966..66ca47f 100644 --- a/stix2/test/test_filesystem.py +++ b/stix2/test/test_filesystem.py @@ -4,7 +4,12 @@ import shutil import pytest from stix2 import (Bundle, Campaign, CustomObject, FileSystemSink, - FileSystemSource, FileSystemStore, Filter, properties) + FileSystemSource, FileSystemStore, Filter, Identity, + Indicator, Malware, Relationship, properties) + +from .constants import (CAMPAIGN_ID, CAMPAIGN_KWARGS, IDENTITY_ID, + IDENTITY_KWARGS, INDICATOR_ID, INDICATOR_KWARGS, + MALWARE_ID, MALWARE_KWARGS, RELATIONSHIP_IDS) FS_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "stix2_data") @@ -40,6 +45,25 @@ def fs_sink(): shutil.rmtree(os.path.join(FS_PATH, "campaign"), True) +@pytest.fixture +def rel_fs_store(): + cam = Campaign(id=CAMPAIGN_ID, **CAMPAIGN_KWARGS) + idy = Identity(id=IDENTITY_ID, **IDENTITY_KWARGS) + ind = Indicator(id=INDICATOR_ID, **INDICATOR_KWARGS) + mal = Malware(id=MALWARE_ID, **MALWARE_KWARGS) + rel1 = Relationship(ind, 'indicates', mal, id=RELATIONSHIP_IDS[0]) + rel2 = Relationship(mal, 'targets', idy, id=RELATIONSHIP_IDS[1]) + rel3 = Relationship(cam, 'uses', mal, id=RELATIONSHIP_IDS[2]) + stix_objs = [cam, idy, ind, mal, rel1, rel2, rel3] + fs = FileSystemStore(FS_PATH) + for o in stix_objs: + fs.add(o) + yield fs + + for o in stix_objs: + os.remove(os.path.join(FS_PATH, o.type, o.id + '.json')) + + def test_filesystem_source_nonexistent_folder(): with pytest.raises(ValueError) as excinfo: FileSystemSource('nonexistent-folder') @@ -375,3 +399,50 @@ def test_filesystem_custom_object(fs_store): # remove dir shutil.rmtree(os.path.join(FS_PATH, "x-new-obj"), True) + + +def test_relationships(rel_fs_store): + mal = rel_fs_store.get(MALWARE_ID) + resp = rel_fs_store.relationships(mal) + + assert len(resp) == 3 + assert any(x['id'] == RELATIONSHIP_IDS[0] for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[1] for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[2] for x in resp) + + +def test_relationships_by_type(rel_fs_store): + mal = rel_fs_store.get(MALWARE_ID) + resp = rel_fs_store.relationships(mal, relationship_type='indicates') + + assert len(resp) == 1 + assert resp[0]['id'] == RELATIONSHIP_IDS[0] + + +def test_relationships_by_source(rel_fs_store): + resp = rel_fs_store.relationships(MALWARE_ID, source_only=True) + + assert len(resp) == 1 + assert resp[0]['id'] == RELATIONSHIP_IDS[1] + + +def test_relationships_by_target(rel_fs_store): + resp = rel_fs_store.relationships(MALWARE_ID, target_only=True) + + assert len(resp) == 2 + assert any(x['id'] == RELATIONSHIP_IDS[0] for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[2] for x in resp) + + +def test_relationships_by_target_and_type(rel_fs_store): + resp = rel_fs_store.relationships(MALWARE_ID, relationship_type='uses', target_only=True) + + assert len(resp) == 1 + assert any(x['id'] == RELATIONSHIP_IDS[2] for x in resp) + + +def test_relationships_by_target_and_source(rel_fs_store): + with pytest.raises(ValueError) as excinfo: + rel_fs_store.relationships(MALWARE_ID, target_only=True, source_only=True) + + assert 'not both' in str(excinfo.value) diff --git a/stix2/test/test_memory.py b/stix2/test/test_memory.py index 6b1219e..d26b7d3 100644 --- a/stix2/test/test_memory.py +++ b/stix2/test/test_memory.py @@ -3,10 +3,15 @@ import shutil import pytest -from stix2 import (Bundle, Campaign, CustomObject, Filter, MemorySource, - MemoryStore, properties) +from stix2 import (Bundle, Campaign, CustomObject, Filter, Identity, Indicator, + Malware, MemorySource, MemoryStore, Relationship, + properties) from stix2.sources import make_id +from .constants import (CAMPAIGN_ID, CAMPAIGN_KWARGS, IDENTITY_ID, + IDENTITY_KWARGS, INDICATOR_ID, INDICATOR_KWARGS, + MALWARE_ID, MALWARE_KWARGS, RELATIONSHIP_IDS) + IND1 = { "created": "2017-01-27T13:49:53.935Z", "id": "indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f", @@ -118,6 +123,19 @@ def mem_source(): yield MemorySource(STIX_OBJS1) +@pytest.fixture +def rel_mem_store(): + cam = Campaign(id=CAMPAIGN_ID, **CAMPAIGN_KWARGS) + idy = Identity(id=IDENTITY_ID, **IDENTITY_KWARGS) + ind = Indicator(id=INDICATOR_ID, **INDICATOR_KWARGS) + mal = Malware(id=MALWARE_ID, **MALWARE_KWARGS) + rel1 = Relationship(ind, 'indicates', mal, id=RELATIONSHIP_IDS[0]) + rel2 = Relationship(mal, 'targets', idy, id=RELATIONSHIP_IDS[1]) + rel3 = Relationship(cam, 'uses', mal, id=RELATIONSHIP_IDS[2]) + stix_objs = [cam, idy, ind, mal, rel1, rel2, rel3] + yield MemoryStore(stix_objs) + + def test_memory_source_get(mem_source): resp = mem_source.get("indicator--d81f86b8-975b-bc0b-775e-810c5ad45a4f") assert resp["id"] == "indicator--d81f86b8-975b-bc0b-775e-810c5ad45a4f" @@ -287,3 +305,50 @@ def test_memory_store_custom_object(mem_store): newobj_r = mem_store.get(newobj.id) assert newobj_r.id == newobj.id assert newobj_r.property1 == 'something' + + +def test_relationships(rel_mem_store): + mal = rel_mem_store.get(MALWARE_ID) + resp = rel_mem_store.relationships(mal) + + assert len(resp) == 3 + assert any(x['id'] == RELATIONSHIP_IDS[0] for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[1] for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[2] for x in resp) + + +def test_relationships_by_type(rel_mem_store): + mal = rel_mem_store.get(MALWARE_ID) + resp = rel_mem_store.relationships(mal, relationship_type='indicates') + + assert len(resp) == 1 + assert resp[0]['id'] == RELATIONSHIP_IDS[0] + + +def test_relationships_by_source(rel_mem_store): + resp = rel_mem_store.relationships(MALWARE_ID, source_only=True) + + assert len(resp) == 1 + assert resp[0]['id'] == RELATIONSHIP_IDS[1] + + +def test_relationships_by_target(rel_mem_store): + resp = rel_mem_store.relationships(MALWARE_ID, target_only=True) + + assert len(resp) == 2 + assert any(x['id'] == RELATIONSHIP_IDS[0] for x in resp) + assert any(x['id'] == RELATIONSHIP_IDS[2] for x in resp) + + +def test_relationships_by_target_and_type(rel_mem_store): + resp = rel_mem_store.relationships(MALWARE_ID, relationship_type='uses', target_only=True) + + assert len(resp) == 1 + assert any(x['id'] == RELATIONSHIP_IDS[2] for x in resp) + + +def test_relationships_by_target_and_source(rel_mem_store): + with pytest.raises(ValueError) as excinfo: + rel_mem_store.relationships(MALWARE_ID, target_only=True, source_only=True) + + assert 'not both' in str(excinfo.value)