Skip to content

Commit

Permalink
Merge pull request #94 from emmanvg/pagination-changes
Browse files Browse the repository at this point in the history 
Pagination Changes
  • Loading branch information
@clenk
clenk authored Mar 12, 2021
2 parents 16a0f31 + 5d6ba88 commit c776650
Show file tree
Hide file tree
Showing 6 changed files with 380 additions and 59 deletions.
26 changes: 13 additions & 13 deletions README.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,13 +104,14 @@ A ``Collection`` can also be instantiated directly:
from taxii2client.v20 import Collection, as_pages
collection = Collection('https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116')
collection.get_object('indicator--252c7c11-daf2-42bd-843b-be65edca9f61')
print(collection.get_object('indicator--252c7c11-daf2-42bd-843b-be65edca9f61'))
# For normal (no pagination) requests
collection.get_objects()
collection.get_manifest()
print(collection.get_objects())
print(collection.get_manifest())
# For pagination requests.
# Use *args for other arguments to the call and **kwargs to pass filter information
for bundle in as_pages(collection.get_objects, per_request=50):
print(bundle)
Expand All @@ -119,23 +120,22 @@ A ``Collection`` can also be instantiated directly:
# ---------------------------------------------------------------- #
# Performing TAXII 2.1 Requests
from taxii2client.v21 import Collection
from taxii2client.v21 import Collection, as_pages
collection = Collection('https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116')
collection.get_object('indicator--252c7c11-daf2-42bd-843b-be65edca9f61')
print(collection.get_object('indicator--252c7c11-daf2-42bd-843b-be65edca9f61'))
# For normal (no pagination) requests
collection.get_objects()
collection.get_manifest()
print(collection.get_objects())
print(collection.get_manifest())
# For pagination requests.
envelope = collection.get_objects(limit=50)
while envelope.get("more", False):
envelope = collection.get_objects(limit=50, next=envelope.get("next", ""))
# Use *args for other arguments to the call and **kwargs to pass filter information
for envelope in as_pages(collection.get_objects, per_request=50):
print(envelope)
envelope = collection.get_manifest(limit=50)
while envelope.get("more", False):
envelope = collection.get_manifest(limit=50, next=envelope.get("next", ""))
for manifest_resource in as_pages(collection.get_manifest, per_request=50):
print(manifest_resource)
In addition to the object-specific properties and methods, all classes have a
``refresh()`` method that reloads the URL corresponding to that resource, to
Expand Down
20 changes: 6 additions & 14 deletions taxii2client/common.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import datetime
import logging
import re

import pytz
import requests
Expand Down Expand Up @@ -117,7 +116,7 @@ def _to_json(resp):
Factors out some JSON parse code with error handling, to hopefully improve
error messages.
:param resp: A "requests" library response
:param resp: A requests.Response instance
:return: Parsed JSON.
:raises: InvalidJSONError If JSON parsing failed.
"""
Expand All @@ -130,18 +129,11 @@ def _to_json(resp):
), e)


def _grab_total_items(resp):
"""Extracts the Total elements available on the Endpoint making the request"""
try:
results = re.match(r"^items (\d+)-(\d+)/(\d+)$", resp.headers["Content-Range"])
return int(results.group(2)) - int(results.group(1)) + 1, int(results.group(3))
except (ValueError, IndexError) as e:
six.raise_from(InvalidJSONError(
"Invalid Content-Range was received from " + resp.request.url
), e)
except KeyError:
log.warning("TAXII Server Response did not include 'Content-Range' header - results could be incomplete")
return 0, 0
def _grab_total_items_from_resource(resp):
"""Returns number of objects in bundle/envelope"""
if isinstance(resp, requests.Response):
resp = _to_json(resp)
return len(resp.get("objects", []))


class TokenAuth(requests.auth.AuthBase):
Expand Down
255 changes: 244 additions & 11 deletions taxii2client/test/test_client_v20.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
AccessError, InvalidArgumentsError, InvalidJSONError,
TAXIIServiceException, ValidationError
)
from taxii2client.v20 import ApiRoot, Collection, Server, Status
from taxii2client.v20 import ApiRoot, Collection, Server, Status, as_pages

TAXII_SERVER = "example.com"
DISCOVERY_URL = "https://{}/taxii/".format(TAXII_SERVER)
Expand Down Expand Up @@ -98,23 +98,29 @@
]
}"""


STIX_OBJECT = """
{
"type": "indicator",
"id": "indicator--252c7c11-daf2-42bd-843b-be65edca9f61",
"created": "2016-04-06T20:03:48.000Z",
"modified": "2016-04-06T20:03:48.000Z",
"pattern": "[ file:hashes.MD5 = 'd41d8cd98f00b204e9800998ecf8427e' ]",
"valid_from": "2016-01-01T00:00:00Z"
}
"""


