~ubuntu-branches/ubuntu/trusty/requests/trusty-proposed : revision 11

1

"""

2

Compatibility code to be able to use `cookielib.CookieJar` with requests.

3

4

requests.utils imports from here, so be careful with imports.

5

"""

6

7

import collections

8

from .compat import cookielib, urlparse, Morsel

9

10

try:

11

import threading

12

# grr, pyflakes: this fixes "redefinition of unused 'threading'"

13

threading

14

except ImportError:

15

import dummy_threading as threading

16

17

class MockRequest(object):

18

"""Wraps a `requests.Request` to mimic a `urllib2.Request`.

19

20

The code in `cookielib.CookieJar` expects this interface in order to correctly

21

manage cookie policies, i.e., determine whether a cookie can be set, given the

22

domains of the request and the cookie.

23

24

The original request object is read-only. The client is responsible for collecting

25

the new headers via `get_new_headers()` and interpreting them appropriately. You

26

probably want `get_cookie_header`, defined below.

27

"""

28

29

def __init__(self, request):

30

self._r = request

31

self._new_headers = {}

32

33

def get_type(self):

34

return urlparse(self._r.full_url).scheme

35

36

def get_host(self):

37

return urlparse(self._r.full_url).netloc

38

39

def get_origin_req_host(self):

40

if self._r.response.history:

41

r = self._r.response.history[0]

42

return urlparse(r).netloc

43

else:

44

return self.get_host()

45

46

def get_full_url(self):

47

return self._r.full_url

48

49

def is_unverifiable(self):

50

# unverifiable == redirected

51

return bool(self._r.response.history)

52

53

def has_header(self, name):

54

return name in self._r.headers or name in self._new_headers

55

56

def get_header(self, name, default=None):

57

return self._r.headers.get(name, self._new_headers.get(name, default))

58

59

def add_header(self, key, val):

60

"""cookielib has no legitimate use for this method; add it back if you find one."""

61

raise NotImplementedError("Cookie headers should be added with add_unredirected_header()")

62

63

def add_unredirected_header(self, name, value):

64

self._new_headers[name] = value

65

66

def get_new_headers(self):

67

return self._new_headers

68

69

class MockResponse(object):

70

"""Wraps a `httplib.HTTPMessage` to mimic a `urllib.addinfourl`.

71

72

...what? Basically, expose the parsed HTTP headers from the server response

73

the way `cookielib` expects to see them.

74

"""

75

76

def __init__(self, headers):

77

"""Make a MockResponse for `cookielib` to read.

78

79

:param headers: a httplib.HTTPMessage or analogous carrying the headers

80

"""

81

self._headers = headers

82

83

def info(self):

84

return self._headers

85

86

def getheaders(self, name):

87

self._headers.getheaders(name)

88

89

def extract_cookies_to_jar(jar, request, response):

90

"""Extract the cookies from the response into a CookieJar.

91

92

:param jar: cookielib.CookieJar (not necessarily a RequestsCookieJar)

93

:param request: our own requests.Request object

94

:param response: urllib3.HTTPResponse object

95

"""

96

# the _original_response field is the wrapped httplib.HTTPResponse object,

97

# and in safe mode, it may be None if the request didn't actually complete.

98

# in that case, just skip the cookie extraction.

99

if response._original_response is not None:

100

req = MockRequest(request)

101

# pull out the HTTPMessage with the headers and put it in the mock:

102

res = MockResponse(response._original_response.msg)

103

jar.extract_cookies(res, req)

104

105

def get_cookie_header(jar, request):

106

"""Produce an appropriate Cookie header string to be sent with `request`, or None."""

107

r = MockRequest(request)

108

jar.add_cookie_header(r)

109

return r.get_new_headers().get('Cookie')

110

111

def remove_cookie_by_name(cookiejar, name, domain=None, path=None):

112

"""Unsets a cookie by name, by default over all domains and paths.

113

114

Wraps CookieJar.clear(), is O(n).

115

"""

116

clearables = []

117

for cookie in cookiejar:

118

if cookie.name == name:

119

if domain is None or domain == cookie.domain:

120

if path is None or path == cookie.path:

121

clearables.append((cookie.domain, cookie.path, cookie.name))

122

123

for domain, path, name in clearables:

124

cookiejar.clear(domain, path, name)

125

126

class RequestsCookieJar(cookielib.CookieJar, collections.MutableMapping):

127

