Add support for multiple cookies

There is a failing test TestPluggableHTTP

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
@@ -225,7 +226,7 @@ def f():
     @wraps(request_fun)
     def wrapper(self, *args, **kwargs):
         if self.token is _NoAuthenticationToken and \
-                self.cookie is _NoAuthenticationToken:
+                len(self.cookies) < 1:
             # Not yet logged in.
             if self.autologin and self.username and self.password:
                 # This will throw an uncaught
@@ -428,9 +429,13 @@ 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
+
+        # By default, there are no cookies
+        self.cookies = {}
+        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
 
     # Shared per-context request headers
     @property
@@ -443,8 +448,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 len(self.cookies) > 0:
+            return [("cookie", "%s=%s" % cookie) for cookie in self.cookies.items()]
         elif self.token is _NoAuthenticationToken:
             return []
         else:
@@ -761,10 +766,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.cookies and self.token only, use the cookie
+        if len(self.cookies) > 0 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
@@ -785,10 +790,12 @@ def login(self):
                 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
+                    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")
@@ -801,8 +808,9 @@ def login(self):
                 raise
 
     def logout(self):
-        """Forgets the current session token."""
+        """Forgets the current session token, and cookies."""
         self.token = _NoAuthenticationToken
+        self.cookies = {}
         return self
 
     def _abspath(self, path_segment,
@@ -1030,6 +1038,8 @@ 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 = {}
 
     def delete(self, url, headers=None, **kwargs):
         """Sends a DELETE request to a URL.
@@ -1143,8 +1153,9 @@ def request(self, url, message, **kwargs):
         # Update the cookie with any HTTP request
         for key, value in response.headers:
             if key.lower() == "set-cookie":
-                self.cookie = value
-                break
+                parsed_cookie = Cookie.SimpleCookie(value)
+                for cookie in parsed_cookie.values():
+                    self.cookies[cookie.key] = cookie.coded_value
         return response
 
 

splunklib/client.py

@@ -1902,9 +1902,13 @@ 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 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)
 
         # 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.cookies, {})
         self.assertRaises(AuthenticationError,
                           self.context.get, "/services")
         self.assertRaises(AuthenticationError,
@@ -507,22 +510,61 @@ 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.cookie)
-        self.assertNotEqual(self.context.cookie, binding._NoAuthenticationToken)
-        self.assertEqual(self.context.cookie[:8], "splunkd_")
+        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_")
 
     def test_got_updated_cookie_with_get(self):
-        old_cookie = self.context.cookie
+        old_cookies = self.context.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)
+
+                parsed_cookies = Cookie.SimpleCookie(value)
+                # 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(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.cookies.items():
+                new_context.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())
+
+            for k, v in self.context.cookies.items():
+                self.assertEqual(new_context.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'],

tests/test_index.py

@@ -16,7 +16,7 @@
 
 import testlib
 import logging
-
+import splunklib.client as client
 try:
     import unittest
 except ImportError:
@@ -71,9 +71,9 @@ def test_disable_enable(self):
 
     def test_submit_and_clean(self):
         self.index.refresh()
-        originalCount = int(self.index['totalEventCount'])
+        original_count = int(self.index['totalEventCount'])
         self.index.submit("Hello again!", sourcetype="Boris", host="meep")
-        self.assertEventuallyTrue(lambda: self.totalEventCount() == originalCount+1, timeout=50)
+        self.assertEventuallyTrue(lambda: self.totalEventCount() == original_count+1, timeout=50)
 
         # Cleaning an enabled index on 4.x takes forever, so we disable it.
         # However, cleaning it on 5 requires it to be enabled.
@@ -87,37 +87,73 @@ def test_prefresh(self):
         self.assertEqual(self.index['disabled'], '0') # Index is prefreshed
 
     def test_submit(self):
-        eventCount = int(self.index['totalEventCount'])
+        event_count = int(self.index['totalEventCount'])
         self.assertEqual(self.index['sync'], '0')
         self.assertEqual(self.index['disabled'], '0')
         self.index.submit("Hello again!", sourcetype="Boris", host="meep")
-        self.assertEventuallyTrue(lambda: self.totalEventCount() == eventCount+1, timeout=50)
+        self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=50)
 
     def test_submit_via_attach(self):
-        eventCount = int(self.index['totalEventCount'])
+        event_count = int(self.index['totalEventCount'])
         cn = self.index.attach()
         cn.send("Hello Boris!\r\n")
         cn.close()
-        self.assertEventuallyTrue(lambda: self.totalEventCount() == eventCount+1, timeout=60)
+        self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=60)
 
     def test_submit_via_attached_socket(self):
-        eventCount = int(self.index['totalEventCount'])
+        event_count = int(self.index['totalEventCount'])
         f = self.index.attached_socket
         with f() as sock:
             sock.send('Hello world!\r\n')
-        self.assertEventuallyTrue(lambda: self.totalEventCount() == eventCount+1, timeout=60)
+        self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=60)
+
+    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])
+        service = client.Service(**{"cookie": cookie})
+        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_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.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."
             return
         self.install_app_from_collection("file_to_upload")
 
-        eventCount = int(self.index['totalEventCount'])
+        event_count = int(self.index['totalEventCount'])
 
         path = self.pathInApp("file_to_upload", ["log.txt"])
         self.index.upload(path)
-        self.assertEventuallyTrue(lambda: self.totalEventCount() == eventCount+4, timeout=60)
+        self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+4, timeout=60)
 
 if __name__ == "__main__":
     try:

tests/test_service.py

@@ -152,22 +152,79 @@ 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.cookie)
-        self.assertEquals(self.service.cookie, _NoAuthenticationToken)
+        self.assertIsNotNone(self.service.cookies)
+        self.assertEquals(len(self.service.cookies), 0)
         self.service.login()
-        self.assertIsNotNone(self.service.cookie)
-        self.assertNotEquals(self.service.cookie, _NoAuthenticationToken)
+        self.assertIsNotNone(self.service.cookies)
+        self.assertNotEquals(self.service.cookies, {})
+        self.assertEquals(len(self.service.cookies), 1)
 
     def test_login_with_cookie(self):
         self.service.login()
-        self.assertIsNotNone(self.service.cookie)
-        self.assertNotEqual(self.service.cookie, _NoAuthenticationToken)
+        self.assertIsNotNone(self.service.cookies)
         # Use the cookie from the other service as the only auth param (don't need user/password)
-        service2 = client.Service(**{"cookie": self.service.cookie})
+        service2 = client.Service(**{"cookie": "%s=%s" % self.service.cookies.items()[0]})
         service2.login()
-        self.assertEqual(service2.cookie, self.service.cookie)
-        self.assertEqual(service2.cookie[:8], "splunkd_")
+        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(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)
+        service2.login()
+        self.assertEquals(service2.cookies, {'bad': 'cookie'})
+
+        # Should get an error with a bad cookie
+        try:
+            service2.apps.get()
+            self.fail()
+        except AuthenticationError as ae:
+            self.assertEqual(ae.message, "Request failed: Session is not logged in.")
+
+    def test_login_fails_with_no_cookie(self):
+        service2 = client.Service()
+        self.assertEquals(len(service2.cookies), 0)
+
+        # Should get an error when no authentication method
+        try:
+            service2.login()
+            self.fail()
+        except AuthenticationError as ae:
+            self.assertEqual(ae.message, "Login failed.")
 
+    def test_login_with_multiple_cookies(self):
+        bad_cookie = 'bad=cookie'
+        self.service.login()
+        self.assertIsNotNone(self.service.cookies)
+
+        service2 = client.Service(**{"cookie": bad_cookie})
+        service2.login()
+
+        # Should get an error with a bad cookie
+        try:
+            service2.apps.get()
+            self.fail()
+        except AuthenticationError as ae:
+            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.login()
+            self.assertEqual(service2.apps.get().status, 200)
 
 class TestSettings(testlib.SDKTestCase):
     def test_read_settings(self):