Repository: oasis-open/cti-python-stix2 (python)

Issue:
TAXIICollectionSource.query() only returns the first page
I noticed that the TAXIICollectionSource methods query(), relationships(), ... only operates with the first page returned by the associated TAXII 2.0 server.

I did tests with a local Medallion server v1.1.0 which have a default max_page_size of 100.

By calling TAXIICollectionSource.query() I can get only 100/10000 objects. Also, when calling relationships(stix2_id) the results are limited to those 100 objects.

The only solution I found is increasing that max_page_size, so all objects are returned in one call.





Pull request description:
Support Pagination in TAXII DataStore
Proposed solution to solve the pagination issues.

Closes #479 

Patch:
diff --git a/setup.py b/setup.py
index 2fc5d70..397b98f 100644
--- a/setup.py
+++ b/setup.py
@@ -60,7 +60,7 @@ setup(
         'Bug Tracker': 'https://github.com/oasis-open/cti-python-stix2/issues/',
     },
     extras_require={
-        'taxii': ['taxii2-client>=2.2.1'],
+        'taxii': ['taxii2-client>=2.3.0'],
         'semantic': ['haversine', 'rapidfuzz'],
     },
 )
diff --git a/stix2/datastore/taxii.py b/stix2/datastore/taxii.py
index 41d1e54..9ad6df9 100644
--- a/stix2/datastore/taxii.py
+++ b/stix2/datastore/taxii.py
@@ -12,6 +12,8 @@ from stix2.parsing import parse
 from stix2.utils import deduplicate
 
 try:
+    from taxii2client import v20 as tcv20
+    from taxii2client import v21 as tcv21
     from taxii2client.exceptions import ValidationError
     _taxii2_client = True
 except ImportError:
@@ -33,9 +35,12 @@ class TAXIICollectionStore(DataStoreMixin):
             side(retrieving data) and False for TAXIICollectionSink
             side(pushing data). However, when parameter is supplied, it will
             be applied to both TAXIICollectionSource/Sink.
+        items_per_page (int): How many STIX objects to request per call
+            to TAXII Server. The value can be tuned, but servers may override
+            if their internal limit is surpassed. Used by TAXIICollectionSource
 
     """
-    def __init__(self, collection, allow_custom=None):
+    def __init__(self, collection, allow_custom=None, items_per_page=5000):
         if allow_custom is None:
             allow_custom_source = True
             allow_custom_sink = False
@@ -43,7 +48,7 @@ class TAXIICollectionStore(DataStoreMixin):
             allow_custom_sink = allow_custom_source = allow_custom
 
         super(TAXIICollectionStore, self).__init__(
-            source=TAXIICollectionSource(collection, allow_custom=allow_custom_source),
+            source=TAXIICollectionSource(collection, allow_custom=allow_custom_source, items_per_page=items_per_page),
             sink=TAXIICollectionSink(collection, allow_custom=allow_custom_sink),
         )
 
@@ -144,9 +149,12 @@ class TAXIICollectionSource(DataSource):
         collection (taxii2.Collection): TAXII Collection instance
         allow_custom (bool): Whether to allow custom STIX content to be
             added to the FileSystemSink. Default: True
+        items_per_page (int): How many STIX objects to request per call
+            to TAXII Server. The value can be tuned, but servers may override
+            if their internal limit is surpassed.
 
     """
-    def __init__(self, collection, allow_custom=True):
+    def __init__(self, collection, allow_custom=True, items_per_page=5000):
         super(TAXIICollectionSource, self).__init__()
         if not _taxii2_client:
             raise ImportError("taxii2client library is required for usage of TAXIICollectionSource")
@@ -167,6 +175,7 @@ class TAXIICollectionSource(DataSource):
             )
 
         self.allow_custom = allow_custom
