Add Malware object with required fields.

stix2/__init__.py

@@ -7,14 +7,19 @@ import pytz
 
 
 def format_datetime(dt):
+    # TODO: how to handle naive datetime
+
     # 1. Convert to UTC
     # 2. Format in isoformat
     # 3. Strip off "+00:00"
     # 4. Add "Z"
+
+    # TODO: how to handle timestamps with subsecond 0's
     return dt.astimezone(pytz.utc).isoformat()[:-6] + "Z"
 
 
-class Indicator(collections.Mapping):
+class _STIXBase(collections.Mapping):
+    """Base class for STIX object types"""
 
     def __getitem__(self, key):
         return self._inner[key]
@@ -32,7 +37,10 @@ class Indicator(collections.Mapping):
     def __setattr__(self, name, value):
         if name != '_inner':
             raise ValueError("Cannot modify properties after creation.")
-        super(Indicator, self).__setattr__(name, value)
+        super(_STIXBase, self).__setattr__(name, value)
+
+
+class Indicator(_STIXBase):
 
     def __init__(self, type='indicator', id=None, created=None, modified=None,
                  labels=None, pattern=None, valid_from=None):
@@ -67,7 +75,7 @@ class Indicator(collections.Mapping):
 
         self._inner = {
             'type': type,
-            'id': id or now,
+            'id': id,
             'created': created or now,
             'modified': modified or now,
             'labels': labels,
@@ -86,3 +94,55 @@ class Indicator(collections.Mapping):
             'pattern': self['pattern'],
             'valid_from': format_datetime(self['valid_from']),
         }, indent=4, sort_keys=True, separators=(",", ": "))  # Don't include spaces after commas.
+
+
+class Malware(_STIXBase):
+
+    def __init__(self, type='malware', id=None, created=None, modified=None,
+                 labels=None, name=None):
+        # TODO:
+        # - created_by_ref
+        # - revoked
+        # - external_references
+        # - object_marking_refs
+        # - granular_markings
+
+        # - description
+        # - kill_chain_phases
+
+        if not created or not modified:
+            now = datetime.now(tz=pytz.UTC)
+
+        if type != 'malware':
+            raise ValueError("Malware must have type='malware'.")
+
+        if not id:
+            id = 'malware--' + str(uuid.uuid4())
+        if not id.startswith('malware--'):
+            raise ValueError("Malware id values must begin with 'malware--'.")
+
+        if not labels:
+            raise ValueError("Missing required field for Malware: 'labels'.")
+
+        if not name:
+            raise ValueError("Missing required field for Malware: 'name'.")
+
+        self._inner = {
+            'type': type,
+            'id': id,
+            'created': created or now,
+            'modified': modified or now,
+            'labels': labels,
+            'name': name,
+        }
+
+    def __str__(self):
+        # TODO: put keys in specific order. Probably need custom JSON encoder.
+        return json.dumps({
+            'type': self['type'],
+            'id': self['id'],
+            'created': format_datetime(self['created']),
+            'modified': format_datetime(self['modified']),
+            'labels': self['labels'],
+            'name': self['name'],
+        }, indent=4, sort_keys=True, separators=(",", ": "))  # Don't include spaces after commas.

stix2/test/test_stix2.py

@@ -21,7 +21,7 @@ def test_timestamp_formatting(dt, timestamp):
     assert stix2.format_datetime(dt) == timestamp
 
 
