Skip to content

Commit 47d9e0e

Browse files
author
Shakeel Mohamed
committed
Refactor cookie authentication functionality + tests
1 parent 71e2878 commit 47d9e0e

File tree

5 files changed

+119
-93
lines changed

5 files changed

+119
-93
lines changed

‎splunklib/binding.py

Lines changed: 58 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,47 @@ def new_f(*args, **kwargs):
6868
return new_f
6969

7070

71+
def parse_cookies(cookie_str, dictionary):
72+
"""Tries to parse any key-value pairs of cookies in a string,
73+
then updates the the dictionary with any key-value pairs found.
74+
75+
**Example**::
76+
77+
parse_cookies('my=value', {})
78+
# Now the following is True
79+
dictionary['my'] == 'value'
80+
81+
:param cookie_str: A string containing "key=value" pairs from an HTTP "Set-Cookie" header
82+
:type msg: ``str``
83+
:param dictionary: A dictionary to update with any found key-value pairs
84+
:type msg: ``dict``
85+
"""
86+
parsed_cookie = Cookie.SimpleCookie(cookie_str)
87+
for cookie in parsed_cookie.values():
88+
dictionary[cookie.key] = cookie.coded_value
89+
90+
def make_cookie_header(cookies):
91+
# TODO: verify
92+
"""
93+
Takes a list of 2-tuples of key-value pairs of
94+
cookies, and returns a valid HTTP ``Cookie``
95+
header.
96+
97+
**Example**::
98+
99+
header = make_cookie_header([("key", "value"), ("key_2", "value_2")])
100+
# Now the following is True
101+
header == "key=value; key_2=value_2"
102+
103+
:param cookies:
104+
:return: ``str` An HTTP header cookie string.
105+
:rtype: ``str``
106+
"""
107+
header = ""
108+
for key, value in cookies:
109+
header += "%s=%s; " % (key, value)
110+
return header
111+
71112
# Singleton values to eschew None
72113
class _NoAuthenticationToken(object):
73114
"""The value stored in a :class:`Context` or :class:`splunklib.client.Service`
@@ -226,7 +267,7 @@ def f():
226267
@wraps(request_fun)
227268
def wrapper(self, *args, **kwargs):
228269
if self.token is _NoAuthenticationToken and \
229-
len(self.cookies) < 1:
270+
len(self.http.cookies) < 1:
230271
# Not yet logged in.
231272
if self.autologin and self.username and self.password:
232273
# This will throw an uncaught
@@ -430,12 +471,9 @@ def __init__(self, handler=None, **kwargs):
430471
self.password = kwargs.get("password", "")
431472
self.autologin = kwargs.get("autologin", False)
432473

433-
# By default, there are no cookies
434-
self.cookies = {}
474+
# Store any cookies in the self.http.cookies dict
435475
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
476+
parse_cookies(kwargs["cookie"], self.http.cookies)
439477

440478
# Shared per-context request headers
441479
@property
@@ -448,8 +486,8 @@ def _auth_headers(self):
448486
449487
:returns: A list of 2-tuples containing key and value
450488
"""
451-
if len(self.cookies) > 0:
452-
return [("cookie", "%s=%s" % cookie) for cookie in self.cookies.items()]
489+
if len(self.http.cookies) > 0:
490+
return [("Cookie", make_cookie_header(self.http.cookies.items()))]
453491
elif self.token is _NoAuthenticationToken:
454492
return []
455493
else:
@@ -766,8 +804,8 @@ def login(self):
766804
c = binding.Context(...).login()
767805
# Then issue requests...
768806
"""
769-
# If self.cookies and self.token only, use the cookie
770-
if len(self.cookies) > 0 and \
807+
808+
if len(self.http.cookies) > 0 and \
771809
(not self.username and not self.password):
772810
# If we were passed session cookie(s), but no username or
773811
# password, then login is a nop, since we're automatically
@@ -789,14 +827,6 @@ def login(self):
789827
password=self.password,
790828
cookie="1") # In Splunk 6.2+, passing "cookie=1" will return the "set-cookie" header
791829

792-
# Store the cookie
793-
794-
for key, value in response.headers:
795-
if key.lower() == "set-cookie":
796-
parsed_cookies = Cookie.SimpleCookie(value)
797-
for cookie in parsed_cookies.values():
798-
self.cookies[cookie.key] = cookie.coded_value
799-
800830
body = response.body.read()
801831
session = XML(body).findtext("./sessionKey")
802832
self.token = "Splunk %s" % session
@@ -810,7 +840,7 @@ def login(self):
810840
def logout(self):
811841
"""Forgets the current session token, and cookies."""
812842
self.token = _NoAuthenticationToken
813-
self.cookies = {}
843+
self.http.cookies = {}
814844
return self
815845

816846
def _abspath(self, path_segment,
@@ -1151,11 +1181,16 @@ def request(self, url, message, **kwargs):
11511181
raise HTTPError(response)
11521182

11531183
# Update the cookie with any HTTP request
1154-
for key, value in response.headers:
1184+
# Initially, assume list of 2-tuples
1185+
key_value_tuples = response.headers
1186+
# If response.headers is a dict, get the key-value pairs as 2-tuples
1187+
# this is the case when using urllib2
1188+
if isinstance(response.headers, dict):
1189+
key_value_tuples = response.headers.items()
1190+
for key, value in key_value_tuples:
11551191
if key.lower() == "set-cookie":
1156-
parsed_cookie = Cookie.SimpleCookie(value)
1157-
for cookie in parsed_cookie.values():
1158-
self.cookies[cookie.key] = cookie.coded_value
1192+
parse_cookies(value, self.cookies)
1193+
11591194
return response
11601195

11611196

‎splunklib/client.py

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@
6767
import socket
6868
import contextlib
6969

70-
from binding import Context, HTTPError, AuthenticationError, namespace, UrlEncoded, _encode, _NoAuthenticationToken
70+
from binding import Context, HTTPError, AuthenticationError, namespace, UrlEncoded, _encode, make_cookie_header
7171
from data import record
7272
import data
7373

@@ -1903,12 +1903,8 @@ def attach(self, host=None, source=None, sourcetype=None):
19031903
cookie_or_auth_header = "Authorization: %s\r\n" % self.service.token
19041904

19051905
# 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)
1906+
if len(self.service.http.cookies) > 0:
1907+
cookie_or_auth_header = "Cookie: %s\r\n" % make_cookie_header(self.service.http.cookies.items())
19121908

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

‎tests/test_binding.py

Lines changed: 19 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -478,7 +478,7 @@ def test_logout(self):
478478
self.assertEqual(response.status, 200)
479479
self.context.logout()
480480
self.assertEqual(self.context.token, binding._NoAuthenticationToken)
481-
self.assertEqual(self.context.cookies, {})
481+
self.assertEqual(self.context.http.cookies, {})
482482
self.assertRaises(AuthenticationError,
483483
self.context.get, "/services")
484484
self.assertRaises(AuthenticationError,
@@ -506,32 +506,31 @@ def test_cookie_in_auth_headers(self):
506506
self.assertNotEqual(self.context._auth_headers, [])
507507
self.assertEqual(len(self.context._auth_headers), 1)
508508
self.assertEqual(len(self.context._auth_headers), 1)
509-
self.assertEqual(self.context._auth_headers[0][0], "cookie")
509+
self.assertEqual(self.context._auth_headers[0][0], "Cookie")
510510
self.assertEqual(self.context._auth_headers[0][1][:8], "splunkd_")
511511

512512
def test_got_cookie_on_connect(self):
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_")
513+
self.assertIsNotNone(self.context.http.cookies)
514+
self.assertNotEqual(self.context.http.cookies, {})
515+
self.assertEqual(len(self.context.http.cookies), 1)
516+
self.assertEqual(self.context.http.cookies.keys()[0][:8], "splunkd_")
517517

518518
def test_got_updated_cookie_with_get(self):
519-
old_cookies = self.context.cookies
519+
old_cookies = self.context.http.cookies
520520
resp = self.context.get("apps/local")
521521
found = False
522522
for key, value in resp.headers:
523523
if key.lower() == "set-cookie":
524524
found = True
525525
self.assertEqual(value[:8], "splunkd_")
526526

527-
parsed_cookies = Cookie.SimpleCookie(value)
527+
new_cookies = {}
528+
binding.parse_cookies(value, new_cookies)
528529
# We're only expecting 1 in this scenario
529530
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])
531+
self.assertTrue(len(new_cookies.values()), 1)
532+
self.assertEqual(old_cookies, new_cookies)
533+
self.assertEqual(new_cookies.values()[0], old_cookies.values()[0])
535534
self.assertTrue(found)
536535

537536
def test_login_fails_with_bad_cookie(self):
@@ -553,15 +552,15 @@ def test_login_with_multiple_cookies(self):
553552
except AuthenticationError as ae:
554553
self.assertEqual(ae.message, "Request failed: Session is not logged in.")
555554
# Bring in a valid cookie now
556-
for key, value in self.context.cookies.items():
557-
new_context.cookies[key] = value
555+
for key, value in self.context.http.cookies.items():
556+
new_context.http.cookies[key] = value
558557

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())
558+
self.assertEqual(len(new_context.http.cookies), 2)
559+
self.assertTrue('bad' in new_context.http.cookies.keys())
560+
self.assertTrue('cookie' in new_context.http.cookies.values())
562561

563-
for k, v in self.context.cookies.items():
564-
self.assertEqual(new_context.cookies[k], v)
562+
for k, v in self.context.http.cookies.items():
563+
self.assertEqual(new_context.http.cookies[k], v)
565564

566565
self.assertEqual(new_context.get("apps/local").status, 200)
567566

‎tests/test_index.py

Lines changed: 2 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ def test_submit_via_attached_socket(self):
110110
def test_submit_via_attach_with_cookie_header(self):
111111
event_count = int(self.service.indexes[self.index_name]['totalEventCount'])
112112

113-
cookie = "%s=%s" % (self.service.cookies.items()[0])
113+
cookie = "%s=%s" % (self.service.http.cookies.items()[0])
114114
service = client.Service(**{"cookie": cookie})
115115
service.login()
116116
cn = service.indexes[self.index_name].attach()
@@ -121,28 +121,13 @@ def test_submit_via_attach_with_cookie_header(self):
121121
def test_submit_via_attach_with_multiple_cookie_headers(self):
122122
event_count = int(self.service.indexes[self.index_name]['totalEventCount'])
123123
service = client.Service(**{"cookie": 'a bad cookie'})
124-
service.cookies.update(self.service.cookies)
124+
service.http.cookies.update(self.service.http.cookies)
125125
service.login()
126126
cn = service.indexes[self.index_name].attach()
127127
cn.send("Hello Boris!\r\n")
128128
cn.close()
129129
self.assertEventuallyTrue(lambda: self.totalEventCount() == event_count+1, timeout=60)
130130

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)
145-
146131
def test_upload(self):
147132
if not self.app_collection_installed():
148133
print "Test requires sdk-app-collection. Skipping."

‎tests/test_service.py

Lines changed: 37 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -152,32 +152,32 @@ 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.cookies)
156-
self.assertEquals(len(self.service.cookies), 0)
155+
self.assertIsNotNone(self.service.http.cookies)
156+
self.assertEquals(len(self.service.http.cookies), 0)
157157
self.service.login()
158-
self.assertIsNotNone(self.service.cookies)
159-
self.assertNotEquals(self.service.cookies, {})
160-
self.assertEquals(len(self.service.cookies), 1)
158+
self.assertIsNotNone(self.service.http.cookies)
159+
self.assertNotEquals(self.service.http.cookies, {})
160+
self.assertEquals(len(self.service.http.cookies), 1)
161161

162162
def test_login_with_cookie(self):
163163
self.service.login()
164-
self.assertIsNotNone(self.service.cookies)
164+
self.assertIsNotNone(self.service.http.cookies)
165165
# Use the cookie from the other service as the only auth param (don't need user/password)
166-
service2 = client.Service(**{"cookie": "%s=%s" % self.service.cookies.items()[0]})
166+
service2 = client.Service(**{"cookie": "%s=%s" % self.service.http.cookies.items()[0]})
167167
service2.login()
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_")
168+
self.assertEqual(len(service2.http.cookies), 1)
169+
self.assertEqual(service2.http.cookies, self.service.http.cookies)
170+
self.assertEqual(len(service2.http.cookies), len(self.service.http.cookies))
171+
self.assertEqual(service2.http.cookies.keys()[0][:8], "splunkd_")
172172
self.assertEqual(service2.apps.get().status, 200)
173173

174174
def test_login_fails_with_bad_cookie(self):
175175
bad_cookie = {'bad': 'cookie'}
176176
service2 = client.Service()
177-
self.assertEquals(len(service2.cookies), 0)
178-
service2.cookies.update(bad_cookie)
177+
self.assertEquals(len(service2.http.cookies), 0)
178+
service2.http.cookies.update(bad_cookie)
179179
service2.login()
180-
self.assertEquals(service2.cookies, {'bad': 'cookie'})
180+
self.assertEquals(service2.http.cookies, {'bad': 'cookie'})
181181

182182
# Should get an error with a bad cookie
183183
try:
@@ -188,7 +188,7 @@ def test_login_fails_with_bad_cookie(self):
188188

189189
def test_login_fails_with_no_cookie(self):
190190
service2 = client.Service()
191-
self.assertEquals(len(service2.cookies), 0)
191+
self.assertEquals(len(service2.http.cookies), 0)
192192

193193
# Should get an error when no authentication method
194194
try:
@@ -197,10 +197,21 @@ def test_login_fails_with_no_cookie(self):
197197
except AuthenticationError as ae:
198198
self.assertEqual(ae.message, "Login failed.")
199199

200+
def test_login_with_multiple_cookie_headers(self):
201+
cookies = {
202+
'bad': 'cookie',
203+
'something_else': 'bad'
204+
}
205+
self.service.logout()
206+
self.service.http.cookies.update(cookies)
207+
208+
self.service.login()
209+
self.assertEqual(self.service.apps.get().status, 200)
210+
200211
def test_login_with_multiple_cookies(self):
201212
bad_cookie = 'bad=cookie'
202213
self.service.login()
203-
self.assertIsNotNone(self.service.cookies)
214+
self.assertIsNotNone(self.service.http.cookies)
204215

205216
service2 = client.Service(**{"cookie": bad_cookie})
206217
service2.login()
@@ -213,16 +224,16 @@ def test_login_with_multiple_cookies(self):
213224
self.assertEqual(ae.message, "Request failed: Session is not logged in.")
214225

215226
# 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())
227+
service2.http.cookies.update(self.service.http.cookies)
228+
229+
self.assertEqual(len(service2.http.cookies), 2)
230+
self.service.http.cookies.update({'bad': 'cookie'})
231+
self.assertEqual(service2.http.cookies, self.service.http.cookies)
232+
self.assertEqual(len(service2.http.cookies), 2)
233+
self.assertEqual(service2.http.cookies.keys()[1][:8], "splunkd_")
234+
self.assertTrue('bad' in service2.http.cookies.keys())
235+
self.assertEqual(service2.http.cookies['bad'], 'cookie')
236+
self.assertEqual(self.service.http.cookies.items(), service2.http.cookies.items())
226237
service2.login()
227238
self.assertEqual(service2.apps.get().status, 200)
228239

0 commit comments

Comments
 (0)