Add Relationship class with required fields.

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'],
+        }

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)