From 6761d1fdfce1e1792e22d26b69b869dca9b531b0 Mon Sep 17 00:00:00 2001
From: Greg Back <gback@mitre.org>
Date: Tue, 17 Jan 2017 13:58:19 -0800
Subject: [PATCH] Add required fields to Indicator.

---
 setup.py                 |  5 ++++
 stix2/__init__.py        | 51 ++++++++++++++++++++++++++++++++++++++--
 stix2/test/test_stix2.py | 48 +++++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index f72753b..22ee0ac 100644
--- a/setup.py
+++ b/setup.py
@@ -1,10 +1,15 @@
 #!/usr/bin/env python
 from setuptools import setup, find_packages
 
+install_requires = [
+    'pytz',
+]
+
 setup(
     name='stix2',
     description="Produce and consume STIX 2 JSON content",
     version='0.0.1',
     packages=find_packages(),
+    install_requires=install_requires,
     keywords="stix stix2 json cti cyber threat intelligence",
 )
diff --git a/stix2/__init__.py b/stix2/__init__.py
index 74aeca4..678374a 100644
--- a/stix2/__init__.py
+++ b/stix2/__init__.py
@@ -1,7 +1,54 @@
+from datetime import datetime
+import json
 import uuid
 
+import pytz
+
+
+def format_datetime(dt):
+    # 1. Convert to UTC
+    # 2. Format in isoformat
+    # 3. Strip off "+00:00"
+    # 4. Add "Z"
+    return dt.astimezone(pytz.utc).isoformat()[:-6] + "Z"
+
+# REQUIRED (all):
+# - type
+# - id
+# - created
+# - modified
+
 
 class Indicator:
+    # REQUIRED (Indicator):
+    # - type
+    # - labels
+    # - pattern
+    # - valid_from
+    required = ['']
 
-    def __init__(self):
-        self.id = "indicator--" + str(uuid.uuid4())
+    def __init__(self, type='indicator', id=None, created=None, modified=None,
+                 labels=None, pattern=None, valid_from=None):
+        now = datetime.now(tz=pytz.UTC)
+
+        self.type = type
+        if not id:
+            id = "indicator--" + str(uuid.uuid4())
+        self.id = id
+        self.created = created or now
+        self.modified = modified or now
+        self.labels = labels
+        self.pattern = pattern
+        self.valid_from = valid_from or now
+
+    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,
+            'pattern': self.pattern,
+            'valid_from': format_datetime(self.valid_from),
+        }, indent=4, sort_keys=True, separators=(",", ": "))  # Don't include spaces after commas.
diff --git a/stix2/test/test_stix2.py b/stix2/test/test_stix2.py
index 0ca98f9..82b7e42 100644
--- a/stix2/test/test_stix2.py
+++ b/stix2/test/test_stix2.py
@@ -1,8 +1,56 @@
 """Tests for the stix2 library"""
 
+import datetime
+
+import pytest
+import pytz
+
 import stix2
 
+amsterdam = pytz.timezone('Europe/Amsterdam')
+eastern = pytz.timezone('US/Eastern')
+
+
+@pytest.mark.parametrize('dt,timestamp', [
+    (datetime.datetime(2017, 1, 1, tzinfo=pytz.utc), '2017-01-01T00:00:00Z'),
+    (amsterdam.localize(datetime.datetime(2017, 1, 1)), '2016-12-31T23:00:00Z'),
+    (eastern.localize(datetime.datetime(2017, 1, 1, 12, 34, 56)), '2017-01-01T17:34:56Z'),
+    (eastern.localize(datetime.datetime(2017, 7, 1)), '2017-07-01T04:00:00Z'),
+])
+def test_timestamp_formatting(dt, timestamp):
+    assert stix2.format_datetime(dt) == timestamp
+
 
 def test_basic_indicator():
     indicator = stix2.Indicator()
     assert indicator.id.startswith("indicator")
+
+
+EXPECTED = """{
+    "created": "2017-01-01T00:00:00Z",
+    "id": "indicator--01234567-89ab-cdef-0123-456789abcdef",
+    "labels": [
+        "malicious-activity"
+    ],
+    "modified": "2017-01-01T00:00:00Z",
+    "pattern": "[file:hashes.MD5 = 'd41d8cd98f00b204e9800998ecf8427e']",
+    "type": "indicator",
+    "valid_from": "1970-01-01T00:00:00Z"
+}"""
+
+
+def test_indicator_with_all_required_fields():
+    now = datetime.datetime(2017, 1, 1, 0, 0, 0, tzinfo=pytz.utc)
+    epoch = datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=pytz.utc)
+
+    indicator = stix2.Indicator(
+        type="indicator",
+        id="indicator--01234567-89ab-cdef-0123-456789abcdef",
+        created=now,
+        modified=now,
+        labels=['malicious-activity'],
+        pattern="[file:hashes.MD5 = 'd41d8cd98f00b204e9800998ecf8427e']",
+        valid_from=epoch,
+    )
+
+    assert str(indicator) == EXPECTED