Update package structure

stix2/__init__.py

@@ -19,27 +19,11 @@
 
 # flake8: noqa
 
-from . import exceptions
+from . import exceptions, v20
-from .common import (TLP_AMBER, TLP_GREEN, TLP_RED, TLP_WHITE, CustomMarking,
+from .core import Bundle, _collect_stix2_obj_maps, _register_type, parse
-                     ExternalReference, GranularMarking, KillChainPhase,
-                     MarkingDefinition, StatementMarking, TLPMarking)
-from .core import Bundle, _register_type, parse
 from .environment import Environment, ObjectFactory
 from .markings import (add_markings, clear_markings, get_markings, is_marked,
                        remove_markings, set_markings)
-from .observables import (URL, AlternateDataStream, ArchiveExt, Artifact,
-                          AutonomousSystem, CustomExtension, CustomObservable,
-                          Directory, DomainName, EmailAddress, EmailMessage,
-                          EmailMIMEComponent, File, HTTPRequestExt, ICMPExt,
-                          IPv4Address, IPv6Address, MACAddress, Mutex,
-                          NetworkTraffic, NTFSExt, PDFExt, Process,
-                          RasterImageExt, SocketExt, Software, TCPExt,
-                          UNIXAccountExt, UserAccount, WindowsPEBinaryExt,
-                          WindowsPEOptionalHeaderType, WindowsPESection,
-                          WindowsProcessExt, WindowsRegistryKey,
-                          WindowsRegistryValueType, WindowsServiceExt,
-                          X509Certificate, X509V3ExtenstionsType,
-                          parse_observable)
 from .patterns import (AndBooleanExpression, AndObservationExpression,
                        BasicObjectPathComponent, EqualityComparisonExpression,
                        FloatConstant, FollowedByObservationExpression,
@@ -58,9 +42,6 @@ from .patterns import (AndBooleanExpression, AndObservationExpression,
                        ReferenceObjectPathComponent, RepeatQualifier,
                        StartStopQualifier, StringConstant, TimestampConstant,
                        WithinQualifier)
-from .sdo import (AttackPattern, Campaign, CourseOfAction, CustomObject,
-                  Identity, Indicator, IntrusionSet, Malware, ObservedData,
-                  Report, ThreatActor, Tool, Vulnerability)
 from .sources import CompositeDataSource
 from .sources.filesystem import (FileSystemSink, FileSystemSource,
                                  FileSystemStore)
@@ -68,6 +49,8 @@ from .sources.filters import Filter
 from .sources.memory import MemorySink, MemorySource, MemoryStore
 from .sources.taxii import (TAXIICollectionSink, TAXIICollectionSource,
                             TAXIICollectionStore)
-from .sro import Relationship, Sighting
 from .utils import get_dict, new_version, revoke
+from .v20 import *  # This import should always be the latest STIX 2.X version
 from .version import __version__
+
+_collect_stix2_obj_maps()

stix2/core.py

@@ -1,15 +1,12 @@
 """STIX 2.0 Objects that are neither SDOs nor SROs."""
 
 from collections import OrderedDict
+import importlib
+import pkgutil
 
 from . import exceptions
 from .base import _STIXBase
-from .common import MarkingDefinition
 from .properties import IDProperty, ListProperty, Property, TypeProperty
-from .sdo import (AttackPattern, Campaign, CourseOfAction, Identity, Indicator,
-                  IntrusionSet, Malware, ObservedData, Report, ThreatActor,
-                  Tool, Vulnerability)
-from .sro import Relationship, Sighting
 from .utils import get_dict
 
 
@@ -62,37 +59,30 @@ class Bundle(_STIXBase):
         super(Bundle, self).__init__(**kwargs)
 
 
-OBJ_MAP = {
+STIX2_OBJ_MAPS = {}
-    'attack-pattern': AttackPattern,
-    'bundle': Bundle,
-    'campaign': Campaign,
-    'course-of-action': CourseOfAction,
-    'identity': Identity,
-    'indicator': Indicator,
-    'intrusion-set': IntrusionSet,
-    'malware': Malware,
-    'marking-definition': MarkingDefinition,
-    'observed-data': ObservedData,
-    'report': Report,
-    'relationship': Relationship,
-    'threat-actor': ThreatActor,
-    'tool': Tool,
-    'sighting': Sighting,
-    'vulnerability': Vulnerability,
-}
 
 
-def parse(data, allow_custom=False):
+def parse(data, allow_custom=False, version=None):
     """Deserialize a string or file-like object into a STIX object.
 
     Args:
         data (str, dict, file-like object): The STIX 2 content to be parsed.
-        allow_custom (bool): Whether to allow custom properties or not. Default: False.
+        allow_custom (bool): Whether to allow custom properties or not.
+            Default: False.
+        version (str): Which STIX2 version to use. (e.g. "2.0", "2.1"). If
+            None, use latest version.
 
     Returns:
         An instantiated Python STIX object.
 
     """
+    if not version:
+        # Use latest version
+        OBJ_MAP = STIX2_OBJ_MAPS[sorted(STIX2_OBJ_MAPS.keys())[-1]]
+    else:
+        v = 'v' + version.replace('.', '')
+        OBJ_MAP = STIX2_OBJ_MAPS[v]
+
     obj = get_dict(data)
 
     if 'type' not in obj:
@@ -105,8 +95,34 @@ def parse(data, allow_custom=False):
     return obj_class(allow_custom=allow_custom, **obj)
 
 
-def _register_type(new_type):
+def _register_type(new_type, version=None):
     """Register a custom STIX Object type.
 
+    Args:
+        new_type (class): A class to register in the Object map.
+        version (str): Which STIX2 version to use. (e.g. "2.0", "2.1"). If
+            None, use latest version.
     """
+    if not version:
+        # Use latest version
+        OBJ_MAP = STIX2_OBJ_MAPS[sorted(STIX2_OBJ_MAPS.keys())[-1]]
+    else:
+        v = 'v' + version.replace('.', '')
+        OBJ_MAP = STIX2_OBJ_MAPS[v]
+
     OBJ_MAP[new_type._type] = new_type
+
+
+def _collect_stix2_obj_maps():
+    """Navigate the package once and retrieve all OBJ_MAP dicts for each v2X
+    package."""
+    if not STIX2_OBJ_MAPS:
+        top_level_module = importlib.import_module('stix2')
+        path = top_level_module.__path__
+        prefix = str(top_level_module.__name__) + '.'
+
+        for module_loader, name, is_pkg in pkgutil.walk_packages(path=path,
+                                                                 prefix=prefix):
+            if name.startswith('stix2.v2') and is_pkg:
+                mod = importlib.import_module(name, top_level_module)
+                STIX2_OBJ_MAPS[name.split('.')[-1]] = mod.OBJ_MAP

stix2/test/test_bundle.py

@@ -132,8 +132,9 @@ def test_create_bundle_invalid(indicator, malware, relationship):
     assert excinfo.value.reason == 'This property may not contain a Bundle object'
 
 
-def test_parse_bundle():
+@pytest.mark.parametrize("version", ["2.0"])
-    bundle = stix2.parse(EXPECTED_BUNDLE)
+def test_parse_bundle(version):
+    bundle = stix2.parse(EXPECTED_BUNDLE, version=version)
 
     assert bundle.type == "bundle"
     assert bundle.id.startswith("bundle--")

stix2/test/test_properties.py

@@ -1,8 +1,7 @@
 import pytest
 
-from stix2 import TCPExt
+from stix2 import EmailMIMEComponent, ExtensionsProperty, TCPExt
 from stix2.exceptions import AtLeastOnePropertyError, DictionaryKeyError
-from stix2.observables import EmailMIMEComponent, ExtensionsProperty
 from stix2.properties import (BinaryProperty, BooleanProperty,
                               DictionaryProperty, EmbeddedObjectProperty,
                               EnumProperty, FloatProperty, HashesProperty,

stix2/v20/__init__.py

@@ -0,0 +1,43 @@
+
+# flake8: noqa
+
+from ..core import Bundle
+from .common import (TLP_AMBER, TLP_GREEN, TLP_RED, TLP_WHITE, CustomMarking,
+                     ExternalReference, GranularMarking, KillChainPhase,
+                     MarkingDefinition, StatementMarking, TLPMarking)
+from .observables import (URL, AlternateDataStream, ArchiveExt, Artifact,
+                          AutonomousSystem, CustomExtension, CustomObservable,
+                          Directory, DomainName, EmailAddress, EmailMessage,
+                          EmailMIMEComponent, ExtensionsProperty, File,
+                          HTTPRequestExt, ICMPExt, IPv4Address, IPv6Address,
+                          MACAddress, Mutex, NetworkTraffic, NTFSExt, PDFExt,
+                          Process, RasterImageExt, SocketExt, Software, TCPExt,
+                          UNIXAccountExt, UserAccount, WindowsPEBinaryExt,
+                          WindowsPEOptionalHeaderType, WindowsPESection,
+                          WindowsProcessExt, WindowsRegistryKey,
+                          WindowsRegistryValueType, WindowsServiceExt,
+                          X509Certificate, X509V3ExtenstionsType,
+                          parse_observable)
+from .sdo import (AttackPattern, Campaign, CourseOfAction, Identity, Indicator,
+                  IntrusionSet, Malware, ObservedData, Report, ThreatActor,
+                  Tool, Vulnerability)
+from .sro import Relationship, Sighting
+
+OBJ_MAP = {
+    'attack-pattern': AttackPattern,
+    'bundle': Bundle,
+    'campaign': Campaign,
+    'course-of-action': CourseOfAction,
+    'identity': Identity,
+    'indicator': Indicator,
+    'intrusion-set': IntrusionSet,
+    'malware': Malware,
+    'marking-definition': MarkingDefinition,
+    'observed-data': ObservedData,
+    'report': Report,
+    'relationship': Relationship,
+    'threat-actor': ThreatActor,
+    'tool': Tool,
+    'sighting': Sighting,
+    'vulnerability': Vulnerability,
+}

stix2/v20/common.py

@@ -2,12 +2,12 @@
 
 from collections import OrderedDict
 
-from .base import _STIXBase
+from ..base import _STIXBase
-from .markings import _MarkingsMixin
+from ..markings import _MarkingsMixin
-from .properties import (HashesProperty, IDProperty, ListProperty, Property,
+from ..properties import (HashesProperty, IDProperty, ListProperty, Property,
                           ReferenceProperty, SelectorProperty, StringProperty,
                           TimestampProperty, TypeProperty)
-from .utils import NOW, get_dict
+from ..utils import NOW, get_dict
 
 
 class ExternalReference(_STIXBase):

stix2/v20/observables.py

@@ -7,15 +7,15 @@ Observable and do not have a ``_type`` attribute.
 
 from collections import OrderedDict
 
-from .base import _Extension, _Observable, _STIXBase
+from ..base import _Extension, _Observable, _STIXBase
-from .exceptions import (AtLeastOnePropertyError, DependentPropertiesError,
+from ..exceptions import (AtLeastOnePropertyError, DependentPropertiesError,
                           ParseError)
-from .properties import (BinaryProperty, BooleanProperty, DictionaryProperty,
+from ..properties import (BinaryProperty, BooleanProperty, DictionaryProperty,
                           EmbeddedObjectProperty, EnumProperty, FloatProperty,
                           HashesProperty, HexProperty, IntegerProperty,
                           ListProperty, ObjectReferenceProperty, Property,
                           StringProperty, TimestampProperty, TypeProperty)
-from .utils import get_dict
+from ..utils import get_dict
 
 
 class ObservableProperty(Property):

stix2/v20/sdo.py

@@ -4,14 +4,14 @@ from collections import OrderedDict
 
 import stix2
 
-from .base import _STIXBase
+from ..base import _STIXBase
-from .common import ExternalReference, GranularMarking, KillChainPhase
+from ..markings import _MarkingsMixin
-from .markings import _MarkingsMixin
+from ..properties import (BooleanProperty, IDProperty, IntegerProperty,
-from .observables import ObservableProperty
-from .properties import (BooleanProperty, IDProperty, IntegerProperty,
                           ListProperty, PatternProperty, ReferenceProperty,
                           StringProperty, TimestampProperty, TypeProperty)
-from .utils import NOW
+from ..utils import NOW
+from .common import ExternalReference, GranularMarking, KillChainPhase
+from .observables import ObservableProperty
 
 
 class STIXDomainObject(_STIXBase, _MarkingsMixin):
@@ -358,7 +358,7 @@ def CustomObject(type='x-custom-type', properties=None):
                         return
                     raise e
 
-        stix2._register_type(_Custom)
+        stix2._register_type(_Custom, version="2.0")
         return _Custom
 
     return custom_builder

stix2/v20/sro.py

@@ -2,13 +2,13 @@
 
 from collections import OrderedDict
 
-from .base import _STIXBase
+from ..base import _STIXBase
-from .common import ExternalReference, GranularMarking
+from ..markings import _MarkingsMixin
-from .markings import _MarkingsMixin
+from ..properties import (BooleanProperty, IDProperty, IntegerProperty,
-from .properties import (BooleanProperty, IDProperty, IntegerProperty,
                           ListProperty, ReferenceProperty, StringProperty,
                           TimestampProperty, TypeProperty)
-from .utils import NOW
+from ..utils import NOW
+from .common import ExternalReference, GranularMarking
 
 
 class STIXRelationshipObject(_STIXBase, _MarkingsMixin):