WIP: Allow custom observables, extensions

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.

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)