From 1a75d830bb63d3c8c88caea7c1c92d1fdffb13d1 Mon Sep 17 00:00:00 2001
From: clenk <clenk@mitre.org>
Date: Wed, 3 May 2017 18:19:30 -0400
Subject: [PATCH] Add Autonomous System

---
 stix2/__init__.py                |  3 ++-
 stix2/observables.py             | 12 +++++++++++-
 stix2/properties.py              | 11 +++++++++++
 stix2/test/test_observed_data.py | 18 ++++++++++++++++++
 4 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/stix2/__init__.py b/stix2/__init__.py
index 1ec1711..763c4d4 100644
--- a/stix2/__init__.py
+++ b/stix2/__init__.py
@@ -3,7 +3,7 @@
 # flake8: noqa
 
 from .bundle import Bundle
-from .observables import Artifact, File
+from .observables import Artifact, AutonomousSystem, File
 from .other import ExternalReference, KillChainPhase, MarkingDefinition, \
     GranularMarking, StatementMarking, TLPMarking
 from .sdo import AttackPattern, Campaign, CourseOfAction, Identity, Indicator, \
@@ -34,6 +34,7 @@ OBJ_MAP = {
 
 OBJ_MAP_OBSERVABLE = {
     'artifact': Artifact,
+    'autonomous-system': AutonomousSystem,
     'file': File,
 }
 
diff --git a/stix2/observables.py b/stix2/observables.py
index 33eded2..4e72cae 100644
--- a/stix2/observables.py
+++ b/stix2/observables.py
@@ -5,7 +5,7 @@ from .base import Observable
 #                          HashesProperty, HexProperty, IDProperty,
 #                          IntegerProperty, ListProperty, ReferenceProperty,
 #                          StringProperty, TimestampProperty, TypeProperty)
-from .properties import BinaryProperty, HashesProperty, StringProperty, TypeProperty
+from .properties import BinaryProperty, HashesProperty, IntegerProperty, StringProperty, TypeProperty
 
 
 class Artifact(Observable):
@@ -19,6 +19,16 @@ class Artifact(Observable):
     }
 
 
+class AutonomousSystem(Observable):
+    _type = 'autonomous-system'
+    _properties = {
+        'type': TypeProperty(_type),
+        'number': IntegerProperty(),
+        'name': StringProperty(),
+        'rir': StringProperty(),
+    }
+
+
 class File(Observable):
     _type = 'file'
     _properties = {
diff --git a/stix2/properties.py b/stix2/properties.py
index 4b4da7a..5aae467 100644
--- a/stix2/properties.py
+++ b/stix2/properties.py
@@ -335,3 +335,14 @@ class SelectorProperty(Property):
         if not SELECTOR_REGEX.match(value):
             raise ValueError("must adhere to selector syntax.")
         return value
+
+
+class ObjectReferenceProperty(Property):
+    def _init(self, valid_refs=None):
+        self.valid_refs = valid_refs
+        super(ObjectReferenceProperty, self).__init__()
+
+    def clean(self, value):
+        if value not in self.valid_refs:
+            raise ValueError("must refer to observable objects in the same "
+                             "Observable Objects container.")
diff --git a/stix2/test/test_observed_data.py b/stix2/test/test_observed_data.py
index 154b395..5472597 100644
--- a/stix2/test/test_observed_data.py
+++ b/stix2/test/test_observed_data.py
@@ -115,4 +115,22 @@ def test_parse_artifact_invalid(data):
     with pytest.raises(ValueError):
         stix2.parse(odata_str)
 
+
+@pytest.mark.parametrize("data", [
+    """"0": {
+        "type": "autonomous-system",
+        "number": 15139,
+        "name": "Slime Industries",
+        "rir": "ARIN"
+    }""",
+])
+def test_parse_autonomous_system_valid(data):
+    odata_str = re.compile('"objects".+\},', re.DOTALL).sub('"objects": { %s },' % data, EXPECTED)
+    odata = stix2.parse(odata_str)
+    assert odata.objects["0"].type == "autonomous-system"
+    assert odata.objects["0"].number == 15139
+    assert odata.objects["0"].name == "Slime Industries"
+    assert odata.objects["0"].rir == "ARIN"
+
+
 # TODO: Add other examples