Repository: oasis-open/cti-taxii-client

Issue:
Provide access to raw TAXII responses
This can probably be behind some sort of "private" interface (like `_raw`), but it would be helpful to have access to the raw JSON bodies on returned responses.

Discussion:

Patch:
diff --git a/taxii2client/__init__.py b/taxii2client/__init__.py
index b4b0646..5c359cf 100644
--- a/taxii2client/__init__.py
+++ b/taxii2client/__init__.py
@@ -178,7 +178,7 @@ class Status(_TAXIIEndpoint):
     # aren't other endpoints to call on the Status object.
 
     def __init__(self, url, conn=None, user=None, password=None, verify=True,
-                 **kwargs):
+                 status_info=None):
         """Create an API root resource endpoint.
 
         Args:
@@ -187,11 +187,16 @@ class Status(_TAXIIEndpoint):
             password (str): password for authentication (optional)
             conn (_HTTPConnection): reuse connection object, as an alternative
                 to providing username/password
+            status_info (dict): Parsed JSON representing a response from the
+                status endpoint, if already known.  If not given, the
+                endpoint will be queried. (optional)
 
         """
         super(Status, self).__init__(url, conn, user, password, verify)
-        if kwargs:
-            self._populate_fields(**kwargs)
+        self.__raw = None
+        if status_info:
+            self._populate_fields(**status_info)
+            self.__raw = status_info
         else:
             self.refresh()
 
@@ -200,9 +205,14 @@ class Status(_TAXIIEndpoint):
 
     __bool__ = __nonzero__
 
+    @property
+    def _raw(self):
+        """Get the "raw" status response (parsed JSON)."""
+        return self.__raw
+
     def refresh(self, accept=MEDIA_TYPE_TAXII_V20):
         """Updates Status information"""
-        response = self._conn.get(self.url, accept=accept)
+        response = self.__raw = self._conn.get(self.url, accept=accept)
         self._populate_fields(**response)
 
     def wait_until_final(self, poll_interval=1, timeout=60):
@@ -317,7 +327,7 @@ class Collection(_TAXIIEndpoint):
     """
 
     def __init__(self, url, conn=None, user=None, password=None, verify=True,
-                 **kwargs):
+                 collection_info=None):
         """
         Initialize a new Collection.  Either user/password or conn may be
         given, but not both.  The latter is intended for internal use, when
@@ -334,20 +344,22 @@ class Collection(_TAXIIEndpoint):
                 case it must be a path to a CA bundle to use. Defaults to
                 `True` (optional)
             conn (_HTTPConnection): A connection to reuse (optional)
-            kwargs: Collection metadata, if known in advance (optional)
+            collection_info: Collection metadata, if known in advance (optional)
 
         """
 
         super(Collection, self).__init__(url, conn, user, password, verify)
 
         self._loaded = False
+        self.__raw = None
 
         # Since the API Root "Get Collections" endpoint returns information on
         # all collections as a list, it's possible that we can create multiple
         # Collection objects from a single HTTPS request, and not need to call
         # `refresh` for each one.
-        if kwargs:
-            self._populate_fields(**kwargs)
+        if collection_info:
+            self._populate_fields(**collection_info)
+            self.__raw = collection_info
             self._loaded = True
 
     @property
@@ -384,6 +396,12 @@ class Collection(_TAXIIEndpoint):
     def objects_url(self):
         return self.url + "objects/"
 
+    @property
+    def _raw(self):
+        """Get the "raw" collection information response (parsed JSON)."""
+        self._ensure_loaded()
+        return self.__raw
+
     def _populate_fields(self, id=None, title=None, description=None,
                          can_read=None, can_write=None, media_types=None):
         self._id = id  # required
@@ -434,7 +452,7 @@ class Collection(_TAXIIEndpoint):
 
     def refresh(self, accept=MEDIA_TYPE_TAXII_V20):
         """Update Collection information"""
-        response = self._conn.get(self.url, accept=accept)
+        response = self.__raw = self._conn.get(self.url, accept=accept)
         self._populate_fields(**response)
         self._loaded = True
 
@@ -514,7 +532,8 @@ class Collection(_TAXIIEndpoint):
             "../../status/{}".format(status_json["id"])
         )
 
-        status = Status(url=status_url, conn=self._conn, **status_json)
+        status = Status(url=status_url, conn=self._conn,
+                        status_info=status_json)
 
         if not wait_for_completion or status.status == "complete":
             return status
@@ -564,6 +583,7 @@ class ApiRoot(_TAXIIEndpoint):
 
         self._loaded_collections = False
         self._loaded_information = False
+        self.__raw = None
 
     @property
     def collections(self):
@@ -591,6 +611,12 @@ class ApiRoot(_TAXIIEndpoint):
         self._ensure_loaded_information()
         return self._max_content_length
 
+    @property
+    def _raw(self):
+        """Get the "raw" API root information response (parsed JSON)."""
+        self._ensure_loaded_information()
+        return self.__raw
+
     def _ensure_loaded_information(self):
         if not self._loaded_information:
             self.refresh_information()
@@ -620,7 +646,7 @@ class ApiRoot(_TAXIIEndpoint):
 
         This invokes the ``Get API Root Information`` endpoint.
         """
