diff --git a/stix2/__init__.py b/stix2/__init__.py index 5a515e8..f881f96 100644 --- a/stix2/__init__.py +++ b/stix2/__init__.py @@ -180,3 +180,74 @@ class Malware(_STIXBase): 'labels': self['labels'], 'name': self['name'], } + + +class Relationship(_STIXBase): + + _properties = [ + 'type', + 'id', + 'created', + 'modified', + 'relationship_type', + 'source_ref', + 'target_ref', + ] + + def __init__(self, **kwargs): + # TODO: + # - created_by_ref + # - revoked + # - external_references + # - object_marking_refs + # - granular_markings + + # - description + + # TODO: do we care about the performance penalty of creating this + # if we won't need it? + now = datetime.now(tz=pytz.UTC) + + if not kwargs.get('type'): + kwargs['type'] = 'relationship' + if kwargs['type'] != 'relationship': + raise ValueError("Relationship must have type='relationship'.") + + if not kwargs.get('id'): + kwargs['id'] = 'relationship--' + str(uuid.uuid4()) + if not kwargs['id'].startswith('relationship--'): + raise ValueError("Relationship id values must begin with 'relationship--'.") + + if not kwargs.get('relationship_type'): + raise ValueError("Missing required field for Relationship: 'relationship_type'.") + + if not kwargs.get('source_ref'): + raise ValueError("Missing required field for Relationship: 'source_ref'.") + + if not kwargs.get('target_ref'): + raise ValueError("Missing required field for Relationship: 'target_ref'.") + + extra_kwargs = list(set(kwargs.keys()) - set(self._properties)) + if extra_kwargs: + raise TypeError("unexpected keyword arguments: " + str(extra_kwargs)) + + self._inner = { + 'type': kwargs['type'], + 'id': kwargs['id'], + 'created': kwargs.get('created', now), + 'modified': kwargs.get('modified', now), + 'relationship_type': kwargs['relationship_type'], + 'source_ref': kwargs['source_ref'], + 'target_ref': kwargs['target_ref'], + } + + def _dict(self): + return { + 'type': self['type'], + 'id': self['id'], + 'created': format_datetime(self['created']), + 'modified': format_datetime(self['modified']), + 'relationship_type': self['relationship_type'], + 'source_ref': self['source_ref'], + 'target_ref': self['target_ref'], + } diff --git a/stix2/test/test_stix2.py b/stix2/test/test_stix2.py index b9423c7..02e721d 100644 --- a/stix2/test/test_stix2.py +++ b/stix2/test/test_stix2.py @@ -34,14 +34,33 @@ MALWARE_KWARGS = dict( ) +# Minimum required args for a Relationship instance +RELATIONSHIP_KWARGS = dict( + relationship_type="indicates", + source_ref="indicator--01234567-89ab-cdef-0123-456789abcdef", + target_ref="malware--fedcba98-7654-3210-fedc-ba9876543210", +) + + @pytest.fixture def indicator(): - return stix2.Indicator(**INDICATOR_KWARGS) + return stix2.Indicator( + id="indicator--01234567-89ab-cdef-0123-456789abcdef", + **INDICATOR_KWARGS + ) @pytest.fixture def malware(): - return stix2.Malware(**MALWARE_KWARGS) + return stix2.Malware( + id="malware--fedcba98-7654-3210-fedc-ba9876543210", + **MALWARE_KWARGS + ) + + +@pytest.fixture +def relationship(): + return stix2.Relationship(**RELATIONSHIP_KWARGS) EXPECTED_INDICATOR = """{ @@ -119,7 +138,7 @@ def test_indicator_required_field_pattern(): assert "Missing required field for Indicator: 'pattern'." in str(excinfo) -def test_cannot_assign_to_attributes(indicator): +def test_cannot_assign_to_indicator_attributes(indicator): with pytest.raises(ValueError) as excinfo: indicator.valid_from = datetime.datetime.now() @@ -202,8 +221,108 @@ def test_malware_required_field_name(): assert "Missing required field for Malware: 'name'." in str(excinfo) -def test_cannot_assign_to_attributes(malware): +def test_cannot_assign_to_malware_attributes(malware): with pytest.raises(ValueError) as excinfo: malware.name = "Cryptolocker II" assert "Cannot modify properties after creation." in str(excinfo) + + +def test_invalid_kwarg_to_malware(): + with pytest.raises(TypeError) as excinfo: + malware = stix2.Malware(my_custom_property="foo", **MALWARE_KWARGS) + assert "unexpected keyword arguments: ['my_custom_property']" in str(excinfo) + + +EXPECTED_RELATIONSHIP = """{ + "created": "2016-04-06T20:06:37Z", + "id": "relationship--44298a74-ba52-4f0c-87a3-1824e67d7fad", + "modified": "2016-04-06T20:06:37Z", + "relationship_type": "indicates", + "source_ref": "indicator--01234567-89ab-cdef-0123-456789abcdef", + "target_ref": "malware--fedcba98-7654-3210-fedc-ba9876543210", + "type": "relationship" +}""" + + +def test_relationship_all_required_fields(indicator, malware): + now = datetime.datetime(2016, 4, 6, 20, 6, 37, tzinfo=pytz.utc) + + relationship = stix2.Relationship( + type='relationship', + id='relationship--44298a74-ba52-4f0c-87a3-1824e67d7fad', + created=now, + modified=now, + relationship_type='indicates', + source_ref=indicator.id, + target_ref=malware.id, + ) + assert str(relationship) == EXPECTED_RELATIONSHIP + + +def test_relationship_autogenerated_fields(relationship): + assert relationship.type == 'relationship' + assert relationship.id.startswith('relationship--') + assert relationship.created is not None + assert relationship.modified is not None + assert relationship.relationship_type == 'indicates' + assert relationship.source_ref == "indicator--01234567-89ab-cdef-0123-456789abcdef" + assert relationship.target_ref == "malware--fedcba98-7654-3210-fedc-ba9876543210" + + assert relationship['type'] == 'relationship' + assert relationship['id'].startswith('relationship--') + assert relationship['created'] is not None + assert relationship['modified'] is not None + assert relationship['relationship_type'] == 'indicates' + assert relationship['source_ref'] == "indicator--01234567-89ab-cdef-0123-456789abcdef" + assert relationship['target_ref'] == "malware--fedcba98-7654-3210-fedc-ba9876543210" + + +def test_relationship_type_must_be_relationship(): + with pytest.raises(ValueError) as excinfo: + relationship = stix2.Relationship(type='xxx') + + assert "Relationship must have type='relationship'." in str(excinfo) + + +def test_relationship_id_must_start_with_relationship(): + with pytest.raises(ValueError) as excinfo: + relationship = stix2.Relationship(id='my-prefix--') + + assert "Relationship id values must begin with 'relationship--'." in str(excinfo) + + +def test_relationship_required_field_relationship_type(): + with pytest.raises(ValueError) as excinfo: + relationship = stix2.Relationship() + assert "Missing required field for Relationship: 'relationship_type'." in str(excinfo) + + +def test_relationship_required_field_source_ref(): + with pytest.raises(ValueError) as excinfo: + # relationship_type is checked first, so make sure that is provided + relationship = stix2.Relationship(relationship_type='indicates') + assert "Missing required field for Relationship: 'source_ref'." in str(excinfo) + + +def test_relationship_required_field_target_ref(): + with pytest.raises(ValueError) as excinfo: + # relationship_type and source_ref are checked first, so make sure those are provided + relationship = stix2.Relationship( + relationship_type='indicates', + source_ref='indicator--01234567-89ab-cdef-0123-456789abcdef' + ) + assert "Missing required field for Relationship: 'target_ref'." in str(excinfo) + + +def test_cannot_assign_to_relationship_attributes(relationship): + with pytest.raises(ValueError) as excinfo: + relationship.relationship_type = "derived-from" + + assert "Cannot modify properties after creation." in str(excinfo) + + +def test_invalid_kwarg_to_relationship(): + with pytest.raises(TypeError) as excinfo: + relationship = stix2.Relationship(my_custom_property="foo", **RELATIONSHIP_KWARGS) + assert "unexpected keyword arguments: ['my_custom_property']" in str(excinfo)