MISP
/
cti-python-stix2
mirror of https://github.com/MISP/cti-python-stix2
2
0
Fork 0
Code Issues Releases Wiki Activity

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.
master
Chris Lenk 2020-04-02 14:15:45 -04:00
parent d33adbcc71
commit 14540c0ea1
4 changed files with 68 additions and 72 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

85
stix2/custom.py
View File

@ -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):

47
stix2/parsing.py
View File

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

4
stix2/test/v20/test_markings.py
View File

@ -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(

4
stix2/test/v21/test_markings.py
View File

@ -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(