diff --git a/stix2/custom.py b/stix2/custom.py index 387906b..f749b04 100644 --- a/stix2/custom.py +++ b/stix2/custom.py @@ -1,5 +1,4 @@ from collections import OrderedDict -import re import six @@ -8,22 +7,29 @@ from .parsing import ( _register_marking, _register_object, _register_observable, _register_observable_extension, ) -from .utils import PREFIX_21_REGEX, get_class_hierarchy_names + + +def _get_properties_dict(properties): + try: + return OrderedDict(properties) + except TypeError as e: + six.raise_from( + ValueError( + "properties must be dict-like, e.g. a list " + "containing tuples. For example, " + "[('property1', IntegerProperty())]", + ), + e, + ) def _custom_object_builder(cls, type, properties, version, base_class): + prop_dict = _get_properties_dict(properties) + class _CustomObject(cls, base_class): - if not properties or not isinstance(properties, list): - raise ValueError("Must supply a list, containing tuples. For example, [('property1', IntegerProperty())]") - - if version == "2.1": - for prop_name, prop in properties: - if not re.match(PREFIX_21_REGEX, prop_name): - raise ValueError("Property name '%s' must begin with an alpha character" % prop_name) - _type = type - _properties = OrderedDict(properties) + _properties = prop_dict def __init__(self, **kwargs): base_class.__init__(self, **kwargs) @@ -34,14 +40,12 @@ def _custom_object_builder(cls, type, properties, version, base_class): def _custom_marking_builder(cls, type, properties, version, base_class): + prop_dict = _get_properties_dict(properties) class _CustomMarking(cls, base_class): - if not properties or not isinstance(properties, list): - raise ValueError("Must supply a list, containing tuples. For example, [('property1', IntegerProperty())]") - _type = type - _properties = OrderedDict(properties) + _properties = prop_dict def __init__(self, **kwargs): base_class.__init__(self, **kwargs) @@ -55,44 +59,12 @@ def _custom_observable_builder(cls, type, properties, version, base_class, id_co if id_contrib_props is None: id_contrib_props = [] + prop_dict = _get_properties_dict(properties) + class _CustomObservable(cls, base_class): - if not properties or not isinstance(properties, list): - raise ValueError("Must supply a list, containing tuples. For example, [('property1', IntegerProperty())]") - - if version == "2.0": - # If using STIX2.0, check properties ending in "_ref/s" are ObjectReferenceProperties - for prop_name, prop in properties: - if prop_name.endswith('_ref') and ('ObjectReferenceProperty' not in get_class_hierarchy_names(prop)): - raise ValueError( - "'%s' is named like an object reference property but " - "is not an ObjectReferenceProperty." % prop_name, - ) - elif (prop_name.endswith('_refs') and ('ListProperty' not in get_class_hierarchy_names(prop) or - 'ObjectReferenceProperty' not in get_class_hierarchy_names(prop.contained))): - raise ValueError( - "'%s' is named like an object reference list property but " - "is not a ListProperty containing ObjectReferenceProperty." % prop_name, - ) - else: - # If using STIX2.1 (or newer...), check properties ending in "_ref/s" are ReferenceProperties - for prop_name, prop in properties: - if not re.match(PREFIX_21_REGEX, prop_name): - raise ValueError("Property name '%s' must begin with an alpha character." % prop_name) - elif prop_name.endswith('_ref') and ('ReferenceProperty' not in get_class_hierarchy_names(prop)): - raise ValueError( - "'%s' is named like a reference property but " - "is not a ReferenceProperty." % prop_name, - ) - elif (prop_name.endswith('_refs') and ('ListProperty' not in get_class_hierarchy_names(prop) or - 'ReferenceProperty' not in get_class_hierarchy_names(prop.contained))): - raise ValueError( - "'%s' is named like a reference list property but " - "is not a ListProperty containing ReferenceProperty." % prop_name, - ) - _type = type - _properties = OrderedDict(properties) + _properties = prop_dict if version != '2.0': _id_contributing_properties = id_contrib_props @@ -105,18 +77,7 @@ def _custom_observable_builder(cls, type, properties, version, base_class, id_co def _custom_extension_builder(cls, observable, type, properties, version, base_class): - - try: - prop_dict = OrderedDict(properties) - except TypeError as e: - six.raise_from( - ValueError( - "Extension properties must be dict-like, e.g. a list " - "containing tuples. For example, " - "[('property1', IntegerProperty())]", - ), - e, - ) + prop_dict = _get_properties_dict(properties) class _CustomExtension(cls, base_class): diff --git a/stix2/parsing.py b/stix2/parsing.py index db0c693..baa3b7b 100644 --- a/stix2/parsing.py +++ b/stix2/parsing.py @@ -9,7 +9,7 @@ import stix2 from .base import _DomainObject, _Observable from .exceptions import DuplicateRegistrationError, ParseError -from .utils import PREFIX_21_REGEX, _get_dict +from .utils import PREFIX_21_REGEX, _get_dict, get_class_hierarchy_names STIX2_OBJ_MAPS = {} @@ -192,7 +192,7 @@ def parse_observable(data, _valid_refs=None, allow_custom=False, version=None): return obj_class(allow_custom=allow_custom, **obj) -def _register_object(new_type, version=None): +def _register_object(new_type, version=stix2.DEFAULT_VERSION): """Register a custom STIX Object type. Args: @@ -213,6 +213,13 @@ def _register_object(new_type, version=None): new_type.__name__, ) + properties = new_type._properties + + if version == "2.1": + for prop_name, prop in properties.items(): + if not re.match(PREFIX_21_REGEX, prop_name): + raise ValueError("Property name '%s' must begin with an alpha character" % prop_name) + if version: v = 'v' + version.replace('.', '') else: @@ -225,7 +232,7 @@ def _register_object(new_type, version=None): OBJ_MAP[new_type._type] = new_type -def _register_marking(new_marking, version=None): +def _register_marking(new_marking, version=stix2.DEFAULT_VERSION): """Register a custom STIX Marking Definition type. Args: @@ -257,7 +264,7 @@ def _register_marking(new_marking, version=None): OBJ_MAP_MARKING[mark_type] = new_marking -def _register_observable(new_observable, version=None): +def _register_observable(new_observable, version=stix2.DEFAULT_VERSION): """Register a custom STIX Cyber Observable type. Args: @@ -266,6 +273,38 @@ def _register_observable(new_observable, version=None): None, use latest version. """ + properties = new_observable._properties + + if version == "2.0": + # If using STIX2.0, check properties ending in "_ref/s" are ObjectReferenceProperties + for prop_name, prop in properties.items(): + if prop_name.endswith('_ref') and ('ObjectReferenceProperty' not in get_class_hierarchy_names(prop)): + raise ValueError( + "'%s' is named like an object reference property but " + "is not an ObjectReferenceProperty." % prop_name, + ) + elif (prop_name.endswith('_refs') and ('ListProperty' not in get_class_hierarchy_names(prop) or + 'ObjectReferenceProperty' not in get_class_hierarchy_names(prop.contained))): + raise ValueError( + "'%s' is named like an object reference list property but " + "is not a ListProperty containing ObjectReferenceProperty." % prop_name, + ) + else: + # If using STIX2.1 (or newer...), check properties ending in "_ref/s" are ReferenceProperties + for prop_name, prop in properties.items(): + if not re.match(PREFIX_21_REGEX, prop_name): + raise ValueError("Property name '%s' must begin with an alpha character." % prop_name) + elif prop_name.endswith('_ref') and ('ReferenceProperty' not in get_class_hierarchy_names(prop)): + raise ValueError( + "'%s' is named like a reference property but " + "is not a ReferenceProperty." % prop_name, + ) + elif (prop_name.endswith('_refs') and ('ListProperty' not in get_class_hierarchy_names(prop) or + 'ReferenceProperty' not in get_class_hierarchy_names(prop.contained))): + raise ValueError( + "'%s' is named like a reference list property but " + "is not a ListProperty containing ReferenceProperty." % prop_name, + ) if version: v = 'v' + version.replace('.', '') diff --git a/stix2/test/v20/test_markings.py b/stix2/test/v20/test_markings.py index b34ef43..0753e86 100644 --- a/stix2/test/v20/test_markings.py +++ b/stix2/test/v20/test_markings.py @@ -249,14 +249,12 @@ def test_not_registered_marking_raises_exception(): def test_marking_wrong_type_construction(): - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError): # Test passing wrong type for properties. @stix2.v20.CustomMarking('x-new-marking-type2', ("a", "b")) class NewObject3(object): pass - assert str(excinfo.value) == "Must supply a list, containing tuples. For example, [('property1', IntegerProperty())]" - def test_campaign_add_markings(): campaign = stix2.v20.Campaign( diff --git a/stix2/test/v21/test_markings.py b/stix2/test/v21/test_markings.py index a2fca51..b0a0d09 100644 --- a/stix2/test/v21/test_markings.py +++ b/stix2/test/v21/test_markings.py @@ -277,14 +277,12 @@ def test_not_registered_marking_raises_exception(): def test_marking_wrong_type_construction(): - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError): # Test passing wrong type for properties. @stix2.v21.CustomMarking('x-new-marking-type2', ("a", "b")) class NewObject3(object): pass - assert str(excinfo.value) == "Must supply a list, containing tuples. For example, [('property1', IntegerProperty())]" - def test_campaign_add_markings(): campaign = stix2.v21.Campaign(