"""Compatibility class; is a cookielib.CookieJar, but exposes a dict interface.

128

129

This is the CookieJar we create by default for requests and sessions that

130

don't specify one, since some clients may expect response.cookies and

131

session.cookies to support dict operations.

132

133

Don't use the dict interface internally; it's just for compatibility with

134

with external client code. All `requests` code should work out of the box

135

with externally provided instances of CookieJar, e.g., LWPCookieJar and

136

FileCookieJar.

137

138

Caution: dictionary operations that are normally O(1) may be O(n).

139

140

Unlike a regular CookieJar, this class is pickleable.

141

"""

142

143

def get(self, name, domain=None, path=None, default=None):

144

try:

145

return self._find(name, domain, path)

146

except KeyError:

147

return default

148

149

def set(self, name, value, **kwargs):

150

# support client code that unsets cookies by assignment of a None value:

151

if value is None:

152

remove_cookie_by_name(self, name, domain=kwargs.get('domain'), path=kwargs.get('path'))

153

return

154

155

if isinstance(value, Morsel):

156

c = morsel_to_cookie(value)

157

else:

158

c = create_cookie(name, value, **kwargs)

159

self.set_cookie(c)

160

return c

161

162

def __getitem__(self, name):

163

return self._find(name)

164

165

def __setitem__(self, name, value):

166

self.set(name, value)

167

168

def __delitem__(self, name):

169

remove_cookie_by_name(self, name)

170

171

def _find(self, name, domain=None, path=None):

172

for cookie in iter(self):

173

if cookie.name == name:

174

if domain is None or cookie.domain == domain:

175

if path is None or cookie.path == path:

176

return cookie.value

177

178

raise KeyError('name=%r, domain=%r, path=%r' % (name, domain, path))

179

180

def __getstate__(self):

181

state = self.__dict__.copy()

182

# remove the unpickleable RLock object

183

state.pop('_cookies_lock')

184

return state

185

186

def __setstate__(self, state):

187

self.__dict__.update(state)

188

if '_cookies_lock' not in self.__dict__:

189

self._cookies_lock = threading.RLock()

190

191

def copy(self):

192

"""We're probably better off forbidding this."""

193

raise NotImplementedError

194

195

def create_cookie(name, value, **kwargs):

196

"""Make a cookie from underspecified parameters.

197

198

By default, the pair of `name` and `value` will be set for the domain ''

199

and sent on every request (this is sometimes called a "supercookie").

200

"""

201

result = dict(

202

version=0,

203

name=name,

204

value=value,

205

port=None,

206

domain='',

207

path='/',

208

secure=False,

209

expires=None,

210

discard=True,

211

comment=None,

212

comment_url=None,

213

rest={'HttpOnly': None},

214

rfc2109=False,

215

)

216

217

badargs = set(kwargs) - set(result)

218

if badargs:

219

err = 'create_cookie() got unexpected keyword arguments: %s'

220

raise TypeError(err % list(badargs))

221

222

result.update(kwargs)

223

result['port_specified'] = bool(result['port'])

224

result['domain_specified'] = bool(result['domain'])

225

result['domain_initial_dot'] = result['domain'].startswith('.')

226

result['path_specified'] = bool(result['path'])

227

228

return cookielib.Cookie(**result)

229

230

def morsel_to_cookie(morsel):

231

"""Convert a Morsel object into a Cookie containing the one k/v pair."""

232

c = create_cookie(

233

name=morsel.key,

234

value=morsel.value,

235

version=morsel['version'] or 0,

236

port=None,

237

port_specified=False,

238

domain=morsel['domain'],

239

domain_specified=bool(morsel['domain']),

240

domain_initial_dot=morsel['domain'].startswith('.'),

241

path=morsel['path'],

242

path_specified=bool(morsel['path']),

243

secure=bool(morsel['secure']),

244

expires=morsel['max-age'] or morsel['expires'],

245

discard=False,

246

comment=morsel['comment'],

247

comment_url=bool(morsel['comment']),

248

rest={'HttpOnly': morsel['httponly']},

249

rfc2109=False,

250

)

251

return c

252

253

def cookiejar_from_dict(cookie_dict, cookiejar=None):

254

"""Returns a CookieJar from a key/value dictionary.

255

256

:param cookie_dict: Dict of key/values to insert into CookieJar.

257

"""

258

if cookiejar is None:

259

cookiejar = RequestsCookieJar()

260

261

if cookie_dict is not None:

262

for name in cookie_dict:

263

cookiejar.set_cookie(create_cookie(name, cookie_dict[name]))

264

return cookiejar