Skip to content

Commit 7db4adf

Browse files
author
Shakeel Mohamed
committed
Cosmetic refactor for cookie auth functionality
1 parent ab2aff9 commit 7db4adf

File tree

5 files changed

+111
-67
lines changed

5 files changed

+111
-67
lines changed

‎splunklib/binding.py

Lines changed: 34 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -68,13 +68,13 @@ def new_f(*args, **kwargs):
6868
return new_f
6969

7070

71-
def parse_cookies(cookie_str, dictionary):
71+
def _parse_cookies(cookie_str, dictionary):
7272
"""Tries to parse any key-value pairs of cookies in a string,
7373
then updates the the dictionary with any key-value pairs found.
7474
7575
**Example**::
76-
77-
parse_cookies('my=value', {})
76+
dictionary = {}
77+
_parse_cookies('my=value', dictionary)
7878
# Now the following is True
7979
dictionary['my'] == 'value'
8080
@@ -87,15 +87,16 @@ def parse_cookies(cookie_str, dictionary):
8787
for cookie in parsed_cookie.values():
8888
dictionary[cookie.key] = cookie.coded_value
8989

90-
def make_cookie_header(cookies):
90+
91+
def _make_cookie_header(cookies):
9192
"""
9293
Takes a list of 2-tuples of key-value pairs of
9394
cookies, and returns a valid HTTP ``Cookie``
9495
header.
9596
9697
**Example**::
9798
98-
header = make_cookie_header([("key", "value"), ("key_2", "value_2")])
99+
header = _make_cookie_header([("key", "value"), ("key_2", "value_2")])
99100
# Now the following is True
100101
header == "key=value; key_2=value_2"
101102
@@ -104,10 +105,7 @@ def make_cookie_header(cookies):
104105
:return: ``str` An HTTP header cookie string.
105106
:rtype: ``str``
106107
"""
107-
header = ""
108-
for key, value in cookies:
109-
header += "%s=%s; " % (key, value)
110-
return header
108+
return "; ".join("%s=%s" % (key, value) for key, value in cookies)
111109

112110
# Singleton values to eschew None
113111
class _NoAuthenticationToken(object):
@@ -267,7 +265,7 @@ def f():
267265
@wraps(request_fun)
268266
def wrapper(self, *args, **kwargs):
269267
if self.token is _NoAuthenticationToken and \
270-
len(self.http.cookies) < 1:
268+
len(self.get_cookies()) < 1:
271269
# Not yet logged in.
272270
if self.autologin and self.username and self.password:
273271
# This will throw an uncaught
@@ -471,9 +469,26 @@ def __init__(self, handler=None, **kwargs):
471469
self.password = kwargs.get("password", "")
472470
self.autologin = kwargs.get("autologin", False)
473471

474-
# Store any cookies in the self.http.cookies dict
472+
# Store any cookies in the self.http._cookies dict
475473
if kwargs.has_key("cookie") and kwargs['cookie'] not in [None, _NoAuthenticationToken]:
476-
parse_cookies(kwargs["cookie"], self.http.cookies)
474+
_parse_cookies(kwargs["cookie"], self.http._cookies)
475+
476+
def get_cookies(self):
477+
"""Gets the dictionary of cookies from the ``HttpLib`` member of this instance.
478+
479+
:return: Dictionary of cookies stored on the ``self.http``.
480+
:rtype: ``dict``
481+
"""
482+
return self.http._cookies
483+
484+
def has_cookies(self):
485+
"""Returns true if the ``HttpLib` member of this instance has at least
486+
one cookie stored.
487+
488+
:return: ``True`` if there is at least one cookie, else ``False``
489+
:rtype: ``bool``
490+
"""
491+
return len(self.get_cookies()) > 0
477492

478493
# Shared per-context request headers
479494
@property
@@ -486,8 +501,8 @@ def _auth_headers(self):
486501
487502
:returns: A list of 2-tuples containing key and value
488503
"""
489-
if len(self.http.cookies) > 0:
490-
return [("Cookie", make_cookie_header(self.http.cookies.items()))]
504+
if self.has_cookies():
505+
return [("Cookie", _make_cookie_header(self.get_cookies().items()))]
491506
elif self.token is _NoAuthenticationToken:
492507
return []
493508
else:
@@ -805,7 +820,7 @@ def login(self):
805820
# Then issue requests...
806821
"""
807822

808-
if len(self.http.cookies) > 0 and \
823+
if self.has_cookies() and \
809824
(not self.username and not self.password):
810825
# If we were passed session cookie(s), but no username or
811826
# password, then login is a nop, since we're automatically
@@ -840,7 +855,7 @@ def login(self):
840855
def logout(self):
841856
"""Forgets the current session token, and cookies."""
842857
self.token = _NoAuthenticationToken
843-
self.http.cookies = {}
858+
self.http._cookies = {}
844859
return self
845860

846861
def _abspath(self, path_segment,
@@ -1068,8 +1083,7 @@ class HttpLib(object):
10681083
"""
10691084
def __init__(self, custom_handler=None):
10701085
self.handler = handler() if custom_handler is None else custom_handler
1071-
if not hasattr(self, 'cookies'):
1072-
self.cookies = {}
1086+
self._cookies = {}
10731087

