creating native DataSourceError exception class to wrap taxii client and server errors

stix2/datastore/__init__.py

@@ -24,15 +24,35 @@ def make_id():
     return str(uuid.uuid4())
 
 
+class DataSourceError(Exception):
+    """General DataSource error instance, used primarily for wrapping
+    lower level errors
+
+    Args:
+        message (str): error message
+        root_exception (Exception): Exception instance of root exception
+            in the case that DataSourceError is wrapping a lower level or
+            other exception
+    """
+    def __init__(self, message, root_exception=None):
+        self.message = message
+        self.root_exception = root_exception
+
+    def __str__(self):
+        if self.root_exception:
+            self.message = "{} \"{}\"".format(self.message, self.root_exception)
+        return self.message
+
+
 class DataStoreMixin(object):
     """Provides mechanisms for storing and retrieving STIX data. The specific
     behavior can be customized by subclasses.
 
     Args:
         source (DataSource): An existing DataSource to use
-             as this DataStore's DataSource component
+            as this DataStore's DataSource component
         sink (DataSink): An existing DataSink to use
-             as this DataStore's DataSink component
+            as this DataStore's DataSink component
 
     Attributes:
         id (str): A unique UUIDv4 to identify this DataStore.
@@ -129,7 +149,7 @@ class DataStoreMixin(object):
             obj (STIX object OR dict OR str): The STIX object (or its ID) whose
                 relationships will be looked up.
             relationship_type (str): Only retrieve Relationships of this type.
-                If None, all relationships will be returned, regardless of type.
+            If None, all relationships will be returned, regardless of type.
             source_only (bool): Only retrieve Relationships for which this
                 object is the source_ref. Default: False.
             target_only (bool): Only retrieve Relationships for which this

stix2/datastore/taxii.py

@@ -5,7 +5,8 @@ from requests.exceptions import HTTPError
 
 from stix2.base import _STIXBase
 from stix2.core import Bundle, parse
-from stix2.datastore import DataSink, DataSource, DataStoreMixin
+from stix2.datastore import (DataSink, DataSource, DataSourceError,
+                             DataStoreMixin)
 from stix2.datastore.filters import Filter, FilterSet, apply_common_filters
 from stix2.utils import deduplicate
 
@@ -65,18 +66,12 @@ class TAXIICollectionSink(DataSink):
             if collection.can_write:
                 self.collection = collection
             else:
-                raise ValueError("The TAXII Collection object provided does not have write access"
+                raise DataSourceError("The TAXII Collection object provided does not have write access"
-                                 " to the underlying linked Collection resource")
+                                      " to the underlying linked Collection resource")
 
         except (HTTPError, ValidationError) as e:
-            added_context = ("The underlying TAXII Collection resource defined in the supplied TAXII"
+            raise DataSourceError("The underlying TAXII Collection resource defined in the supplied TAXII"
-                             " Collection object provided could not be reached.")
+                                  " Collection object provided could not be reached. Receved error:", e)
-            if not e.args:
-                e.args = (added_context,)
-            else:
-                e.args = (added_context,) + e.args
-
-            raise
 
         self.allow_custom = allow_custom
 
@@ -144,18 +139,12 @@ class TAXIICollectionSource(DataSource):
             if collection.can_read:
                 self.collection = collection
             else:
-                raise ValueError("The TAXII Collection object provided does not have read access"
+                raise DataSourceError("The TAXII Collection object provided does not have read access"
-                                 " to the underlying linked Collection resource")
+                                      " to the underlying linked Collection resource")
 
         except (HTTPError, ValidationError) as e:
-            added_context = ("The underlying TAXII Collection resource defined in the supplied TAXII"
+            raise DataSourceError("The underlying TAXII Collection resource defined in the supplied TAXII"
-                             " Collection object provided could not be reached.")
+                                  " Collection object provided could not be reached. Recieved error:", e)
-            if not e.args:
-                e.args = (added_context,)
-            else:
-                e.args = (added_context,) + e.args
-
-            raise
 
         self.allow_custom = allow_custom
 
@@ -190,12 +179,12 @@ class TAXIICollectionSource(DataSource):
             stix_objs = self.collection.get_object(stix_id)["objects"]
             stix_obj = list(apply_common_filters(stix_objs, query))
 
-        except HTTPError as err:
+        except HTTPError as e:
-            if err.response.status_code == 404:
+            if e.response.status_code == 404:
                 # if resource not found or access is denied from TAXII server, return None
                 stix_obj = []
             else:
-                raise
+                raise DataSourceError("TAXII Collection resource returned error", e)
 
         if len(stix_obj):
             stix_obj = parse(stix_obj[0], allow_custom=self.allow_custom, version=version)
@@ -286,15 +275,9 @@ class TAXIICollectionSource(DataSource):
         except HTTPError as e:
             # if resources not found or access is denied from TAXII server, return empty list
             if e.response.status_code == 404:
-                added_context = ("The requested STIX objects for the TAXII Collection resource defined in"
+                raise DataSourceError("The requested STIX objects for the TAXII Collection resource defined in"
-                                 " the supplied TAXII Collection object is either not found or access is"
+                                      " the supplied TAXII Collection object are either not found or access is"
-                                 " denied. Received error: ")
+                                      " denied. Received error: ", e)
-            if not e.args:
-                e.args = (added_context,)
-            else:
-                e.args = (added_context,) + e.args
-
-            raise
 
         # parse python STIX objects from the STIX object dicts
         stix_objs = [parse(stix_obj_dict, allow_custom=self.allow_custom, version=version) for stix_obj_dict in all_data]

stix2/test/test_datastore_taxii.py

@@ -2,12 +2,12 @@ import json
 
 from medallion.filters.basic_filter import BasicFilter
 import pytest
-from requests.exceptions import HTTPError
 from requests.models import Response
 from taxii2client import Collection, _filter_kwargs_to_query_params
 
 from stix2 import (Bundle, TAXIICollectionSink, TAXIICollectionSource,
                    TAXIICollectionStore, ThreatActor)
+from stix2.datastore import DataSourceError
 from stix2.datastore.filters import Filter
 
 COLLECTION_URL = 'https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/'
@@ -326,7 +326,7 @@ def test_get_all_versions(collection):
 def test_can_read_error(collection_no_rw_access):
     """create a TAXIICOllectionSource with a taxii2client.Collection
     instance that does not have read access, check ValueError exception is raised"""
-    with pytest.raises(ValueError) as excinfo:
+    with pytest.raises(DataSourceError) as excinfo:
         TAXIICollectionSource(collection_no_rw_access)
     assert "Collection object provided does not have read access" in str(excinfo.value)
 
@@ -334,7 +334,7 @@ def test_can_read_error(collection_no_rw_access):
 def test_can_write_error(collection_no_rw_access):
     """create a TAXIICOllectionSink with a taxii2client.Collection
     instance that does not have write access, check ValueError exception is raised"""
-    with pytest.raises(ValueError) as excinfo:
+    with pytest.raises(DataSourceError) as excinfo:
         TAXIICollectionSink(collection_no_rw_access)
     assert "Collection object provided does not have write access" in str(excinfo.value)
 
@@ -359,9 +359,9 @@ def test_all_versions_404(collection):
     """ a TAXIICollectionSource.all_version() call that recieves an HTTP 404
     response code from the taxii2client should be returned as an exception"""
     ds = TAXIICollectionStore(collection)
-    with pytest.raises(HTTPError) as excinfo:
+    with pytest.raises(DataSourceError) as excinfo:
         ds.all_versions("indicator--1")
-    assert "is either not found or access is denied" in str(excinfo.value)
+    assert "are either not found or access is denied" in str(excinfo.value)
     assert "404" in str(excinfo.value)
 
 
@@ -371,7 +371,7 @@ def test_query_404(collection):
     ds = TAXIICollectionStore(collection)
     query = [Filter("type", "=", "malware")]
 
-    with pytest.raises(HTTPError) as excinfo:
+    with pytest.raises(DataSourceError) as excinfo:
         ds.query(query=query)
-    assert "is either not found or access is denied" in str(excinfo.value)
+    assert "are either not found or access is denied" in str(excinfo.value)
     assert "404" in str(excinfo.value)