Refactor cookie authentication functionality + tests

Author: Shakeel Mohamed

splunklib/binding.py

@@ -68,6 +68,47 @@ 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**::
+
+        parse_cookies('my=value', {})
+        # Now the following is True
+        dictionary['my'] == 'value'
+
+    :param cookie_str: A string containing "key=value" pairs from an HTTP "Set-Cookie" header
+    :type msg: ``str``
+    :param dictionary: A dictionary to update with any found key-value pairs
+    :type msg: ``dict``
+    """
+    parsed_cookie = Cookie.SimpleCookie(cookie_str)
+    for cookie in parsed_cookie.values():
+        dictionary[cookie.key] = cookie.coded_value
+
+def make_cookie_header(cookies):
+    # TODO: verify
+    """
+    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:
+    :return: ``str` An HTTP header cookie string.
+    :rtype: ``str``
+    """
+    header = ""
+    for key, value in cookies:
+        header += "%s=%s; " % (key, value)
+    return header
+
 # Singleton values to eschew None
 class _NoAuthenticationToken(object):
     """The value stored in a :class:`Context` or :class:`splunklib.client.Service`
@@ -226,7 +267,7 @@ def f():
     @wraps(request_fun)
     def wrapper(self, *args, **kwargs):
         if self.token is _NoAuthenticationToken and \
-                len(self.cookies) < 1:
+                len(self.http.cookies) < 1:
             # Not yet logged in.
             if self.autologin and self.username and self.password:
                 # This will throw an uncaught
@@ -430,12 +471,9 @@ def __init__(self, handler=None, **kwargs):
         self.password = kwargs.get("password", "")
         self.autologin = kwargs.get("autologin", False)
 
-        # By default, there are no cookies
-        self.cookies = {}
+        # Store any cookies in the self.http.cookies dict
         if kwargs.has_key("cookie") and kwargs['cookie'] not in [None, _NoAuthenticationToken]:
-            parsed_cookie = Cookie.SimpleCookie(kwargs.get(("cookie")))
-            for cookie in parsed_cookie.values():
-                self.cookies[cookie.key] = cookie.coded_value
+            parse_cookies(kwargs["cookie"], self.http.cookies)
 
     # Shared per-context request headers
     @property
@@ -448,8 +486,8 @@ def _auth_headers(self):
 
         :returns: A list of 2-tuples containing key and value
         """
-        if len(self.cookies) > 0:
-            return [("cookie", "%s=%s" % cookie) for cookie in self.cookies.items()]
+        if len(self.http.cookies) > 0:
+            return [("Cookie", make_cookie_header(self.http.cookies.items()))]
         elif self.token is _NoAuthenticationToken:
             return []
         else:
@@ -766,8 +804,8 @@ def login(self):
             c = binding.Context(...).login()
             # Then issue requests...
         """
-        # If self.cookies and self.token only, use the cookie
-        if len(self.cookies) > 0 and \
+
+        if len(self.http.cookies) > 0 and \
                 (not self.username and not self.password):
             # If we were passed session cookie(s), but no username or
             # password, then login is a nop, since we're automatically
@@ -789,14 +827,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":
-                    parsed_cookies = Cookie.SimpleCookie(value)
-                    for cookie in parsed_cookies.values():
-                        self.cookies[cookie.key] = cookie.coded_value
-
             body = response.body.read()
             session = XML(body).findtext("./sessionKey")
             self.token = "Splunk %s" % session
@@ -810,7 +840,7 @@ def login(self):
     def logout(self):
         """Forgets the current session token, and cookies."""
         self.token = _NoAuthenticationToken
-        self.cookies = {}
+        self.http.cookies = {}
         return self
 
     def _abspath(self, path_segment,
@@ -1151,11 +1181,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":
-                parsed_cookie = Cookie.SimpleCookie(value)
-                for cookie in parsed_cookie.values():
-                    self.cookies[cookie.key] = cookie.coded_value
+                parse_cookies(value, self.cookies)
+
         return response
 
 

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
 
@@ -1903,12 +1903,8 @@ def attach(self, host=None, source=None, sourcetype=None):
         cookie_or_auth_header = "Authorization: %s\r\n" % self.service.token
 
         # If we have cookie(s), use them instead of "Authorization: ..."
-        if len(self.service.cookies) > 0:
-            cookies = []
-            for cookie in self.service.cookies.items():
-                cookies.append("Cookie: %s=%s\r\n" % cookie)
-            if len(cookies) > 0:
-                cookie_or_auth_header = "".join(cookies)
+        if len(self.service.http.cookies) > 0:
+            cookie_or_auth_header = "Cookie: %s\r\n" % make_cookie_header(self.service.http.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

@@ -478,7 +478,7 @@ def test_logout(self):
         self.assertEqual(response.status, 200)
         self.context.logout()
         self.assertEqual(self.context.token, binding._NoAuthenticationToken)
-        self.assertEqual(self.context.cookies, {})
+        self.assertEqual(self.context.http.cookies, {})
         self.assertRaises(AuthenticationError,
                           self.context.get, "/services")
         self.assertRaises(AuthenticationError,
@@ -506,32 +506,31 @@ 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.cookies)
-        self.assertNotEqual(self.context.cookies, {})
-        self.assertEqual(len(self.context.cookies), 1)
-        self.assertEqual(self.context.cookies.keys()[0][:8], "splunkd_")
+        self.assertIsNotNone(self.context.http.cookies)
+        self.assertNotEqual(self.context.http.cookies, {})
+        self.assertEqual(len(self.context.http.cookies), 1)
+        self.assertEqual(self.context.http.cookies.keys()[0][:8], "splunkd_")
 
     def test_got_updated_cookie_with_get(self):
-        old_cookies = self.context.cookies
+        old_cookies = self.context.http.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_")
 
-                parsed_cookies = Cookie.SimpleCookie(value)
+                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(parsed_cookies.values()), 1)
-                parsed_cookie = parsed_cookies.values()[0]
-
-                self.assertEqual(parsed_cookie.key, old_cookies.keys()[0])
-                self.assertEqual(parsed_cookie.coded_value, old_cookies.values()[0])
+                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):
@@ -553,15 +552,15 @@ def test_login_with_multiple_cookies(self):
         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.cookies.items():
-                new_context.cookies[key] = value
+            for key, value in self.context.http.cookies.items():
+                new_context.http.cookies[key] = value
 
-            self.assertEqual(len(new_context.cookies), 2)
-            self.assertTrue('bad' in new_context.cookies.keys())
-            self.assertTrue('cookie' in new_context.cookies.values())
+            self.assertEqual(len(new_context.http.cookies), 2)
+            self.assertTrue('bad' in new_context.http.cookies.keys())
+            self.assertTrue('cookie' in new_context.http.cookies.values())
 
-            for k, v in self.context.cookies.items():
-                self.assertEqual(new_context.cookies[k], v)
+            for k, v in self.context.http.cookies.items():
+                self.assertEqual(new_context.http.cookies[k], v)
 
             self.assertEqual(new_context.get("apps/local").status, 200)
 

tests/test_index.py

@@ -110,7 +110,7 @@ def test_submit_via_attached_socket(self):
     def test_submit_via_attach_with_cookie_header(self):
         event_count = int(self.service.indexes[self.index_name]['totalEventCount'])
 
-        cookie = "%s=%s" % (self.service.cookies.items()[0])
+        cookie = "%s=%s" % (self.service.http.cookies.items()[0])
         service = client.Service(**{"cookie": cookie})
         service.login()
         cn = service.indexes[self.index_name].attach()
@@ -121,28 +121,13 @@ def test_submit_via_attach_with_cookie_header(self):
     def test_submit_via_attach_with_multiple_cookie_headers(self):
         event_count = int(self.service.indexes[self.index_name]['totalEventCount'])
         service = client.Service(**{"cookie": 'a bad cookie'})
-        service.cookies.update(self.service.cookies)
+        service.http.cookies.update(self.service.http.cookies)
         service.login()
         cn = service.indexes[self.index_name].attach()
         cn.send("Hello Boris!\r\n")
         cn.close()
         self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=60)
 
-    def test_login_with_multiple_cookie_headers(self):
-        event_count = int(self.index['totalEventCount'])
-
-        cookies = {}
-        cookies['bad'] = 'cookie'
-        cookies['something_else'] = 'bad'
-        self.service.logout()
-        self.service.cookies.update(cookies)
-
-        self.service.login()
-        cn = self.service.indexes[self.index_name].attach()
-        cn.send("Hello Boris!\r\n")
-        cn.close()
-        self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=60)
-
     def test_upload(self):
         if not self.app_collection_installed():
             print "Test requires sdk-app-collection. Skipping."

tests/test_service.py

@@ -152,32 +152,32 @@ def setUp(self):
             self.skipTest("Skipping cookie-auth tests, running in %d.%d.%d, this feature was added in 6.2+" % splver)
 
     def test_login_and_store_cookie(self):
-        self.assertIsNotNone(self.service.cookies)
-        self.assertEquals(len(self.service.cookies), 0)
+        self.assertIsNotNone(self.service.http.cookies)
+        self.assertEquals(len(self.service.http.cookies), 0)
         self.service.login()
-        self.assertIsNotNone(self.service.cookies)
-        self.assertNotEquals(self.service.cookies, {})
-        self.assertEquals(len(self.service.cookies), 1)
+        self.assertIsNotNone(self.service.http.cookies)
+        self.assertNotEquals(self.service.http.cookies, {})
+        self.assertEquals(len(self.service.http.cookies), 1)
 
     def test_login_with_cookie(self):
         self.service.login()
-        self.assertIsNotNone(self.service.cookies)
+        self.assertIsNotNone(self.service.http.cookies)
         # Use the cookie from the other service as the only auth param (don't need user/password)
-        service2 = client.Service(**{"cookie": "%s=%s" % self.service.cookies.items()[0]})
+        service2 = client.Service(**{"cookie": "%s=%s" % self.service.http.cookies.items()[0]})
         service2.login()
-        self.assertEqual(len(service2.cookies), 1)
-        self.assertEqual(service2.cookies, self.service.cookies)
-        self.assertEqual(len(service2.cookies), len(self.service.cookies))
-        self.assertEqual(service2.cookies.keys()[0][:8], "splunkd_")
+        self.assertEqual(len(service2.http.cookies), 1)
+        self.assertEqual(service2.http.cookies, self.service.http.cookies)
+        self.assertEqual(len(service2.http.cookies), len(self.service.http.cookies))
+        self.assertEqual(service2.http.cookies.keys()[0][:8], "splunkd_")
         self.assertEqual(service2.apps.get().status, 200)
 
     def test_login_fails_with_bad_cookie(self):
         bad_cookie = {'bad': 'cookie'}
         service2 = client.Service()
-        self.assertEquals(len(service2.cookies), 0)
-        service2.cookies.update(bad_cookie)
+        self.assertEquals(len(service2.http.cookies), 0)
+        service2.http.cookies.update(bad_cookie)
         service2.login()
-        self.assertEquals(service2.cookies, {'bad': 'cookie'})
+        self.assertEquals(service2.http.cookies, {'bad': 'cookie'})
 
         # Should get an error with a bad cookie
         try:
@@ -188,7 +188,7 @@ def test_login_fails_with_bad_cookie(self):
 
     def test_login_fails_with_no_cookie(self):
         service2 = client.Service()
-        self.assertEquals(len(service2.cookies), 0)
+        self.assertEquals(len(service2.http.cookies), 0)
 
         # Should get an error when no authentication method
         try:
@@ -197,10 +197,21 @@ def test_login_fails_with_no_cookie(self):
         except AuthenticationError as ae:
             self.assertEqual(ae.message, "Login failed.")
 
+    def test_login_with_multiple_cookie_headers(self):
+        cookies = {
+            'bad': 'cookie',
+            'something_else': 'bad'
+        }
+        self.service.logout()
+        self.service.http.cookies.update(cookies)
+
+        self.service.login()
+        self.assertEqual(self.service.apps.get().status, 200)
+
     def test_login_with_multiple_cookies(self):
         bad_cookie = 'bad=cookie'
         self.service.login()
-        self.assertIsNotNone(self.service.cookies)
+        self.assertIsNotNone(self.service.http.cookies)
 
         service2 = client.Service(**{"cookie": bad_cookie})
         service2.login()
@@ -213,16 +224,16 @@ def test_login_with_multiple_cookies(self):
             self.assertEqual(ae.message, "Request failed: Session is not logged in.")
 
             # Add on valid cookies, and try to use all of them
-            service2.cookies.update(self.service.cookies)
-
-            self.assertEqual(len(service2.cookies), 2)
-            self.service.cookies.update({'bad': 'cookie'})
-            self.assertEqual(service2.cookies, self.service.cookies)
-            self.assertEqual(len(service2.cookies), 2)
-            self.assertEqual(service2.cookies.keys()[1][:8], "splunkd_")
-            self.assertTrue('bad' in service2.cookies.keys())
-            self.assertEqual(service2.cookies['bad'], 'cookie')
-            self.assertEqual(self.service.cookies.items(), service2.cookies.items())
+            service2.http.cookies.update(self.service.http.cookies)
+
+            self.assertEqual(len(service2.http.cookies), 2)
+            self.service.http.cookies.update({'bad': 'cookie'})
+            self.assertEqual(service2.http.cookies, self.service.http.cookies)
+            self.assertEqual(len(service2.http.cookies), 2)
+            self.assertEqual(service2.http.cookies.keys()[1][:8], "splunkd_")
+            self.assertTrue('bad' in service2.http.cookies.keys())
+            self.assertEqual(service2.http.cookies['bad'], 'cookie')
+            self.assertEqual(self.service.http.cookies.items(), service2.http.cookies.items())
             service2.login()
             self.assertEqual(service2.apps.get().status, 200)