From 2645bf2c717429a40077703b7e4c603df01ef948 Mon Sep 17 00:00:00 2001 From: Greg Back Date: Fri, 24 Feb 2017 10:46:21 -0600 Subject: [PATCH] Convert 'type' to a new Property class. --- stix2/bundle.py | 5 ++--- stix2/common.py | 8 +------- stix2/properties.py | 14 +++++++++++++- stix2/sdo.py | 14 +++++++++++++- stix2/sro.py | 3 ++- stix2/test/test_properties.py | 20 ++++++++++++++++---- stix2/test/test_stix2.py | 8 ++++---- 7 files changed, 51 insertions(+), 21 deletions(-) diff --git a/stix2/bundle.py b/stix2/bundle.py index 8723640..6ee62c6 100644 --- a/stix2/bundle.py +++ b/stix2/bundle.py @@ -1,15 +1,14 @@ """STIX 2 Bundle object""" from .base import _STIXBase -from .common import TYPE_PROPERTY -from .properties import IDProperty +from .properties import IDProperty, TypeProperty class Bundle(_STIXBase): _type = 'bundle' _properties = { - 'type': TYPE_PROPERTY, + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'spec_version': { 'fixed': "2.0", diff --git a/stix2/common.py b/stix2/common.py index 1198832..be417af 100644 --- a/stix2/common.py +++ b/stix2/common.py @@ -4,11 +4,6 @@ import re from .base import _STIXBase from .utils import NOW -TYPE_PROPERTY = { - 'default': (lambda x: x._type), - 'validate': (lambda x, val: val == x._type) -} - ref_regex = ("^[a-z][a-z-]+[a-z]--[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}" "-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$") @@ -23,8 +18,7 @@ BOOL_PROPERTY = { } COMMON_PROPERTIES = { - 'type': TYPE_PROPERTY, - # 'id' should be defined on each individual type + # 'type' and 'id' should be defined on each individual type 'created': { 'default': NOW, }, diff --git a/stix2/properties.py b/stix2/properties.py index 0a2341e..b3f6ca4 100644 --- a/stix2/properties.py +++ b/stix2/properties.py @@ -44,10 +44,16 @@ class Property(object): lambdas cannot raise their own exceptions. """ + def _default_validate(self, value): + if value != self._fixed_value: + raise ValueError("must equal '{0}'.".format(self._fixed_value)) + return value + def __init__(self, required=False, fixed=None, clean=None, validate=None, default=None): self.required = required if fixed: - self.validate = lambda x: x == fixed + self._fixed_value = fixed + self.validate = self._default_validate self.default = lambda: fixed if clean: self.clean = clean @@ -97,6 +103,12 @@ class List(Property): return [self.contained(x) for x in value] +class TypeProperty(Property): + + def __init__(self, type): + super(TypeProperty, self).__init__(fixed=type) + + class IDProperty(Property): def __init__(self, type): diff --git a/stix2/sdo.py b/stix2/sdo.py index d6ac7de..35977b4 100644 --- a/stix2/sdo.py +++ b/stix2/sdo.py @@ -2,7 +2,7 @@ from .base import _STIXBase from .common import COMMON_PROPERTIES -from .properties import IDProperty +from .properties import IDProperty, TypeProperty from .utils import NOW @@ -11,6 +11,7 @@ class AttackPattern(_STIXBase): _type = 'attack-pattern' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'name': { 'required': True, @@ -37,6 +38,7 @@ class Campaign(_STIXBase): _type = 'campaign' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'name': { 'required': True, @@ -69,6 +71,7 @@ class CourseOfAction(_STIXBase): _type = 'course-of-action' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'name': { 'required': True, @@ -93,6 +96,7 @@ class Identity(_STIXBase): _type = 'identity' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'name': { 'required': True, @@ -125,6 +129,7 @@ class Indicator(_STIXBase): _type = 'indicator' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'labels': { 'required': True, @@ -161,6 +166,7 @@ class IntrusionSet(_STIXBase): _type = 'intrusion-set' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'name': { 'required': True, @@ -199,6 +205,7 @@ class Malware(_STIXBase): _type = 'malware' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'labels': { 'required': True, @@ -228,6 +235,7 @@ class ObservedData(_STIXBase): _type = 'observed-data' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'first_observed': {}, 'last_observed': {}, @@ -255,6 +263,7 @@ class Report(_STIXBase): _type = 'report' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'labels': { 'required': True, @@ -286,6 +295,7 @@ class ThreatActor(_STIXBase): _type = 'threat-actor' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'labels': { 'required': True, @@ -329,6 +339,7 @@ class Tool(_STIXBase): _type = 'tool' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'labels': { 'required': True, @@ -360,6 +371,7 @@ class Vulnerability(_STIXBase): _type = 'vulnerability' _properties = COMMON_PROPERTIES.copy() _properties.update({ + 'type': TypeProperty(_type), 'id': IDProperty(_type), 'name': { 'required': True, diff --git a/stix2/sro.py b/stix2/sro.py index 0226011..10340be 100644 --- a/stix2/sro.py +++ b/stix2/sro.py @@ -2,7 +2,7 @@ from .base import _STIXBase from .common import COMMON_PROPERTIES -from .properties import IDProperty +from .properties import IDProperty, TypeProperty class Relationship(_STIXBase): @@ -11,6 +11,7 @@ class Relationship(_STIXBase): _properties = COMMON_PROPERTIES.copy() _properties.update({ 'id': IDProperty(_type), + 'type': TypeProperty(_type), 'relationship_type': { 'required': True, }, diff --git a/stix2/test/test_properties.py b/stix2/test/test_properties.py index 3a793a4..82f6bc6 100644 --- a/stix2/test/test_properties.py +++ b/stix2/test/test_properties.py @@ -1,6 +1,6 @@ import pytest -from stix2.properties import Property, IDProperty +from stix2.properties import Property, IDProperty, TypeProperty def test_property(): @@ -39,11 +39,23 @@ def test_default_field(): def test_fixed_property(): p = Property(fixed="2.0") - assert p.validate("2.0") is True - assert p.validate("x") is False - assert p.validate(2.0) is False + assert p.validate("2.0") + with pytest.raises(ValueError): + assert p.validate("x") is False + with pytest.raises(ValueError): + assert p.validate(2.0) is False assert p.default() == "2.0" + assert p.validate(p.default()) + + +def test_type_property(): + prop = TypeProperty('my-type') + + assert prop.validate('my-type') + with pytest.raises(ValueError): + prop.validate('not-my-type') + assert prop.validate(prop.default()) def test_id_property(): diff --git a/stix2/test/test_stix2.py b/stix2/test/test_stix2.py index 40a8f9b..bbd030e 100644 --- a/stix2/test/test_stix2.py +++ b/stix2/test/test_stix2.py @@ -154,7 +154,7 @@ def test_indicator_type_must_be_indicator(): with pytest.raises(ValueError) as excinfo: indicator = stix2.Indicator(type='xxx', **INDICATOR_KWARGS) - assert str(excinfo.value) == "Indicator must have type='indicator'." + assert str(excinfo.value) == "Invalid value for Indicator 'type': must equal 'indicator'." def test_indicator_id_must_start_with_indicator(): @@ -255,7 +255,7 @@ def test_malware_type_must_be_malware(): with pytest.raises(ValueError) as excinfo: malware = stix2.Malware(type='xxx', **MALWARE_KWARGS) - assert str(excinfo.value) == "Malware must have type='malware'." + assert str(excinfo.value) == "Invalid value for Malware 'type': must equal 'malware'." def test_malware_id_must_start_with_malware(): @@ -338,7 +338,7 @@ def test_relationship_type_must_be_relationship(): with pytest.raises(ValueError) as excinfo: relationship = stix2.Relationship(type='xxx', **RELATIONSHIP_KWARGS) - assert str(excinfo.value) == "Relationship must have type='relationship'." + assert str(excinfo.value) == "Invalid value for Relationship 'type': must equal 'relationship'." def test_relationship_id_must_start_with_relationship(): @@ -457,7 +457,7 @@ def test_bundle_with_wrong_type(): with pytest.raises(ValueError) as excinfo: bundle = stix2.Bundle(type="not-a-bundle") - assert str(excinfo.value) == "Bundle must have type='bundle'." + assert str(excinfo.value) == "Invalid value for Bundle 'type': must equal 'bundle'." def test_bundle_id_must_start_with_bundle():