From 22c435168883d4cd6a2b37beb10a1fe5c05f4768 Mon Sep 17 00:00:00 2001
From: Emmanuelle Vargas-Gonzalez <emmanuelle@mitre.org>
Date: Tue, 10 Nov 2020 17:08:51 -0500
Subject: [PATCH] flatten the extensions map, remove enclosing_type from
 ExtensionsProperty

other minor changes
---
 stix2/properties.py   | 16 +++++++++-------
 stix2/v20/__init__.py | 32 ++++++++++++--------------------
 stix2/v21/__init__.py | 33 ++++++++++++---------------------
 stix2/workbench.py    |  2 +-
 4 files changed, 34 insertions(+), 49 deletions(-)

diff --git a/stix2/properties.py b/stix2/properties.py
index 1ca2dbe..cea664f 100644
--- a/stix2/properties.py
+++ b/stix2/properties.py
@@ -24,8 +24,8 @@ try:
 except ImportError:
     from collections import Mapping
 
-TYPE_REGEX = re.compile(r'^\-?[a-z0-9]+(-[a-z0-9]+)*\-?$')
-TYPE_21_REGEX = re.compile(r'^([a-z][a-z0-9]*)+(-[a-z0-9]+)*\-?$')
+TYPE_REGEX = re.compile(r'^-?[a-z0-9]+(-[a-z0-9]+)*-?$')
+TYPE_21_REGEX = re.compile(r'^([a-z][a-z0-9]*)+([a-z0-9-]+)*-?$')
 ERROR_INVALID_ID = (
     "not a valid STIX identifier, must match <object-type>--<UUID>: {}"
 )
@@ -638,9 +638,8 @@ class ExtensionsProperty(DictionaryProperty):
     """Property for representing extensions on Observable objects.
     """
 
-    def __init__(self, spec_version=stix2.DEFAULT_VERSION, allow_custom=False, enclosing_type=None, required=False):
+    def __init__(self, spec_version=stix2.DEFAULT_VERSION, allow_custom=False, required=False):
         self.allow_custom = allow_custom
-        self.enclosing_type = enclosing_type
         super(ExtensionsProperty, self).__init__(spec_version=spec_version, required=required)
 
     def clean(self, value):
@@ -655,10 +654,10 @@ class ExtensionsProperty(DictionaryProperty):
 
         v = 'v' + self.spec_version.replace('.', '')
 
-        specific_type_map = STIX2_OBJ_MAPS[v]['observable-extensions'].get(self.enclosing_type, {})
+        extension_type_map = STIX2_OBJ_MAPS[v].get('extensions', {})
         for key, subvalue in dictified.items():
-            if key in specific_type_map:
-                cls = specific_type_map[key]
+            if key in extension_type_map:
+                cls = extension_type_map[key]
                 if type(subvalue) is dict:
                     if self.allow_custom:
                         subvalue['allow_custom'] = True
@@ -673,6 +672,9 @@ class ExtensionsProperty(DictionaryProperty):
             else:
                 if self.allow_custom:
                     dictified[key] = subvalue
+                elif key.startswith('stix-extension--'):
+                    _validate_id(key, '2.1', 'stix-extension')
+                    dictified[key] = subvalue
                 else:
                     raise CustomContentError("Can't parse unknown extension type: {}".format(key))
         return dictified
diff --git a/stix2/v20/__init__.py b/stix2/v20/__init__.py
index 002343b..85609e7 100644
--- a/stix2/v20/__init__.py
+++ b/stix2/v20/__init__.py
@@ -82,26 +82,18 @@ OBJ_MAP_OBSERVABLE = {
 }
 
 EXT_MAP = {
-    'file': {
-        'archive-ext': ArchiveExt,
-        'ntfs-ext': NTFSExt,
-        'pdf-ext': PDFExt,
-        'raster-image-ext': RasterImageExt,
-        'windows-pebinary-ext': WindowsPEBinaryExt,
-    },
-    'network-traffic': {
-        'http-request-ext': HTTPRequestExt,
-        'icmp-ext': ICMPExt,
-        'socket-ext': SocketExt,
-        'tcp-ext': TCPExt,
-    },
-    'process': {
-        'windows-process-ext': WindowsProcessExt,
-        'windows-service-ext': WindowsServiceExt,
-    },
-    'user-account': {
-        'unix-account-ext': UNIXAccountExt,
-    },
+    'archive-ext': ArchiveExt,
+    'ntfs-ext': NTFSExt,
+    'pdf-ext': PDFExt,
+    'raster-image-ext': RasterImageExt,
+    'windows-pebinary-ext': WindowsPEBinaryExt,
+    'http-request-ext': HTTPRequestExt,
+    'icmp-ext': ICMPExt,
+    'socket-ext': SocketExt,
+    'tcp-ext': TCPExt,
+    'windows-process-ext': WindowsProcessExt,
+    'windows-service-ext': WindowsServiceExt,
+    'unix-account-ext': UNIXAccountExt,
 }
 
 
diff --git a/stix2/v21/__init__.py b/stix2/v21/__init__.py
index 856ff13..81ab980 100644
--- a/stix2/v21/__init__.py
+++ b/stix2/v21/__init__.py
@@ -91,29 +91,20 @@ OBJ_MAP_OBSERVABLE = {
 }
 
 EXT_MAP = {
-    'file': {
-        'archive-ext': ArchiveExt,
-        'ntfs-ext': NTFSExt,
-        'pdf-ext': PDFExt,
-        'raster-image-ext': RasterImageExt,
-        'windows-pebinary-ext': WindowsPEBinaryExt,
-    },
-    'network-traffic': {
-        'http-request-ext': HTTPRequestExt,
-        'icmp-ext': ICMPExt,
-        'socket-ext': SocketExt,
-        'tcp-ext': TCPExt,
-    },
-    'process': {
-        'windows-process-ext': WindowsProcessExt,
-        'windows-service-ext': WindowsServiceExt,
-    },
-    'user-account': {
-        'unix-account-ext': UNIXAccountExt,
-    },
+    'archive-ext': ArchiveExt,
+    'ntfs-ext': NTFSExt,
+    'pdf-ext': PDFExt,
+    'raster-image-ext': RasterImageExt,
+    'windows-pebinary-ext': WindowsPEBinaryExt,
+    'http-request-ext': HTTPRequestExt,
+    'icmp-ext': ICMPExt,
+    'socket-ext': SocketExt,
+    'tcp-ext': TCPExt,
+    'windows-process-ext': WindowsProcessExt,
+    'windows-service-ext': WindowsServiceExt,
+    'unix-account-ext': UNIXAccountExt,
 }
 
-
 # Ensure star-imports from this module get the right symbols.  "base" is a
 # known problem, since there are multiple modules with that name and one can
 # accidentally overwrite another.
diff --git a/stix2/workbench.py b/stix2/workbench.py
index 3724bdb..ff5d680 100644
--- a/stix2/workbench.py
+++ b/stix2/workbench.py
@@ -55,7 +55,7 @@ from . import (  # noqa: F401
     WindowsPEBinaryExt, WindowsPEOptionalHeaderType,
     WindowsPESection, WindowsProcessExt, WindowsRegistryKey,
     WindowsRegistryValueType, WindowsServiceExt, X509Certificate,
-    X509V3ExtenstionsType
+    X509V3ExtensionsType
 )
 from .datastore.filters import FilterSet