From a3313bc08fb754877bbcab53e9193c835a439e34 Mon Sep 17 00:00:00 2001 From: = Date: Wed, 16 May 2018 13:23:50 -0400 Subject: [PATCH] creating native DataSourceError exception class to wrap taxii client and server errors --- stix2/datastore/__init__.py | 26 ++++++++++++++-- stix2/datastore/taxii.py | 49 ++++++++++-------------------- stix2/test/test_datastore_taxii.py | 14 ++++----- 3 files changed, 46 insertions(+), 43 deletions(-) diff --git a/stix2/datastore/__init__.py b/stix2/datastore/__init__.py index 7fdf515..ac03959 100644 --- a/stix2/datastore/__init__.py +++ b/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 diff --git a/stix2/datastore/taxii.py b/stix2/datastore/taxii.py index 4a047a8..88897a9 100644 --- a/stix2/datastore/taxii.py +++ b/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" - " to the underlying linked Collection resource") + raise DataSourceError("The TAXII Collection object provided does not have write access" + " to the underlying linked Collection resource") except (HTTPError, ValidationError) as e: - added_context = ("The underlying TAXII Collection resource defined in the supplied TAXII" - " Collection object provided could not be reached.") - if not e.args: - e.args = (added_context,) - else: - e.args = (added_context,) + e.args - - raise + raise DataSourceError("The underlying TAXII Collection resource defined in the supplied TAXII" + " Collection object provided could not be reached. Receved error:", e) 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" - " to the underlying linked Collection resource") + raise DataSourceError("The TAXII Collection object provided does not have read access" + " to the underlying linked Collection resource") except (HTTPError, ValidationError) as e: - added_context = ("The underlying TAXII Collection resource defined in the supplied TAXII" - " Collection object provided could not be reached.") - if not e.args: - e.args = (added_context,) - else: - e.args = (added_context,) + e.args - - raise + raise DataSourceError("The underlying TAXII Collection resource defined in the supplied TAXII" + " Collection object provided could not be reached. Recieved error:", e) 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: - if err.response.status_code == 404: + except HTTPError as e: + 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" - " the supplied TAXII Collection object is either not found or access is" - " denied. Received error: ") - if not e.args: - e.args = (added_context,) - else: - e.args = (added_context,) + e.args - - raise + raise DataSourceError("The requested STIX objects for the TAXII Collection resource defined in" + " the supplied TAXII Collection object are either not found or access is" + " denied. Received error: ", e) # 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] diff --git a/stix2/test/test_datastore_taxii.py b/stix2/test/test_datastore_taxii.py index a153ce8..78a3284 100644 --- a/stix2/test/test_datastore_taxii.py +++ b/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)