-        response = self._conn.get(self.url, accept=accept)
+        response = self.__raw = self._conn.get(self.url, accept=accept)
 
         self._title = response.get("title")  # required
         self._description = response.get("description")  # optional
@@ -641,7 +667,8 @@ class ApiRoot(_TAXIIEndpoint):
         self._collections = []
         for item in response.get("collections", []):  # optional
             collection_url = url + item["id"] + "/"
-            collection = Collection(collection_url, conn=self._conn, **item)
+            collection = Collection(collection_url, conn=self._conn,
+                                    collection_info=item)
             self._collections.append(collection)
 
         self._loaded_collections = True
@@ -649,7 +676,7 @@ class ApiRoot(_TAXIIEndpoint):
     def get_status(self, status_id, accept=MEDIA_TYPE_TAXII_V20):
         status_url = self.url + "status/" + status_id + "/"
         response = self._conn.get(status_url, accept=accept)
-        return Status(status_url, conn=self._conn, **response)
+        return Status(status_url, conn=self._conn, status_info=response)
 
 
 class Server(_TAXIIEndpoint):
@@ -685,6 +712,7 @@ class Server(_TAXIIEndpoint):
         self._password = password
         self._verify = verify
         self._loaded = False
+        self.__raw = None
 
     @property
     def title(self):
@@ -711,6 +739,12 @@ class Server(_TAXIIEndpoint):
         self._ensure_loaded()
         return self._api_roots
 
+    @property
+    def _raw(self):
+        """Get the "raw" server discovery response (parsed JSON)."""
+        self._ensure_loaded()
+        return self.__raw
+
     def _ensure_loaded(self):
         if not self._loaded:
             self.refresh()
@@ -724,7 +758,8 @@ class Server(_TAXIIEndpoint):
 
     def refresh(self):
         """Update the Server information and list of API Roots"""
-        response = self._conn.get(self.url, accept=MEDIA_TYPE_TAXII_V20)
+        response = self.__raw = self._conn.get(self.url,
+                                               accept=MEDIA_TYPE_TAXII_V20)
 
         self._title = response.get("title")  # required
         self._description = response.get("description")  # optional


Test patch:
diff --git a/taxii2client/test/test_client.py b/taxii2client/test/test_client.py
index 5d70679..0505ec4 100644
--- a/taxii2client/test/test_client.py
+++ b/taxii2client/test/test_client.py
@@ -272,11 +272,21 @@ def set_discovery_response(response):
                   content_type=MEDIA_TYPE_TAXII_V20)
 
 
+def set_collections_response():
+    responses.add(responses.GET, COLLECTIONS_URL, COLLECTIONS_RESPONSE,
+                  status=200, content_type=MEDIA_TYPE_TAXII_V20)
+
+
 def set_collection_response(url=COLLECTION_URL, response=COLLECTION_RESPONSE):
     responses.add(responses.GET, url, response, status=200,
                   content_type=MEDIA_TYPE_TAXII_V20)
 
 
+def set_status_response():
+    responses.add(responses.GET, STATUS_URL, STATUS_RESPONSE,
+                  status=200, content_type=MEDIA_TYPE_TAXII_V20)
+
+
 @responses.activate
 def test_server_discovery(server):
     set_discovery_response(DISCOVERY_RESPONSE)
@@ -296,6 +306,9 @@ def test_server_discovery(server):
     assert api_root._loaded_information is False
     assert api_root._loaded_collections is False
 
