1# -*- coding: utf-8 -*- 2 3""" 4Compatibility code to be able to use `cookielib.CookieJar` with requests. 5 6requests.utils imports from here, so be careful with imports. 7""" 8 9import copy 10import time 11import calendar 12import collections 13from .compat import cookielib, urlparse, urlunparse, Morsel 14 15try: 16 import threading 17 # grr, pyflakes: this fixes "redefinition of unused 'threading'" 18 threading 19except ImportError: 20 import dummy_threading as threading 21 22 23class MockRequest(object): 24 """Wraps a `requests.Request` to mimic a `urllib2.Request`. 25 26 The code in `cookielib.CookieJar` expects this interface in order to correctly 27 manage cookie policies, i.e., determine whether a cookie can be set, given the 28 domains of the request and the cookie. 29 30 The original request object is read-only. The client is responsible for collecting 31 the new headers via `get_new_headers()` and interpreting them appropriately. You 32 probably want `get_cookie_header`, defined below. 33 """ 34 35 def __init__(self, request): 36 self._r = request 37 self._new_headers = {} 38 self.type = urlparse(self._r.url).scheme 39 40 def get_type(self): 41 return self.type 42 43 def get_host(self): 44 return urlparse(self._r.url).netloc 45 46 def get_origin_req_host(self): 47 return self.get_host() 48 49 def get_full_url(self): 50 # Only return the response's URL if the user hadn't set the Host 51 # header 52 if not self._r.headers.get('Host'): 53 return self._r.url 54 # If they did set it, retrieve it and reconstruct the expected domain 55 host = self._r.headers['Host'] 56 parsed = urlparse(self._r.url) 57 # Reconstruct the URL as we expect it 58 return urlunparse([ 59 parsed.scheme, host, parsed.path, parsed.params, parsed.query, 60 parsed.fragment 61 ]) 62 63 def is_unverifiable(self): 64 return True 65 66 def has_header(self, name): 67 return name in self._r.headers or name in self._new_headers 68 69 def get_header(self, name, default=None): 70 return self._r.headers.get(name, self._new_headers.get(name, default)) 71 72 def add_header(self, key, val): 73 """cookielib has no legitimate use for this method; add it back if you find one.""" 74 raise NotImplementedError("Cookie headers should be added with add_unredirected_header()") 75 76 def add_unredirected_header(self, name, value): 77 self._new_headers[name] = value 78 79 def get_new_headers(self): 80 return self._new_headers 81 82 @property 83 def unverifiable(self): 84 return self.is_unverifiable() 85 86 @property 87 def origin_req_host(self): 88 return self.get_origin_req_host() 89 90 @property 91 def host(self): 92 return self.get_host() 93 94 95class MockResponse(object): 96 """Wraps a `httplib.HTTPMessage` to mimic a `urllib.addinfourl`. 97 98 ...what? Basically, expose the parsed HTTP headers from the server response 99 the way `cookielib` expects to see them. 100 """ 101 102 def __init__(self, headers): 103 """Make a MockResponse for `cookielib` to read. 104 105 :param headers: a httplib.HTTPMessage or analogous carrying the headers 106 """ 107 self._headers = headers 108 109 def info(self): 110 return self._headers 111 112 def getheaders(self, name): 113 self._headers.getheaders(name) 114 115 116def extract_cookies_to_jar(jar, request, response): 117 """Extract the cookies from the response into a CookieJar. 118 119 :param jar: cookielib.CookieJar (not necessarily a RequestsCookieJar) 120 :param request: our own requests.Request object 121 :param response: urllib3.HTTPResponse object 122 """ 123 if not (hasattr(response, '_original_response') and 124 response._original_response): 125 return 126 # the _original_response field is the wrapped httplib.HTTPResponse object, 127 req = MockRequest(request) 128 # pull out the HTTPMessage with the headers and put it in the mock: 129 res = MockResponse(response._original_response.msg) 130 jar.extract_cookies(res, req) 131 132 133def get_cookie_header(jar, request): 134 """Produce an appropriate Cookie header string to be sent with `request`, or None.""" 135 r = MockRequest(request) 136 jar.add_cookie_header(r) 137 return r.get_new_headers().get('Cookie') 138 139 140def remove_cookie_by_name(cookiejar, name, domain=None, path=None): 141 """Unsets a cookie by name, by default over all domains and paths. 142 143 Wraps CookieJar.clear(), is O(n). 144 """ 145 clearables = [] 146 for cookie in cookiejar: 147 if cookie.name != name: 148 continue 149 if domain is not None and domain != cookie.domain: 150 continue 151 if path is not None and path != cookie.path: 152 continue 153 clearables.append((cookie.domain, cookie.path, cookie.name)) 154 155 for domain, path, name in clearables: 156 cookiejar.clear(domain, path, name) 157 158 159class CookieConflictError(RuntimeError): 160 """There are two cookies that meet the criteria specified in the cookie jar. 161 Use .get and .set and include domain and path args in order to be more specific.""" 162 163 164class RequestsCookieJar(cookielib.CookieJar, collections.MutableMapping): 165 """Compatibility class; is a cookielib.CookieJar, but exposes a dict 166 interface. 167 168 This is the CookieJar we create by default for requests and sessions that 169 don't specify one, since some clients may expect response.cookies and 170 session.cookies to support dict operations. 171 172 Requests does not use the dict interface internally; it's just for 173 compatibility with external client code. All requests code should work 174 out of the box with externally provided instances of ``CookieJar``, e.g. 175 ``LWPCookieJar`` and ``FileCookieJar``. 176 177 Unlike a regular CookieJar, this class is pickleable. 178 179 .. warning:: dictionary operations that are normally O(1) may be O(n). 180 """ 181 def get(self, name, default=None, domain=None, path=None): 182 """Dict-like get() that also supports optional domain and path args in 183 order to resolve naming collisions from using one cookie jar over 184 multiple domains. 185 186 .. warning:: operation is O(n), not O(1).""" 187 try: 188 return self._find_no_duplicates(name, domain, path) 189 except KeyError: 190 return default 191 192 def set(self, name, value, **kwargs): 193 """Dict-like set() that also supports optional domain and path args in 194 order to resolve naming collisions from using one cookie jar over 195 multiple domains.""" 196 # support client code that unsets cookies by assignment of a None value: 197 if value is None: 198 remove_cookie_by_name(self, name, domain=kwargs.get('domain'), path=kwargs.get('path')) 199 return 200 201 if isinstance(value, Morsel): 202 c = morsel_to_cookie(value) 203 else: 204 c = create_cookie(name, value, **kwargs) 205 self.set_cookie(c) 206 return c 207 208 def iterkeys(self): 209 """Dict-like iterkeys() that returns an iterator of names of cookies 210 from the jar. See itervalues() and iteritems().""" 211 for cookie in iter(self): 212 yield cookie.name 213 214 def keys(self): 215 """Dict-like keys() that returns a list of names of cookies from the 216 jar. See values() and items().""" 217 return list(self.iterkeys()) 218 219 def itervalues(self): 220 """Dict-like itervalues() that returns an iterator of values of cookies 221 from the jar. See iterkeys() and iteritems().""" 222 for cookie in iter(self): 223 yield cookie.value 224 225 def values(self): 226 """Dict-like values() that returns a list of values of cookies from the 227 jar. See keys() and items().""" 228 return list(self.itervalues()) 229 230 def iteritems(self): 231 """Dict-like iteritems() that returns an iterator of name-value tuples 232 from the jar. See iterkeys() and itervalues().""" 233 for cookie in iter(self): 234 yield cookie.name, cookie.value 235 236 def items(self): 237 """Dict-like items() that returns a list of name-value tuples from the 238 jar. See keys() and values(). Allows client-code to call 239 ``dict(RequestsCookieJar)`` and get a vanilla python dict of key value 240 pairs.""" 241 return list(self.iteritems()) 242 243 def list_domains(self): 244 """Utility method to list all the domains in the jar.""" 245 domains = [] 246 for cookie in iter(self): 247 if cookie.domain not in domains: 248 domains.append(cookie.domain) 249 return domains 250 251 def list_paths(self): 252 """Utility method to list all the paths in the jar.""" 253 paths = [] 254 for cookie in iter(self): 255 if cookie.path not in paths: 256 paths.append(cookie.path) 257 return paths 258 259 def multiple_domains(self): 260 """Returns True if there are multiple domains in the jar. 261 Returns False otherwise.""" 262 domains = [] 263 for cookie in iter(self): 264 if cookie.domain is not None and cookie.domain in domains: 265 return True 266 domains.append(cookie.domain) 267 return False # there is only one domain in jar 268 269 def get_dict(self, domain=None, path=None): 270 """Takes as an argument an optional domain and path and returns a plain 271 old Python dict of name-value pairs of cookies that meet the 272 requirements.""" 273 dictionary = {} 274 for cookie in iter(self): 275 if (domain is None or cookie.domain == domain) and (path is None 276 or cookie.path == path): 277 dictionary[cookie.name] = cookie.value 278 return dictionary 279 280 def __getitem__(self, name): 281 """Dict-like __getitem__() for compatibility with client code. Throws 282 exception if there are more than one cookie with name. In that case, 283 use the more explicit get() method instead. 284 285 .. warning:: operation is O(n), not O(1).""" 286 287 return self._find_no_duplicates(name) 288 289 def __setitem__(self, name, value): 290 """Dict-like __setitem__ for compatibility with client code. Throws 291 exception if there is already a cookie of that name in the jar. In that 292 case, use the more explicit set() method instead.""" 293 294 self.set(name, value) 295 296 def __delitem__(self, name): 297 """Deletes a cookie given a name. Wraps ``cookielib.CookieJar``'s 298 ``remove_cookie_by_name()``.""" 299 remove_cookie_by_name(self, name) 300 301 def set_cookie(self, cookie, *args, **kwargs): 302 if hasattr(cookie.value, 'startswith') and cookie.value.startswith('"') and cookie.value.endswith('"'): 303 cookie.value = cookie.value.replace('\\"', '') 304 return super(RequestsCookieJar, self).set_cookie(cookie, *args, **kwargs) 305 306 def update(self, other): 307 """Updates this jar with cookies from another CookieJar or dict-like""" 308 if isinstance(other, cookielib.CookieJar): 309 for cookie in other: 310 self.set_cookie(copy.copy(cookie)) 311 else: 312 super(RequestsCookieJar, self).update(other) 313 314 def _find(self, name, domain=None, path=None): 315 """Requests uses this method internally to get cookie values. Takes as 316 args name and optional domain and path. Returns a cookie.value. If 317 there are conflicting cookies, _find arbitrarily chooses one. See 318 _find_no_duplicates if you want an exception thrown if there are 319 conflicting cookies.""" 320 for cookie in iter(self): 321 if cookie.name == name: 322 if domain is None or cookie.domain == domain: 323 if path is None or cookie.path == path: 324 return cookie.value 325 326 raise KeyError('name=%r, domain=%r, path=%r' % (name, domain, path)) 327 328 def _find_no_duplicates(self, name, domain=None, path=None): 329 """Both ``__get_item__`` and ``get`` call this function: it's never 330 used elsewhere in Requests. Takes as args name and optional domain and 331 path. Returns a cookie.value. Throws KeyError if cookie is not found 332 and CookieConflictError if there are multiple cookies that match name 333 and optionally domain and path.""" 334 toReturn = None 335 for cookie in iter(self): 336 if cookie.name == name: 337 if domain is None or cookie.domain == domain: 338 if path is None or cookie.path == path: 339 if toReturn is not None: # if there are multiple cookies that meet passed in criteria 340 raise CookieConflictError('There are multiple cookies with name, %r' % (name)) 341 toReturn = cookie.value # we will eventually return this as long as no cookie conflict 342 343 if toReturn: 344 return toReturn 345 raise KeyError('name=%r, domain=%r, path=%r' % (name, domain, path)) 346 347 def __getstate__(self): 348 """Unlike a normal CookieJar, this class is pickleable.""" 349 state = self.__dict__.copy() 350 # remove the unpickleable RLock object 351 state.pop('_cookies_lock') 352 return state 353 354 def __setstate__(self, state): 355 """Unlike a normal CookieJar, this class is pickleable.""" 356 self.__dict__.update(state) 357 if '_cookies_lock' not in self.__dict__: 358 self._cookies_lock = threading.RLock() 359 360 def copy(self): 361 """Return a copy of this RequestsCookieJar.""" 362 new_cj = RequestsCookieJar() 363 new_cj.update(self) 364 return new_cj 365 366 367def _copy_cookie_jar(jar): 368 if jar is None: 369 return None 370 371 if hasattr(jar, 'copy'): 372 # We're dealing with an instance of RequestsCookieJar 373 return jar.copy() 374 # We're dealing with a generic CookieJar instance 375 new_jar = copy.copy(jar) 376 new_jar.clear() 377 for cookie in jar: 378 new_jar.set_cookie(copy.copy(cookie)) 379 return new_jar 380 381 382def create_cookie(name, value, **kwargs): 383 """Make a cookie from underspecified parameters. 384 385 By default, the pair of `name` and `value` will be set for the domain '' 386 and sent on every request (this is sometimes called a "supercookie"). 387 """ 388 result = dict( 389 version=0, 390 name=name, 391 value=value, 392 port=None, 393 domain='', 394 path='/', 395 secure=False, 396 expires=None, 397 discard=True, 398 comment=None, 399 comment_url=None, 400 rest={'HttpOnly': None}, 401 rfc2109=False,) 402 403 badargs = set(kwargs) - set(result) 404 if badargs: 405 err = 'create_cookie() got unexpected keyword arguments: %s' 406 raise TypeError(err % list(badargs)) 407 408 result.update(kwargs) 409 result['port_specified'] = bool(result['port']) 410 result['domain_specified'] = bool(result['domain']) 411 result['domain_initial_dot'] = result['domain'].startswith('.') 412 result['path_specified'] = bool(result['path']) 413 414 return cookielib.Cookie(**result) 415 416 417def morsel_to_cookie(morsel): 418 """Convert a Morsel object into a Cookie containing the one k/v pair.""" 419 420 expires = None 421 if morsel['max-age']: 422 try: 423 expires = int(time.time() + int(morsel['max-age'])) 424 except ValueError: 425 raise TypeError('max-age: %s must be integer' % morsel['max-age']) 426 elif morsel['expires']: 427 time_template = '%a, %d-%b-%Y %H:%M:%S GMT' 428 expires = calendar.timegm( 429 time.strptime(morsel['expires'], time_template) 430 ) 431 return create_cookie( 432 comment=morsel['comment'], 433 comment_url=bool(morsel['comment']), 434 discard=False, 435 domain=morsel['domain'], 436 expires=expires, 437 name=morsel.key, 438 path=morsel['path'], 439 port=None, 440 rest={'HttpOnly': morsel['httponly']}, 441 rfc2109=False, 442 secure=bool(morsel['secure']), 443 value=morsel.value, 444 version=morsel['version'] or 0, 445 ) 446 447 448def cookiejar_from_dict(cookie_dict, cookiejar=None, overwrite=True): 449 """Returns a CookieJar from a key/value dictionary. 450 451 :param cookie_dict: Dict of key/values to insert into CookieJar. 452 :param cookiejar: (optional) A cookiejar to add the cookies to. 453 :param overwrite: (optional) If False, will not replace cookies 454 already in the jar with new ones. 455 """ 456 if cookiejar is None: 457 cookiejar = RequestsCookieJar() 458 459 if cookie_dict is not None: 460 names_from_jar = [cookie.name for cookie in cookiejar] 461 for name in cookie_dict: 462 if overwrite or (name not in names_from_jar): 463 cookiejar.set_cookie(create_cookie(name, cookie_dict[name])) 464 465 return cookiejar 466 467 468def merge_cookies(cookiejar, cookies): 469 """Add cookies to cookiejar and returns a merged CookieJar. 470 471 :param cookiejar: CookieJar object to add the cookies to. 472 :param cookies: Dictionary or CookieJar object to be added. 473 """ 474 if not isinstance(cookiejar, cookielib.CookieJar): 475 raise ValueError('You can only merge into CookieJar') 476 477 if isinstance(cookies, dict): 478 cookiejar = cookiejar_from_dict( 479 cookies, cookiejar=cookiejar, overwrite=False) 480 elif isinstance(cookies, cookielib.CookieJar): 481 try: 482 cookiejar.update(cookies) 483 except AttributeError: 484 for cookie_in_jar in cookies: 485 cookiejar.set_cookie(cookie_in_jar) 486 487 return cookiejar 488