From b633fd37856ccd15257d7d6bae3223411a90e9ef Mon Sep 17 00:00:00 2001 From: Chris Lenk Date: Tue, 10 Apr 2018 12:54:27 -0400 Subject: [PATCH] WIP: Allow custom observables, extensions --- stix2/test/test_custom.py | 31 +++++++++++++++++++++++++++++++ stix2/v20/observables.py | 15 +++++++++++---- 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/stix2/test/test_custom.py b/stix2/test/test_custom.py index cc8b32b..a50819b 100644 --- a/stix2/test/test_custom.py +++ b/stix2/test/test_custom.py @@ -363,6 +363,7 @@ def test_parse_custom_observable_object(): }""" nt = stix2.parse_observable(nt_string, []) + assert isinstance(nt, stix2.core._STIXBase) assert nt.property1 == 'something' @@ -376,6 +377,32 @@ def test_parse_unregistered_custom_observable_object(): stix2.parse_observable(nt_string) assert "Can't parse unknown observable type" in str(excinfo.value) + parsed_custom = stix2.parse_observable(nt_string, allow_custom=True) + assert parsed_custom['property1'] == 'something' + with pytest.raises(AttributeError) as excinfo: + assert parsed_custom.property1 == 'something' + assert not isinstance(parsed_custom, stix2.core._STIXBase) + + +def test_parse_observed_data_with_custom_observable(): + input_str = """{ + "type": "observed-data", + "id": "observed-data--dc20c4ca-a2a3-4090-a5d5-9558c3af4758", + "created": "2016-04-06T19:58:16.000Z", + "modified": "2016-04-06T19:58:16.000Z", + "first_observed": "2015-12-21T19:00:00Z", + "last_observed": "2015-12-21T19:00:00Z", + "number_observed": 1, + "objects": { + "0": { + "type": "x-foobar-observable", + "property1": "something" + } + } + }""" + parsed = stix2.parse(input_str, allow_custom=True) + assert parsed.objects['0']['property1'] == 'something' + def test_parse_invalid_custom_observable_object(): nt_string = """{ @@ -585,6 +612,10 @@ def test_parse_observable_with_unregistered_custom_extension(): stix2.parse_observable(input_str) assert "Can't parse Unknown extension type" in str(excinfo.value) + parsed_ob = stix2.parse_observable(input_str, allow_custom=True) + assert parsed_ob['extensions']['x-foobar-ext']['property1'] == 'foo' + assert not isinstance(parsed_ob['extensions']['x-foobar-ext'], stix2.core._STIXBase) + def test_register_custom_object(): # Not the way to register custom object. diff --git a/stix2/v20/observables.py b/stix2/v20/observables.py index 83600b0..4449527 100644 --- a/stix2/v20/observables.py +++ b/stix2/v20/observables.py @@ -923,15 +923,22 @@ def parse_observable(data, _valid_refs=None, allow_custom=False): try: obj_class = OBJ_MAP_OBSERVABLE[obj['type']] except KeyError: + if allow_custom: + # flag allows for unknown custom objects too, but will not + # be parsed into STIX observable object, just returned as is + return obj raise ParseError("Can't parse unknown observable type '%s'! For custom observables, " "use the CustomObservable decorator." % 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 ParseError("Can't parse Unknown extension type '%s' for observable type '%s'!" % (name, obj['type'])) - ext_class = EXT_MAP[obj['type']][name] - obj['extensions'][name] = ext_class(allow_custom=allow_custom, **obj['extensions'][name]) + try: + ext_class = EXT_MAP[obj['type']][name] + except KeyError: + if not allow_custom: + raise ParseError("Can't parse Unknown extension type '%s' for observable type '%s'!" % (name, obj['type'])) + else: # extension was found + obj['extensions'][name] = ext_class(allow_custom=allow_custom, **obj['extensions'][name]) return obj_class(allow_custom=allow_custom, **obj)