+    discovery_dict = json.loads(DISCOVERY_RESPONSE)
+    assert server._raw == discovery_dict
+
 
 @responses.activate
 def test_minimal_discovery_response(server):
@@ -389,8 +402,7 @@ def test_api_root_no_max_content_length(api_root):
 
 @responses.activate
 def test_api_root(api_root):
-    responses.add(responses.GET, API_ROOT_URL, API_ROOT_RESPONSE,
-                  status=200, content_type=MEDIA_TYPE_TAXII_V20)
+    set_api_root_response(API_ROOT_RESPONSE)
 
     assert api_root._loaded_information is False
     assert api_root.title == "Malware Research Group"
@@ -399,11 +411,13 @@ def test_api_root(api_root):
     assert api_root.versions == ["taxii-2.0"]
     assert api_root.max_content_length == 9765625
 
+    apiroot_dict = json.loads(API_ROOT_RESPONSE)
+    assert api_root._raw == apiroot_dict
+
 
 @responses.activate
 def test_api_root_collections(api_root):
-    responses.add(responses.GET, COLLECTIONS_URL, COLLECTIONS_RESPONSE, status=200,
-                  content_type=MEDIA_TYPE_TAXII_V20)
+    set_collections_response()
 
     assert api_root._loaded_collections is False
     assert len(api_root.collections) == 2
@@ -420,6 +434,9 @@ def test_api_root_collections(api_root):
     assert coll.can_write is False
     assert coll.media_types == [MEDIA_TYPE_STIX_V20]
 
+    collection_dict = json.loads(COLLECTION_RESPONSE)
+    assert coll._raw == collection_dict
+
 
 @responses.activate
 def test_collection(collection):
@@ -433,6 +450,9 @@ def test_collection(collection):
     assert collection.can_write is False
     assert collection.media_types == [MEDIA_TYPE_STIX_V20]
 
+    collection_dict = json.loads(COLLECTION_RESPONSE)
+    assert collection._raw == collection_dict
+
 
 def test_collection_unexpected_kwarg():
     with pytest.raises(TypeError):
@@ -481,6 +501,9 @@ def test_add_object_to_collection(writable_collection):
     assert status.failure_count == 0
     assert status.pending_count == 0
 
+    status_dict = json.loads(ADD_OBJECTS_RESPONSE)
+    assert status._raw == status_dict
+
 
 @responses.activate
 def test_add_object_to_collection_dict(writable_collection):
@@ -536,9 +559,8 @@ def test_get_manifest(collection):
 
 
 @responses.activate
-def test_get_status(api_root):
-    responses.add(responses.GET, STATUS_URL, STATUS_RESPONSE,
-                  status=200, content_type=MEDIA_TYPE_TAXII_V20)
+def test_get_status(api_root, status_dict):
+    set_status_response()
 
     status = api_root.get_status(STATUS_ID)
 
@@ -550,6 +572,17 @@ def test_get_status(api_root):
     assert status.pending_count == 2
     assert len(status.pendings) == 2
 
+    assert status._raw == status_dict
+
+
+@responses.activate
+def test_status_raw(status_dict):
+    """Test Status object created directly (not obtained via ApiRoot),
+    and _raw property."""
+    set_status_response()
+    status = Status(STATUS_URL)
+    assert status_dict == status._raw
+
 
 @responses.activate
 def test_content_type_valid(collection):
@@ -675,7 +708,8 @@ def test_status_missing_id_property(status_dict):
     with pytest.raises(ValidationError) as excinfo:
         status_dict.pop("id")
         Status("https://example.com/api1/status/12345678-1234-1234-1234-123456789012/",
-               user="foo", password="bar", verify=False, **status_dict)
+               user="foo", password="bar", verify=False,
+               status_info=status_dict)
 
     assert "No 'id' in Status for request 'https://example.com/api1/status/12345678-1234-1234-1234-123456789012/'" == str(excinfo.value)
 
@@ -684,7 +718,8 @@ def test_status_missing_status_property(status_dict):
     with pytest.raises(ValidationError) as excinfo:
         status_dict.pop("status")
         Status("https://example.com/api1/status/12345678-1234-1234-1234-123456789012/",
-               user="foo", password="bar", verify=False, **status_dict)
+               user="foo", password="bar", verify=False,
+               status_info=status_dict)
 
     assert "No 'status' in Status for request 'https://example.com/api1/status/12345678-1234-1234-1234-123456789012/'" == str(excinfo.value)
 
