diff --git a/stix2/__init__.py b/stix2/__init__.py index 573e9f8..bf48a48 100644 --- a/stix2/__init__.py +++ b/stix2/__init__.py @@ -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) diff --git a/stix2/test/test_stix2.py b/stix2/test/test_stix2.py index e8c1e21..591be64 100644 --- a/stix2/test/test_stix2.py +++ b/stix2/test/test_stix2.py @@ -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