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