10741088
def delete(self, url, headers=None, **kwargs):
10751089
"""Sends a DELETE request to a URL.
@@ -1189,7 +1203,7 @@ def request(self, url, message, **kwargs):
11891203
key_value_tuples = response.headers.items()
11901204
for key, value in key_value_tuples:
11911205
if key.lower() == "set-cookie":
1192-
parse_cookies(value, self.cookies)
1206+
_parse_cookies(value, self._cookies)
11931207

11941208
return response
11951209

@@ -1301,7 +1315,7 @@ def request(url, message, **kwargs):
13011315
"Content-Length": str(len(body)),
13021316
"Host": host,
13031317
"User-Agent": "splunk-sdk-python/0.1",
1304-
"Accept": "*/*"
1318+
"Accept": "*/*",
13051319
} # defaults
13061320
for key, value in message["headers"]:
13071321
head[key] = value

‎splunklib/client.py

Lines changed: 3 additions & 3 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, make_cookie_header
70+
from binding import Context, HTTPError, AuthenticationError, namespace, UrlEncoded, _encode, _make_cookie_header
7171
from data import record
7272
import data
7373

@@ -1908,8 +1908,8 @@ def attach(self, host=None, source=None, sourcetype=None):
19081908
cookie_or_auth_header = "Authorization: %s\r\n" % self.service.token
19091909

19101910
# If we have cookie(s), use them instead of "Authorization: ..."
1911-
if len(self.service.http.cookies) > 0:
1912-
cookie_or_auth_header = "Cookie: %s\r\n" % make_cookie_header(self.service.http.cookies.items())
1911+
if self.service.has_cookies():
1912+
cookie_or_auth_header = "Cookie: %s\r\n" % _make_cookie_header(self.service.get_cookies().items())
19131913

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

‎tests/test_binding.py

Lines changed: 32 additions & 14 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.http.cookies, {})
481+
self.assertEqual(self.context.get_cookies(), {})
482482
self.assertRaises(AuthenticationError,
483483
self.context.get, "/services")
484484
self.assertRaises(AuthenticationError,
@@ -510,13 +510,31 @@ def test_cookie_in_auth_headers(self):
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.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_")
513+
self.assertIsNotNone(self.context.get_cookies())
514+
self.assertNotEqual(self.context.get_cookies(), {})
515+
self.assertEqual(len(self.context.get_cookies()), 1)
516+
self.assertEqual(self.context.get_cookies().keys()[0][:8], "splunkd_")
517+
518+
def test_cookie_with_autologin(self):
519+
self.context.autologin = True
520+
self.assertEqual(self.context.get("/services").status, 200)
521+
self.assertTrue(self.context.has_cookies())
522+
self.context.logout()
523+
self.assertFalse(self.context.has_cookies())
524+
self.assertEqual(self.context.get("/services").status, 200)
525+
self.assertTrue(self.context.has_cookies())
526+
527+
def test_cookie_without_autologin(self):
528+
self.context.autologin = False
529+
self.assertEqual(self.context.get("/services").status, 200)
530+
self.assertTrue(self.context.has_cookies())
531+
self.context.logout()
532+
self.assertFalse(self.context.has_cookies())
533+
self.assertRaises(AuthenticationError,
534+
self.context.get, "/services")
517535

518536
def test_got_updated_cookie_with_get(self):
519-
old_cookies = self.context.http.cookies
537+
old_cookies = self.context.get_cookies()
520538
resp = self.context.get("apps/local")
521539
found = False
522540
for key, value in resp.headers:
@@ -525,7 +543,7 @@ def test_got_updated_cookie_with_get(self):
525543
self.assertEqual(value[:8], "splunkd_")
526544

527545
new_cookies = {}
528-
binding.parse_cookies(value, new_cookies)
546+
binding._parse_cookies(value, new_cookies)
529547
# We're only expecting 1 in this scenario
530548
self.assertEqual(len(old_cookies), 1)
531549
self.assertTrue(len(new_cookies.values()), 1)
@@ -552,15 +570,15 @@ def test_login_with_multiple_cookies(self):
552570
except AuthenticationError as ae:
553571
self.assertEqual(ae.message, "Request failed: Session is not logged in.")
554572
# Bring in a valid cookie now
555-
for key, value in self.context.http.cookies.items():
556-
new_context.http.cookies[key] = value
573+
for key, value in self.context.get_cookies().items():
574+
new_context.get_cookies()[key] = value
557575

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())
576+
self.assertEqual(len(new_context.get_cookies()), 2)
577+
self.assertTrue('bad' in new_context.get_cookies().keys())
578+
self.assertTrue('cookie' in new_context.get_cookies().values())
561579

562-
for k, v in self.context.http.cookies.items():
563-
self.assertEqual(new_context.http.cookies[k], v)
580+
for k, v in self.context.get_cookies().items():
581+
self.assertEqual(new_context.get_cookies()[k], v)
564582

565583
self.assertEqual(new_context.get("apps/local").status, 200)
566584

‎tests/test_index.py

Lines changed: 2 additions & 2 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.http.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,7 +121,7 @@ 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.http.cookies.update(self.service.http.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")

‎tests/test_service.py

Lines changed: 40 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
import splunklib.client as client
2121
from splunklib.client import AuthenticationError
2222
from splunklib.client import Service
23-
from splunklib.binding import HTTPError, _NoAuthenticationToken
23+
from splunklib.binding import HTTPError
2424

2525

2626
class ServiceTestCase(testlib.SDKTestCase):
@@ -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.http.cookies)
156-
self.assertEquals(len(self.service.http.cookies), 0)
155+
self.assertIsNotNone(self.service.get_cookies())
156+
self.assertEquals(len(self.service.get_cookies()), 0)
157157
self.service.login()
158-
self.assertIsNotNone(self.service.http.cookies)
159-
self.assertNotEquals(self.service.http.cookies, {})
160-
self.assertEquals(len(self.service.http.cookies), 1)
158+
self.assertIsNotNone(self.service.get_cookies())
159+
self.assertNotEquals(self.service.get_cookies(), {})
160+
self.assertEquals(len(self.service.get_cookies()), 1)
161161

162162
def test_login_with_cookie(self):
163163
self.service.login()
164-
self.assertIsNotNone(self.service.http.cookies)
164+
self.assertIsNotNone(self.service.get_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.http.cookies.items()[0]})
166+
service2 = client.Service(**{"cookie": "%s=%s" % self.service.get_cookies().items()[0]})
167167
service2.login()
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_")
168+
self.assertEqual(len(service2.get_cookies()), 1)
169+
self.assertEqual(service2.get_cookies(), self.service.get_cookies())
170+
self.assertEqual(len(service2.get_cookies()), len(self.service.get_cookies()))
171+
self.assertEqual(service2.get_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.http.cookies), 0)
178-
service2.http.cookies.update(bad_cookie)
177+
self.assertEquals(len(service2.get_cookies()), 0)
178+
service2.get_cookies().update(bad_cookie)
179179
service2.login()
180-
self.assertEquals(service2.http.cookies, {'bad': 'cookie'})
180+
self.assertEquals(service2.get_cookies(), {'bad': 'cookie'})
181181

182182
# Should get an error with a bad cookie
183183
try:
@@ -186,9 +186,21 @@ def test_login_fails_with_bad_cookie(self):
186186
except AuthenticationError as ae:
187187
self.assertEqual(ae.message, "Request failed: Session is not logged in.")
188188

189+
def test_autologin_with_cookie(self):
190+
self.service.login()
191+
self.assertTrue(self.service.has_cookies())
192+
service = client.connect(
193+
autologin=True,
194+
cookie="%s=%s" % self.service.get_cookies().items()[0],
195+
**self.opts.kwargs)
196+
self.assertTrue(service.has_cookies())
197+
self.service.restart(timeout=120)
198+
reader = service.jobs.oneshot("search index=internal | head 1")
199+
self.assertIsNotNone(reader)
200+
189201
def test_login_fails_with_no_cookie(self):
190202
service2 = client.Service()
191-
self.assertEquals(len(service2.http.cookies), 0)
203+
self.assertEquals(len(service2.get_cookies()), 0)
192204

193205
# Should get an error when no authentication method
194206
try:
@@ -203,15 +215,15 @@ def test_login_with_multiple_cookie_headers(self):
203215
'something_else': 'bad'
204216
}
205217
self.service.logout()
206-
self.service.http.cookies.update(cookies)
218+
self.service.get_cookies().update(cookies)
207219

208220
self.service.login()
209221
self.assertEqual(self.service.apps.get().status, 200)
210222

211223
def test_login_with_multiple_cookies(self):
212224
bad_cookie = 'bad=cookie'
213225
self.service.login()
214-
self.assertIsNotNone(self.service.http.cookies)
226+
self.assertIsNotNone(self.service.get_cookies())
215227

216228
service2 = client.Service(**{"cookie": bad_cookie})
217229
service2.login()
@@ -224,16 +236,16 @@ def test_login_with_multiple_cookies(self):
224236
self.assertEqual(ae.message, "Request failed: Session is not logged in.")
225237

226238
# Add on valid cookies, and try to use all of them
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())
239+
service2.get_cookies().update(self.service.get_cookies())
240+
241+
self.assertEqual(len(service2.get_cookies()), 2)
242+
self.service.get_cookies().update({'bad': 'cookie'})
243+
self.assertEqual(service2.get_cookies(), self.service.get_cookies())
244+
self.assertEqual(len(service2.get_cookies()), 2)
245+
self.assertEqual(service2.get_cookies().keys()[1][:8], "splunkd_")
246+
self.assertTrue('bad' in service2.get_cookies().keys())
247+
self.assertEqual(service2.get_cookies()['bad'], 'cookie')
248+
self.assertEqual(self.service.get_cookies().items(), service2.get_cookies().items())
237249
service2.login()
238250
self.assertEqual(service2.apps.get().status, 200)
239251

0 commit comments

Comments
 (0)