From c911cff97fe9d89855c83e805f0702f2d9bdc02e Mon Sep 17 00:00:00 2001 From: "Desai, Kartikey H" Date: Fri, 27 Mar 2020 14:58:18 -0400 Subject: [PATCH] Add duplicate checking to markings and observable extensions, and fix some tests and add some tests. Fixes #363 --- stix2/core.py | 13 +++++-- stix2/exceptions.py | 14 +++++--- stix2/test/v20/test_core.py | 53 ---------------------------- stix2/test/v20/test_custom.py | 65 ++++++++++++++++++++++++++++++++--- stix2/test/v21/test_core.py | 43 ----------------------- stix2/test/v21/test_custom.py | 56 ++++++++++++++++++++++++++---- 6 files changed, 129 insertions(+), 115 deletions(-) diff --git a/stix2/core.py b/stix2/core.py index ffe77f2..2f8de3e 100644 --- a/stix2/core.py +++ b/stix2/core.py @@ -8,7 +8,7 @@ import re import stix2 from .base import _Observable, _STIXBase -from .exceptions import DuplicateObjectRegistrationError, ParseError +from .exceptions import DuplicateRegistrationError, ParseError from .markings import _MarkingsMixin from .utils import SCO21_EXT_REGEX, TYPE_REGEX, _get_dict @@ -218,7 +218,7 @@ def _register_object(new_type, version=None): OBJ_MAP = STIX2_OBJ_MAPS[v]['objects'] if new_type._type in OBJ_MAP.keys(): - raise DuplicateObjectRegistrationError("An object with type '%s' already exists and cannot be created again." % new_type._type) + raise DuplicateRegistrationError("STIX Object", new_type._type) OBJ_MAP[new_type._type] = new_type @@ -238,6 +238,8 @@ def _register_marking(new_marking, version=None): v = 'v' + stix2.DEFAULT_VERSION.replace('.', '') OBJ_MAP_MARKING = STIX2_OBJ_MAPS[v]['markings'] + if new_marking._type in OBJ_MAP_MARKING.keys(): + raise DuplicateRegistrationError("STIX Marking", new_marking._type) OBJ_MAP_MARKING[new_marking._type] = new_marking @@ -258,7 +260,7 @@ def _register_observable(new_observable, version=None): OBJ_MAP_OBSERVABLE = STIX2_OBJ_MAPS[v]['observables'] if new_observable._type in OBJ_MAP_OBSERVABLE.keys(): - raise DuplicateObjectRegistrationError("An observable with type '%s' already exists and cannot be created again." % new_observable._type) + raise DuplicateRegistrationError("Cyber Observable", new_observable._type) OBJ_MAP_OBSERVABLE[new_observable._type] = new_observable @@ -323,6 +325,11 @@ def _register_observable_extension( EXT_MAP = STIX2_OBJ_MAPS[v]['observable-extensions'] try: + try: + if ext_type in EXT_MAP[observable_type].keys(): + raise DuplicateRegistrationError("Observable Extension", ext_type) + except AttributeError: + pass EXT_MAP[observable_type][ext_type] = new_extension except KeyError: if observable_type not in OBJ_MAP_OBSERVABLE: diff --git a/stix2/exceptions.py b/stix2/exceptions.py index f65f48b..edcc352 100644 --- a/stix2/exceptions.py +++ b/stix2/exceptions.py @@ -235,8 +235,14 @@ class STIXDeprecationWarning(DeprecationWarning): pass -class DuplicateObjectRegistrationError(STIXError): - """An object (or observable) with the same type as an existing object (or observable) is being registered""" +class DuplicateRegistrationError(STIXError): + """A STIX object with the same type as an existing object is being registered""" - def __init__(self, msg): - super(DuplicateObjectRegistrationError, self).__init__(msg) + def __init__(self, obj_type, reg_obj_type): + super(DuplicateRegistrationError, self).__init__() + self.obj_type = obj_type + self.reg_obj_type = reg_obj_type + + def __str__(self): + msg = "A(n) {0} with type '{1}' already exists and cannot be registered again" + return msg.format(self.obj_type, self.reg_obj_type) diff --git a/stix2/test/v20/test_core.py b/stix2/test/v20/test_core.py index 2808ba5..080c6f6 100644 --- a/stix2/test/v20/test_core.py +++ b/stix2/test/v20/test_core.py @@ -3,8 +3,6 @@ import pytest import stix2 from stix2 import core, exceptions -from .constants import IDENTITY_ID - BUNDLE = { "type": "bundle", "spec_version": "2.0", @@ -74,54 +72,3 @@ def test_register_marking_with_version(): assert stix2.v20.TLP_WHITE.definition._type in core.STIX2_OBJ_MAPS[v]['markings'] assert v in str(stix2.v20.TLP_WHITE.__class__) - - -@pytest.mark.xfail(reason="The default version is no longer 2.0", condition=stix2.DEFAULT_VERSION != "2.0") -def test_register_marking_with_no_version(): - # Uses default version (2.0 in this case) - core._register_marking(stix2.v20.TLP_WHITE.__class__) - v = 'v20' - - assert stix2.v20.TLP_WHITE.definition._type in core.STIX2_OBJ_MAPS[v]['markings'] - assert v in str(stix2.v20.TLP_WHITE.__class__) - - -def test_register_observable_extension_with_version(): - observed_data = stix2.v20.ObservedData( - id="observed-data--b67d30ff-02ac-498a-92f9-32f845f448cf", - created_by_ref=IDENTITY_ID, - created="2016-04-06T19:58:16.000Z", - modified="2016-04-06T19:58:16.000Z", - first_observed="2015-12-21T19:00:00Z", - last_observed="2015-12-21T19:00:00Z", - number_observed=50, - objects={ - "0": { - "name": "foo.exe", - "type": "file", - "extensions": { - "ntfs-ext": { - "alternate_data_streams": [ - { - "name": "second.stream", - "size": 25536, - }, - ], - }, - }, - }, - "1": { - "type": "directory", - "path": "/usr/home", - "contains_refs": ["0"], - }, - }, - ) - core._register_observable_extension(observed_data.objects['0'], observed_data.objects['0'].extensions['ntfs-ext'].__class__, version='2.0') - v = 'v20' - - assert observed_data.objects['0'].type in core.STIX2_OBJ_MAPS[v]['observables'] - assert v in str(observed_data.objects['0'].__class__) - - assert observed_data.objects['0'].extensions['ntfs-ext']._type in core.STIX2_OBJ_MAPS[v]['observable-extensions']['file'] - assert v in str(observed_data.objects['0'].extensions['ntfs-ext'].__class__) diff --git a/stix2/test/v20/test_custom.py b/stix2/test/v20/test_custom.py index 19cb28a..4aee66e 100644 --- a/stix2/test/v20/test_custom.py +++ b/stix2/test/v20/test_custom.py @@ -4,7 +4,7 @@ import stix2 from stix2 import core import stix2.v20 -from ...exceptions import DuplicateObjectRegistrationError, InvalidValueError +from ...exceptions import DuplicateRegistrationError, InvalidValueError from .constants import FAKE_TIME, IDENTITY_ID, MARKING_DEFINITION_ID IDENTITY_CUSTOM_PROP = stix2.v20.Identity( @@ -1040,7 +1040,7 @@ def test_register_custom_object_with_version(): def test_register_duplicate_object_with_version(): - with pytest.raises(DuplicateObjectRegistrationError) as excinfo: + with pytest.raises(DuplicateRegistrationError) as excinfo: @stix2.v20.CustomObject( 'x-new-type-2', [ ('property1', stix2.properties.StringProperty()), @@ -1049,7 +1049,7 @@ def test_register_duplicate_object_with_version(): ) class NewType2(object): pass - assert "An object with type 'x-new-type-2' already exists and cannot be created again." in str(excinfo.value) + assert "cannot be registered again" in str(excinfo.value) @stix2.v20.CustomObservable( @@ -1069,7 +1069,7 @@ def test_register_observable_with_version(): def test_register_duplicate_observable_with_version(): - with pytest.raises(DuplicateObjectRegistrationError) as excinfo: + with pytest.raises(DuplicateRegistrationError) as excinfo: @stix2.v20.CustomObservable( 'x-new-observable-2', [ ('property1', stix2.properties.StringProperty()), @@ -1077,4 +1077,59 @@ def test_register_duplicate_observable_with_version(): ) class NewObservable2(object): pass - assert "An observable with type 'x-new-observable-2' already exists and cannot be created again." in str(excinfo.value) + assert "cannot be registered again" in str(excinfo.value) + + +@pytest.mark.xfail(reason="The default version is no longer 2.0", condition=stix2.DEFAULT_VERSION != "2.0") +def test_register_marking_with_no_version(): + @stix2.v20.CustomMarking( + 'x-new-obj-2', [ + ('property1', stix2.properties.StringProperty(required=True)), + ], + ) + class NewObj2(): + pass + v = 'v20' + + no = NewObj2(property1='something') + assert no._type in core.STIX2_OBJ_MAPS[v]['markings'] + + +def test_register_observable_extension_with_version(): + @stix2.v20.CustomExtension( + stix2.v20.UserAccount, 'some-extension-2', [ + ('keys', stix2.properties.StringProperty(required=True)), + ], + ) + class SomeCustomExtension2: + pass + + v = 'v20' + example = SomeCustomExtension2(keys='test123') + + assert example._type in core.STIX2_OBJ_MAPS[v]['observable-extensions']['user-account'] + + +def test_register_duplicate_observable_extension(): + with pytest.raises(DuplicateRegistrationError) as excinfo: + @stix2.v20.CustomExtension( + stix2.v20.UserAccount, 'some-extension-2', [ + ('property1', stix2.properties.StringProperty(required=True)), + ('property2', stix2.properties.IntegerProperty()), + ], + ) + class NewExtension2(): + pass + assert "cannot be registered again" in str(excinfo.value) + + +def test_register_duplicate_marking(): + with pytest.raises(DuplicateRegistrationError) as excinfo: + @stix2.v20.CustomMarking( + 'x-new-obj-2', [ + ('property1', stix2.properties.StringProperty(required=True)), + ], + ) + class NewObj2(): + pass + assert "cannot be registered again" in str(excinfo.value) diff --git a/stix2/test/v21/test_core.py b/stix2/test/v21/test_core.py index f9f9de6..78e492a 100644 --- a/stix2/test/v21/test_core.py +++ b/stix2/test/v21/test_core.py @@ -3,8 +3,6 @@ import pytest import stix2 from stix2 import core, exceptions -from .constants import IDENTITY_ID, OBSERVED_DATA_ID - BUNDLE = { "type": "bundle", "id": "bundle--00000000-0000-4000-8000-000000000007", @@ -89,44 +87,3 @@ def test_register_marking_with_no_version(): assert stix2.v21.TLP_WHITE.definition._type in core.STIX2_OBJ_MAPS[v]['markings'] assert v in str(stix2.v21.TLP_WHITE.__class__) - - -def test_register_observable_extension_with_default_version(): - observed_data = stix2.v21.ObservedData( - id=OBSERVED_DATA_ID, - created_by_ref=IDENTITY_ID, - created="2016-04-06T19:58:16.000Z", - modified="2016-04-06T19:58:16.000Z", - first_observed="2015-12-21T19:00:00Z", - last_observed="2015-12-21T19:00:00Z", - number_observed=50, - objects={ - "0": { - "name": "foo.exe", - "type": "file", - "extensions": { - "ntfs-ext": { - "alternate_data_streams": [ - { - "name": "second.stream", - "size": 25536, - }, - ], - }, - }, - }, - "1": { - "type": "directory", - "path": "/usr/home", - "contains_refs": ["file--420bc087-8b53-5ae9-8210-20d27d5e96c8"], - }, - }, - ) - core._register_observable_extension(observed_data.objects['0'], observed_data.objects['0'].extensions['ntfs-ext'].__class__) - v = 'v21' - - assert observed_data.objects['0'].type in core.STIX2_OBJ_MAPS[v]['observables'] - assert v in str(observed_data.objects['0'].__class__) - - assert observed_data.objects['0'].extensions['ntfs-ext']._type in core.STIX2_OBJ_MAPS[v]['observable-extensions']['file'] - assert v in str(observed_data.objects['0'].extensions['ntfs-ext'].__class__) diff --git a/stix2/test/v21/test_custom.py b/stix2/test/v21/test_custom.py index 92cacf3..4522463 100644 --- a/stix2/test/v21/test_custom.py +++ b/stix2/test/v21/test_custom.py @@ -6,7 +6,7 @@ import stix2 import stix2.base import stix2.v21 -from ...exceptions import DuplicateObjectRegistrationError, InvalidValueError +from ...exceptions import DuplicateRegistrationError, InvalidValueError from .constants import FAKE_TIME, IDENTITY_ID, MARKING_DEFINITION_ID IDENTITY_CUSTOM_PROP = stix2.v21.Identity( @@ -1092,7 +1092,7 @@ def test_register_custom_object_with_version(): def test_register_duplicate_object_with_version(): - with pytest.raises(DuplicateObjectRegistrationError) as excinfo: + with pytest.raises(DuplicateRegistrationError) as excinfo: @stix2.v21.CustomObject( 'x-new-type-2', [ ('property1', stix2.properties.StringProperty()), @@ -1101,7 +1101,7 @@ def test_register_duplicate_object_with_version(): ) class NewType2(object): pass - assert "An object with type 'x-new-type-2' already exists and cannot be created again." in str(excinfo.value) + assert "cannot be registered again" in str(excinfo.value) @stix2.v21.CustomObservable( @@ -1113,15 +1113,15 @@ class NewObservable3(object): pass -def test_register_observable_with_version(): +def test_register_observable(): custom_obs = NewObservable3(property1="Test Observable") v = 'v21' assert custom_obs.type in stix2.core.STIX2_OBJ_MAPS[v]['observables'] -def test_register_duplicate_observable_with_version(): - with pytest.raises(DuplicateObjectRegistrationError) as excinfo: +def test_register_duplicate_observable(): + with pytest.raises(DuplicateRegistrationError) as excinfo: @stix2.v21.CustomObservable( 'x-new-observable-2', [ ('property1', stix2.properties.StringProperty()), @@ -1129,4 +1129,46 @@ def test_register_duplicate_observable_with_version(): ) class NewObservable2(object): pass - assert "An observable with type 'x-new-observable-2' already exists and cannot be created again." in str(excinfo.value) + assert "cannot be registered again" in str(excinfo.value) + + +def test_register_observable_custom_extension(): + @stix2.v21.CustomExtension( + stix2.v21.DomainName, 'x-new-2-ext', [ + ('property1', stix2.properties.StringProperty(required=True)), + ('property2', stix2.properties.IntegerProperty()), + ], + ) + class NewExtension2(): + pass + + example = NewExtension2(property1="Hi there") + v = 'v21' + + assert 'domain-name' in stix2.core.STIX2_OBJ_MAPS[v]['observables'] + assert example._type in stix2.core.STIX2_OBJ_MAPS[v]['observable-extensions']['domain-name'] + + +def test_register_duplicate_observable_extension(): + with pytest.raises(DuplicateRegistrationError) as excinfo: + @stix2.v21.CustomExtension( + stix2.v21.DomainName, 'x-new-2-ext', [ + ('property1', stix2.properties.StringProperty(required=True)), + ('property2', stix2.properties.IntegerProperty()), + ], + ) + class NewExtension2(): + pass + assert "cannot be registered again" in str(excinfo.value) + + +def test_register_duplicate_marking(): + with pytest.raises(DuplicateRegistrationError) as excinfo: + @stix2.v21.CustomMarking( + 'x-new-obj', [ + ('property1', stix2.properties.StringProperty(required=True)), + ], + ) + class NewObj2(): + pass + assert "cannot be registered again" in str(excinfo.value)