diff --git a/stix2/markings/__init__.py b/stix2/markings/__init__.py index f543f34..f158723 100644 --- a/stix2/markings/__init__.py +++ b/stix2/markings/__init__.py @@ -63,9 +63,9 @@ def set_markings(obj, selectors, marking): """ if selectors is None: - object_markings.set_markings(obj, marking) + return object_markings.set_markings(obj, marking) else: - granular_markings.set_markings(obj, selectors, marking) + return granular_markings.set_markings(obj, selectors, marking) def remove_markings(obj, selectors, marking): @@ -89,9 +89,9 @@ def remove_markings(obj, selectors, marking): """ if selectors is None: - object_markings.remove_markings(obj, marking) + return object_markings.remove_markings(obj, marking) else: - granular_markings.remove_markings(obj, selectors, marking) + return granular_markings.remove_markings(obj, selectors, marking) def add_markings(obj, selectors, marking): @@ -134,9 +134,9 @@ def clear_markings(obj, selectors): """ if selectors is None: - object_markings.clear_markings(obj) + return object_markings.clear_markings(obj) else: - granular_markings.clear_markings(obj, selectors) + return granular_markings.clear_markings(obj, selectors) def is_marked(obj, selectors, marking=None, inherited=False, descendants=False): diff --git a/stix2/markings/object_markings.py b/stix2/markings/object_markings.py index 9d2b2bc..c0d6d72 100644 --- a/stix2/markings/object_markings.py +++ b/stix2/markings/object_markings.py @@ -63,8 +63,11 @@ def remove_markings(obj, marking): raise AssertionError("Unable to remove Object Level Marking(s) from " "internal collection. Marking(s) not found...") - return obj.new_version(object_marking_refs=[x for x in object_markings - if x not in marking]) + new_markings = [x for x in object_markings if x not in marking] + if new_markings: + return obj.new_version(object_marking_refs=new_markings) + else: + return obj.new_version(object_marking_refs=None) def set_markings(obj, marking): diff --git a/stix2/properties.py b/stix2/properties.py index f63ec8b..69d2309 100644 --- a/stix2/properties.py +++ b/stix2/properties.py @@ -311,6 +311,7 @@ class ReferenceProperty(Property): def clean(self, value): if isinstance(value, _STIXBase): value = value.id + value = str(value) if self.type: if not value.startswith(self.type): raise ValueError("must start with '{0}'.".format(self.type)) diff --git a/stix2/test/constants.py b/stix2/test/constants.py index 958120b..343488f 100644 --- a/stix2/test/constants.py +++ b/stix2/test/constants.py @@ -20,6 +20,15 @@ TOOL_ID = "tool--8e2e2d2b-17d4-4cbf-938f-98ee46b3cd3f" SIGHTING_ID = "sighting--bfbc19db-ec35-4e45-beed-f8bde2a772fb" VULNERABILITY_ID = "vulnerability--0c7b5b88-8ff7-4a4d-aa9d-feb398cd0061" +MARKING_IDS = [ + "marking-definition--613f2e26-407d-48c7-9eca-b8e91df99dc9", + "marking-definition--443eb5c3-a76c-4a0a-8caa-e93998e7bc09", + "marking-definition--57fcd772-9c1d-41b0-8d1f-3d47713415d9", + "marking-definition--462bf1a6-03d2-419c-b74e-eee2238b2de4", + "marking-definition--68520ae2-fefe-43a9-84ee-2c2a934d2c7d", + "marking-definition--2802dfb1-1019-40a8-8848-68d0ec0e417f", +] + # Minimum required args for an Identity instance IDENTITY_KWARGS = dict( name="John Smith", diff --git a/stix2/test/test_object_markings.py b/stix2/test/test_object_markings.py index 15f35a4..f2d0374 100644 --- a/stix2/test/test_object_markings.py +++ b/stix2/test/test_object_markings.py @@ -1,10 +1,20 @@ import pytest -from stix2 import markings +from stix2 import Malware, exceptions, markings + +from .constants import FAKE_TIME, MALWARE_ID, MARKING_IDS +from .constants import MALWARE_KWARGS as MALWARE_KWARGS_CONST """Tests for the Data Markings API.""" +MALWARE_KWARGS = MALWARE_KWARGS_CONST.copy() +MALWARE_KWARGS.update({ + 'id': MALWARE_ID, + 'created': FAKE_TIME, + 'modified': FAKE_TIME, +}) + def test_add_markings_one_marking(): before = { @@ -15,10 +25,10 @@ def test_add_markings_one_marking(): after = { "title": "test title", "description": "test description", - "object_marking_refs": ["marking-definition--1"] + "object_marking_refs": [MARKING_IDS[0]] } - markings.add_markings(before, None, "marking-definition--1") + markings.add_markings(before, None, MARKING_IDS[0]) assert before == after @@ -32,41 +42,38 @@ def test_add_markings_multiple_marking(): after = { "title": "test title", "description": "test description", - "object_marking_refs": ["marking-definition--1", "marking-definition--2"] + "object_marking_refs": [MARKING_IDS[0], MARKING_IDS[1]] } - markings.add_markings(before, None, ["marking-definition--1", "marking-definition--2"]) + markings.add_markings(before, None, [MARKING_IDS[0], MARKING_IDS[1]]) for m in before["object_marking_refs"]: assert m in after["object_marking_refs"] def test_add_markings_combination(): - before = { - "title": "test title", - "description": "test description" - } - - after = { - "title": "test title", - "description": "test description", - "object_marking_refs": ["marking-definition--1", "marking-definition--2"], - "granular_markings": [ + before = Malware( + **MALWARE_KWARGS + ) + after = Malware( + object_marking_refs=[MARKING_IDS[0], MARKING_IDS[1]], + granular_markings=[ { - "selectors": ["title"], - "marking_ref": "marking-definition--3" + "selectors": ["labels"], + "marking_ref": MARKING_IDS[2] }, { - "selectors": ["description"], - "marking_ref": "marking-definition--4" + "selectors": ["name"], + "marking_ref": MARKING_IDS[3] }, - ] - } + ], + **MALWARE_KWARGS + ) - markings.add_markings(before, None, "marking-definition--1") - markings.add_markings(before, None, "marking-definition--2") - markings.add_markings(before, "title", "marking-definition--3") - markings.add_markings(before, "description", "marking-definition--4") + before = markings.add_markings(before, None, MARKING_IDS[0]) + before = markings.add_markings(before, None, MARKING_IDS[1]) + before = markings.add_markings(before, "labels", MARKING_IDS[2]) + before = markings.add_markings(before, "name", MARKING_IDS[3]) for m in before["granular_markings"]: assert m in after["granular_markings"] @@ -79,15 +86,14 @@ def test_add_markings_combination(): ([""]), (""), ([]), - (["marking-definition--1", 456]) + ([MARKING_IDS[0], 456]) ]) def test_add_markings_bad_markings(data): - before = { - "title": "test title", - "description": "test description" - } - with pytest.raises(AssertionError): - markings.add_markings(before, None, data) + before = Malware( + **MALWARE_KWARGS + ) + with pytest.raises(exceptions.InvalidValueError): + before = markings.add_markings(before, None, data) assert "object_marking_refs" not in before @@ -240,65 +246,58 @@ def test_get_markings_object_and_granular_combinations(data): def test_remove_markings_object_level(): - after = { - "title": "test title", - "description": "test description" - } + before = Malware( + object_marking_refs=[MARKING_IDS[0]], + **MALWARE_KWARGS + ) + after = Malware( + **MALWARE_KWARGS + ) - before = { - "title": "test title", - "description": "test description", - "object_marking_refs": ["marking-definition--1"] - } + before = markings.remove_markings(before, None, MARKING_IDS[0]) - markings.remove_markings(before, None, "marking-definition--1") - - assert before == after + assert 'object_marking_refs' not in before + assert 'object_marking_refs' not in after def test_remove_markings_multiple(): - after = { - "title": "test title", - "description": "test description", - "object_marking_refs": ["marking-definition--2"] - } + before = Malware( + object_marking_refs=[MARKING_IDS[0], MARKING_IDS[1], MARKING_IDS[2]], + **MALWARE_KWARGS + ) + after = Malware( + object_marking_refs=[MARKING_IDS[1]], + **MALWARE_KWARGS + ) - before = { - "title": "test title", - "description": "test description", - "object_marking_refs": ["marking-definition--1", "marking-definition--2", "marking-definition--3"] - } + before = markings.remove_markings(before, None, [MARKING_IDS[0], MARKING_IDS[2]]) - markings.remove_markings(before, None, ["marking-definition--1", "marking-definition--3"]) - - assert before == after + assert before['object_marking_refs'] == after['object_marking_refs'] def test_remove_markings_bad_markings(): before = { "title": "test title", "description": "test description", - "object_marking_refs": ["marking-definition--1", "marking-definition--2", "marking-definition--3"] + "object_marking_refs": [MARKING_IDS[0], MARKING_IDS[1], MARKING_IDS[2]] } with pytest.raises(AssertionError): - markings.remove_markings(before, None, ["marking-definition--5"]) + markings.remove_markings(before, None, [MARKING_IDS[4]]) def test_clear_markings(): - after = { - "title": "test title", - "description": "test description" - } + before = Malware( + object_marking_refs=[MARKING_IDS[0], MARKING_IDS[1], MARKING_IDS[2]], + **MALWARE_KWARGS + ) + after = Malware( + **MALWARE_KWARGS + ) - before = { - "title": "test title", - "description": "test description", - "object_marking_refs": ["marking-definition--1", "marking-definition--2", "marking-definition--3"] - } + before = markings.clear_markings(before, None) - markings.clear_markings(before, None) - - assert before == after + assert 'object_marking_refs' not in before + assert 'object_marking_refs' not in after def test_is_marked_object_and_granular_combinations(): @@ -442,23 +441,21 @@ def test_is_marked_object_and_granular_combinations(): def test_set_marking(): - before = { - "title": "test title", - "description": "test description", - "object_marking_refs": ["marking-definition--1", "marking-definition--2", "marking-definition--3"] - } - after = { - "title": "test title", - "description": "test description", - "object_marking_refs": ["marking-definition--7", "marking-definition--9"] - } + before = Malware( + object_marking_refs=[MARKING_IDS[0], MARKING_IDS[1], MARKING_IDS[2]], + **MALWARE_KWARGS + ) + after = Malware( + object_marking_refs=[MARKING_IDS[4], MARKING_IDS[5]], + **MALWARE_KWARGS + ) - markings.set_markings(before, None, ["marking-definition--7", "marking-definition--9"]) + before = markings.set_markings(before, None, [MARKING_IDS[4], MARKING_IDS[5]]) for m in before["object_marking_refs"]: - assert m in ["marking-definition--7", "marking-definition--9"] + assert m in [MARKING_IDS[4], MARKING_IDS[5]] - assert ["marking-definition--1", "marking-definition--2", "marking-definition--3"] not in before["object_marking_refs"] + assert [MARKING_IDS[0], MARKING_IDS[1], MARKING_IDS[2]] not in before["object_marking_refs"] for x in before["object_marking_refs"]: assert x in after["object_marking_refs"] @@ -468,20 +465,18 @@ def test_set_marking(): ([]), ([""]), (""), - (["marking-definition--7", 687]) + ([MARKING_IDS[4], 687]) ]) def test_set_marking_bad_input(data): - before = { - "description": "test description", - "title": "foo", - "object_marking_refs": ["marking-definition--1"] - } - after = { - "description": "test description", - "title": "foo", - "object_marking_refs": ["marking-definition--1"] - } - with pytest.raises(AssertionError): - markings.set_markings(before, None, data) + before = Malware( + object_marking_refs=[MARKING_IDS[0]], + **MALWARE_KWARGS + ) + after = Malware( + object_marking_refs=[MARKING_IDS[0]], + **MALWARE_KWARGS + ) + with pytest.raises(exceptions.InvalidValueError): + before = markings.set_markings(before, None, data) assert before == after