From 3b297c17b548b660e61ad56dad49c7bf767e2cd9 Mon Sep 17 00:00:00 2001
From: Chris Lenk <clenk@mitre.org>
Date: Thu, 1 Nov 2018 17:23:55 -0400
Subject: [PATCH] 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.
---
 stix2/test/test_custom.py     | 27 ++++++++++++++++++++-------
 stix2/test/test_properties.py |  2 +-
 stix2/v20/observables.py      | 33 +++++++++++++++------------------
 3 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/stix2/test/test_custom.py b/stix2/test/test_custom.py
index 18ad03a..18b8f18 100644
--- a/stix2/test/test_custom.py
+++ b/stix2/test/test_custom.py
@@ -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)
 
diff --git a/stix2/test/test_properties.py b/stix2/test/test_properties.py
index 1106bcb..19419bb 100644
--- a/stix2/test/test_properties.py
+++ b/stix2/test/test_properties.py
@@ -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():
diff --git a/stix2/v20/observables.py b/stix2/v20/observables.py
index 36def28..faa5868 100644
--- a/stix2/v20/observables.py
+++ b/stix2/v20/observables.py
@@ -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