# This bundle is used as the response to get_objects(), and also the bundle
# POST'ed with add_objects().
STIX_BUNDLE = """{
STIX_BUNDLE = f"""{{
"type": "bundle",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"spec_version": "2.0",
"objects": [
{
"type": "indicator",
"id": "indicator--252c7c11-daf2-42bd-843b-be65edca9f61",
"created": "2016-04-06T20:03:48.000Z",
"modified": "2016-04-06T20:03:48.000Z",
"pattern": "[ file:hashes.MD5 = 'd41d8cd98f00b204e9800998ecf8427e' ]",
"valid_from": "2016-01-01T00:00:00Z"
}
{STIX_OBJECT}
]
}"""
}}"""
GET_OBJECTS_RESPONSE = STIX_BUNDLE
# get_object() still returns a bundle. In this case, the bundle has only one
# object (the correct one.)
Expand Down Expand Up @@ -489,6 +495,233 @@ def test_collection_unexpected_kwarg():
Collection(url="", conn=None, foo="bar")


@responses.activate
def test_get_collection_objects_paged_1(collection):
obj_return = []
for x in range(0, 50):
obj_return.append(json.loads(STIX_OBJECT))

responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[:10]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items 0-9/50'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[10:20]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items 10-19/50'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[20:30]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items 20-29/50'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[30:40]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items 30-39/50'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[40:50]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items 40-49/50'})
response = []

for bundle in as_pages(collection.get_objects, per_request=10):
response.extend(bundle.get("objects", []))

assert len(response) == 50


@responses.activate
def test_get_collection_objects_paged_2(collection):
obj_return = []
for x in range(0, 50):
obj_return.append(json.loads(STIX_OBJECT))

responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[:10]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items 0-9/*'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[10:20]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items 10-19/*'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[20:30]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items 20-29/*'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[30:40]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items 30-39/*'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[40:50]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items 40-49/*'})
responses.add(responses.GET, GET_OBJECTS_URL, "",
status=406, content_type=MEDIA_TYPE_STIX_V20)
response = []

# The status errors (any 400-500) will make it stop
with pytest.raises(requests.exceptions.HTTPError):
for bundle in as_pages(collection.get_objects, per_request=10):
response.extend(bundle.get("objects", []))

assert len(response) == 50


@responses.activate
def test_get_collection_objects_paged_3(collection):
obj_return = []
for x in range(0, 50):
obj_return.append(json.loads(STIX_OBJECT))

responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[:10]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items */50'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[10:20]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items */50'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[20:30]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items */50'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[30:40]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items */50'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[40:50]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items */50'})
response = []

for bundle in as_pages(collection.get_objects, per_request=10):
response.extend(bundle.get("objects", []))

assert len(response) == 50


@responses.activate
def test_get_collection_objects_paged_4(collection):
obj_return = []
for x in range(0, 50):
obj_return.append(json.loads(STIX_OBJECT))

responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[:10]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items */*'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[10:20]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items */*'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[20:30]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items */*'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[30:40]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items */*'})
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[40:50]}),
status=200, content_type=MEDIA_TYPE_STIX_V20,
headers={'Content-Range': 'items */*'})
responses.add(responses.GET, GET_OBJECTS_URL, "",
status=406, content_type=MEDIA_TYPE_STIX_V20)
response = []

# The status errors (any 400-500) will make it stop
with pytest.raises(requests.exceptions.HTTPError):
for bundle in as_pages(collection.get_objects, per_request=10):
response.extend(bundle.get("objects", []))

assert len(response) == 50


@responses.activate
def test_get_collection_objects_paged_5(collection):
obj_return = []
for x in range(0, 50):
obj_return.append(json.loads(STIX_OBJECT))

responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[:10]}),
status=200, content_type=MEDIA_TYPE_STIX_V20)
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[10:20]}),
status=200, content_type=MEDIA_TYPE_STIX_V20)
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[20:30]}),
status=200, content_type=MEDIA_TYPE_STIX_V20)
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[30:40]}),
status=200, content_type=MEDIA_TYPE_STIX_V20)
responses.add(responses.GET, GET_OBJECTS_URL,
json.dumps({"type": "bundle", "spec_version": "2.0",
"id": "bundle--5d0092c5-5f74-4287-9642-33f4c354e56d",
"objects": obj_return[40:50]}),
status=200, content_type=MEDIA_TYPE_STIX_V20)
responses.add(responses.GET, GET_OBJECTS_URL, "",
status=406, content_type=MEDIA_TYPE_STIX_V20)
response = []

# The status errors (any 400-500) will make it stop
with pytest.raises(requests.exceptions.HTTPError):
for bundle in as_pages(collection.get_objects, per_request=10):
response.extend(bundle.get("objects", []))

assert len(response) == 50


@responses.activate
def test_get_collection_objects(collection):
responses.add(responses.GET, GET_OBJECTS_URL, GET_OBJECTS_RESPONSE,
Expand Down
Loading

0 comments on commit c776650

Please sign in to comment.