Skip to content

Commit 71e2878

Browse files
author
Shakeel Mohamed
committed
Add support for multiple cookies
There is a failing test TestPluggableHTTP
1 parent 44140cb commit 71e2878

File tree

5 files changed

+193
-43
lines changed

5 files changed

+193
-43
lines changed

‎splunklib/binding.py

Lines changed: 25 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import urllib
3232
import io
3333
import sys
34+
import Cookie
3435

3536
from datetime import datetime
3637
from functools import wraps
@@ -225,7 +226,7 @@ def f():
225226
@wraps(request_fun)
226227
def wrapper(self, *args, **kwargs):
227228
if self.token is _NoAuthenticationToken and \
228-
self.cookie is _NoAuthenticationToken:
229+
len(self.cookies) < 1:
229230
# Not yet logged in.
230231
if self.autologin and self.username and self.password:
231232
# This will throw an uncaught
@@ -428,9 +429,13 @@ def __init__(self, handler=None, **kwargs):
428429
self.username = kwargs.get("username", "")
429430
self.password = kwargs.get("password", "")
430431
self.autologin = kwargs.get("autologin", False)
431-
self.cookie = kwargs.get("cookie", _NoAuthenticationToken)
432-
if self.cookie is None: # In case someone explicitly passes cookie=None
433-
self.cookie = _NoAuthenticationToken
432+
433+
# By default, there are no cookies
434+
self.cookies = {}
435+
if kwargs.has_key("cookie") and kwargs['cookie'] not in [None, _NoAuthenticationToken]:
436+
parsed_cookie = Cookie.SimpleCookie(kwargs.get(("cookie")))
437+
for cookie in parsed_cookie.values():
438+
self.cookies[cookie.key] = cookie.coded_value
434439

435440
# Shared per-context request headers
436441
@property
@@ -443,8 +448,8 @@ def _auth_headers(self):
443448
444449
:returns: A list of 2-tuples containing key and value
445450
"""
446-
if self.cookie is not _NoAuthenticationToken:
447-
return [("cookie", self.cookie)]
451+
if len(self.cookies) > 0:
452+
return [("cookie", "%s=%s" % cookie) for cookie in self.cookies.items()]
448453
elif self.token is _NoAuthenticationToken:
449454
return []
450455
else:
@@ -761,10 +766,10 @@ def login(self):
761766
c = binding.Context(...).login()
762767
# Then issue requests...
763768
"""
764-
# If self.cookie and self.token only, use the cookie
765-
if self.cookie is not _NoAuthenticationToken and \
769+
# If self.cookies and self.token only, use the cookie
770+
if len(self.cookies) > 0 and \
766771
(not self.username and not self.password):
767-
# If we were passed a session cookie, but no username or
772+
# If we were passed session cookie(s), but no username or
768773
# password, then login is a nop, since we're automatically
769774
# logged in.
770775
return
@@ -785,10 +790,12 @@ def login(self):
785790
cookie="1") # In Splunk 6.2+, passing "cookie=1" will return the "set-cookie" header
786791

787792
# Store the cookie
793+
788794
for key, value in response.headers:
789795
if key.lower() == "set-cookie":
790-
self.cookie = value
791-
break
796+
parsed_cookies = Cookie.SimpleCookie(value)
797+
for cookie in parsed_cookies.values():
798+
self.cookies[cookie.key] = cookie.coded_value
792799

793800
body = response.body.read()
794801
session = XML(body).findtext("./sessionKey")
@@ -801,8 +808,9 @@ def login(self):
801808
raise
802809

803810
def logout(self):
804-
"""Forgets the current session token."""
811+
"""Forgets the current session token, and cookies."""
805812
self.token = _NoAuthenticationToken
813+
self.cookies = {}
806814
return self
807815

