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