diff --git a/stix2/__init__.py b/stix2/__init__.py index f3dfd20..fa015f4 100644 --- a/stix2/__init__.py +++ b/stix2/__init__.py @@ -95,46 +95,58 @@ EXT_MAP = { } -def parse(data): - """Deserialize a string or file-like object into a STIX object""" +def parse(data, allow_custom=False): + """Deserialize a string or file-like object into a STIX object. + + Args: + data: The STIX 2 string to be parsed. + allow_custom (bool): Whether to allow custom properties or not. Default: False. + + Returns: + An instantiated Python STIX object. + """ obj = get_dict(data) if 'type' not in obj: - # TODO parse external references, kill chain phases, and granular markings - pass - else: - try: - obj_class = OBJ_MAP[obj['type']] - except KeyError: - # TODO handle custom objects - raise ValueError("Can't parse unknown object type '%s'!" % obj['type']) - return obj_class(**obj) + raise exceptions.ParseError("Can't parse object with no 'type' property: %s" % str(obj)) - return obj + try: + obj_class = OBJ_MAP[obj['type']] + except KeyError: + # TODO handle custom objects + raise exceptions.ParseError("Can't parse unknown object type '%s'!" % obj['type']) + return obj_class(allow_custom=allow_custom, **obj) -def parse_observable(data, _valid_refs): - """Deserialize a string or file-like object into a STIX Cyber Observable - object. +def parse_observable(data, _valid_refs, allow_custom=False): + """Deserialize a string or file-like object into a STIX Cyber Observable object. + + Args: + data: The STIX 2 string to be parsed. + _valid_refs: A list of object references valid for the scope of the object being parsed. + allow_custom: Whether to allow custom properties or not. Default: False. + + Returns: + An instantiated Python STIX Cyber Observable object. """ obj = get_dict(data) obj['_valid_refs'] = _valid_refs if 'type' not in obj: - raise ValueError("'type' is a required property!") + raise exceptions.ParseError("Can't parse object with no 'type' property: %s" % str(obj)) try: obj_class = OBJ_MAP_OBSERVABLE[obj['type']] except KeyError: # TODO handle custom observable objects - raise ValueError("Can't parse unknown object type '%s'!" % obj['type']) + raise exceptions.ParseError("Can't parse unknown object type '%s'!" % obj['type']) if 'extensions' in obj and obj['type'] in EXT_MAP: for name, ext in obj['extensions'].items(): if name not in EXT_MAP[obj['type']]: - raise ValueError("Can't parse Unknown extension type '%s' for object type '%s'!" % (name, obj['type'])) + raise exceptions.ParseError("Can't parse Unknown extension type '%s' for object type '%s'!" % (name, obj['type'])) ext_class = EXT_MAP[obj['type']][name] - obj['extensions'][name] = ext_class(**obj['extensions'][name]) + obj['extensions'][name] = ext_class(allow_custom=allow_custom, **obj['extensions'][name]) - return obj_class(**obj) + return obj_class(allow_custom=allow_custom, **obj) diff --git a/stix2/base.py b/stix2/base.py index d3cb072..cb17f11 100644 --- a/stix2/base.py +++ b/stix2/base.py @@ -83,7 +83,7 @@ class _STIXBase(collections.Mapping): # TODO: check selectors pass - def __init__(self, **kwargs): + def __init__(self, allow_custom=False, **kwargs): cls = self.__class__ # Use the same timestamp for any auto-generated datetimes @@ -93,9 +93,10 @@ class _STIXBase(collections.Mapping): custom_props = kwargs.pop('custom_properties', {}) if custom_props and not isinstance(custom_props, dict): raise ValueError("'custom_properties' must be a dictionary") - extra_kwargs = list(set(kwargs) - set(cls._properties)) - if extra_kwargs: - raise ExtraPropertiesError(cls, extra_kwargs) + if not allow_custom: + extra_kwargs = list(set(kwargs) - set(cls._properties)) + if extra_kwargs: + raise ExtraPropertiesError(cls, extra_kwargs) # Remove any keyword arguments whose value is None setting_kwargs = {} diff --git a/stix2/exceptions.py b/stix2/exceptions.py index 3043047..ef47dd0 100644 --- a/stix2/exceptions.py +++ b/stix2/exceptions.py @@ -150,3 +150,10 @@ class RevokeError(STIXError, ValueError): return "Cannot revoke an already revoked object." else: return "Cannot create a new version of a revoked object." + + +class ParseError(STIXError, ValueError): + """Could not parse object""" + + def __init__(self, msg): + super(ParseError, self).__init__(msg) diff --git a/stix2/test/test_identity.py b/stix2/test/test_identity.py index e4a67e6..edc3400 100644 --- a/stix2/test/test_identity.py +++ b/stix2/test/test_identity.py @@ -52,6 +52,16 @@ def test_parse_identity(data): 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", @@ -77,4 +87,48 @@ def test_identity_custom_property_invalid(): 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(): + with pytest.raises(stix2.exceptions.ParseError): + stix2.parse(""" + { + "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" + }""") + # TODO: Add other examples