Add EnumProperty, use it in WindowsRegistryValueType

stix2/__init__.py

@@ -5,10 +5,11 @@
 from . import exceptions
 from .bundle import Bundle
 from .observables import (URL, Artifact, AutonomousSystem, Directory,
-                          DomainName, EmailAddress, EmailMessage, File,
+                          DomainName, EmailAddress, EmailMessage,
-                          IPv4Address, IPv6Address, MACAddress, Mutex,
+                          EmailMIMEComponent, File, IPv4Address, IPv6Address,
-                          NetworkTraffic, Process, Software, UserAccount,
+                          MACAddress, Mutex, NetworkTraffic, Process, Software,
-                          WindowsRegistryKey, X509Certificate)
+                          UserAccount, WindowsRegistryKey,
+                          WindowsRegistryValueType, X509Certificate)
 from .other import (ExternalReference, GranularMarking, KillChainPhase,
                     MarkingDefinition, StatementMarking, TLPMarking)
 from .sdo import (AttackPattern, Campaign, CourseOfAction, Identity, Indicator,

stix2/observables.py

@@ -8,8 +8,8 @@ and do not have a '_type' attribute.
 from .base import _Observable, _STIXBase
 from .exceptions import ObjectConstraintError
 from .properties import (BinaryProperty, BooleanProperty, DictionaryProperty,
-                         EmbeddedObjectProperty, HashesProperty, HexProperty,
+                         EmbeddedObjectProperty, EnumProperty, HashesProperty,
-                         IntegerProperty, ListProperty,
+                         HexProperty, IntegerProperty, ListProperty,
                          ObjectReferenceProperty, Property, StringProperty,
                          TimestampProperty, TypeProperty)
 
@@ -263,7 +263,21 @@ class WindowsRegistryValueType(_STIXBase):
     _properties = {
         'name': StringProperty(required=True),
         'data': StringProperty(),
-        'data_type': Property()
+        'data_type': EnumProperty([
+            'REG_NONE',
+            'REG_SZ',
+            'REG_EXPAND_SZ',
+            'REG_BINARY',
+            'REG_DWORD',
+            'REG_DWORD_BIG_ENDIAN',
+            'REG_LINK',
+            'REG_MULTI_SZ',
+            'REG_RESOURCE_LIST',
+            'REG_FULL_RESOURCE_DESCRIPTION',
+            'REG_RESOURCE_REQUIREMENTS_LIST',
+            'REG_QWORD',
+            'REG_INVALID_TYPE',
+        ]),
     }
 
 
@@ -279,6 +293,11 @@ class WindowsRegistryKey(_Observable):
         'number_of_subkeys': IntegerProperty(),
     }
 
+    @property
+    def values(self):
+        # Needed because 'values' is a property on collections.Mapping objects
+        return self._inner['values']
+
 
 class X509Certificate(_Observable):
     _type = 'x509-certificate'

stix2/properties.py

@@ -143,6 +143,7 @@ class StringProperty(Property):
 
 
 class TypeProperty(Property):
+
     def __init__(self, type):
         super(TypeProperty, self).__init__(fixed=type)
 
@@ -310,6 +311,7 @@ REF_REGEX = re.compile("^[a-z][a-z-]+[a-z]--[0-9a-fA-F]{8}-[0-9a-fA-F]{4}"
 
 
 class ReferenceProperty(Property):
+
     def __init__(self, required=False, type=None):
         """
         references sometimes must be to a specific object type
@@ -332,6 +334,7 @@ SELECTOR_REGEX = re.compile("^[a-z0-9_-]{3,250}(\\.(\\[\\d+\\]|[a-z0-9_-]{1,250}
 
 
 class SelectorProperty(Property):
+
     def __init__(self, type=None):
         # ignore type
         super(SelectorProperty, self).__init__()
@@ -347,6 +350,7 @@ class ObjectReferenceProperty(StringProperty):
 
 
 class EmbeddedObjectProperty(Property):
+
     def __init__(self, type, required=False):
         self.type = type
         super(EmbeddedObjectProperty, self).__init__(required, type=type)
@@ -357,3 +361,18 @@ class EmbeddedObjectProperty(Property):
         elif not isinstance(value, self.type):
             raise ValueError("must be of type %s." % self.type.__name__)
         return value
+
+
+class EnumProperty(StringProperty):
+
+    def __init__(self, allowed, **kwargs):
+        if type(allowed) is not list:
+            allowed = list(allowed)
+        self.allowed = allowed
+        super(EnumProperty, self).__init__(**kwargs)
+
+    def clean(self, value):
+        value = super(EnumProperty, self).clean(value)
+        if value not in self.allowed:
+            raise ValueError("value '%s' is not valid for this enumeration." % value)
+        return self.string_type(value)

stix2/test/test_observed_data.py

@@ -399,3 +399,20 @@ def test_software_example():
     assert s.cpe == "cpe:2.3:a:microsoft:word:2000:*:*:*:*:*:*:*"
     assert s.version == "2002"
     assert s.vendor == "Microsoft"
+
+
+def test_windows_registry_key_example():
+    with pytest.raises(ValueError):
+        v = stix2.WindowsRegistryValueType(name="Foo",
+                                           data="qwerty",
+                                           data_type="string")
+
+    v = stix2.WindowsRegistryValueType(name="Foo",
+                                       data="qwerty",
+                                       data_type="REG_SZ")
+    w = stix2.WindowsRegistryKey(key="hkey_local_machine\\system\\bar\\foo",
+                                 values=[v])
+    assert w.key == "hkey_local_machine\\system\\bar\\foo"
+    assert w.values[0].name == "Foo"
+    assert w.values[0].data == "qwerty"
+    assert w.values[0].data_type == "REG_SZ"

stix2/test/test_properties.py

@@ -4,9 +4,9 @@ from stix2.exceptions import DictionaryKeyError
 from stix2.observables import EmailMIMEComponent
 from stix2.properties import (BinaryProperty, BooleanProperty,
                               DictionaryProperty, EmbeddedObjectProperty,
-                              HashesProperty, HexProperty, IDProperty,
+                              EnumProperty, HashesProperty, HexProperty,
-                              IntegerProperty, ListProperty, Property,
+                              IDProperty, IntegerProperty, ListProperty,
-                              ReferenceProperty, StringProperty,
+                              Property, ReferenceProperty, StringProperty,
                               TimestampProperty, TypeProperty)
 
 from .constants import FAKE_TIME
@@ -247,3 +247,11 @@ def test_embedded_property():
 
     with pytest.raises(ValueError):
         emb_prop.clean("string")
+
+
+def test_enum_property():
+    enum_prop = EnumProperty(['a', 'b', 'c'])
+    assert enum_prop.clean('b')
+
+    with pytest.raises(ValueError):
+        enum_prop.clean('z')