808816
def _abspath(self, path_segment,
@@ -1030,6 +1038,8 @@ class HttpLib(object):
10301038
"""
10311039
def __init__(self, custom_handler=None):
10321040
self.handler = handler() if custom_handler is None else custom_handler
1041+
if not hasattr(self, 'cookies'):
1042+
self.cookies = {}
10331043

10341044
def delete(self, url, headers=None, **kwargs):
10351045
"""Sends a DELETE request to a URL.
@@ -1143,8 +1153,9 @@ def request(self, url, message, **kwargs):
11431153
# Update the cookie with any HTTP request
11441154
for key, value in response.headers:
11451155
if key.lower() == "set-cookie":
1146-
self.cookie = value
1147-
break
1156+
parsed_cookie = Cookie.SimpleCookie(value)
1157+
for cookie in parsed_cookie.values():
1158+
self.cookies[cookie.key] = cookie.coded_value
11481159
return response
11491160

11501161

‎splunklib/client.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1902,9 +1902,13 @@ def attach(self, host=None, source=None, sourcetype=None):
19021902

19031903
cookie_or_auth_header = "Authorization: %s\r\n" % self.service.token
19041904

1905-
# If we have a cookie, use it instead of "Authorization: ..."
1906-
if self.service.cookie is not _NoAuthenticationToken:
1907-
cookie_or_auth_header = "Cookie: %s\r\n" % self.service.cookie
1905+
# If we have cookie(s), use them instead of "Authorization: ..."
1906+
if len(self.service.cookies) > 0:
1907+
cookies = []
1908+
for cookie in self.service.cookies.items():
1909+
cookies.append("Cookie: %s=%s\r\n" % cookie)
1910+
if len(cookies) > 0:
1911+
cookie_or_auth_header = "".join(cookies)
19081912

19091913
# Since we need to stream to the index connection, we have to keep
19101914
# the connection open and use the Splunk extension headers to note

‎tests/test_binding.py

Lines changed: 48 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import socket
2727
import sys
2828
import ssl
29+
import Cookie
2930

3031
import splunklib.binding as binding
3132
from splunklib.binding import HTTPError, AuthenticationError, UrlEncoded
@@ -476,6 +477,8 @@ def test_logout(self):
476477
response = self.context.get("/services")
477478
self.assertEqual(response.status, 200)
478479
self.context.logout()
480+
self.assertEqual(self.context.token, binding._NoAuthenticationToken)
481+
self.assertEqual(self.context.cookies, {})
479482
self.assertRaises(AuthenticationError,
480483
self.context.get, "/services")
481484
self.assertRaises(AuthenticationError,
@@ -507,22 +510,61 @@ def test_cookie_in_auth_headers(self):
507510
self.assertEqual(self.context._auth_headers[0][1][:8], "splunkd_")
508511

509512
def test_got_cookie_on_connect(self):
510-
self.assertIsNotNone(self.context.cookie)
511-
self.assertNotEqual(self.context.cookie, binding._NoAuthenticationToken)
512-
self.assertEqual(self.context.cookie[:8], "splunkd_")
513+
self.assertIsNotNone(self.context.cookies)
514+
self.assertNotEqual(self.context.cookies, {})
515+
self.assertEqual(len(self.context.cookies), 1)
516+
self.assertEqual(self.context.cookies.keys()[0][:8], "splunkd_")
513517

514518
def test_got_updated_cookie_with_get(self):
515-
old_cookie = self.context.cookie
519+
old_cookies = self.context.cookies
516520
resp = self.context.get("apps/local")
517521
found = False
518522
for key, value in resp.headers:
519523
if key.lower() == "set-cookie":
520524
found = True
521525
self.assertEqual(value[:8], "splunkd_")
522-
# It's unlikely that the cookie will change during this short test
523-
self.assertEqual(value, old_cookie)
526+
527+
parsed_cookies = Cookie.SimpleCookie(value)
528+
# We're only expecting 1 in this scenario
529+
self.assertEqual(len(old_cookies), 1)
530+
self.assertTrue(len(parsed_cookies.values()), 1)
531+
parsed_cookie = parsed_cookies.values()[0]
532+
533+
self.assertEqual(parsed_cookie.key, old_cookies.keys()[0])
534+
self.assertEqual(parsed_cookie.coded_value, old_cookies.values()[0])
524535
self.assertTrue(found)
525536

537+
def test_login_fails_with_bad_cookie(self):
538+
new_context = binding.connect(**{"cookie": "bad=cookie"})
539+
# We should get an error if using a bad cookie
540+
try:
541+
new_context.get("apps/local")
542+
self.fail()
543+
except AuthenticationError as ae:
544+
self.assertEqual(ae.message, "Request failed: Session is not logged in.")
545+
546+
def test_login_with_multiple_cookies(self):
547+
bad_cookie = 'bad=cookie'
548+
new_context = binding.connect(**{"cookie": bad_cookie})
549+
# We should get an error if using a bad cookie
550+
try:
551+
new_context.get("apps/local")
552+
self.fail()
553+
except AuthenticationError as ae:
554+
self.assertEqual(ae.message, "Request failed: Session is not logged in.")
555+
# Bring in a valid cookie now
556+
for key, value in self.context.cookies.items():
557+
new_context.cookies[key] = value
558+
559+
self.assertEqual(len(new_context.cookies), 2)
560+
self.assertTrue('bad' in new_context.cookies.keys())
561+
self.assertTrue('cookie' in new_context.cookies.values())
562+
563+
for k, v in self.context.cookies.items():
564+
self.assertEqual(new_context.cookies[k], v)
565+
566+
self.assertEqual(new_context.get("apps/local").status, 200)
567+
526568
def test_login_fails_without_cookie_or_token(self):
527569
opts = {
528570
'host': self.opts.kwargs['host'],

‎tests/test_index.py

Lines changed: 47 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616

1717
import testlib
1818
import logging
19-
19+
import splunklib.client as client
2020
try:
2121
import unittest
2222
except ImportError:
@@ -71,9 +71,9 @@ def test_disable_enable(self):
7171

7272
def test_submit_and_clean(self):
7373
self.index.refresh()
74-
originalCount = int(self.index['totalEventCount'])
74+
original_count = int(self.index['totalEventCount'])
7575
self.index.submit("Hello again!", sourcetype="Boris", host="meep")
76-
self.assertEventuallyTrue(lambda: self.totalEventCount() == originalCount+1, timeout=50)
76+
self.assertEventuallyTrue(lambda: self.totalEventCount() == original_count+1, timeout=50)
7777

7878
# Cleaning an enabled index on 4.x takes forever, so we disable it.
7979
# However, cleaning it on 5 requires it to be enabled.
@@ -87,37 +87,73 @@ def test_prefresh(self):
8787
self.assertEqual(self.index['disabled'], '0') # Index is prefreshed
8888

8989
def test_submit(self):
90-
eventCount = int(self.index['totalEventCount'])
90+
event_count = int(self.index['totalEventCount'])
9191
self.assertEqual(self.index['sync'], '0')
9292
self.assertEqual(self.index['disabled'], '0')
9393
self.index.submit("Hello again!", sourcetype="Boris", host="meep")
94-
self.assertEventuallyTrue(lambda: self.totalEventCount() == eventCount+1, timeout=50)
94+
self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=50)
9595

9696
def test_submit_via_attach(self):
97-
eventCount = int(self.index['totalEventCount'])
97+
event_count = int(self.index['totalEventCount'])
9898
cn = self.index.attach()
9999
cn.send("Hello Boris!\r\n")
100100
cn.close()
101-
self.assertEventuallyTrue(lambda: self.totalEventCount() == eventCount+1, timeout=60)
101+
self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=60)
102102

103103
def test_submit_via_attached_socket(self):
104-
eventCount = int(self.index['totalEventCount'])
104+
event_count = int(self.index['totalEventCount'])
105105
f = self.index.attached_socket
106106
with f() as sock:
107107
sock.send('Hello world!\r\n')
108-
self.assertEventuallyTrue(lambda: self.totalEventCount() == eventCount+1, timeout=60)
108+
self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=60)
109+
110+
def test_submit_via_attach_with_cookie_header(self):
111+
event_count = int(self.service.indexes[self.index_name]['totalEventCount'])
112+
113+
cookie = "%s=%s" % (self.service.cookies.items()[0])
114+
service = client.Service(**{"cookie": cookie})
115+
service.login()
116+
cn = service.indexes[self.index_name].attach()
117+
cn.send("Hello Boris!\r\n")
118+
cn.close()
119+
self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=60)
120+
121+
def test_submit_via_attach_with_multiple_cookie_headers(self):
122+
event_count = int(self.service.indexes[self.index_name]['totalEventCount'])
123+
service = client.Service(**{"cookie": 'a bad cookie'})
124+
service.cookies.update(self.service.cookies)
125+
service.login()
126+
cn = service.indexes[self.index_name].attach()
127+
cn.send("Hello Boris!\r\n")
128+
cn.close()
129+
self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=60)
130+
131+
def test_login_with_multiple_cookie_headers(self):
132+
event_count = int(self.index['totalEventCount'])
133+
134+
cookies = {}
135+
cookies['bad'] = 'cookie'
136+
cookies['something_else'] = 'bad'
137+
self.service.logout()
138+
self.service.cookies.update(cookies)
139+
140+
self.service.login()
141+
cn = self.service.indexes[self.index_name].attach()
142+
cn.send("Hello Boris!\r\n")
143+
cn.close()
144+
self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=60)
109145

