Allow creating relationships from objects, not just IDs.
parent
e683acbf48
commit
fd548a5f41
|
@ -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:
|
||||||
|
|
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue