Merge pull request splunk#209 from abroglesc/additional-headers-get-method

shakeelmohamed · web-flow · commit 711b8953d8eb · 2018-06-28T10:28:02.000-07:00
Additional headers on every request and fix unicode decoding issue
diff --git a/splunklib/binding.py b/splunklib/binding.py
@@ -25,30 +25,30 @@
 """
 
 from __future__ import absolute_import
+
+import io
 import logging
 import socket
 import ssl
-from io import BytesIO
-
-from splunklib.six.moves import urllib
-import io
 import sys
-
 from base64 import b64encode
+from contextlib import contextmanager
 from datetime import datetime
 from functools import wraps
+from io import BytesIO
+from xml.etree.ElementTree import XML
+
+from splunklib import six
 from splunklib.six import StringIO
+from splunklib.six.moves import urllib
 
-from contextlib import contextmanager
+from .data import record
 
-from xml.etree.ElementTree import XML
-from splunklib import six
 try:
     from xml.etree.ElementTree import ParseError
 except ImportError as e:
     from xml.parsers.expat import ExpatError as ParseError
 
-from .data import record
 
 __all__ = [
     "AuthenticationError",
@@ -449,6 +449,8 @@ class Context(object):
     :type username: ``string``
     :param password: The password for the Splunk account.
     :type password: ``string``
+    :param headers: List of extra HTTP headers to send (optional).
+    :type headers: ``list`` of 2-tuples.
     :param handler: The HTTP request handler (optional).
     :returns: A ``Context`` instance.
 
@@ -478,6 +480,7 @@ def __init__(self, handler=None, **kwargs):
         self.password = kwargs.get("password", "")
         self.basic = kwargs.get("basic", False)
         self.autologin = kwargs.get("autologin", False)
+        self.additional_headers = kwargs.get("headers", [])
 
         # Store any cookies in the self.http._cookies dict
         if "cookie" in kwargs and kwargs['cookie'] not in [None, _NoAuthenticationToken]:
@@ -613,7 +616,7 @@ def delete(self, path_segment, owner=None, app=None, sharing=None, **query):
 
     @_authentication
     @_log_duration
-    def get(self, path_segment, owner=None, app=None, sharing=None, **query):
+    def get(self, path_segment, owner=None, app=None, headers=None, sharing=None, **query):
         """Performs a GET operation from the REST path segment with the given
         namespace and query.
 
@@ -636,6 +639,8 @@ def get(self, path_segment, owner=None, app=None, sharing=None, **query):
         :type owner: ``string``
         :param app: The app context of the namespace (optional).
         :type app: ``string``
+        :param headers: List of extra HTTP headers to send (optional).
+        :type headers: ``list`` of 2-tuples.
         :param sharing: The sharing mode of the namespace (optional).
         :type sharing: ``string``
         :param query: All other keyword arguments, which are used as query
@@ -663,10 +668,14 @@ def get(self, path_segment, owner=None, app=None, sharing=None, **query):
             c.logout()
             c.get('apps/local') # raises AuthenticationError
         """
+        if headers is None:
+            headers = []
+
         path = self.authority + self._abspath(path_segment, owner=owner,
                                               app=app, sharing=sharing)
         logging.debug("GET request to %s (body: %s)", path, repr(query))
-        response = self.http.get(path, self._auth_headers, **query)
+        all_headers = headers + self.additional_headers + self._auth_headers
+        response = self.http.get(path, all_headers, **query)
         return response
 
     @_authentication
@@ -738,7 +747,7 @@ def post(self, path_segment, owner=None, app=None, sharing=None, headers=None, *
 
         path = self.authority + self._abspath(path_segment, owner=owner, app=app, sharing=sharing)
         logging.debug("POST request to %s (body: %s)", path, repr(query))
-        all_headers = headers + self._auth_headers
+        all_headers = headers + self.additional_headers + self._auth_headers
         response = self.http.post(path, all_headers, **query)
         return response
 
@@ -804,7 +813,7 @@ def request(self, path_segment, method="GET", headers=None, body="",
         path = self.authority \
             + self._abspath(path_segment, owner=owner,
                             app=app, sharing=sharing)
-        all_headers = headers + self._auth_headers
+        all_headers = headers + self.additional_headers + self._auth_headers
         logging.debug("%s request to %s (headers: %s, body: %s)",
                       method, path, str(all_headers), repr(body))
         response = self.http.request(path,
@@ -858,6 +867,7 @@ def login(self):
                 self.authority + self._abspath("/services/auth/login"),
                 username=self.username,
                 password=self.password,
+                headers=self.additional_headers,
                 cookie="1") # In Splunk 6.2+, passing "cookie=1" will return the "set-cookie" header
 
             body = response.body.read()
@@ -968,6 +978,8 @@ def connect(**kwargs):
     :type username: ``string``
     :param password: The password for the Splunk account.
     :type password: ``string``
+    :param headers: List of extra HTTP headers to send (optional).
+    :type headers: ``list`` of 2-tuples.
     :param autologin: When ``True``, automatically tries to log in again if the
         session terminates.
     :type autologin: ``Boolean``
@@ -1190,7 +1202,7 @@ def post(self, url, headers=None, **kwargs):
         # to support the receivers/stream endpoint.
         if 'body' in kwargs:
             # We only use application/x-www-form-urlencoded if there is no other
-            # Content-Type header present. This can happen in cases where we 
+            # Content-Type header present. This can happen in cases where we
             # send requests as application/json, e.g. for KV Store.
             if len([x for x in headers if x[0].lower() == "content-type"]) == 0:
                 headers.append(("Content-Type", "application/x-www-form-urlencoded"))
diff --git a/splunklib/client.py b/splunklib/client.py
@@ -58,19 +58,22 @@
     my_app.package()  # Creates a compressed package of this application
 """
 
+import contextlib
 import datetime
 import json
-from splunklib.six.moves import urllib
 import logging
-from time import sleep
-from datetime import datetime, timedelta
 import socket
-import contextlib
+from datetime import datetime, timedelta
+from time import sleep
 
 from splunklib import six
-from .binding import Context, HTTPError, AuthenticationError, namespace, UrlEncoded, _encode, _make_cookie_header, _NoAuthenticationToken
-from .data import record
+from splunklib.six.moves import urllib
+
 from . import data
+from .binding import (AuthenticationError, Context, HTTPError, UrlEncoded,
+                      _encode, _make_cookie_header, _NoAuthenticationToken,
+                      namespace)
+from .data import record
 
 __all__ = [
     "connect",
@@ -193,8 +196,11 @@ def _path(base, name):
 
 
 # Load an atom record from the body of the given response
+# this will ultimately be sent to an xml ElementTree so we
+# should use the xmlcharrefreplace option
 def _load_atom(response, match=None):
-    return data.load(response.body.read().decode('utf-8'), match)
+    return data.load(response.body.read()
+                     .decode('utf-8', 'xmlcharrefreplace'), match)
 
 
 # Load an array of atom entries from the body of the given response
@@ -557,7 +563,7 @@ def restart(self, timeout=None):
         # This message will be deleted once the server actually restarts.
         self.messages.create(name="restart_required", **msg)
         result = self.post("/services/server/control/restart")
-        if timeout is None: 
+        if timeout is None:
             return result
         start = datetime.now()
         diff = timedelta(seconds=timeout)
@@ -1631,7 +1637,7 @@ def get(self, name="", owner=None, app=None, sharing=None, **query):
                 and ``status``
 
         Example:
-        
+
         import splunklib.client
             s = client.service(...)
             saved_searches = s.saved_searches
@@ -1685,7 +1691,7 @@ def __getitem__(self, key):
         # The superclass implementation is designed for collections that contain
         # entities. This collection (Configurations) contains collections
         # (ConfigurationFile).
-        # 
+        #
         # The configurations endpoint returns multiple entities when we ask for a single file.
         # This screws up the default implementation of __getitem__ from Collection, which thinks
         # that multiple entities means a name collision, so we have to override it here.
@@ -1749,9 +1755,9 @@ class Stanza(Entity):
     """This class contains a single configuration stanza."""
 
     def submit(self, stanza):
-        """Adds keys to the current configuration stanza as a 
+        """Adds keys to the current configuration stanza as a
         dictionary of key-value pairs.
-        
+
         :param stanza: A dictionary of key-value pairs for the stanza.
         :type stanza: ``dict``
         :return: The :class:`Stanza` object.
@@ -1962,7 +1968,7 @@ def attach(self, host=None, source=None, sourcetype=None):
                    cookie_or_auth_header.encode('utf-8'),
                    b"X-Splunk-Input-Mode: Streaming\r\n",
                    b"\r\n"]
-        
+
         for h in headers:
             sock.write(h)
         return sock
@@ -3695,13 +3701,13 @@ def batch_find(self, *dbqueries):
 
         :param dbqueries: Array of individual queries as dictionaries
         :type dbqueries: ``array`` of ``dict``
-        
+
         :return: Results of each query
         :rtype: ``array`` of ``array``
         """
-        if len(dbqueries) < 1: 
+        if len(dbqueries) < 1:
             raise Exception('Must have at least one query.')
-        
+
         data = json.dumps(dbqueries)
 
         return json.loads(self._post('batch_find', headers=KVStoreCollectionData.JSON_HEADER, body=data).body.read().decode('utf-8'))
@@ -3712,13 +3718,13 @@ def batch_save(self, *documents):
 
         :param documents: Array of documents to save as dictionaries
         :type documents: ``array`` of ``dict``
-        
+
         :return: Results of update operation as overall stats
         :rtype: ``dict``
         """
-        if len(documents) < 1: 
+        if len(documents) < 1:
             raise Exception('Must have at least one document.')
-        
+
         data = json.dumps(documents)
 
         return json.loads(self._post('batch_save', headers=KVStoreCollectionData.JSON_HEADER, body=data).body.read().decode('utf-8'))
