Add support for required fields with no default values.

stix2.1
Greg Back 2017-02-02 08:53:46 -06:00
parent 67c3311672
commit 675a29dbfb
1 changed files with 33 additions and 33 deletions

View File

@ -58,6 +58,10 @@ class _STIXBase(collections.Mapping):
for prop_name, prop_metadata in cls._properties.items(): for prop_name, prop_metadata in cls._properties.items():
if prop_name not in kwargs: if prop_name not in kwargs:
if prop_metadata.get('required'):
msg = "Missing required field for {type}: '{field}'."
raise ValueError(msg.format(type=class_name,
field=prop_name))
if prop_metadata.get('default'): if prop_metadata.get('default'):
kwargs[prop_name] = prop_metadata['default'](cls) kwargs[prop_name] = prop_metadata['default'](cls)
elif prop_metadata.get('fixed'): elif prop_metadata.get('fixed'):
@ -145,8 +149,12 @@ class Indicator(_STIXBase):
_type = 'indicator' _type = 'indicator'
_properties = COMMON_PROPERTIES.copy() _properties = COMMON_PROPERTIES.copy()
_properties.update({ _properties.update({
'labels': {}, 'labels': {
'pattern': {}, 'required': True,
},
'pattern': {
'required': True,
},
'valid_from': {}, 'valid_from': {},
}) })
@ -170,12 +178,6 @@ class Indicator(_STIXBase):
# TODO: remove once we check all the fields in the right order # TODO: remove once we check all the fields in the right order
kwargs = self._check_kwargs(**kwargs) kwargs = self._check_kwargs(**kwargs)
if not kwargs.get('labels'):
raise ValueError("Missing required field for Indicator: 'labels'.")
if not kwargs.get('pattern'):
raise ValueError("Missing required field for Indicator: 'pattern'.")
kwargs.update({ kwargs.update({
'created': kwargs.get('created', now), 'created': kwargs.get('created', now),
'modified': kwargs.get('modified', now), 'modified': kwargs.get('modified', now),
@ -189,8 +191,12 @@ class Malware(_STIXBase):
_type = 'malware' _type = 'malware'
_properties = COMMON_PROPERTIES.copy() _properties = COMMON_PROPERTIES.copy()
_properties.update({ _properties.update({
'labels': {}, 'labels': {
'name': {}, 'required': True,
},
'name': {
'required': True,
},
}) })
def __init__(self, **kwargs): def __init__(self, **kwargs):
@ -211,12 +217,6 @@ class Malware(_STIXBase):
# TODO: remove once we check all the fields in the right order # TODO: remove once we check all the fields in the right order
kwargs = self._check_kwargs(**kwargs) kwargs = self._check_kwargs(**kwargs)
if not kwargs.get('labels'):
raise ValueError("Missing required field for Malware: 'labels'.")
if not kwargs.get('name'):
raise ValueError("Missing required field for Malware: 'name'.")
kwargs.update({ kwargs.update({
'created': kwargs.get('created', now), 'created': kwargs.get('created', now),
'modified': kwargs.get('modified', now), 'modified': kwargs.get('modified', now),
@ -229,9 +229,15 @@ class Relationship(_STIXBase):
_type = 'relationship' _type = 'relationship'
_properties = COMMON_PROPERTIES.copy() _properties = COMMON_PROPERTIES.copy()
_properties.update({ _properties.update({
'relationship_type': {}, 'relationship_type': {
'source_ref': {}, 'required': True,
'target_ref': {}, },
'source_ref': {
'required': True,
},
'target_ref': {
'required': True,
},
}) })
# Explicitly define the first three kwargs to make readable Relationship declarations. # Explicitly define the first three kwargs to make readable Relationship declarations.
@ -246,9 +252,7 @@ class Relationship(_STIXBase):
# - description # - description
# TODO: remove once we check all the fields in the right order # Allow (source_ref, relationship_type, target_ref) as positional args.
kwargs = self._check_kwargs(**kwargs)
if source_ref and not kwargs.get('source_ref'): if source_ref and not kwargs.get('source_ref'):
kwargs['source_ref'] = source_ref kwargs['source_ref'] = source_ref
if relationship_type and not kwargs.get('relationship_type'): if relationship_type and not kwargs.get('relationship_type'):
@ -256,28 +260,24 @@ class Relationship(_STIXBase):
if target_ref and not kwargs.get('target_ref'): if target_ref and not kwargs.get('target_ref'):
kwargs['target_ref'] = target_ref kwargs['target_ref'] = target_ref
# TODO: remove once we check all the fields in the right order
kwargs = self._check_kwargs(**kwargs)
# TODO: do we care about the performance penalty of creating this # TODO: do we care about the performance penalty of creating this
# if we won't need it? # if we won't need it?
now = datetime.datetime.now(tz=pytz.UTC) now = datetime.datetime.now(tz=pytz.UTC)
if not kwargs.get('relationship_type'): # If actual STIX objects (vs. just the IDs) are passed in, extract the
raise ValueError("Missing required field for Relationship: 'relationship_type'.") # ID values to use in the Relationship object.
if kwargs.get('source_ref') and isinstance(kwargs['source_ref'], _STIXBase):
if not kwargs.get('source_ref'):
raise ValueError("Missing required field for Relationship: 'source_ref'.")
elif isinstance(kwargs['source_ref'], _STIXBase):
kwargs['source_ref'] = kwargs['source_ref'].id kwargs['source_ref'] = kwargs['source_ref'].id
if not kwargs.get('target_ref'): if kwargs.get('target_ref') and isinstance(kwargs['target_ref'], _STIXBase):
raise ValueError("Missing required field for Relationship: 'target_ref'.")
elif isinstance(kwargs['target_ref'], _STIXBase):
kwargs['target_ref'] = kwargs['target_ref'].id kwargs['target_ref'] = kwargs['target_ref'].id
kwargs.update({ kwargs.update({
'created': kwargs.get('created', now), 'created': kwargs.get('created', now),
'modified': kwargs.get('modified', now), 'modified': kwargs.get('modified', now),
'relationship_type': kwargs['relationship_type'],
'target_ref': kwargs['target_ref'],
}) })
super(Relationship, self).__init__(**kwargs) super(Relationship, self).__init__(**kwargs)