Added the option of always adding a particular header to a request. Also fixed unicode xml decoding error

Author: abroglesc

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",
@@ -478,6 +478,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]:
@@ -671,7 +672,7 @@ def get(self, path_segment, owner=None, app=None, headers=None, sharing=None, **
         path = self.authority + self._abspath(path_segment, owner=owner,
                                               app=app, sharing=sharing)
         logging.debug("GET 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.get(path, all_headers, **query)
         return response
 
@@ -744,7 +745,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
 
@@ -810,7 +811,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,
@@ -864,6 +865,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()

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'))