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

Add custom STIX Object types

stix2.1
clenk 2017-06-12 16:15:12 -04:00
parent 8f1ae4e6d3
commit bcfb13f23c
4 changed files with 132 additions and 69 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
stix2/__init__.py
View File

@ -150,3 +150,10 @@ def parse_observable(data, _valid_refs, allow_custom=False):
obj['extensions'][name] = ext_class(allow_custom=allow_custom, **obj['extensions'][name]) obj['extensions'][name] = ext_class(allow_custom=allow_custom, **obj['extensions'][name])
return obj_class(allow_custom=allow_custom, **obj) return obj_class(allow_custom=allow_custom, **obj)
def _register_type(new_type):
"""Register a custom STIX Object type.
"""
OBJ_MAP[new_type._type] = new_type

21
stix2/sdo.py
View File

@ -1,5 +1,7 @@
"""STIX 2.0 Domain Objects""" """STIX 2.0 Domain Objects"""
import stix2
from .base import _STIXBase from .base import _STIXBase
from .common import COMMON_PROPERTIES from .common import COMMON_PROPERTIES
from .other import KillChainPhase from .other import KillChainPhase
@ -190,3 +192,22 @@ class Vulnerability(_STIXBase):
'name': StringProperty(required=True), 'name': StringProperty(required=True),
'description': StringProperty(), 'description': StringProperty(),
}) })
def CustomObject(type='x-custom-type', properties={}):
"""Custom STIX Object type decorator"""
def custom_builder(cls):
class _Custom(_STIXBase):
_type = type
_properties = COMMON_PROPERTIES.copy()
_properties.update({
'id': IDProperty(_type),
'type': TypeProperty(_type),
})
_properties.update(properties)
stix2._register_type(_Custom)
return _Custom
return custom_builder

104
stix2/test/test_custom.py Normal file
View File

@ -0,0 +1,104 @@
import pytest
import stix2
def test_identity_custom_property():
with pytest.raises(ValueError):
stix2.Identity(
id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
created="2015-12-21T19:59:11Z",
modified="2015-12-21T19:59:11Z",
name="John Smith",
identity_class="individual",
custom_properties="foobar",
)
identity = stix2.Identity(
id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
created="2015-12-21T19:59:11Z",
modified="2015-12-21T19:59:11Z",
name="John Smith",
identity_class="individual",
custom_properties={
"foo": "bar",
},
)
assert identity.foo == "bar"
def test_identity_custom_property_invalid():
with pytest.raises(stix2.exceptions.ExtraPropertiesError):
stix2.Identity(
id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
created="2015-12-21T19:59:11Z",
modified="2015-12-21T19:59:11Z",
name="John Smith",
identity_class="individual",
x_foo="bar",
)
def test_identity_custom_property_allowed():
identity = stix2.Identity(
id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
created="2015-12-21T19:59:11Z",
modified="2015-12-21T19:59:11Z",
name="John Smith",
identity_class="individual",
x_foo="bar",
allow_custom=True,
)
assert identity.x_foo == "bar"
@pytest.mark.parametrize("data", [
"""{
"type": "identity",
"id": "identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
"created": "2015-12-21T19:59:11Z",
"modified": "2015-12-21T19:59:11Z",
"name": "John Smith",
"identity_class": "individual",
"foo": "bar"
}""",
])
def test_parse_identity_custom_property(data):
with pytest.raises(stix2.exceptions.ExtraPropertiesError):
identity = stix2.parse(data)
identity = stix2.parse(data, allow_custom=True)
assert identity.foo == "bar"
@stix2.sdo.CustomObject('new-type', {
'property1': stix2.properties.StringProperty(required=True),
'property2': stix2.properties.IntegerProperty(),
})
class NewType():
def __init__(self, property2=None, **kwargs):
if property2 < 10:
raise ValueError("'property2' is too small.")
def test_custom_object_type():
nt = NewType(property1='something')
assert nt.property1 == 'something'
with pytest.raises(stix2.exceptions.MissingPropertiesError):
NewType(property2=42)
with pytest.raises(ValueError):
NewType(property2=4)
def test_parse_custom_object_type():
nt_string = """{
"type": "new-type",
"created": "2015-12-21T19:59:11Z",
"property1": "something"
}"""
nt = stix2.parse(nt_string)
assert nt.property1 == 'something'

69
stix2/test/test_identity.py
View File

@ -51,75 +51,6 @@ def test_parse_identity(data):
assert identity.name == "John Smith" assert identity.name == "John Smith"
def test_identity_custom_property():
with pytest.raises(ValueError):
stix2.Identity(
id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
created="2015-12-21T19:59:11Z",
modified="2015-12-21T19:59:11Z",
name="John Smith",
identity_class="individual",
custom_properties="foobar",
)
identity = stix2.Identity(
id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
created="2015-12-21T19:59:11Z",
modified="2015-12-21T19:59:11Z",
name="John Smith",
identity_class="individual",
custom_properties={
"foo": "bar",
},
)
assert identity.foo == "bar"
def test_identity_custom_property_invalid():
with pytest.raises(stix2.exceptions.ExtraPropertiesError):
stix2.Identity(
id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
created="2015-12-21T19:59:11Z",
modified="2015-12-21T19:59:11Z",
name="John Smith",
identity_class="individual",
x_foo="bar",
)
def test_identity_custom_property_allowed():
identity = stix2.Identity(
id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
created="2015-12-21T19:59:11Z",
modified="2015-12-21T19:59:11Z",
name="John Smith",
identity_class="individual",
x_foo="bar",
allow_custom=True,
)
assert identity.x_foo == "bar"
@pytest.mark.parametrize("data", [
"""{
"type": "identity",
"id": "identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
"created": "2015-12-21T19:59:11Z",
"modified": "2015-12-21T19:59:11Z",
"name": "John Smith",
"identity_class": "individual",
"foo": "bar"
}""",
])
def test_parse_identity_custom_property(data):
with pytest.raises(stix2.exceptions.ExtraPropertiesError):
identity = stix2.parse(data)
identity = stix2.parse(data, allow_custom=True)
assert identity.foo == "bar"
def test_parse_no_type(): def test_parse_no_type():
with pytest.raises(stix2.exceptions.ParseError): with pytest.raises(stix2.exceptions.ParseError):
stix2.parse(""" stix2.parse("""