cookiejar.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190
  1. from requests.cookies import RequestsCookieJar as _RequestsCookieJar
  2. from http.cookiejar import Cookie as _Cookie
  3. _cookie_attrs = dict(
  4. version=0, name="", value="",
  5. port=None, domain='', path='/',
  6. secure=False, expires=None, discard=False,
  7. comment=None, comment_url=None,
  8. rfc2109=False,
  9. )
  10. _bool_attrs = (
  11. ('port_specified', lambda c: bool(c['port'])),
  12. ('domain_specified', lambda c: bool(c['domain'])),
  13. ('domain_initial_dot', lambda c: c['domain'].startswith('.')),
  14. ('path_specified', lambda c: bool(c['path'])),
  15. )
  16. class ClientCookieJar(_RequestsCookieJar):
  17. """
  18. Custom CookieJar that can be saved to/loaded from dicts
  19. """
  20. # TODO: bring this into line with RequestsCookieJar (it was formerly inheriting http.cookielib.CookieJar)
  21. @staticmethod
  22. def cookie_to_dict(cookie):
  23. dct = {}
  24. for attr in _cookie_attrs:
  25. val = getattr(cookie, attr)
  26. if val == _cookie_attrs[attr]:
  27. # don't store default values
  28. continue
  29. dct[attr] = getattr(cookie, attr)
  30. if cookie._rest:
  31. dct['rest'] = cookie._rest
  32. return dct
  33. @staticmethod
  34. def cookie_from_dict(dct):
  35. """
  36. Constructs a cookie from a dict.
  37. Fills in any missing parameters from a set of defaults.
  38. This method was based on Requests' "create_cookie"
  39. function, originally written by Miguel Turner (dhagrow):
  40. https://github.com/dhagrow/requests/blame/develop/requests/packages/oreos/cookiejar.py#L126
  41. """
  42. if 'name' not in dct or 'value' not in dct:
  43. raise TypeError('Cookie dictionary must contain name and value')
  44. cookie_kwargs = _cookie_attrs.copy()
  45. cookie_kwargs['rest'] = {}
  46. extra_args = set(dct) - set(cookie_kwargs)
  47. if extra_args:
  48. err = 'Unexpected keys in Cookie dictionary: {}'
  49. raise TypeError(err.format(sorted(extra_args)))
  50. cookie_kwargs.update(dct)
  51. for key, func in _bool_attrs:
  52. cookie_kwargs[key] = func(cookie_kwargs)
  53. return _Cookie(**cookie_kwargs)
  54. def __init__(self, cookies=None, policy=None):
  55. _RequestsCookieJar.__init__(self, policy)
  56. if cookies:
  57. self.update(cookies)
  58. def update(self, cookies):
  59. """
  60. Updates from a dictionary of cookies, optionally nested by domain and/or path
  61. May raise TypeError if cookie_dict is invalid or contains unexpected keys
  62. """
  63. self._update(cookies, domain=None, path=None)
  64. def _update(self, cookies, domain=None, path=None):
  65. if not cookies:
  66. return
  67. for key in cookies:
  68. # will fail if path or cookie.name == 'name'
  69. # as is, this check allows mixed nesting
  70. # e.g. cookies and domains at the same level
  71. if 'name' not in cookies[key]:
  72. if domain is not None:
  73. if path is not None:
  74. err = 'No Cookies found in dictionary'
  75. raise TypeError(err)
  76. else:
  77. self._update(cookies[key], domain=domain, path=key)
  78. else:
  79. self._update(cookies[key], domain=domain)
  80. else:
  81. self.set_cookie(self.cookie_from_dict(cookies[key]))
  82. @property
  83. def expires_earliest(self):
  84. # if len(self) > 0:
  85. # # sometimes a cookie has no expiration?
  86. # return min([cookie.expires for cookie in self if cookie.expires])
  87. # return None
  88. # Compatibility Note: the default argument was added to min() in Python 3.4
  89. return min([(cookie.expires or 0) for cookie in self], default=None)
  90. def to_dict(self, ignore_domain=False, ignore_path=False):
  91. """
  92. Returns a dict representation of the CookieJar
  93. If more than one domain exists, or more than one path in
  94. each domain, cookies will be nested under their respective
  95. domain/path. Otherwise all cookies will be stored at the
  96. topmost level.
  97. Nesting can be disabled with ignore_domain and ignore_path
  98. Examples:
  99. One domain, one path:
  100. {
  101. cookie1.name: {key: val, ...},
  102. cookie2.name: {key: val, ...},
  103. ...
  104. }
  105. Multiple domains, one path per domain:
  106. {
  107. domain1: {
  108. cookie1.name: {key: val, ...},
  109. ...
  110. },
  111. domain2: {
  112. cookie1.name: {key: val, ...},
  113. ...
  114. },
  115. ...
  116. }
  117. One domain, multiple paths:
  118. {
  119. path1: {
  120. cookie1.name: {key: val, ...},
  121. ...
  122. },
  123. path2: {
  124. cookie1.name: {key: val, ...},
  125. ...
  126. },
  127. ...
  128. }
  129. Multiple domains, multiple paths per domain:
  130. {
  131. domain1: {
  132. path1: {
  133. cookie1.name: {key: val, ...},
  134. ...
  135. },
  136. ...
  137. },
  138. ...
  139. }
  140. set_cookies_from_dict can handle any of the above variants.
  141. """
  142. target = cookie_dict = {}
  143. if not ignore_domain and len(self._cookies) > 1:
  144. nest_domain = True
  145. else:
  146. nest_domain = False
  147. for domain in self._cookies:
  148. if nest_domain:
  149. target = cookie_dict[domain] = {}
  150. if not ignore_path and len(self._cookies[domain]) > 1:
  151. nest_path = True
  152. else:
  153. nest_path = False
  154. for path in self._cookies[domain]:
  155. if nest_path:
  156. target = target[path] = {}
  157. for name in self._cookies[domain][path]:
  158. cookie = self._cookies[domain][path][name]
  159. target[name] = self.cookie_to_dict(cookie)
  160. return cookie_dict