@@ -693,7 +728,8 @@ def test_status_missing_total_count_property(status_dict):
     with pytest.raises(ValidationError) as excinfo:
         status_dict.pop("total_count")
         Status("https://example.com/api1/status/12345678-1234-1234-1234-123456789012/",
-               user="foo", password="bar", verify=False, **status_dict)
+               user="foo", password="bar", verify=False,
+               status_info=status_dict)
 
     assert "No 'total_count' in Status for request 'https://example.com/api1/status/12345678-1234-1234-1234-123456789012/'" == str(excinfo.value)
 
@@ -702,7 +738,8 @@ def test_status_missing_success_count_property(status_dict):
     with pytest.raises(ValidationError) as excinfo:
         status_dict.pop("success_count")
         Status("https://example.com/api1/status/12345678-1234-1234-1234-123456789012/",
-               user="foo", password="bar", verify=False, **status_dict)
+               user="foo", password="bar", verify=False,
+               status_info=status_dict)
 
     assert "No 'success_count' in Status for request 'https://example.com/api1/status/12345678-1234-1234-1234-123456789012/'" == str(excinfo.value)
 
@@ -711,7 +748,8 @@ def test_status_missing_failure_count_property(status_dict):
     with pytest.raises(ValidationError) as excinfo:
         status_dict.pop("failure_count")
         Status("https://example.com/api1/status/12345678-1234-1234-1234-123456789012/",
-               user="foo", password="bar", verify=False, **status_dict)
+               user="foo", password="bar", verify=False,
+               status_info=status_dict)
 
     assert "No 'failure_count' in Status for request 'https://example.com/api1/status/12345678-1234-1234-1234-123456789012/'" == str(excinfo.value)
 
@@ -720,7 +758,8 @@ def test_status_missing_pending_count_property(status_dict):
     with pytest.raises(ValidationError) as excinfo:
         status_dict.pop("pending_count")
         Status("https://example.com/api1/status/12345678-1234-1234-1234-123456789012/",
-               user="foo", password="bar", verify=False, **status_dict)
+               user="foo", password="bar", verify=False,
+               status_info=status_dict)
 
     assert "No 'pending_count' in Status for request 'https://example.com/api1/status/12345678-1234-1234-1234-123456789012/'" == str(excinfo.value)
 
@@ -729,7 +768,8 @@ def test_collection_missing_id_property(collection_dict):
     with pytest.raises(ValidationError) as excinfo:
         collection_dict.pop("id")
         Collection("https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/",
-                   user="foo", password="bar", verify=False, **collection_dict)
+                   user="foo", password="bar", verify=False,
+                   collection_info=collection_dict)
 
     assert "No 'id' in Collection for request 'https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/'" == str(excinfo.value)
 
@@ -738,7 +778,8 @@ def test_collection_missing_title_property(collection_dict):
     with pytest.raises(ValidationError) as excinfo:
         collection_dict.pop("title")
         Collection("https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/",
-                   user="foo", password="bar", verify=False, **collection_dict)
+                   user="foo", password="bar", verify=False,
+                   collection_info=collection_dict)
 
     assert "No 'title' in Collection for request 'https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/'" == str(excinfo.value)
 
@@ -747,7 +788,8 @@ def test_collection_missing_can_read_property(collection_dict):
     with pytest.raises(ValidationError) as excinfo:
         collection_dict.pop("can_read")
         Collection("https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/",
-                   user="foo", password="bar", verify=False, **collection_dict)
+                   user="foo", password="bar", verify=False,
+                   collection_info=collection_dict)
 
     assert "No 'can_read' in Collection for request 'https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/'" == str(excinfo.value)
 
@@ -756,6 +798,7 @@ def test_collection_missing_can_write_property(collection_dict):
     with pytest.raises(ValidationError) as excinfo:
         collection_dict.pop("can_write")
         Collection("https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/",
-                   user="foo", password="bar", verify=False, **collection_dict)
+                   user="foo", password="bar", verify=False,
+                   collection_info=collection_dict)
 
     assert "No 'can_write' in Collection for request 'https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/'" == str(excinfo.value)