Repository: oasis-open/cti-taxii-client

Issue:
Allow users to specify the user agent string...
...and/or use a default that is specific to the taxii client.

Comments:

Patch:
diff --git a/docs/conf.py b/docs/conf.py
index 47930be..70f3419 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -33,7 +33,6 @@
 # ones.
 extensions = [
     'sphinx-prompt',
-    'nbsphinx',
     'sphinx.ext.autodoc',
     'sphinx.ext.autosummary',
     'sphinx.ext.napoleon',
diff --git a/setup.py b/setup.py
index 132e987..2ca235c 100644
--- a/setup.py
+++ b/setup.py
@@ -47,6 +47,7 @@ setup(
         ],
         'docs': [
             'sphinx',
+            'sphinx-prompt',
         ]
     }
 )
diff --git a/taxii2client/__init__.py b/taxii2client/__init__.py
index e001bb6..36e5929 100644
--- a/taxii2client/__init__.py
+++ b/taxii2client/__init__.py
@@ -8,6 +8,7 @@ import time
 
 import pytz
 import requests
+import requests.structures  # is this public API?
 import six
 import six.moves.urllib.parse as urlparse
 
@@ -15,6 +16,7 @@ __version__ = '0.3.1'
 
 MEDIA_TYPE_STIX_V20 = "application/vnd.oasis.stix+json; version=2.0"
 MEDIA_TYPE_TAXII_V20 = "application/vnd.oasis.taxii+json; version=2.0"
+DEFAULT_USER_AGENT = "taxii2-client/" + __version__
 
 
 class TAXIIServiceException(Exception):
@@ -217,7 +219,8 @@ class Status(_TAXIIEndpoint):
 
     def refresh(self, accept=MEDIA_TYPE_TAXII_V20):
         """Updates Status information"""
-        response = self.__raw = self._conn.get(self.url, accept=accept)
+        response = self.__raw = self._conn.get(self.url,
+                                               headers={"Accept": accept})
         self._populate_fields(**response)
 
     def wait_until_final(self, poll_interval=1, timeout=60):
@@ -457,7 +460,8 @@ class Collection(_TAXIIEndpoint):
 
     def refresh(self, accept=MEDIA_TYPE_TAXII_V20):
         """Update Collection information"""
-        response = self.__raw = self._conn.get(self.url, accept=accept)
+        response = self.__raw = self._conn.get(self.url,
+                                               headers={"Accept": accept})
         self._populate_fields(**response)
         self._loaded = True
 
@@ -465,7 +469,7 @@ class Collection(_TAXIIEndpoint):
         """Implement the ``Get Objects`` endpoint (section 5.3)"""
         self._verify_can_read()
         query_params = _filter_kwargs_to_query_params(filter_kwargs)