+        self.items_per_page = items_per_page
 
     def get(self, stix_id, version=None, _composite_filters=None):
         """Retrieve STIX object from local/remote STIX Collection
@@ -286,8 +295,12 @@ class TAXIICollectionSource(DataSource):
         taxii_filters_dict = dict((f.property, f.value) for f in taxii_filters)
 
         # query TAXII collection
+        all_data = []
         try:
-            all_data = self.collection.get_objects(**taxii_filters_dict).get('objects', [])
+            paged_request = tcv21.as_pages if isinstance(self.collection, tcv21.Collection) else tcv20.as_pages
+
+            for resource in paged_request(self.collection.get_objects, per_request=self.items_per_page, **taxii_filters_dict):
+                all_data.extend(resource.get("objects", []))
 
             # deduplicate data (before filtering as reduces wasted filtering)
             all_data = deduplicate(all_data)


Test patch:
diff --git a/stix2/test/v20/test_datastore_taxii.py b/stix2/test/v20/test_datastore_taxii.py
index 0b21981..075f0a3 100644
--- a/stix2/test/v20/test_datastore_taxii.py
+++ b/stix2/test/v20/test_datastore_taxii.py
@@ -3,9 +3,8 @@ import json
 from medallion.filters.basic_filter import BasicFilter
 import pytest
 from requests.models import Response
-import six
 from taxii2client.common import _filter_kwargs_to_query_params
-from taxii2client.v20 import Collection
+from taxii2client.v20 import MEDIA_TYPE_STIX_V20, Collection
 
 import stix2
 from stix2.datastore import DataSourceError
@@ -27,7 +26,7 @@ class MockTAXIICollectionEndpoint(Collection):
 
     def add_objects(self, bundle):
         self._verify_can_write()
-        if isinstance(bundle, six.string_types):
+        if isinstance(bundle, str):
             bundle = json.loads(bundle)
         for object in bundle.get("objects", []):
             self.objects.append(object)
@@ -35,12 +34,12 @@ class MockTAXIICollectionEndpoint(Collection):
                 {
                     "date_added": get_timestamp(),
                     "id": object["id"],
-                    "media_type": "application/stix+json;version=2.1",
+                    "media_type": "application/stix+json;version=2.0",
                     "version": object.get("modified", object.get("created", get_timestamp())),
                 },
             )
 
-    def get_objects(self, **filter_kwargs):
+    def get_objects(self, accept=MEDIA_TYPE_STIX_V20, start=0, per_request=0, **filter_kwargs):
         self._verify_can_read()
         query_params = _filter_kwargs_to_query_params(filter_kwargs)
         assert isinstance(query_params, dict)
@@ -52,7 +51,12 @@ class MockTAXIICollectionEndpoint(Collection):
             100,
         )[0]
         if objs:
-            return stix2.v20.Bundle(objects=objs)
+            resp = Response()
+            resp.status_code = 200
+            resp.headers["Content-Range"] = f"items 0-{len(objs)}/{len(objs)}"
+            resp.encoding = "utf-8"
+            resp._content = bytes(stix2.v20.Bundle(objects=objs).serialize(ensure_ascii=False), resp.encoding)
+            return resp
         else:
             resp = Response()
             resp.status_code = 404
diff --git a/stix2/test/v21/test_datastore_taxii.py b/stix2/test/v21/test_datastore_taxii.py
index 92ae6dc..62ddadc 100644
--- a/stix2/test/v21/test_datastore_taxii.py
+++ b/stix2/test/v21/test_datastore_taxii.py
@@ -3,7 +3,6 @@ import json
 from medallion.filters.basic_filter import BasicFilter
 import pytest
 from requests.models import Response
-import six
 from taxii2client.common import _filter_kwargs_to_query_params
 from taxii2client.v21 import Collection
 
@@ -27,16 +26,16 @@ class MockTAXIICollectionEndpoint(Collection):
 
     def add_objects(self, bundle):
         self._verify_can_write()
-        if isinstance(bundle, six.string_types):
+        if isinstance(bundle, str):
             bundle = json.loads(bundle)
-        for object in bundle.get("objects", []):
-            self.objects.append(object)
+        for obj in bundle.get("objects", []):
+            self.objects.append(obj)
             self.manifests.append(
                 {
                     "date_added": get_timestamp(),
-                    "id": object["id"],
+                    "id": obj["id"],
                     "media_type": "application/stix+json;version=2.1",
-                    "version": object.get("modified", object.get("created", get_timestamp())),
+                    "version": obj.get("modified", obj.get("created", get_timestamp())),
                 },
             )
 
@@ -52,7 +51,10 @@ class MockTAXIICollectionEndpoint(Collection):
             100,
         )[0]
         if objs:
-            return stix2.v21.Bundle(objects=objs)
+            return {
+                "objects": objs,
+                "more": False,
+            }
         else:
             resp = Response()
             resp.status_code = 404
@@ -76,7 +78,10 @@ class MockTAXIICollectionEndpoint(Collection):
         else:
             filtered_objects = []
         if filtered_objects:
-            return stix2.v21.Bundle(objects=filtered_objects)
+            return {
+                "objects": filtered_objects,
+                "more": False,
+            }
         else:
             resp = Response()
             resp.status_code = 404