Cosmetic refactor for cookie auth functionality

Author: Shakeel Mohamed

splunklib/binding.py

@@ -68,13 +68,13 @@ def new_f(*args, **kwargs):
     return new_f
 
 
-def parse_cookies(cookie_str, dictionary):
+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', {})
+        dictionary = {}
+        _parse_cookies('my=value', dictionary)
         # Now the following is True
         dictionary['my'] == 'value'
 
@@ -87,15 +87,16 @@ def parse_cookies(cookie_str, dictionary):
     for cookie in parsed_cookie.values():
         dictionary[cookie.key] = cookie.coded_value
 
-def make_cookie_header(cookies):
+
+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")])
+        header = _make_cookie_header([("key", "value"), ("key_2", "value_2")])
         # Now the following is True
         header == "key=value; key_2=value_2"
 
@@ -104,10 +105,7 @@ def make_cookie_header(cookies):
     :return: ``str` An HTTP header cookie string.
     :rtype: ``str``
     """
-    header = ""
-    for key, value in cookies:
-        header += "%s=%s; " % (key, value)
-    return header
+    return "; ".join("%s=%s" % (key, value) for key, value in cookies)
 
 # Singleton values to eschew None
 class _NoAuthenticationToken(object):
@@ -267,7 +265,7 @@ def f():
     @wraps(request_fun)
     def wrapper(self, *args, **kwargs):
         if self.token is _NoAuthenticationToken and \
-                len(self.http.cookies) < 1:
+                len(self.get_cookies()) < 1:
             # Not yet logged in.
             if self.autologin and self.username and self.password:
                 # This will throw an uncaught
@@ -471,9 +469,26 @@ def __init__(self, handler=None, **kwargs):
         self.password = kwargs.get("password", "")
         self.autologin = kwargs.get("autologin", False)
 
-        # Store any cookies in the self.http.cookies dict
+        # 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)
+            _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
@@ -486,8 +501,8 @@ def _auth_headers(self):
 
         :returns: A list of 2-tuples containing key and value
         """
-        if len(self.http.cookies) > 0:
-            return [("Cookie", make_cookie_header(self.http.cookies.items()))]
+        if self.has_cookies():
+            return [("Cookie", _make_cookie_header(self.get_cookies().items()))]
         elif self.token is _NoAuthenticationToken:
             return []
         else:
@@ -805,7 +820,7 @@ def login(self):
             # Then issue requests...
         """
 
-        if len(self.http.cookies) > 0 and \
+        if self.has_cookies() 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
@@ -840,7 +855,7 @@ def login(self):
     def logout(self):
         """Forgets the current session token, and cookies."""
         self.token = _NoAuthenticationToken
-        self.http.cookies = {}
+        self.http._cookies = {}
         return self
 
     def _abspath(self, path_segment,
@@ -1068,8 +1083,7 @@ class HttpLib(object):
     """
     def __init__(self, custom_handler=None):
         self.handler = handler() if custom_handler is None else custom_handler
-        if not hasattr(self, 'cookies'):
-            self.cookies = {}
+        self._cookies = {}
 
     def delete(self, url, headers=None, **kwargs):
         """Sends a DELETE request to a URL.
@@ -1189,7 +1203,7 @@ def request(self, url, message, **kwargs):
             key_value_tuples = response.headers.items()
         for key, value in key_value_tuples:
             if key.lower() == "set-cookie":
-                parse_cookies(value, self.cookies)
+                _parse_cookies(value, self._cookies)
 
         return response
 
@@ -1301,7 +1315,7 @@ def request(url, message, **kwargs):
             "Content-Length": str(len(body)),
             "Host": host,
             "User-Agent": "splunk-sdk-python/0.1",
-            "Accept": "*/*"
+            "Accept": "*/*",
         } # 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, make_cookie_header
+from binding import Context, HTTPError, AuthenticationError, namespace, UrlEncoded, _encode, _make_cookie_header
 from data import record
 import data
 
@@ -1908,8 +1908,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.http.cookies) > 0:
-            cookie_or_auth_header = "Cookie: %s\r\n" % make_cookie_header(self.service.http.cookies.items())
+        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

@@ -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.http.cookies, {})
+        self.assertEqual(self.context.get_cookies(), {})
         self.assertRaises(AuthenticationError,
                           self.context.get, "/services")
         self.assertRaises(AuthenticationError,
@@ -510,13 +510,31 @@ def test_cookie_in_auth_headers(self):
         self.assertEqual(self.context._auth_headers[0][1][:8], "splunkd_")
 
     def test_got_cookie_on_connect(self):
-        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_")
+        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_cookies = self.context.http.cookies
+        old_cookies = self.context.get_cookies()
         resp = self.context.get("apps/local")
         found = False
         for key, value in resp.headers:
@@ -525,7 +543,7 @@ def test_got_updated_cookie_with_get(self):
                 self.assertEqual(value[:8], "splunkd_")
 
                 new_cookies = {}
-                binding.parse_cookies(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(new_cookies.values()), 1)
@@ -552,15 +570,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.http.cookies.items():
-                new_context.http.cookies[key] = value
+            for key, value in self.context.get_cookies().items():
+                new_context.get_cookies()[key] = value
 
-            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())
+            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.http.cookies.items():
-                self.assertEqual(new_context.http.cookies[k], v)
+            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)
 

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.http.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,7 +121,7 @@ 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.http.cookies.update(self.service.http.cookies)
+        service.http._cookies.update(self.service.http._cookies)
         service.login()
         cn = service.indexes[self.index_name].attach()
         cn.send("Hello Boris!\r\n")

tests/test_service.py

@@ -20,7 +20,7 @@
 import splunklib.client as client
 from splunklib.client import AuthenticationError
 from splunklib.client import Service
-from splunklib.binding import HTTPError, _NoAuthenticationToken
+from splunklib.binding import HTTPError
 
 
 class ServiceTestCase(testlib.SDKTestCase):
@@ -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.http.cookies)
-        self.assertEquals(len(self.service.http.cookies), 0)
+        self.assertIsNotNone(self.service.get_cookies())
+        self.assertEquals(len(self.service.get_cookies()), 0)
         self.service.login()
-        self.assertIsNotNone(self.service.http.cookies)
-        self.assertNotEquals(self.service.http.cookies, {})
-        self.assertEquals(len(self.service.http.cookies), 1)
+        self.assertIsNotNone(self.service.get_cookies())
+        self.assertNotEquals(self.service.get_cookies(), {})
+        self.assertEquals(len(self.service.get_cookies()), 1)
 
     def test_login_with_cookie(self):
         self.service.login()
-        self.assertIsNotNone(self.service.http.cookies)
+        self.assertIsNotNone(self.service.get_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.http.cookies.items()[0]})
+        service2 = client.Service(**{"cookie": "%s=%s" % self.service.get_cookies().items()[0]})
         service2.login()
-        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(len(service2.get_cookies()), 1)
+        self.assertEqual(service2.get_cookies(), self.service.get_cookies())
+        self.assertEqual(len(service2.get_cookies()), len(self.service.get_cookies()))
+        self.assertEqual(service2.get_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.http.cookies), 0)
-        service2.http.cookies.update(bad_cookie)
+        self.assertEquals(len(service2.get_cookies()), 0)
+        service2.get_cookies().update(bad_cookie)
         service2.login()
-        self.assertEquals(service2.http.cookies, {'bad': 'cookie'})
+        self.assertEquals(service2.get_cookies(), {'bad': 'cookie'})
 
         # Should get an error with a bad cookie
         try:
@@ -186,9 +186,21 @@ def test_login_fails_with_bad_cookie(self):
         except AuthenticationError as ae:
             self.assertEqual(ae.message, "Request failed: Session is not logged in.")
 
+    def test_autologin_with_cookie(self):
+        self.service.login()
+        self.assertTrue(self.service.has_cookies())
+        service = client.connect(
+            autologin=True,
+            cookie="%s=%s" % self.service.get_cookies().items()[0],
+            **self.opts.kwargs)
+        self.assertTrue(service.has_cookies())
+        self.service.restart(timeout=120)
+        reader = service.jobs.oneshot("search index=internal | head 1")
+        self.assertIsNotNone(reader)
+
     def test_login_fails_with_no_cookie(self):
         service2 = client.Service()
-        self.assertEquals(len(service2.http.cookies), 0)
+        self.assertEquals(len(service2.get_cookies()), 0)
 
         # Should get an error when no authentication method
         try:
@@ -203,15 +215,15 @@ def test_login_with_multiple_cookie_headers(self):
             'something_else': 'bad'
         }
         self.service.logout()
-        self.service.http.cookies.update(cookies)
+        self.service.get_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.http.cookies)
+        self.assertIsNotNone(self.service.get_cookies())
 
         service2 = client.Service(**{"cookie": bad_cookie})
         service2.login()
@@ -224,16 +236,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.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.get_cookies().update(self.service.get_cookies())
+
+            self.assertEqual(len(service2.get_cookies()), 2)
+            self.service.get_cookies().update({'bad': 'cookie'})
+            self.assertEqual(service2.get_cookies(), self.service.get_cookies())
+            self.assertEqual(len(service2.get_cookies()), 2)
+            self.assertEqual(service2.get_cookies().keys()[1][:8], "splunkd_")
+            self.assertTrue('bad' in service2.get_cookies().keys())
+            self.assertEqual(service2.get_cookies()['bad'], 'cookie')
+            self.assertEqual(self.service.get_cookies().items(), service2.get_cookies().items())
             service2.login()
             self.assertEqual(service2.apps.get().status, 200)