110146
def test_upload(self):
111147
if not self.app_collection_installed():
112148
print "Test requires sdk-app-collection. Skipping."
113149
return
114150
self.install_app_from_collection("file_to_upload")
115151

116-
eventCount = int(self.index['totalEventCount'])
152+
event_count = int(self.index['totalEventCount'])
117153

118154
path = self.pathInApp("file_to_upload", ["log.txt"])
119155
self.index.upload(path)
120-
self.assertEventuallyTrue(lambda: self.totalEventCount() == eventCount+4, timeout=60)
156+
self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+4, timeout=60)
121157

122158
if __name__ == "__main__":
123159
try:

‎tests/test_service.py

Lines changed: 66 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -152,22 +152,79 @@ def setUp(self):
152152
self.skipTest("Skipping cookie-auth tests, running in %d.%d.%d, this feature was added in 6.2+" % splver)
153153

154154
def test_login_and_store_cookie(self):
155-
self.assertIsNotNone(self.service.cookie)
156-
self.assertEquals(self.service.cookie, _NoAuthenticationToken)
155+
self.assertIsNotNone(self.service.cookies)
156+
self.assertEquals(len(self.service.cookies), 0)
157157
self.service.login()
158-
self.assertIsNotNone(self.service.cookie)
159-
self.assertNotEquals(self.service.cookie, _NoAuthenticationToken)
158+
self.assertIsNotNone(self.service.cookies)
159+
self.assertNotEquals(self.service.cookies, {})
160+
self.assertEquals(len(self.service.cookies), 1)
160161