-EXPECTED = """{
+EXPECTED_INDICATOR = """{
     "created": "2017-01-01T00:00:00Z",
     "id": "indicator--01234567-89ab-cdef-0123-456789abcdef",
     "labels": [
@@ -48,18 +48,18 @@ def test_indicator_with_all_required_fields():
         valid_from=epoch,
     )
 
-    assert str(indicator) == EXPECTED
+    assert str(indicator) == EXPECTED_INDICATOR
 
 
-# Minimum required args for an indicator
-KWARGS = dict(
+# Minimum required args for an Indicator instance
+INDICATOR_KWARGS = dict(
     labels=['malicious-activity'],
     pattern="[file:hashes.MD5 = 'd41d8cd98f00b204e9800998ecf8427e']",
 )
 
 
 def test_indicator_autogenerated_fields():
-    indicator = stix2.Indicator(**KWARGS)
+    indicator = stix2.Indicator(**INDICATOR_KWARGS)
     assert indicator.type == 'indicator'
     assert indicator.id.startswith('indicator--')
     assert indicator.created is not None
@@ -105,9 +105,96 @@ def test_indicator_required_field_pattern():
 
 
 def test_cannot_assign_to_attributes():
-    indicator = stix2.Indicator(**KWARGS)
+    indicator = stix2.Indicator(**INDICATOR_KWARGS)
 
     with pytest.raises(ValueError) as excinfo:
         indicator.valid_from = datetime.datetime.now()
 
     assert "Cannot modify properties after creation." in str(excinfo)
+
+
+EXPECTED_MALWARE = """{
+    "created": "2016-05-12T08:17:27Z",
+    "id": "malware--0c7b5b88-8ff7-4a4d-aa9d-feb398cd0061",
+    "labels": [
+        "ransomware"
+    ],
+    "modified": "2016-05-12T08:17:27Z",
+    "name": "Cryptolocker",
+    "type": "malware"
+}"""
+
+
+def test_malware_with_all_required_fields():
+    now = datetime.datetime(2016, 5, 12, 8, 17, 27, tzinfo=pytz.utc)
+
+    malware = stix2.Malware(
+        type="malware",
+        id="malware--0c7b5b88-8ff7-4a4d-aa9d-feb398cd0061",
+        created=now,
+        modified=now,
+        labels=["ransomware"],
+        name="Cryptolocker",
+    )
+
+    assert str(malware) == EXPECTED_MALWARE
+
+
+# Minimum required args for a Malware instance
+MALWARE_KWARGS = dict(
+    labels=['ransomware'],
+    name="Cryptolocker",
+)
+
+
+def test_malware_autogenerated_fields():
+    malware = stix2.Malware(**MALWARE_KWARGS)
+    assert malware.type == 'malware'
+    assert malware.id.startswith('malware--')
+    assert malware.created is not None
+    assert malware.modified is not None
+    assert malware.labels == ['ransomware']
+    assert malware.name == "Cryptolocker"
+
+    assert malware['type'] == 'malware'
+    assert malware['id'].startswith('malware--')
+    assert malware['created'] is not None
+    assert malware['modified'] is not None
+    assert malware['labels'] == ['ransomware']
+    assert malware['name'] == "Cryptolocker"
+
+
+def test_malware_type_must_be_malware():
+    with pytest.raises(ValueError) as excinfo:
+        malware = stix2.Malware(type='xxx')
+
+    assert "Malware must have type='malware'." in str(excinfo)
+
+
+def test_malware_id_must_start_with_malware():
+    with pytest.raises(ValueError) as excinfo:
+        malware = stix2.Malware(id='my-prefix--')
+
+    assert "Malware id values must begin with 'malware--'." in str(excinfo)
+
+
+def test_malware_required_field_labels():
+    with pytest.raises(ValueError) as excinfo:
+        malware = stix2.Malware()
+    assert "Missing required field for Malware: 'labels'." in str(excinfo)
+
+
+def test_malware_required_field_name():
+    with pytest.raises(ValueError) as excinfo:
+        # Label is checked first, so make sure that is provided
+        malware = stix2.Malware(labels=['ransomware'])
+    assert "Missing required field for Malware: 'name'." in str(excinfo)
+
+
+def test_cannot_assign_to_attributes():
+    malware = stix2.Malware(**MALWARE_KWARGS)
+
+    with pytest.raises(ValueError) as excinfo:
+        malware.name = "Cryptolocker II"
+
+    assert "Cannot modify properties after creation." in str(excinfo)