diff --git a/stix2/datastore/__init__.py b/stix2/datastore/__init__.py index 7fdf515..c2963e2 100644 --- a/stix2/datastore/__init__.py +++ b/stix2/datastore/__init__.py @@ -24,15 +24,36 @@ 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: + return "{} \"{}\"".format(self.message, self.root_exception) + else: + 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. diff --git a/stix2/datastore/taxii.py b/stix2/datastore/taxii.py index 872a510..c815e12 100644 --- a/stix2/datastore/taxii.py +++ b/stix2/datastore/taxii.py @@ -5,10 +5,18 @@ 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 +try: + from taxii2client import ValidationError + _taxii2_client = True +except ImportError: + _taxii2_client = False + + TAXII_FILTERS = ['added_after', 'id', 'type', 'version'] @@ -51,7 +59,20 @@ class TAXIICollectionSink(DataSink): """ def __init__(self, collection, allow_custom=False): super(TAXIICollectionSink, self).__init__() - self.collection = collection + if not _taxii2_client: + raise ImportError("taxii2client library is required for usage of TAXIICollectionSink") + + try: + if collection.can_write: + self.collection = collection + else: + raise DataSourceError("The TAXII Collection object provided does not have write access" + " to the underlying linked Collection resource") + + except (HTTPError, ValidationError) as e: + 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 def add(self, stix_data, version=None): @@ -111,7 +132,20 @@ class TAXIICollectionSource(DataSource): """ def __init__(self, collection, allow_custom=True): super(TAXIICollectionSource, self).__init__() - self.collection = collection + if not _taxii2_client: + raise ImportError("taxii2client library is required for usage of TAXIICollectionSource") + + try: + if collection.can_read: + self.collection = collection + else: + raise DataSourceError("The TAXII Collection object provided does not have read access" + " to the underlying linked Collection resource") + + except (HTTPError, ValidationError) as e: + 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 def get(self, stix_id, version=None, _composite_filters=None): @@ -145,9 +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: - # if resource not found or access is denied from TAXII server, return None - stix_obj = [] + 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 DataSourceError("TAXII Collection resource returned error", e) if len(stix_obj): stix_obj = parse(stix_obj[0], allow_custom=self.allow_custom, version=version) @@ -235,9 +272,12 @@ class TAXIICollectionSource(DataSource): query.remove(taxii_filters) all_data = list(apply_common_filters(all_data, query)) - except HTTPError: + except HTTPError as e: # if resources not found or access is denied from TAXII server, return empty list - all_data = [] + if e.response.status_code == 404: + 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 098f944..8675378 100644 --- a/stix2/test/test_datastore_taxii.py +++ b/stix2/test/test_datastore_taxii.py @@ -2,10 +2,12 @@ import json from medallion.filters.basic_filter import BasicFilter import pytest +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/' @@ -36,7 +38,12 @@ class MockTAXIICollectionEndpoint(Collection): ("id", "type", "version"), [] ) - return Bundle(objects=objs) + if objs: + return Bundle(objects=objs) + else: + resp = Response() + resp.status_code = 404 + resp.raise_for_status() def get_object(self, id, version=None): self._verify_can_read() @@ -51,7 +58,12 @@ class MockTAXIICollectionEndpoint(Collection): ("version",), [] ) - return Bundle(objects=objs) + if objs: + return Bundle(objects=objs) + else: + resp = Response() + resp.status_code = 404 + resp.raise_for_status() @pytest.fixture @@ -71,6 +83,23 @@ def collection(stix_objs1): return mock +@pytest.fixture +def collection_no_rw_access(stix_objs1): + mock = MockTAXIICollectionEndpoint(COLLECTION_URL, **{ + "id": "91a7b528-80eb-42ed-a74d-c6fbd5a26116", + "title": "Not writeable or readable Collection", + "description": "This collection is a dropbox for submitting indicators", + "can_read": False, + "can_write": False, + "media_types": [ + "application/vnd.oasis.stix+json; version=2.0" + ] + }) + + mock.objects.extend(stix_objs1) + return mock + + def test_ds_taxii(collection): ds = TAXIICollectionSource(collection) assert ds.collection is not None @@ -292,3 +321,71 @@ def test_get_all_versions(collection): indicators = ds.all_versions('indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f') # There are 3 indicators but 2 share the same 'modified' timestamp assert len(indicators) == 2 + + +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(DataSourceError) as excinfo: + TAXIICollectionSource(collection_no_rw_access) + assert "Collection object provided does not have read access" in str(excinfo.value) + + +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(DataSourceError) as excinfo: + TAXIICollectionSink(collection_no_rw_access) + assert "Collection object provided does not have write access" in str(excinfo.value) + + +def test_get_404(): + """a TAXIICollectionSource.get() call that receives an HTTP 404 response + code from the taxii2client should be be returned as None. + + TAXII spec states that a TAXII server can return a 404 for nonexistent + resources or lack of access. Decided that None is acceptable reponse + to imply that state of the TAXII endpoint. + """ + + class TAXIICollection404(): + can_read = True + + def get_object(self, id, version=None): + resp = Response() + resp.status_code = 404 + resp.raise_for_status() + + ds = TAXIICollectionSource(TAXIICollection404()) + + # this will raise 404 from mock TAXII Client but TAXIICollectionStore + # should handle gracefully and return None + stix_obj = ds.get("indicator--1") + assert stix_obj is None + + +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(DataSourceError) as excinfo: + ds.all_versions("indicator--1") + assert "are either not found or access is denied" in str(excinfo.value) + assert "404" in str(excinfo.value) + + +def test_query_404(collection): + """ a TAXIICollectionSource.query() call that recieves an HTTP 404 + response code from the taxii2client should be returned as an exception""" + + ds = TAXIICollectionStore(collection) + query = [Filter("type", "=", "malware")] + + with pytest.raises(DataSourceError) as excinfo: + ds.query(query=query) + assert "are either not found or access is denied" in str(excinfo.value) + assert "404" in str(excinfo.value)