This repository was archived by the owner on May 22, 2021. It is now read-only.
Repository navigation
Expand file tree
/
Copy pathsession.py
More file actions
153 lines (121 loc) · 4.89 KB
/
Copy pathsession.py
File metadata and controls
153 lines (121 loc) · 4.89 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
# -*- coding: utf-8 -*-
import requests
from collections import Callable
from . import __version__
from logging import getLogger
from contextlib import contextmanager
__url_cache__ = {}
__logs__ = getLogger(__package__)
def requires_2fa(response):
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 GitHubSession(requests.Session):
auth = None
__attrs__ = requests.Session.__attrs__ + ['base_url', 'two_factor_auth_cb']
def __init__(self):
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 = (username, password)
# Disable token authentication
self.headers.pop('Authorization', None)
def build_url(self, *args, **kwargs):
"""Builds 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):
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):
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):
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):
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.headers.update({
'Authorization': 'token {0}'.format(token)
})
# Unset username/password so we stop sending them
self.auth = None
@contextmanager
def temporary_basic_auth(self, *auth):
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