Fix marking definition extension unit test to use a plain
old CustomExtension. Add a unit test for a custom extension-based marking definition where the extension is of type "toplevel-property-extension".pull/1/head
parent
d27e9e6e55
commit
34e82e489f
|
@ -24,6 +24,33 @@ IDENTITY_CUSTOM_PROP = stix2.v21.Identity(
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@contextlib.contextmanager
|
||||||
|
def _register_extension(ext, props):
|
||||||
|
"""
|
||||||
|
A contextmanager useful for registering an extension and then ensuring
|
||||||
|
it gets unregistered again. A random extension-definition STIX ID is
|
||||||
|
generated for the extension and yielded as the contextmanager's value.
|
||||||
|
|
||||||
|
:param ext: The class which would normally be decorated with the
|
||||||
|
CustomExtension decorator.
|
||||||
|
:param props: Properties as would normally be passed into the
|
||||||
|
CustomExtension decorator.
|
||||||
|
"""
|
||||||
|
|
||||||
|
ext_def_id = "extension-definition--" + str(uuid.uuid4())
|
||||||
|
|
||||||
|
stix2.v21.CustomExtension(
|
||||||
|
ext_def_id,
|
||||||
|
props,
|
||||||
|
)(ext)
|
||||||
|
|
||||||
|
try:
|
||||||
|
yield ext_def_id
|
||||||
|
finally:
|
||||||
|
# "unregister" the extension
|
||||||
|
del stix2.registry.STIX2_OBJ_MAPS["2.1"]["extensions"][ext_def_id]
|
||||||
|
|
||||||
|
|
||||||
def test_identity_custom_property():
|
def test_identity_custom_property():
|
||||||
identity = stix2.v21.Identity(
|
identity = stix2.v21.Identity(
|
||||||
id=IDENTITY_ID,
|
id=IDENTITY_ID,
|
||||||
|
@ -1645,54 +1672,64 @@ def test_registered_new_extension_sco_allow_custom_false():
|
||||||
|
|
||||||
|
|
||||||
def test_registered_new_extension_marking_allow_custom_false():
|
def test_registered_new_extension_marking_allow_custom_false():
|
||||||
@stix2.v21.CustomMarking(
|
|
||||||
'my-favorite-marking', [
|
|
||||||
('some_marking_field', stix2.properties.StringProperty(required=True)),
|
|
||||||
], 'extension-definition--a932fcc6-e032-176c-126f-cb970a5a1fff',
|
|
||||||
)
|
|
||||||
class MyFavMarking:
|
|
||||||
pass
|
|
||||||
|
|
||||||
my_favorite_marking = {
|
class MyFavMarking:
|
||||||
'type': 'marking-definition',
|
extension_type = "property-extension"
|
||||||
'spec_version': '2.1',
|
|
||||||
'id': 'marking-definition--f9dbe89c-0030-4a9d-8b78-0dcd0a0de874',
|
props = {
|
||||||
'name': 'This is the name of my favorite Marking',
|
'some_marking_field': stix2.properties.StringProperty(required=True),
|
||||||
'extensions': {
|
|
||||||
'extension-definition--a932fcc6-e032-176c-126f-cb970a5a1fff': {
|
|
||||||
'extension_type': 'property-extension',
|
|
||||||
'some_marking_field': 'value',
|
|
||||||
},
|
|
||||||
},
|
|
||||||
}
|
}
|
||||||
|
|
||||||
marking_object = stix2.parse(my_favorite_marking)
|
with _register_extension(MyFavMarking, props) as ext_def_id:
|
||||||
assert isinstance(marking_object, stix2.v21.MarkingDefinition)
|
|
||||||
assert isinstance(
|
|
||||||
marking_object.extensions['extension-definition--a932fcc6-e032-176c-126f-cb970a5a1fff'],
|
|
||||||
stix2.v21.EXT_MAP['extension-definition--a932fcc6-e032-176c-126f-cb970a5a1fff'],
|
|
||||||
)
|
|
||||||
|
|
||||||
marking_serialized = marking_object.serialize(sort_keys=True)
|
my_favorite_marking = {
|
||||||
assert '"extensions": {"extension-definition--a932fcc6-e032-176c-126f-cb970a5a1fff": ' \
|
'type': 'marking-definition',
|
||||||
'{"extension_type": "property-extension", "some_marking_field": "value"}}' in marking_serialized
|
'spec_version': '2.1',
|
||||||
|
'id': 'marking-definition--f9dbe89c-0030-4a9d-8b78-0dcd0a0de874',
|
||||||
|
'name': 'This is the name of my favorite Marking',
|
||||||
|
'extensions': {
|
||||||
|
ext_def_id: {
|
||||||
|
'extension_type': 'property-extension',
|
||||||
|
'some_marking_field': 'value',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
marking_object = stix2.parse(my_favorite_marking)
|
||||||
|
assert isinstance(marking_object, stix2.v21.MarkingDefinition)
|
||||||
|
assert isinstance(
|
||||||
|
marking_object.extensions[ext_def_id],
|
||||||
|
stix2.v21.EXT_MAP[ext_def_id],
|
||||||
|
)
|
||||||
|
|
||||||
|
marking_serialized = marking_object.serialize(sort_keys=True)
|
||||||
|
assert '"extensions": {{"{}": ' \
|
||||||
|
'{{"extension_type": "property-extension", "some_marking_field": "value"}}}}'.format(ext_def_id) in marking_serialized
|
||||||
|
|
||||||
|
|
||||||
@contextlib.contextmanager
|
def test_custom_marking_toplevel_properties():
|
||||||
def _register_extension(ext, props):
|
class CustomMarking:
|
||||||
|
extension_type = "toplevel-property-extension"
|
||||||
|
|
||||||
ext_def_id = "extension-definition--" + str(uuid.uuid4())
|
props = {
|
||||||
|
"foo": stix2.properties.StringProperty(required=True)
|
||||||
|
}
|
||||||
|
|
||||||
stix2.v21.CustomExtension(
|
with _register_extension(CustomMarking, props) as ext_def_id:
|
||||||
ext_def_id,
|
|
||||||
props,
|
|
||||||
)(ext)
|
|
||||||
|
|
||||||
try:
|
marking_dict = {
|
||||||
yield ext_def_id
|
"type": "marking-definition",
|
||||||
finally:
|
"spec_version": "2.1",
|
||||||
# "unregister" the extension
|
"foo": "hello",
|
||||||
del stix2.registry.STIX2_OBJ_MAPS["2.1"]["extensions"][ext_def_id]
|
"extensions": {
|
||||||
|
ext_def_id: {
|
||||||
|
"extension_type": "toplevel-property-extension"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
marking = stix2.parse(marking_dict)
|
||||||
|
assert marking.foo == "hello"
|
||||||
|
|
||||||
|
|
||||||
def test_nested_ext_prop_meta():
|
def test_nested_ext_prop_meta():
|
||||||
|
|
Loading…
Reference in New Issue