Merge branch 'feature/cookie-auth' into bug/namespace-for-server-endpoints

Author: Shakeel Mohamed

splunklib/binding.py

@@ -31,6 +31,7 @@
 import urllib
 import io
 import sys
+import Cookie
 
 from datetime import datetime
 from functools import wraps
@@ -67,6 +68,45 @@ def new_f(*args, **kwargs):
     return new_f
 
 
+def _parse_cookies(cookie_str, dictionary):
+    """Tries to parse any key-value pairs of cookies in a string,
+    then updates the the dictionary with any key-value pairs found.
+
+    **Example**::
+        dictionary = {}
+        _parse_cookies('my=value', dictionary)
+        # Now the following is True
+        dictionary['my'] == 'value'
+
+    :param cookie_str: A string containing "key=value" pairs from an HTTP "Set-Cookie" header.
+    :type cookie_str: ``str``
+    :param dictionary: A dictionary to update with any found key-value pairs.
+    :type dictionary: ``dict``
+    """
+    parsed_cookie = Cookie.SimpleCookie(cookie_str)
+    for cookie in parsed_cookie.values():
+        dictionary[cookie.key] = cookie.coded_value
+
+
+def _make_cookie_header(cookies):
+    """
+    Takes a list of 2-tuples of key-value pairs of
+    cookies, and returns a valid HTTP ``Cookie``
+    header.
+
+    **Example**::
+
+        header = _make_cookie_header([("key", "value"), ("key_2", "value_2")])
+        # Now the following is True
+        header == "key=value; key_2=value_2"
+
+    :param cookies: A list of 2-tuples of cookie key-value pairs.
+    :type cookies: ``list`` of 2-tuples
+    :return: ``str` An HTTP header cookie string.
+    :rtype: ``str``
+    """
+    return "; ".join("%s=%s" % (key, value) for key, value in cookies)
+
 # Singleton values to eschew None
 class _NoAuthenticationToken(object):
     """The value stored in a :class:`Context` or :class:`splunklib.client.Service`
@@ -225,7 +265,7 @@ def f():
     @wraps(request_fun)
     def wrapper(self, *args, **kwargs):
         if self.token is _NoAuthenticationToken and \
-                self.cookie is _NoAuthenticationToken:
+                not self.has_cookies():
             # Not yet logged in.
             if self.autologin and self.username and self.password:
                 # This will throw an uncaught
@@ -428,9 +468,27 @@ def __init__(self, handler=None, **kwargs):
         self.username = kwargs.get("username", "")
         self.password = kwargs.get("password", "")
         self.autologin = kwargs.get("autologin", False)
-        self.cookie = kwargs.get("cookie", _NoAuthenticationToken)
-        if self.cookie is None: # In case someone explicitly passes cookie=None
-            self.cookie = _NoAuthenticationToken
+
+        # Store any cookies in the self.http._cookies dict
+        if kwargs.has_key("cookie") and kwargs['cookie'] not in [None, _NoAuthenticationToken]:
+            _parse_cookies(kwargs["cookie"], self.http._cookies)
+
+    def get_cookies(self):
+        """Gets the dictionary of cookies from the ``HttpLib`` member of this instance.
+
+        :return: Dictionary of cookies stored on the ``self.http``.
+        :rtype: ``dict``
+        """
+        return self.http._cookies
+
+    def has_cookies(self):
+        """Returns true if the ``HttpLib` member of this instance has at least
+        one cookie stored.
+
+        :return: ``True`` if there is at least one cookie, else ``False``
+        :rtype: ``bool``
+        """
+        return len(self.get_cookies()) > 0
 
     # Shared per-context request headers
     @property
@@ -443,8 +501,8 @@ def _auth_headers(self):
 
         :returns: A list of 2-tuples containing key and value
         """
-        if self.cookie is not _NoAuthenticationToken:
-            return [("cookie", self.cookie)]
+        if self.has_cookies():
+            return [("Cookie", _make_cookie_header(self.get_cookies().items()))]
         elif self.token is _NoAuthenticationToken:
             return []
         else:
@@ -761,10 +819,10 @@ def login(self):
             c = binding.Context(...).login()
             # Then issue requests...
         """
-        # If self.cookie and self.token only, use the cookie
-        if self.cookie is not _NoAuthenticationToken and \
+
+        if self.has_cookies() and \
                 (not self.username and not self.password):
-            # If we were passed a session cookie, but no username or
+            # If we were passed session cookie(s), but no username or
             # password, then login is a nop, since we're automatically
             # logged in.
             return
@@ -784,12 +842,6 @@ def login(self):
                 password=self.password,
                 cookie="1") # In Splunk 6.2+, passing "cookie=1" will return the "set-cookie" header
 
-            # Store the cookie
-            for key, value in response.headers:
-                if key.lower() == "set-cookie":
-                    self.cookie = value
-                    break
-
             body = response.body.read()
             session = XML(body).findtext("./sessionKey")
             self.token = "Splunk %s" % session
@@ -801,8 +853,9 @@ def login(self):
                 raise
 
     def logout(self):
-        """Forgets the current session token."""
+        """Forgets the current session token, and cookies."""
         self.token = _NoAuthenticationToken
+        self.http._cookies = {}
         return self
 
     def _abspath(self, path_segment,
@@ -889,6 +942,9 @@ def connect(**kwargs):
     :param token: The current session token (optional). Session tokens can be
         shared across multiple service instances.
     :type token: ``string``
+    :param cookie: A session cookie. When provided, you don't need to call :meth:`login`.
+        This parameter is only supported for Splunk 6.2+.
+    :type cookie: ``string``
     :param username: The Splunk account username, which is used to
         authenticate the Splunk instance.
     :type username: ``string``
@@ -1030,6 +1086,7 @@ class HttpLib(object):
     """
     def __init__(self, custom_handler=None):
         self.handler = handler() if custom_handler is None else custom_handler
+        self._cookies = {}
 
     def delete(self, url, headers=None, **kwargs):
         """Sends a DELETE request to a URL.
@@ -1141,10 +1198,16 @@ def request(self, url, message, **kwargs):
             raise HTTPError(response)
 
         # Update the cookie with any HTTP request
-        for key, value in response.headers:
+        # Initially, assume list of 2-tuples
+        key_value_tuples = response.headers
+        # If response.headers is a dict, get the key-value pairs as 2-tuples
+        # this is the case when using urllib2
+        if isinstance(response.headers, dict):
+            key_value_tuples = response.headers.items()
+        for key, value in key_value_tuples:
             if key.lower() == "set-cookie":
-                self.cookie = value
-                break
+                _parse_cookies(value, self._cookies)
+
         return response
 
 
@@ -1256,7 +1319,6 @@ def request(url, message, **kwargs):
             "Host": host,
             "User-Agent": "splunk-sdk-python/0.1",
             "Accept": "*/*",
-            "Cookie": "1"
         } # defaults
         for key, value in message["headers"]:
             head[key] = value

splunklib/client.py

@@ -67,7 +67,7 @@
 import socket
 import contextlib
 
-from binding import Context, HTTPError, AuthenticationError, namespace, UrlEncoded, _encode, _NoAuthenticationToken
+from binding import Context, HTTPError, AuthenticationError, namespace, UrlEncoded, _encode, _make_cookie_header
 from data import record
 import data
 
@@ -262,7 +262,6 @@ def _parse_atom_metadata(content):
 
     return record({'access': access, 'fields': fields})
 
-
 # kwargs: scheme, host, port, app, owner, username, password
 def connect(**kwargs):
     """This function connects and logs in to a Splunk instance.
@@ -285,6 +284,9 @@ def connect(**kwargs):
     :param `token`: The current session token (optional). Session tokens can be
                     shared across multiple service instances.
     :type token: ``string``
+    :param cookie: A session cookie. When provided, you don't need to call :meth:`login`.
+        This parameter is only supported for Splunk 6.2+.
+    :type cookie: ``string``
     :param autologin: When ``True``, automatically tries to log in again if the
         session terminates.
     :type autologin: ``boolean``
@@ -302,7 +304,9 @@ def connect(**kwargs):
         a = s.apps["my_app"]
         ...
     """
-    return Service(**kwargs).login()
+    s = Service(**kwargs)
+    s.login()
+    return s
 
 
 # In preparation for adding Storm support, we added an
@@ -345,6 +349,9 @@ class Service(_BaseService):
     :param `token`: The current session token (optional). Session tokens can be
                     shared across multiple service instances.
     :type token: ``string``
+    :param cookie: A session cookie. When provided, you don't need to call :meth:`login`.
+        This parameter is only supported for Splunk 6.2+.
+    :type cookie: ``string``
     :param `username`: The Splunk account username, which is used to
                        authenticate the Splunk instance.
     :type username: ``string``
@@ -362,6 +369,8 @@ class Service(_BaseService):
         s = client.connect(username="boris", password="natasha")
         # Or if you already have a session token
         s = client.Service(token="atg232342aa34324a")
+        # Or if you already have a valid cookie
+        s = client.Service(cookie="splunkd_8089=...")
     """
     def __init__(self, **kwargs):
         super(Service, self).__init__(**kwargs)
@@ -1902,9 +1911,9 @@ def attach(self, host=None, source=None, sourcetype=None):
 
         cookie_or_auth_header = "Authorization: %s\r\n" % self.service.token
 
-        # If we have a cookie, use it instead of "Authorization: ..."
-        if self.service.cookie is not _NoAuthenticationToken:
-            cookie_or_auth_header = "Cookie: %s\r\n" % self.service.cookie
+        # If we have cookie(s), use them instead of "Authorization: ..."
+        if self.service.has_cookies():
+            cookie_or_auth_header = "Cookie: %s\r\n" % _make_cookie_header(self.service.get_cookies().items())
 
         # Since we need to stream to the index connection, we have to keep
         # the connection open and use the Splunk extension headers to note

tests/test_binding.py

@@ -26,6 +26,7 @@
 import socket
 import sys
 import ssl
+import Cookie
 
 import splunklib.binding as binding
 from splunklib.binding import HTTPError, AuthenticationError, UrlEncoded
@@ -476,6 +477,8 @@ def test_logout(self):
         response = self.context.get("/services")
         self.assertEqual(response.status, 200)
         self.context.logout()
+        self.assertEqual(self.context.token, binding._NoAuthenticationToken)
+        self.assertEqual(self.context.get_cookies(), {})
         self.assertRaises(AuthenticationError,
                           self.context.get, "/services")
         self.assertRaises(AuthenticationError,
@@ -503,26 +506,82 @@ def test_cookie_in_auth_headers(self):
         self.assertNotEqual(self.context._auth_headers, [])
         self.assertEqual(len(self.context._auth_headers), 1)
         self.assertEqual(len(self.context._auth_headers), 1)
-        self.assertEqual(self.context._auth_headers[0][0], "cookie")
+        self.assertEqual(self.context._auth_headers[0][0], "Cookie")
         self.assertEqual(self.context._auth_headers[0][1][:8], "splunkd_")
 
     def test_got_cookie_on_connect(self):
-        self.assertIsNotNone(self.context.cookie)
-        self.assertNotEqual(self.context.cookie, binding._NoAuthenticationToken)
-        self.assertEqual(self.context.cookie[:8], "splunkd_")
+        self.assertIsNotNone(self.context.get_cookies())
+        self.assertNotEqual(self.context.get_cookies(), {})
+        self.assertEqual(len(self.context.get_cookies()), 1)
+        self.assertEqual(self.context.get_cookies().keys()[0][:8], "splunkd_")
+
+    def test_cookie_with_autologin(self):
+        self.context.autologin = True
+        self.assertEqual(self.context.get("/services").status, 200)
+        self.assertTrue(self.context.has_cookies())
+        self.context.logout()
+        self.assertFalse(self.context.has_cookies())
+        self.assertEqual(self.context.get("/services").status, 200)
+        self.assertTrue(self.context.has_cookies())
+
+    def test_cookie_without_autologin(self):
+        self.context.autologin = False
+        self.assertEqual(self.context.get("/services").status, 200)
+        self.assertTrue(self.context.has_cookies())
+        self.context.logout()
+        self.assertFalse(self.context.has_cookies())
+        self.assertRaises(AuthenticationError,
+                          self.context.get, "/services")
 
     def test_got_updated_cookie_with_get(self):
-        old_cookie = self.context.cookie
+        old_cookies = self.context.get_cookies()
         resp = self.context.get("apps/local")
         found = False
         for key, value in resp.headers:
             if key.lower() == "set-cookie":
                 found = True
                 self.assertEqual(value[:8], "splunkd_")
-                # It's unlikely that the cookie will change during this short test
-                self.assertEqual(value, old_cookie)
+
+                new_cookies = {}
+                binding._parse_cookies(value, new_cookies)
+                # We're only expecting 1 in this scenario
+                self.assertEqual(len(old_cookies), 1)
+                self.assertTrue(len(new_cookies.values()), 1)
+                self.assertEqual(old_cookies, new_cookies)
+                self.assertEqual(new_cookies.values()[0], old_cookies.values()[0])
         self.assertTrue(found)
 
+    def test_login_fails_with_bad_cookie(self):
+        new_context = binding.connect(**{"cookie": "bad=cookie"})
+        # We should get an error if using a bad cookie
+        try:
+            new_context.get("apps/local")
+            self.fail()
+        except AuthenticationError as ae:
+            self.assertEqual(ae.message, "Request failed: Session is not logged in.")
+
+    def test_login_with_multiple_cookies(self):
+        bad_cookie = 'bad=cookie'
+        new_context = binding.connect(**{"cookie": bad_cookie})
+        # We should get an error if using a bad cookie
+        try:
+            new_context.get("apps/local")
+            self.fail()
+        except AuthenticationError as ae:
+            self.assertEqual(ae.message, "Request failed: Session is not logged in.")
+            # Bring in a valid cookie now
+            for key, value in self.context.get_cookies().items():
+                new_context.get_cookies()[key] = value
+
+            self.assertEqual(len(new_context.get_cookies()), 2)
+            self.assertTrue('bad' in new_context.get_cookies().keys())
+            self.assertTrue('cookie' in new_context.get_cookies().values())
+
+            for k, v in self.context.get_cookies().items():
+                self.assertEqual(new_context.get_cookies()[k], v)
+
+            self.assertEqual(new_context.get("apps/local").status, 200)
+
     def test_login_fails_without_cookie_or_token(self):
         opts = {
             'host': self.opts.kwargs['host'],