From d87718c15cabc17a247791625c719a52eed6290b Mon Sep 17 00:00:00 2001 From: Michael Chisholm Date: Wed, 30 Jun 2021 17:50:00 -0400 Subject: [PATCH] Bug fixes, hackage removal, and some pre-commit stylistic fixes. - Fixed bugged logic in _STIXBase._check_at_least_one_property(), and revamped the code to be simpler and clearer. - Changed custom extension registration to auto-create an "extension_type" property based on the attribute of that name on the custom class, if present. - The custom extension registration change above uncovered what seemed like a bug in a unit test: a custom extension was registered, but it was not given an extension type. The test used the extension as extension_type="property-extension"; this now causes a standard error about an extra property. I fixed the test to assign the custom extension the proper type. --- stix2/base.py | 40 ++++++++++++++++++---------- stix2/properties.py | 6 ++--- stix2/test/v20/test_observed_data.py | 1 + stix2/test/v21/test_custom.py | 8 +++--- stix2/test/v21/test_observed_data.py | 4 +-- stix2/v21/base.py | 11 -------- stix2/v21/common.py | 2 +- stix2/v21/observables.py | 19 ++++++++++++- stix2/v21/vocab.py | 2 +- 9 files changed, 56 insertions(+), 37 deletions(-) diff --git a/stix2/base.py b/stix2/base.py index 4a28d8f..5b88998 100644 --- a/stix2/base.py +++ b/stix2/base.py @@ -78,20 +78,34 @@ class _STIXBase(collections.abc.Mapping): if count > 1 or (at_least_one and count == 0): raise MutuallyExclusivePropertiesError(self.__class__, list_of_properties) - def _check_at_least_one_property(self, list_of_properties=None): - if not list_of_properties: - list_of_properties = sorted(self.__class__._properties.keys()) + def _check_at_least_one_property(self, properties_checked=None): + """ + Check whether one or more of the given properties are present. + + :param properties_checked: An iterable of the names of the properties + of interest, or None to check against a default list. The default + list includes all properties defined on the object, with some + hard-coded exceptions. + :raises AtLeastOnePropertyError: If none of the given properties are + present. + """ + if properties_checked is None: + property_exceptions = {"extensions", "type"} if isinstance(self, _Observable): - props_to_remove = {"type", "id", "defanged", "spec_version"} - else: - props_to_remove = {"type"} + property_exceptions |= {"id", "defanged", "spec_version"} - list_of_properties = [prop for prop in list_of_properties if prop not in props_to_remove] - current_properties = self.properties_populated() - list_of_properties_populated = set(list_of_properties).intersection(current_properties) + properties_checked = self._properties.keys() - property_exceptions - if list_of_properties and (not list_of_properties_populated or list_of_properties_populated == {'extensions'}): - raise AtLeastOnePropertyError(self.__class__, list_of_properties) + elif not isinstance(properties_checked, set): + properties_checked = set(properties_checked) + + if properties_checked: + properties_checked_assigned = properties_checked & self.keys() + + if not properties_checked_assigned: + raise AtLeastOnePropertyError( + self.__class__, properties_checked + ) def _check_properties_dependency(self, list_of_properties, list_of_dependent_properties): failed_dependency_pairs = [] @@ -118,8 +132,6 @@ class _STIXBase(collections.abc.Mapping): raise ValueError("'custom_properties' must be a dictionary") extra_kwargs = kwargs.keys() - self._properties.keys() - if extra_kwargs and issubclass(cls, stix2.v21._Extension): - extra_kwargs = {prop for prop in extra_kwargs if prop != 'extension_type'} if extra_kwargs and not allow_custom: ext_found = False @@ -129,7 +141,7 @@ class _STIXBase(collections.abc.Mapping): for key_id, ext_def in kwargs.get('extensions', {}).items(): if ( key_id.startswith('extension-definition--') and - ext_def.get('extension_type', '') == 'toplevel-property-extension' + ext_def.get('extension_type') == 'toplevel-property-extension' ): ext_found = True break diff --git a/stix2/properties.py b/stix2/properties.py index e63206a..be1c4b6 100644 --- a/stix2/properties.py +++ b/stix2/properties.py @@ -793,8 +793,8 @@ class ExtensionsProperty(DictionaryProperty): else: raise TypeError( "Can't create extension '{}' from {}.".format( - key, type(subvalue) - ) + key, type(subvalue), + ), ) has_custom = has_custom or ext.has_custom @@ -818,7 +818,7 @@ class ExtensionsProperty(DictionaryProperty): if key.startswith('extension-definition--'): _validate_id( - key, self.spec_version, 'extension-definition--' + key, self.spec_version, 'extension-definition--', ) elif allow_custom: has_custom = True diff --git a/stix2/test/v20/test_observed_data.py b/stix2/test/v20/test_observed_data.py index 79ed6c4..196fb46 100644 --- a/stix2/test/v20/test_observed_data.py +++ b/stix2/test/v20/test_observed_data.py @@ -1117,6 +1117,7 @@ def test_process_example_empty_error(): assert excinfo.value.cls == stix2.v20.Process properties_of_process = list(stix2.v20.Process._properties.keys()) properties_of_process.remove("type") + properties_of_process.remove("extensions") assert excinfo.value.properties == sorted(properties_of_process) msg = "At least one of the ({1}) properties for {0} must be populated." msg = msg.format( diff --git a/stix2/test/v21/test_custom.py b/stix2/test/v21/test_custom.py index ff0464b..1b3166f 100644 --- a/stix2/test/v21/test_custom.py +++ b/stix2/test/v21/test_custom.py @@ -1237,9 +1237,9 @@ def test_unregistered_new_style_extension(): "extension-definition--31adb724-a9a4-44b6-8ec2-fd4b181c9507": { "extension-type": "property-extension", "a": 1, - "b": True - } - } + "b": True, + }, + }, } f = stix2.parse(f_dict, allow_custom=False) @@ -1508,7 +1508,7 @@ def test_registered_embedded_extension_passes_with_allow_custom_false(): ], ) class ExtensionFoo1: - pass + extension_type = "property-extension" indicator = stix2.v21.Indicator( id='indicator--e97bfccf-8970-4a3c-9cd1-5b5b97ed5d0c', diff --git a/stix2/test/v21/test_observed_data.py b/stix2/test/v21/test_observed_data.py index 28c40c0..d2ccec4 100644 --- a/stix2/test/v21/test_observed_data.py +++ b/stix2/test/v21/test_observed_data.py @@ -1218,8 +1218,8 @@ def test_process_example_empty_error(): stix2.v21.Process() assert excinfo.value.cls == stix2.v21.Process - properties_of_process = list(stix2.v21.Process._properties.keys()) - properties_of_process = [prop for prop in properties_of_process if prop not in ["type", "id", "defanged", "spec_version"]] + properties_of_process = stix2.v21.Process._properties.keys() + properties_of_process -= {"type", "id", "defanged", "spec_version", "extensions"} assert excinfo.value.properties == sorted(properties_of_process) msg = "At least one of the ({1}) properties for {0} must be populated." msg = msg.format( diff --git a/stix2/v21/base.py b/stix2/v21/base.py index 90b9fd2..3878b79 100644 --- a/stix2/v21/base.py +++ b/stix2/v21/base.py @@ -29,17 +29,6 @@ class _Observable(_Observable, _STIXBase21): class _Extension(_Extension, _STIXBase21): extension_type = None - def __init__(self, **kwargs): - super(_Extension, self).__init__(**kwargs) - if getattr(self, "extension_type", None): - self._inner["extension_type"] = self.extension_type - - def _check_at_least_one_property(self, list_of_properties=None): - new_ext_check = getattr(self, "extension_type", None) - - if new_ext_check is None: - super(_Extension, self)._check_at_least_one_property(list_of_properties=list_of_properties) - class _DomainObject(_DomainObject, _STIXBase21): pass diff --git a/stix2/v21/common.py b/stix2/v21/common.py index 7c9cd70..e62bdae 100644 --- a/stix2/v21/common.py +++ b/stix2/v21/common.py @@ -14,7 +14,7 @@ from ..properties import ( ) from ..utils import NOW, _get_dict from .base import _STIXBase21 -from .vocab import HASHING_ALGORITHM, EXTENSION_TYPE +from .vocab import EXTENSION_TYPE, HASHING_ALGORITHM class ExternalReference(_STIXBase21): diff --git a/stix2/v21/observables.py b/stix2/v21/observables.py index 4e5876d..44b346c 100644 --- a/stix2/v21/observables.py +++ b/stix2/v21/observables.py @@ -6,6 +6,7 @@ _Observable and do not have a ``_type`` attribute. """ from collections import OrderedDict +from collections.abc import Mapping import itertools from ..custom import _custom_extension_builder, _custom_observable_builder @@ -20,7 +21,7 @@ from ..properties import ( from .base import _Extension, _Observable, _STIXBase21 from .common import GranularMarking from .vocab import ( - ACCOUNT_TYPE, ENCRYPTION_ALGORITHM, HASHING_ALGORITHM, + ACCOUNT_TYPE, ENCRYPTION_ALGORITHM, EXTENSION_TYPE, HASHING_ALGORITHM, NETWORK_SOCKET_ADDRESS_FAMILY, NETWORK_SOCKET_TYPE, WINDOWS_INTEGRITY_LEVEL, WINDOWS_PEBINARY_TYPE, WINDOWS_REGISTRY_DATATYPE, WINDOWS_SERVICE_START_TYPE, WINDOWS_SERVICE_STATUS, WINDOWS_SERVICE_TYPE, @@ -901,5 +902,21 @@ def CustomExtension(type='x-custom-observable-ext', properties=None): """Custom STIX Object Extension decorator. """ def wrapper(cls): + + # Auto-create an "extension_type" property from the class attribute, if + # it exists. + extension_type = getattr(cls, "extension_type", None) + if extension_type: + extension_type_prop = EnumProperty( + EXTENSION_TYPE, + required=False, + fixed=extension_type, + ) + + if isinstance(properties, Mapping): + properties["extension_type"] = extension_type_prop + else: + properties.append(("extension_type", extension_type_prop)) + return _custom_extension_builder(cls, type, properties, '2.1', _Extension) return wrapper diff --git a/stix2/v21/vocab.py b/stix2/v21/vocab.py index 9443d0e..303692e 100644 --- a/stix2/v21/vocab.py +++ b/stix2/v21/vocab.py @@ -99,7 +99,7 @@ EXTENSION_TYPE = [ EXTENSION_TYPE_NEW_SCO, EXTENSION_TYPE_NEW_SRO, EXTENSION_TYPE_PROPERTY_EXTENSION, - EXTENSION_TYPE_TOPLEVEL_PROPERTY_EXTENSION + EXTENSION_TYPE_TOPLEVEL_PROPERTY_EXTENSION, ]