Allow shorter syntax for creating relationships.
parent
fd548a5f41
commit
742d9645d6
|
@ -194,7 +194,8 @@ class Relationship(_STIXBase):
|
|||
'target_ref',
|
||||
]
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
def __init__(self, source_ref=None, relationship_type=None, target_ref=None,
|
||||
**kwargs):
|
||||
# TODO:
|
||||
# - created_by_ref
|
||||
# - revoked
|
||||
|
@ -204,6 +205,13 @@ class Relationship(_STIXBase):
|
|||
|
||||
# - description
|
||||
|
||||
if source_ref and not kwargs.get('source_ref'):
|
||||
kwargs['source_ref'] = source_ref
|
||||
if relationship_type and not kwargs.get('relationship_type'):
|
||||
kwargs['relationship_type'] = relationship_type
|
||||
if target_ref and not kwargs.get('target_ref'):
|
||||
kwargs['target_ref'] = target_ref
|
||||
|
||||
# TODO: do we care about the performance penalty of creating this
|
||||
# if we won't need it?
|
||||
now = datetime.now(tz=pytz.UTC)
|
||||
|
|
|
@ -336,3 +336,11 @@ def test_create_relationship_from_objects_rather_than_ids(indicator, malware):
|
|||
assert relationship.relationship_type == 'indicates'
|
||||
assert relationship.source_ref == INDICATOR_ID
|
||||
assert relationship.target_ref == MALWARE_ID
|
||||
|
||||
|
||||
def test_create_relationship_with_positional_args(indicator, malware):
|
||||
relationship = stix2.Relationship(indicator, 'indicates', malware)
|
||||
|
||||
assert relationship.relationship_type == 'indicates'
|
||||
assert relationship.source_ref == INDICATOR_ID
|
||||
assert relationship.target_ref == MALWARE_ID
|
||||
|
|
Loading…
Reference in New Issue