Repository: oasis-open/cti-python-stix2
Pull request #112 (2017-11-20 03:40:28)

Add get_type_from_id utility function
Added get_type_from_id to utils and created unit tests for it.


Patch:
diff --git a/stix2/utils.py b/stix2/utils.py
index f23dbe2..541e6d8 100644
--- a/stix2/utils.py
+++ b/stix2/utils.py
@@ -255,3 +255,7 @@ def get_class_hierarchy_names(obj):
     for cls in obj.__class__.__mro__:
         names.append(cls.__name__)
     return names
+
+
+def get_type_from_id(stix_id):
+    return stix_id.split('--', 1)[0]
diff --git a/stix2/v20/common.py b/stix2/v20/common.py
index ef45060..2d15529 100644
--- a/stix2/v20/common.py
+++ b/stix2/v20/common.py
@@ -145,14 +145,7 @@ def CustomMarking(type='x-custom-marking', properties=None):
 
             def __init__(self, **kwargs):
                 _STIXBase.__init__(self, **kwargs)
-                try:
-                    cls.__init__(self, **kwargs)
-                except (AttributeError, TypeError) as e:
-                    # Don't accidentally catch errors raised in a custom __init__()
-                    if ("has no attribute '__init__'" in str(e) or
-                            str(e) == "object.__init__() takes no parameters"):
-                        return
-                    raise e
+                cls.__init__(self, **kwargs)
 
         _register_marking(_Custom)
         return _Custom


Test patch:
diff --git a/stix2/test/test_custom.py b/stix2/test/test_custom.py
index 7c1832b..92d5d4c 100644
--- a/stix2/test/test_custom.py
+++ b/stix2/test/test_custom.py
@@ -94,28 +94,6 @@ def test_custom_property_in_bundled_object():
     assert '"x_foo": "bar"' in str(bundle)
 
 
-def test_custom_marking_no_init_1():
-    @stix2.CustomMarking('x-new-obj', [
-        ('property1', stix2.properties.StringProperty(required=True)),
-    ])
-    class NewObj():
-        pass
-
-    no = NewObj(property1='something')
-    assert no.property1 == 'something'
-
-
-def test_custom_marking_no_init_2():
-    @stix2.CustomMarking('x-new-obj2', [
-        ('property1', stix2.properties.StringProperty(required=True)),
-    ])
-    class NewObj2(object):
-        pass
-
-    no2 = NewObj2(property1='something')
-    assert no2.property1 == 'something'
-
-
 @stix2.sdo.CustomObject('x-new-type', [
     ('property1', stix2.properties.StringProperty(required=True)),
     ('property2', stix2.properties.IntegerProperty()),
@@ -124,15 +102,6 @@ class NewType(object):
     def __init__(self, property2=None, **kwargs):
         if property2 and property2 < 10:
             raise ValueError("'property2' is too small.")
-        if "property3" in kwargs and not isinstance(kwargs.get("property3"), int):
-            raise TypeError("Must be integer!")
-
-
-def test_custom_object_raises_exception():
-    with pytest.raises(TypeError) as excinfo:
-        NewType(property1='something', property3='something', allow_custom=True)
-
-    assert str(excinfo.value) == "Must be integer!"
 
 
 def test_custom_object_type():
@@ -148,7 +117,7 @@ def test_custom_object_type():
     assert "'property2' is too small." in str(excinfo.value)
 
 
-def test_custom_object_no_init_1():
+def test_custom_object_no_init():
     @stix2.sdo.CustomObject('x-new-obj', [
         ('property1', stix2.properties.StringProperty(required=True)),
     ])
@@ -158,8 +127,6 @@ def test_custom_object_no_init_1():
     no = NewObj(property1='something')
     assert no.property1 == 'something'
 
-
-def test_custom_object_no_init_2():
     @stix2.sdo.CustomObject('x-new-obj2', [
         ('property1', stix2.properties.StringProperty(required=True)),
     ])
@@ -203,36 +170,23 @@ class NewObservable():
     def __init__(self, property2=None, **kwargs):
         if property2 and property2 < 10:
             raise ValueError("'property2' is too small.")
-        if "property3" in kwargs and not isinstance(kwargs.get("property3"), int):
-            raise TypeError("Must be integer!")
 
 
-def test_custom_observable_object_1():
+def test_custom_observable_object():
     no = NewObservable(property1='something')
     assert no.property1 == 'something'
 
-
-def test_custom_observable_object_2():
     with pytest.raises(stix2.exceptions.MissingPropertiesError) as excinfo:
         NewObservable(property2=42)
     assert excinfo.value.properties == ['property1']
     assert "No values for required properties" in str(excinfo.value)
 
-
-def test_custom_observable_object_3():
     with pytest.raises(ValueError) as excinfo:
         NewObservable(property1='something', property2=4)
     assert "'property2' is too small." in str(excinfo.value)
 
 
-def test_custom_observable_raises_exception():
-    with pytest.raises(TypeError) as excinfo:
-        NewObservable(property1='something', property3='something', allow_custom=True)
-
-    assert str(excinfo.value) == "Must be integer!"
-
-
-def test_custom_observable_object_no_init_1():
+def test_custom_observable_object_no_init():
     @stix2.observables.CustomObservable('x-new-observable', [
         ('property1', stix2.properties.StringProperty()),
     ])
@@ -242,8 +196,6 @@ def test_custom_observable_object_no_init_1():
     no = NewObs(property1='something')
     assert no.property1 == 'something'
 
-
-def test_custom_observable_object_no_init_2():
     @stix2.observables.CustomObservable('x-new-obs2', [
         ('property1', stix2.properties.StringProperty()),
     ])
@@ -402,15 +354,6 @@ class NewExtension():
     def __init__(self, property2=None, **kwargs):
         if property2 and property2 < 10:
             raise ValueError("'property2' is too small.")
-        if "property3" in kwargs and not isinstance(kwargs.get("property3"), int):
-            raise TypeError("Must be integer!")
-
-
-def test_custom_extension_raises_exception():
-    with pytest.raises(TypeError) as excinfo:
-        NewExtension(property1='something', property3='something', allow_custom=True)
-
-    assert str(excinfo.value) == "Must be integer!"
 
 
 def test_custom_extension():
@@ -490,7 +433,7 @@ def test_custom_extension_empty_properties():
     assert "'properties' must be a dict!" in str(excinfo.value)
 
 
-def test_custom_extension_no_init_1():
+def test_custom_extension_no_init():
     @stix2.observables.CustomExtension(stix2.DomainName, 'x-new-extension', {
         'property1': stix2.properties.StringProperty(required=True),
     })
@@ -500,8 +443,6 @@ def test_custom_extension_no_init_1():
     ne = NewExt(property1="foobar")
     assert ne.property1 == "foobar"
 
-
-def test_custom_extension_no_init_2():
     @stix2.observables.CustomExtension(stix2.DomainName, 'x-new-ext2', {
         'property1': stix2.properties.StringProperty(required=True),
     })
diff --git a/stix2/test/test_markings.py b/stix2/test/test_markings.py
index d2271f0..456bf92 100644
--- a/stix2/test/test_markings.py
+++ b/stix2/test/test_markings.py
@@ -187,8 +187,7 @@ def test_parse_marking_definition(data):
 ])
 class NewMarking(object):
     def __init__(self, property2=None, **kwargs):
-        if "property3" in kwargs and not isinstance(kwargs.get("property3"), int):
-            raise TypeError("Must be integer!")
+        return
 
 
 def test_registered_custom_marking():
@@ -209,13 +208,6 @@ def test_registered_custom_marking():
     assert marking_def.definition_type == "x-new-marking-type"
 
 
-def test_registered_custom_marking_raises_exception():
-    with pytest.raises(TypeError) as excinfo:
-        NewMarking(property1='something', property3='something', allow_custom=True)
-
-    assert str(excinfo.value) == "Must be integer!"
-
-
 def test_not_registered_marking_raises_exception():
     with pytest.raises(ValueError) as excinfo:
         # Used custom object on purpose to demonstrate a not-registered marking
diff --git a/stix2/test/test_utils.py b/stix2/test/test_utils.py
index c73bcd2..cbe5b0f 100644
--- a/stix2/test/test_utils.py
+++ b/stix2/test/test_utils.py
@@ -74,3 +74,11 @@ def test_get_dict(data):
 def test_get_dict_invalid(data):
     with pytest.raises(ValueError):
         stix2.utils.get_dict(data)
+
+
+@pytest.mark.parametrize('stix_id, typ', [
+    ('malware--d69c8146-ab35-4d50-8382-6fc80e641d43', 'malware'),
+    ('intrusion-set--899ce53f-13a0-479b-a0e4-67d46e241542', 'intrusion-set')
+])
+def test_get_type_from_id(stix_id, typ):
+    assert stix2.utils.get_type_from_id(stix_id) == typ