161162
def test_login_with_cookie(self):
162163
self.service.login()
163-
self.assertIsNotNone(self.service.cookie)
164-
self.assertNotEqual(self.service.cookie, _NoAuthenticationToken)
164+
self.assertIsNotNone(self.service.cookies)
165165
# Use the cookie from the other service as the only auth param (don't need user/password)
166-
service2 = client.Service(**{"cookie": self.service.cookie})
166+
service2 = client.Service(**{"cookie": "%s=%s" % self.service.cookies.items()[0]})
167167
service2.login()
168-
self.assertEqual(service2.cookie, self.service.cookie)
169-
self.assertEqual(service2.cookie[:8], "splunkd_")
168+
self.assertEqual(len(service2.cookies), 1)
169+
self.assertEqual(service2.cookies, self.service.cookies)
170+
self.assertEqual(len(service2.cookies), len(self.service.cookies))
171+
self.assertEqual(service2.cookies.keys()[0][:8], "splunkd_")
172+
self.assertEqual(service2.apps.get().status, 200)
173+
174+
def test_login_fails_with_bad_cookie(self):
175+
bad_cookie = {'bad': 'cookie'}
176+
service2 = client.Service()
177+
self.assertEquals(len(service2.cookies), 0)
178+
service2.cookies.update(bad_cookie)
179+
service2.login()
180+
self.assertEquals(service2.cookies, {'bad': 'cookie'})
181+
182+
# Should get an error with a bad cookie
183+
try:
184+
service2.apps.get()
185+
self.fail()
186+
except AuthenticationError as ae:
187+
self.assertEqual(ae.message, "Request failed: Session is not logged in.")
188+
189+
def test_login_fails_with_no_cookie(self):
190+
service2 = client.Service()
191+
self.assertEquals(len(service2.cookies), 0)
192+
193+
# Should get an error when no authentication method
194+
try:
195+
service2.login()
196+
self.fail()
197+
except AuthenticationError as ae:
198+
self.assertEqual(ae.message, "Login failed.")
170199

200+
def test_login_with_multiple_cookies(self):
201+
bad_cookie = 'bad=cookie'
202+
self.service.login()
203+
self.assertIsNotNone(self.service.cookies)
204+
205+
service2 = client.Service(**{"cookie": bad_cookie})
206+
service2.login()
207+
208+
# Should get an error with a bad cookie
209+
try:
210+
service2.apps.get()
211+
self.fail()
212+
except AuthenticationError as ae:
213+
self.assertEqual(ae.message, "Request failed: Session is not logged in.")
214+
215+
# Add on valid cookies, and try to use all of them
216+
service2.cookies.update(self.service.cookies)
217+
218+
self.assertEqual(len(service2.cookies), 2)
219+
self.service.cookies.update({'bad': 'cookie'})
220+
self.assertEqual(service2.cookies, self.service.cookies)
221+
self.assertEqual(len(service2.cookies), 2)
222+
self.assertEqual(service2.cookies.keys()[1][:8], "splunkd_")
223+
self.assertTrue('bad' in service2.cookies.keys())
224+
self.assertEqual(service2.cookies['bad'], 'cookie')
225+
self.assertEqual(self.service.cookies.items(), service2.cookies.items())
226+
service2.login()
227+
self.assertEqual(service2.apps.get().status, 200)
171228

172229
class TestSettings(testlib.SDKTestCase):
173230
def test_read_settings(self):

0 commit comments

Comments
 (0)