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={}):
|
||||
"""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):
|
||||
|
||||
class _Custom(_STIXBase):
|
||||
class _Custom(cls, _STIXBase):
|
||||
_type = type
|
||||
_properties = COMMON_PROPERTIES.copy()
|
||||
_properties.update({
|
||||
|
@ -207,6 +231,11 @@ def CustomObject(type='x-custom-type', properties={}):
|
|||
'type': TypeProperty(_type),
|
||||
})
|
||||
_properties.update(properties)
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
_STIXBase.__init__(self, **kwargs)
|
||||
cls.__init__(self, **kwargs)
|
||||
|
||||
stix2._register_type(_Custom)
|
||||
return _Custom
|
||||
|
||||
|
|
|
@ -78,7 +78,7 @@ def test_parse_identity_custom_property(data):
|
|||
})
|
||||
class NewType():
|
||||
def __init__(self, property2=None, **kwargs):
|
||||
if property2 < 10:
|
||||
if property2 and property2 < 10:
|
||||
raise ValueError("'property2' is too small.")
|
||||
|
||||
|
||||
|
@ -90,7 +90,7 @@ def test_custom_object_type():
|
|||
NewType(property2=42)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
NewType(property2=4)
|
||||
NewType(property1='something', property2=4)
|
||||
|
||||
|
||||
def test_parse_custom_object_type():
|
||||
|
|
Loading…
Reference in New Issue