Allow adding validation to custom object types
_Custom.__init__() is a solution to recursion errors arising from using a decorator. See https://stackoverflow.com/q/14739809/1013457stix2.1
parent
bcfb13f23c
commit
d4e92dd813
33
stix2/sdo.py
33
stix2/sdo.py
|
@ -195,11 +195,35 @@ class Vulnerability(_STIXBase):
|
||||||
|
|
||||||
|
|
||||||
def CustomObject(type='x-custom-type', properties={}):
|
def CustomObject(type='x-custom-type', properties={}):
|
||||||
"""Custom STIX Object type decorator"""
|
"""Custom STIX Object type decorator
|
||||||
|
|
||||||
|
Example 1:
|
||||||
|
|
||||||
|
@CustomObject('type-name', {
|
||||||
|
'property1': StringProperty(required=True),
|
||||||
|
'property2': IntegerProperty(),
|
||||||
|
})
|
||||||
|
class MyNewObjectType():
|
||||||
|
pass
|
||||||
|
|
||||||
|
Supply an __init__() function to add any special validations to the custom
|
||||||
|
type. Don't call super().__init() though - doing so will cause an error.
|
||||||
|
|
||||||
|
Example 2:
|
||||||
|
|
||||||
|
@CustomObject('type-name', {
|
||||||
|
'property1': StringProperty(required=True),
|
||||||
|
'property2': IntegerProperty(),
|
||||||
|
})
|
||||||
|
class MyNewObjectType():
|
||||||
|
def __init__(self, property2=None, **kwargs):
|
||||||
|
if property2 and property2 < 10:
|
||||||
|
raise ValueError("'property2' is too small.")
|
||||||
|
"""
|
||||||
|
|
||||||
def custom_builder(cls):
|
def custom_builder(cls):
|
||||||
|
|
||||||
class _Custom(_STIXBase):
|
class _Custom(cls, _STIXBase):
|
||||||
_type = type
|
_type = type
|
||||||
_properties = COMMON_PROPERTIES.copy()
|
_properties = COMMON_PROPERTIES.copy()
|
||||||
_properties.update({
|
_properties.update({
|
||||||
|
@ -207,6 +231,11 @@ def CustomObject(type='x-custom-type', properties={}):
|
||||||
'type': TypeProperty(_type),
|
'type': TypeProperty(_type),
|
||||||
})
|
})
|
||||||
_properties.update(properties)
|
_properties.update(properties)
|
||||||
|
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
_STIXBase.__init__(self, **kwargs)
|
||||||
|
cls.__init__(self, **kwargs)
|
||||||
|
|
||||||
stix2._register_type(_Custom)
|
stix2._register_type(_Custom)
|
||||||
return _Custom
|
return _Custom
|
||||||
|
|
||||||
|
|
|
@ -78,7 +78,7 @@ def test_parse_identity_custom_property(data):
|
||||||
})
|
})
|
||||||
class NewType():
|
class NewType():
|
||||||
def __init__(self, property2=None, **kwargs):
|
def __init__(self, property2=None, **kwargs):
|
||||||
if property2 < 10:
|
if property2 and property2 < 10:
|
||||||
raise ValueError("'property2' is too small.")
|
raise ValueError("'property2' is too small.")
|
||||||
|
|
||||||
|
|
||||||
|
@ -90,7 +90,7 @@ def test_custom_object_type():
|
||||||
NewType(property2=42)
|
NewType(property2=42)
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
NewType(property2=4)
|
NewType(property1='something', property2=4)
|
||||||
|
|
||||||
|
|
||||||
def test_parse_custom_object_type():
|
def test_parse_custom_object_type():
|
||||||
|
|
Loading…
Reference in New Issue