Use consistent errors for observable extensions

Whether or not the Observable type is in the EXT_MAP already, using a
custom extension without also using allow_custom=True should result in
the same behavior/error message.
revert-222-multi_version_filesystem_store
Chris Lenk 2018-11-01 17:23:55 -04:00
parent 3084c9f51f
commit 3b297c17b5
3 changed files with 36 additions and 26 deletions

View File

@ -826,9 +826,10 @@ def test_parse_observable_with_custom_extension():
assert parsed.extensions['x-new-ext'].property2 == 12
def test_parse_observable_with_unregistered_custom_extension():
input_str = """{
"type": "domain-name",
@pytest.mark.parametrize("data", [
# URL is not in EXT_MAP
"""{
"type": "url",
"value": "example.com",
"extensions": {
"x-foobar-ext": {
@ -836,13 +837,25 @@ def test_parse_observable_with_unregistered_custom_extension():
"property2": 12
}
}
}"""
}""",
# File is in EXT_MAP
"""{
"type": "file",
"name": "foo.txt",
"extensions": {
"x-foobar-ext": {
"property1": "foo",
"property2": 12
}
}
}""",
])
def test_parse_observable_with_unregistered_custom_extension(data):
with pytest.raises(ValueError) as excinfo:
stix2.parse_observable(input_str)
stix2.parse_observable(data)
assert "Can't parse unknown extension type" in str(excinfo.value)
parsed_ob = stix2.parse_observable(input_str, allow_custom=True)
parsed_ob = stix2.parse_observable(data, allow_custom=True)
assert parsed_ob['extensions']['x-foobar-ext']['property1'] == 'foo'
assert not isinstance(parsed_ob['extensions']['x-foobar-ext'], stix2.core._STIXBase)

View File

@ -412,7 +412,7 @@ def test_extension_property_invalid_type():
'pe_type': 'exe'
}}
)
assert 'no extensions defined' in str(excinfo.value)
assert "Can't parse unknown extension" in str(excinfo.value)
def test_extension_at_least_one_property_constraint():

View File

@ -73,26 +73,23 @@ class ExtensionsProperty(DictionaryProperty):
if dictified == {}:
raise ValueError("The extensions property must contain a non-empty dictionary")
if self.enclosing_type in EXT_MAP:
specific_type_map = EXT_MAP[self.enclosing_type]
for key, subvalue in dictified.items():
if key in specific_type_map:
cls = specific_type_map[key]
if type(subvalue) is dict:
if self.allow_custom:
subvalue['allow_custom'] = True
dictified[key] = cls(**subvalue)
else:
dictified[key] = cls(**subvalue)
elif type(subvalue) is cls:
# If already an instance of an _Extension class, assume it's valid
dictified[key] = subvalue
specific_type_map = EXT_MAP.get(self.enclosing_type, {})
for key, subvalue in dictified.items():
if key in specific_type_map:
cls = specific_type_map[key]
if type(subvalue) is dict:
if self.allow_custom:
subvalue['allow_custom'] = True
dictified[key] = cls(**subvalue)
else:
raise ValueError("Cannot determine extension type.")
dictified[key] = cls(**subvalue)
elif type(subvalue) is cls:
# If already an instance of an _Extension class, assume it's valid
dictified[key] = subvalue
else:
raise CustomContentError("Can't parse unknown extension type: {}".format(key))
else:
raise ValueError("The enclosing type '%s' has no extensions defined" % self.enclosing_type)
raise ValueError("Cannot determine extension type.")
else:
raise CustomContentError("Can't parse unknown extension type: {}".format(key))
return dictified