Clean up TAXII Datastore to use latest TAXII client

.isort.cfg

@@ -1,7 +1,7 @@
 [settings]
 check=1
 diff=1
-known_third_party=dateutil,pytest,pytz,six,requests
+known_third_party=dateutil,pytest,pytz,six,requests,taxii2_client
 known_first_party=stix2
 not_skip=__init__.py
 force_sort_within_sections=1

stix2/sources/taxii.py

@@ -21,67 +21,28 @@ TAXII_FILTERS = ['added_after', 'id', 'type', 'version']
 class TAXIICollectionStore(DataStore):
     """
     """
-    def __init__(self,
+    def __init__(self, collection, name="TAXIICollectionStore"):
-                 taxii_client=None,
+        """
-                 api_root_name=None,
+        Create a new TAXII Collection Data store
-                 collection_id=None,
+
-                 user=None,
+        Args:
-                 password=None,
+            collection (taxii2.Collection): Collection instance
-                 name="TAXIICollectionStore"):
+        """
 
         self.name = name
         self.id = make_id()
-        self.source = TAXIICollectionSource(taxii_client, api_root_name, collection_id, user, password)
+        self.source = TAXIICollectionSource(collection)
-        self.sink = TAXIICollectionSink(taxii_client, api_root_name, collection_id, user, password)
+        self.sink = TAXIICollectionSink(collection)
-
-    # file system sink API calls
-
-    def add(self, stix_objs):
-        return self.sink.add(stix_objs=stix_objs)
-
-    # file sytem source API calls
-
-    def get(self, stix_id):
-        return self.source.get(stix_id=stix_id)
-
-    def all_versions(self, stix_id):
-        return self.source.all_versions(stix_id=stix_id)
-
-    def query(self, query):
-        return self.source.query(query=query)
 
 
 class TAXIICollectionSink(DataSink):
     """
     """
 
-    def __init__(self, taxii_client=None, api_root_name=None, collection_id=None, user=None, password=None, name="TAXIICollectionSink"):
+    def __init__(self, collection, name="TAXIICollectionSink"):
         super(TAXIICollectionSink, self).__init__(name=name)
 
-        self.taxii_client = taxii_client
+        self.collection = collection
-        self.taxii_client.populate_available_information()
-
-        if not api_root_name:
-            raise ValueError("No api_root specified.")
-        else:
-            self.api_root = None
-            for a_r in self.taxii_client.api_roots:
-                if api_root_name == a_r.name:
-                    self.api_root = a_r
-                    break
-            if not self.api_root:
-                raise ValueError("The api_root %s is not found on this taxii server" % api_root_name)
-        if not collection_id:
-            raise ValueError("No collection specified.")
-        else:
-            self.collection = None
-            for c in self.api_root.collections:
-                if c.id_ == collection_id:
-                    self.collection = c
-                    break
-            if not self.collection:
-                raise ValueError("The collection %s is not found on the api_root %s of this taxii server" %
-                                 (collection_id, api_root_name))
 
     def add(self, stix_obj):
         """
@@ -95,53 +56,14 @@ class TAXIICollectionSink(DataSink):
                     spec_version="2.0",
                     type="bundle")
 
-    # utility functions for the current set collection and api root
-    def get_api_root_info(self):
-        """
-        """
-        return self.api_root.get_information()
-
-    def get_api_root_collections(self):
-        """
-        """
-        return self.api_root.get_collections()
-
-    def get_collection_manifest(self):
-        """
-        """
-        return self.collection.get_collection_manifest()
-
 
 class TAXIICollectionSource(DataSource):
     """
     """
-    def __init__(self, taxii_client=None, api_root_name=None, collection_id=None, user=None, password=None, name="TAXIICollectionSourc"):
+    def __init__(self, collection, name="TAXIICollectionSource"):
         super(TAXIICollectionSource, self).__init__(name=name)
 
-        self.taxii_client = taxii_client
+        self.collection = collection
-        self.taxii_client.populate_available_information()
-
-        if not api_root_name:
-            raise ValueError("No api_root specified.")
-        else:
-            self.api_root = None
-            for a_r in self.taxii_client.api_roots:
-                if api_root_name == a_r.name:
-                    self.api_root = a_r
-                    break
-            if not self.api_root:
-                raise ValueError("The api_root %s is not found on this taxii server" % api_root_name)
-        if not collection_id:
-            raise ValueError("No collection specified.")
-        else:
-            self.collection = None
-            for c in self.api_root.collections:
-                if c.id_ == collection_id:
-                    self.collection = c
-                    break
-            if not self.collection:
-                raise ValueError("The collection %s is not found on the api_root %s of this taxii server" %
-                                 (collection_id, api_root_name))
 
     def get(self, stix_id, _composite_filters=None):
         """
@@ -244,54 +166,3 @@ class TAXIICollectionSource(DataSource):
                     taxii_field = "match[" + filter_["field"] + ']'
                     params[taxii_field] = filter_["value"]
         return params
-
-    # utility functions for the current attached collection and api root
-    def get_api_root_info(self):
-        """
-        """
-        return self.api_root.get_information()
-
-    def get_api_root_collections(self):
-        """
-        """
-        return self.api_root.get_collections()
-
-    def get_collection_manifest(self):
-        """
-        """
-        return self.collection.get_collection_manifest()
-
-
-def get_server_api_roots(taxii_client):
-    """
-    """
-    api_root_info = []
-    taxii_client.populate_available_information()
-
-    for api_root in taxii_client.api_roots:
-        api_root_info.append(api_root.information())
-
-    return api_root_info
-
-
-def get_server_collections(taxii_client):
-    """
-    """
-    server_collections = []
-
-    taxii_client.populate_available_information()
-
-    for api_root in taxii_client.api_roots:
-        server_collections.extend(api_root.get_collections())
-
-    return server_collections
-
-
-def get_api_root_collections(taxii_client, api_root_name):
-    """
-    """
-    taxii_client.populate_available_information()
-
-    for api_root in taxii_client.api_roots:
-        if api_root == api_root_name:
-            return api_root.get_collections()

stix2/test/test_data_sources.py

@@ -1,25 +1,37 @@
+import pytest
+from taxii2_client import Collection
+
 from stix2.sources import taxii
 
-
+COLLECTION_URL = 'https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/'
-def test_ds_taxii():
-    ds = taxii.TAXIICollectionSource()
-    assert ds.name == 'TAXII'
 
 
-def test_ds_taxii_name():
+class MockTAXIIClient(object):
-    ds = taxii.TAXIICollectionSource(name='My Data Source Name')
+    """Mock for taxii2_client.TAXIIClient"""
+
+    def get(self):
+        return {}
+
+    def post(self):
+        return {}
+
+
+@pytest.fixture
+def collection():
+    return Collection(COLLECTION_URL, MockTAXIIClient())
+
+
+def test_ds_taxii(collection):
+    ds = taxii.TAXIICollectionSource(collection)
+    assert ds.name == 'TAXIICollectionSource'
+
+
+def test_ds_taxii_name(collection):
+    ds = taxii.TAXIICollectionSource(collection, name='My Data Source Name')
     assert ds.name == "My Data Source Name"
 
 
-def test_ds_params():
+def test_parse_taxii_filters(collection):
-    url = "http://taxii_url.com:5000"
-    creds = {"username": "Wade", "password": "Wilson"}
-    ds = taxii.TAXIICollectionSource(api_root=url, auth=creds)
-    assert ds.taxii_info['api_root']['url'] == url
-    assert ds.taxii_info['auth'] == creds
-
-
-def test_parse_taxii_filters():
     query = [
         {
             "field": "added_after",
@@ -55,14 +67,14 @@ def test_parse_taxii_filters():
         "match[version]": "first"
     }
 
-    ds = taxii.TAXIICollectionSource()
+    ds = taxii.TAXIICollectionSource(collection)
 
     taxii_filters = ds._parse_taxii_filters(query)
 
     assert taxii_filters == expected_params
 
 
-def test_add_get_remove_filter():
+def test_add_get_remove_filter(collection):
 
     class dummy(object):
         x = 4
@@ -114,7 +126,7 @@ def test_add_get_remove_filter():
         "Filter 'value' type is not supported. The type(value) must be python immutable type or dictionary"
     ]
 
-    ds = taxii.TAXIICollectionSource()
+    ds = taxii.TAXIICollectionSource(collection)
     # add
     ids, statuses = ds.add_filter(filters)
 
@@ -156,7 +168,7 @@ def test_add_get_remove_filter():
         assert id_ in ids[:3]
 
 
-def test_apply_common_filters():
+def test_apply_common_filters(collection):
     stix_objs = [
         {
             "created": "2017-01-27T13:49:53.997Z",
@@ -210,7 +222,7 @@ def test_apply_common_filters():
         }
     ]
 
-    ds = taxii.TAXIICollectionSource()
+    ds = taxii.TAXIICollectionSource(collection)
 
     resp = ds.apply_common_filters(stix_objs, [filters[0]])
     ids = [r['id'] for r in resp]
@@ -224,7 +236,7 @@ def test_apply_common_filters():
     assert resp[0]['id'] == stix_objs[0]['id']
 
 
-def test_deduplicate():
+def test_deduplicate(collection):
     stix_objs = [
         {
             "created": "2017-01-27T13:49:53.935Z",
@@ -288,7 +300,7 @@ def test_deduplicate():
         }
     ]
 
-    ds = taxii.TAXIICollectionSource()
+    ds = taxii.TAXIICollectionSource(collection)
     unique = ds.deduplicate(stix_objs)
 
     # Only 3 objects are unique