Repository navigation
Expand file tree
/
Copy pathsession.py
More file actions
287 lines (224 loc) · 9.1 KB
/
Copy pathsession.py
File metadata and controls
287 lines (224 loc) · 9.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
# -*- coding: utf-8 -*-
"""Module containing session and auth logic."""
from collections import Callable
import datetime
from contextlib import contextmanager
from logging import getLogger
import dateutil.parser
import requests
from . import __version__
from . import exceptions as exc
__url_cache__ = {}
__logs__ = getLogger(__package__)
def requires_2fa(response):
"""Determine whether a response requires us to prompt the user for 2FA."""
if (
response.status_code == 401
and "X-GitHub-OTP" in response.headers
and "required" in response.headers["X-GitHub-OTP"]
):
return True
return False
class BasicAuth(requests.auth.HTTPBasicAuth):
"""Sub-class requests's class so we have a nice repr."""
def __repr__(self):
"""Use the username as the representation."""
return "basic {}".format(self.username)
class TokenAuth(requests.auth.AuthBase):
"""Auth class that handles simple tokens."""
header_format_str = "token {}"
def __init__(self, token):
"""Store our token."""
self.token = token
def __repr__(self):
"""Return a nice view of the token in use."""
return "token {}...".format(self.token[:4])
def __ne__(self, other):
"""Test for equality, or the lack thereof."""
return not self == other
def __eq__(self, other):
"""Test for equality, or the lack thereof."""
return self.token == getattr(other, "token", None)
def __call__(self, request):
"""Add the authorization header and format it."""
request.headers["Authorization"] = self.header_format_str.format(
self.token
)
return request
class GitHubSession(requests.Session):
"""Our slightly specialized Session object."""
auth = None
__attrs__ = requests.Session.__attrs__ + [
"base_url",
"two_factor_auth_cb",
]
def __init__(self):
"""Slightly modify how we initialize our session."""
super(GitHubSession, self).__init__()
self.headers.update(
{
# Only accept JSON responses
"Accept": "application/vnd.github.v3.full+json",
# Only accept UTF-8 encoded data
"Accept-Charset": "utf-8",
# Always sending JSON
"Content-Type": "application/json",
# Set our own custom User-Agent string
"User-Agent": "github3.py/{0}".format(__version__),
}
)
self.base_url = "https://api.github.com"
self.two_factor_auth_cb = None
self.request_counter = 0
def basic_auth(self, username, password):
"""Set the Basic Auth credentials on this Session.
:param str username: Your GitHub username
:param str password: Your GitHub password
"""
if not (username and password):
return
self.auth = BasicAuth(username, password)
def build_url(self, *args, **kwargs):
"""Build a new API url from scratch."""
parts = [kwargs.get("base_url") or self.base_url]
parts.extend(args)
parts = [str(p) for p in parts]
key = tuple(parts)
__logs__.info("Building a url from %s", key)
if key not in __url_cache__:
__logs__.info("Missed the cache building the url")
__url_cache__[key] = "/".join(parts)
return __url_cache__[key]
def handle_two_factor_auth(self, args, kwargs):
"""Handler for when the user has 2FA turned on."""
headers = kwargs.pop("headers", {})
headers.update({"X-GitHub-OTP": str(self.two_factor_auth_cb())})
kwargs.update(headers=headers)
return super(GitHubSession, self).request(*args, **kwargs)
def has_auth(self):
"""Check for whether or not the user has authentication configured."""
return self.auth or self.headers.get("Authorization")
def oauth2_auth(self, client_id, client_secret):
"""Use OAuth2 for authentication.
It is suggested you install requests-oauthlib to use this.
:param str client_id: Client ID retrieved from GitHub
:param str client_secret: Client secret retrieved from GitHub
"""
raise NotImplementedError("These features are not implemented yet")
def request(self, *args, **kwargs):
"""Make a request, count it, and handle 2FA if necessary."""
response = super(GitHubSession, self).request(*args, **kwargs)
self.request_counter += 1
if requires_2fa(response) and self.two_factor_auth_cb:
# No need to flatten and re-collect the args in
# handle_two_factor_auth
new_response = self.handle_two_factor_auth(args, kwargs)
new_response.history.append(response)
response = new_response
return response
def retrieve_client_credentials(self):
"""Return the client credentials.
:returns: tuple(client_id, client_secret)
"""
client_id = self.params.get("client_id")
client_secret = self.params.get("client_secret")
return (client_id, client_secret)
def two_factor_auth_callback(self, callback):
"""Register our 2FA callback specified by the user."""
if not callback:
return
if not isinstance(callback, Callable):
raise ValueError("Your callback should be callable")
self.two_factor_auth_cb = callback
def token_auth(self, token):
"""Use an application token for authentication.
:param str token: Application token retrieved from GitHub's
/authorizations endpoint
"""
if not token:
return
self.auth = TokenAuth(token)
def app_bearer_token_auth(self, headers, expire_in):
"""Authenticate as an App to be able to view its metadata."""
if not headers:
return
self.auth = AppBearerTokenAuth(headers, expire_in)
def app_installation_token_auth(self, json):
"""Use an access token generated by an App's installation."""
if not json:
return
self.auth = AppInstallationTokenAuth(
json["token"], json["expires_at"]
)
@contextmanager
def temporary_basic_auth(self, *auth):
"""Allow us to temporarily swap out basic auth credentials."""
old_basic_auth = self.auth
old_token_auth = self.headers.get("Authorization")
self.basic_auth(*auth)
yield
self.auth = old_basic_auth
if old_token_auth:
self.headers["Authorization"] = old_token_auth
@contextmanager
def no_auth(self):
"""Unset authentication temporarily as a context manager."""
old_basic_auth, self.auth = self.auth, None
old_token_auth = self.headers.pop("Authorization", None)
yield
self.auth = old_basic_auth
if old_token_auth:
self.headers["Authorization"] = old_token_auth
def _utcnow():
return datetime.datetime.now(dateutil.tz.UTC)
class AppInstallationTokenAuth(TokenAuth):
"""Use token authentication but throw an exception on expiration."""
def __init__(self, token, expires_at):
"""Set-up our authentication handler."""
super(AppInstallationTokenAuth, self).__init__(token)
self.expires_at_str = expires_at
self.expires_at = dateutil.parser.parse(expires_at)
def __repr__(self):
"""Return a nice view of the token in use."""
return "app installation token {}... expiring at {}".format(
self.token[:4], self.expires_at_str
)
@property
def expired(self):
"""Indicate whether our token is expired or not."""
now = _utcnow()
return now > self.expires_at
def __call__(self, request):
"""Add the authorization header and format it."""
if self.expired:
raise exc.AppInstallationTokenExpired(
"Your app installation token expired at {}".format(
self.expires_at_str
)
)
return super(AppInstallationTokenAuth, self).__call__(request)
class AppBearerTokenAuth(TokenAuth):
"""Use JWT authentication but throw an exception on expiration."""
header_format_str = "Bearer {}"
def __init__(self, token, expire_in):
"""Set-up our authentication handler."""
super(AppBearerTokenAuth, self).__init__(token)
expire_in = datetime.timedelta(seconds=expire_in)
self.expires_at = _utcnow() + expire_in
def __repr__(self):
"""Return a helpful view of the token."""
return "app bearer token {} expiring at {}".format(
self.token[:4], str(self.expires_at)
)
@property
def expired(self):
"""Indicate whether our token is expired or not."""
now = _utcnow()
return now > self.expires_at
def __call__(self, request):
"""Add the authorization header and format it."""
if self.expired:
raise exc.AppTokenExpired(
"Your app token expired at {}".format(str(self.expires_at))
)
return super(AppBearerTokenAuth, self).__call__(request)