-        return self._conn.get(self.objects_url, accept=accept,
+        return self._conn.get(self.objects_url, headers={"Accept": accept},
                               params=query_params)
 
     def get_object(self, obj_id, version=None, accept=MEDIA_TYPE_STIX_V20):
@@ -475,7 +479,7 @@ class Collection(_TAXIIEndpoint):
         query_params = None
         if version:
             query_params = _filter_kwargs_to_query_params({"version": version})
-        return self._conn.get(url, accept=accept,
+        return self._conn.get(url, headers={"Accept": accept},
                               params=query_params)
 
     def add_objects(self, bundle, wait_for_completion=True, poll_interval=1,
@@ -552,7 +556,7 @@ class Collection(_TAXIIEndpoint):
         self._verify_can_read()
         query_params = _filter_kwargs_to_query_params(filter_kwargs)
         return self._conn.get(self.url + "manifest/",
-                              accept=accept,
+                              headers={"Accept": accept},
                               params=query_params)
 
 
@@ -651,7 +655,8 @@ class ApiRoot(_TAXIIEndpoint):
 
         This invokes the ``Get API Root Information`` endpoint.
         """
-        response = self.__raw = self._conn.get(self.url, accept=accept)
+        response = self.__raw = self._conn.get(self.url,
+                                               headers={"Accept": accept})
 
         self._title = response.get("title")  # required
         self._description = response.get("description")  # optional
@@ -667,7 +672,7 @@ class ApiRoot(_TAXIIEndpoint):
         This invokes the ``Get Collections`` endpoint.
         """
         url = self.url + "collections/"
-        response = self._conn.get(url, accept=accept)
+        response = self._conn.get(url, headers={"Accept": accept})
 
         self._collections = []
         for item in response.get("collections", []):  # optional
@@ -680,7 +685,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)
+        response = self._conn.get(status_url, headers={"Accept": accept})
         return Status(status_url, conn=self._conn, status_info=response)
 
 
@@ -763,8 +768,7 @@ class Server(_TAXIIEndpoint):
 
     def refresh(self):
         """Update the Server information and list of API Roots"""
-        response = self.__raw = self._conn.get(self.url,
-                                               accept=MEDIA_TYPE_TAXII_V20)
+        response = self.__raw = self._conn.get(self.url)
 
         self._title = response.get("title")  # required
         self._description = response.get("description")  # optional
@@ -801,17 +805,22 @@ class _HTTPConnection(object):
 
     """
 
-    def __init__(self, user=None, password=None, verify=True, proxies=None):
+    def __init__(self, user=None, password=None, verify=True, proxies=None,
+                 user_agent=DEFAULT_USER_AGENT):
         """Create a connection session.
 
         Args:
             user (str): username for authentication (optional)
             password (str): password for authentication (optional)
             verify (bool): validate the entity credentials. (default: True)
-
+            user_agent (str): A value to use for the User-Agent header in
+                requests.  If not given, use a default value which represents
+                this library.
         """
         self.session = requests.Session()
         self.session.verify = verify
+        # enforce that we always have a connection-default user agent.
+        self.user_agent = user_agent or DEFAULT_USER_AGENT
         if user and password:
             self.session.auth = requests.auth.HTTPBasicAuth(user, password)
         if proxies:
@@ -834,22 +843,27 @@ class _HTTPConnection(object):
              content_type_tokens[0] == 'application/vnd.oasis.stix+json')
         )
 
-    def get(self, url, accept, params=None):
+    def get(self, url, headers=None, params=None):
         """Perform an HTTP GET, using the saved requests.Session and auth info.
+        If "Accept" isn't one of the given headers, a default TAXII mime type is
+        used.  Regardless, the response type is checked against the accept
+        header value, and an exception is raised if they don't match.
 
         Args:
             url (str): URL to retrieve
-            accept (str): media type to include in the ``Accept:`` header. This
-                function checks that the ``Content-Type:`` header on the HTTP
-                response matches this media type.
+            headers (dict): Any other headers to be added to the request.
             params: dictionary or bytes to be sent in the query string for the
                 request. (optional)
 
         """
-        headers = {
-            "Accept": accept
-        }
-        resp = self.session.get(url, headers=headers, params=params)
+
+        merged_headers = self._merge_headers(headers)
+
+        if "Accept" not in merged_headers:
+            merged_headers["Accept"] = MEDIA_TYPE_TAXII_V20
+        accept = merged_headers["Accept"]
+
+        resp = self.session.get(url, headers=merged_headers, params=params)
 
         resp.raise_for_status()
 
@@ -866,8 +880,20 @@ class _HTTPConnection(object):
         URL query parameters, and the given JSON in the request body.  The
         extra query parameters are merged with any which already exist in the
         URL.
+
+        Args:
+            url (str): URL to retrieve
+            headers (dict): Any other headers to be added to the request.
+            params: dictionary or bytes to be sent in the query string for the
+                request. (optional)
+            data: data to post as dictionary, list of tuples, bytes, or
+                file-like object
         """
-        resp = self.session.post(url, headers=headers, params=params, data=data)
+
+        merged_headers = self._merge_headers(headers)
+
+        resp = self.session.post(url, headers=merged_headers, params=params,
+                                 data=data)
         resp.raise_for_status()
         return _to_json(resp)
 
@@ -875,6 +901,42 @@ class _HTTPConnection(object):
         """Closes connections.  This object is no longer usable."""
         self.session.close()
 
+    def _merge_headers(self, call_specific_headers):
+        """
+        Merge headers from different sources together.  Headers passed to the
+        post/get methods have highest priority, then headers associated with
+        the connection object itself have next priority.
+
+        :param call_specific_headers: A header dict from the get/post call, or
+            None (the default for those methods).
+        :return: A key-case-insensitive MutableMapping object which contains
+            the merged headers.  (This doesn't actually return a dict.)
+        """
+
+        # A case-insensitive mapping is necessary here so that there is
+        # predictable behavior.  If a plain dict were used, you'd get keys in
+        # the merged dict which differ only in case.  The requests library
+        # would merge them internally, and it would be unpredictable which key
+        # is chosen for the final set of headers.  Another possible approach
+        # would be to upper/lower-case everything, but this seemed easier.  On
+        # the other hand, I don't know if CaseInsensitiveDict is public API...?
+
+        # First establish defaults
+        merged_headers = requests.structures.CaseInsensitiveDict({
+            "User-Agent": self.user_agent
+        })
+
+        # Then overlay with specifics from post/get methods
+        if call_specific_headers:
+            merged_headers.update(call_specific_headers)
+
+        # Special "User-Agent" header check, to ensure one is always sent.
+        # The call-specific overlay could have null'd out that header.
+        if not merged_headers.get("User-Agent"):
+            merged_headers["User-Agent"] = self.user_agent
+
+        return merged_headers
+
 
 def _to_json(resp):
     """


Test patch:
diff --git a/taxii2client/test/test_client.py b/taxii2client/test/test_client.py
index 958dc6f..d80b2a3 100644
--- a/taxii2client/test/test_client.py
+++ b/taxii2client/test/test_client.py
@@ -6,10 +6,10 @@ import responses
 import six
 
 from taxii2client import (
-    MEDIA_TYPE_STIX_V20, MEDIA_TYPE_TAXII_V20, AccessError, ApiRoot,
-    Collection, InvalidArgumentsError, InvalidJSONError, Server, Status,
-    TAXIIServiceException, ValidationError, _filter_kwargs_to_query_params,
-    _HTTPConnection, _TAXIIEndpoint
+    DEFAULT_USER_AGENT, MEDIA_TYPE_STIX_V20, MEDIA_TYPE_TAXII_V20, AccessError,
+    ApiRoot, Collection, InvalidArgumentsError, InvalidJSONError, Server,
+    Status, TAXIIServiceException, ValidationError,
+    _filter_kwargs_to_query_params, _HTTPConnection, _TAXIIEndpoint
 )
 
 TAXII_SERVER = "example.com"
@@ -697,7 +697,8 @@ def test_valid_content_type_for_connection():
                   content_type=MEDIA_TYPE_TAXII_V20 + "; charset=utf-8")
 
     conn = _HTTPConnection(user="foo", password="bar", verify=False)
-    conn.get("https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/", MEDIA_TYPE_TAXII_V20, None)
+    conn.get("https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/",
+             headers={"Accept": MEDIA_TYPE_TAXII_V20})
 
 
 @responses.activate
@@ -708,7 +709,8 @@ def test_invalid_content_type_for_connection():
 
     with pytest.raises(TAXIIServiceException) as excinfo:
         conn = _HTTPConnection(user="foo", password="bar", verify=False)
-        conn.get("https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/", MEDIA_TYPE_TAXII_V20 + "; charset=utf-8", None)
+        conn.get("https://example.com/api1/collections/91a7b528-80eb-42ed-a74d-c6fbd5a26116/",
+                 headers={"Accept": MEDIA_TYPE_TAXII_V20 + "; charset=utf-8"})
 
     assert ("Unexpected Response. Got Content-Type: 'application/vnd.oasis.taxii+json; "
             "version=2.0' for Accept: 'application/vnd.oasis.taxii+json; version=2.0; "
@@ -813,3 +815,58 @@ def test_collection_missing_can_write_property(collection_dict):
                    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)
+
+
+def test_user_agent_defaulting():
+    conn = _HTTPConnection(user_agent="foo/1.0")
+    headers = conn._merge_headers({})
+
+    # also test key access is case-insensitive
+    assert headers["user-agent"] == "foo/1.0"
+
+
+def test_user_agent_overriding():
+    conn = _HTTPConnection(user_agent="foo/1.0")
+    headers = conn._merge_headers({"User-Agent": "bar/2.0"})
+
+    assert headers["user-agent"] == "bar/2.0"
+
+
+def test_user_agent_enforcing1():
+    conn = _HTTPConnection(user_agent=None)
+    headers = conn._merge_headers({})
+
+    assert headers["user-agent"] == DEFAULT_USER_AGENT
+
+
+def test_user_agent_enforcing2():
+    conn = _HTTPConnection()
+    headers = conn._merge_headers({"User-Agent": None})
+
+    assert headers["user-agent"] == DEFAULT_USER_AGENT
+
+
+def test_user_agent_enforcing3():
+    conn = _HTTPConnection(user_agent=None)
+    headers = conn._merge_headers({"User-Agent": None})
+
+    assert headers["user-agent"] == DEFAULT_USER_AGENT
+
+
+def test_header_merging():
+    conn = _HTTPConnection()
+    headers = conn._merge_headers({"AddedHeader": "addedvalue"})
+
+    assert headers == {
+        "user-agent": DEFAULT_USER_AGENT,
+        "addedheader": "addedvalue"
+    }
+
+
+def test_header_merge_none():
+    conn = _HTTPConnection()
+    headers = conn._merge_headers(None)
+
+    assert headers == {
+        "user-agent": DEFAULT_USER_AGENT
+    }