Clean up _register_* functions

Made them consistent with _register_observable_extension, by:
- moving validation logic there from _custom_*_builder functions
- using a new function for ensuring properties are dict-like
- using the library default spec version instead of None

Fix #371, fix #372, fix #373.

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):
-
+    prop_dict = _get_properties_dict(properties)
-    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,
-        )
 
     class _CustomExtension(cls, base_class):
 

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('.', '')

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(

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(