Allow creating relationships from objects, not just IDs.

stix2.1
Greg Back 2017-01-18 15:14:22 -08:00
parent e683acbf48
commit fd548a5f41
2 changed files with 16 additions and 0 deletions

View File

@ -223,9 +223,13 @@ class Relationship(_STIXBase):
if not kwargs.get('source_ref'): if not kwargs.get('source_ref'):
raise ValueError("Missing required field for Relationship: 'source_ref'.") raise ValueError("Missing required field for Relationship: 'source_ref'.")
elif isinstance(kwargs['source_ref'], _STIXBase):
kwargs['source_ref'] = kwargs['source_ref'].id
if not kwargs.get('target_ref'): if not kwargs.get('target_ref'):
raise ValueError("Missing required field for Relationship: 'target_ref'.") raise ValueError("Missing required field for Relationship: 'target_ref'.")
elif isinstance(kwargs['target_ref'], _STIXBase):
kwargs['target_ref'] = kwargs['target_ref'].id
extra_kwargs = list(set(kwargs.keys()) - set(self._properties)) extra_kwargs = list(set(kwargs.keys()) - set(self._properties))
if extra_kwargs: if extra_kwargs:

View File

@ -324,3 +324,15 @@ def test_invalid_kwarg_to_relationship():
with pytest.raises(TypeError) as excinfo: with pytest.raises(TypeError) as excinfo:
relationship = stix2.Relationship(my_custom_property="foo", **RELATIONSHIP_KWARGS) relationship = stix2.Relationship(my_custom_property="foo", **RELATIONSHIP_KWARGS)
assert "unexpected keyword arguments: ['my_custom_property']" in str(excinfo) assert "unexpected keyword arguments: ['my_custom_property']" in str(excinfo)
def test_create_relationship_from_objects_rather_than_ids(indicator, malware):
relationship = stix2.Relationship(
relationship_type="indicates",
source_ref=indicator,
target_ref=malware,
)
assert relationship.relationship_type == 'indicates'
assert relationship.source_ref == INDICATOR_ID
assert relationship.target_ref == MALWARE_ID