Commit | Line | Data |
---|---|---|
a29d3526 | 1 | #!/usr/bin/env python |
e2c48b2f | 2 | # -*- coding: utf-8 -*- |
a29d3526 | 3 | |
36f274c0 MM |
4 | |
5 | """This module abstracts connection to pod. | |
6 | """ | |
7 | ||
8 | ||
b74513c5 | 9 | import json |
8993810c MM |
10 | import re |
11 | import requests | |
ff3d2ab4 | 12 | import warnings |
8993810c | 13 | |
2cf8467c MM |
14 | from diaspy import errors |
15 | ||
a29d3526 | 16 | |
36f274c0 | 17 | DEBUG = True |
178faa46 MM |
18 | |
19 | ||
a29d3526 | 20 | class Connection(): |
d95ff94a C |
21 | """Object representing connection with the pod. |
22 | """ | |
23 | _token_regex = re.compile(r'name="csrf-token"\s+content="(.*?)"') | |
24 | _userinfo_regex = re.compile(r'window.current_user_attributes = ({.*})') | |
25 | # this is for older version of D* | |
26 | _token_regex_2 = re.compile(r'content="(.*?)"\s+name="csrf-token') | |
27 | _userinfo_regex_2 = re.compile(r'gon.user=({.*?});gon.') | |
28 | _verify_SSL = True | |
29 | ||
8bfe79c7 | 30 | def __init__(self, pod, username, password, schema='https', **requestsKwargs): |
d95ff94a C |
31 | """ |
32 | :param pod: The complete url of the diaspora pod to use. | |
33 | :type pod: str | |
34 | :param username: The username used to log in. | |
35 | :type username: str | |
36 | :param password: The password used to log in. | |
37 | :type password: str | |
8bfe79c7 C |
38 | :param requestsKwargs: default kwargs for requests (proxy, timeout, etc) |
39 | :type requestsKwargs: keyworded arguments | |
d95ff94a C |
40 | """ |
41 | self.pod = pod | |
42 | self._session = requests.Session() | |
43 | self._login_data = {'user[remember_me]': 1, 'utf8': '✓'} | |
44 | self._userdata = {} | |
45 | self._token = '' | |
46 | self._diaspora_session = '' | |
d95ff94a | 47 | self._fetch_token_from = 'stream' |
b9d087bb C |
48 | self._requests_kwargs = {'verify':self._verify_SSL} |
49 | if requestsKwargs: self._requests_kwargs.update(requestsKwargs) | |
50 | ||
8bfe79c7 | 51 | self._camo_enabled = False |
74dbe1b3 | 52 | try: self._setlogin(username, password) |
d95ff94a C |
53 | except requests.exceptions.MissingSchema: |
54 | self.pod = '{0}://{1}'.format(schema, self.pod) | |
55 | warnings.warn('schema was missing') | |
8bfe79c7 | 56 | |
74dbe1b3 C |
57 | try: self._setlogin(username, password) |
58 | except Exception as e: | |
59 | raise errors.LoginError('cannot create login data (caused by: {0})'.format(e)) | |
60 | self._cookies = self._fetchcookies() | |
d95ff94a | 61 | |
8bfe79c7 C |
62 | def __bool__(self): |
63 | if self._token: return True | |
64 | return False | |
65 | ||
d95ff94a C |
66 | def _fetchcookies(self): |
67 | request = self.get('stream') | |
68 | return request.cookies | |
69 | ||
70 | def __repr__(self): | |
71 | """Returns token string. | |
72 | It will be easier to change backend if programs will just use: | |
73 | repr(connection) | |
74 | instead of calling a specified method. | |
75 | """ | |
76 | return self._fetchtoken() | |
77 | ||
8bfe79c7 C |
78 | def requestsKwargs(self): |
79 | """Returns keyworded arguments set to use for all requests. | |
80 | """ | |
81 | return self._requests_kwargs | |
82 | ||
83 | def setRequestsKwargs(self, **requestsKwargs): | |
84 | """Sets keyworded arguments that will be used for earch request. | |
85 | """ | |
86 | self._requests_kwargs = requestsKwargs | |
87 | ||
d95ff94a C |
88 | def get(self, string, headers={}, params={}, direct=False, **kwargs): |
89 | """This method gets data from session. | |
90 | Performs additional checks if needed. | |
91 | ||
92 | Example: | |
93 | To obtain 'foo' from pod one should call `get('foo')`. | |
94 | ||
95 | :param string: URL to get without the pod's URL and slash eg. 'stream'. | |
96 | :type string: str | |
97 | :param direct: if passed as True it will not be expanded | |
98 | :type direct: bool | |
99 | """ | |
100 | if not direct: url = '{0}/{1}'.format(self.pod, string) | |
101 | else: url = string | |
8bfe79c7 | 102 | if not kwargs: kwargs = self._requests_kwargs |
b9d087bb | 103 | return self._session.get(url, params=params, headers=headers, **kwargs) |
d95ff94a C |
104 | |
105 | def tokenFrom(self, location): | |
106 | """Sets location for the *next* fetch of CSRF token. | |
107 | Intended to be used for oneliners like this one: | |
108 | ||
109 | connection.tokenFrom('somewhere').delete(...) | |
110 | ||
111 | where the token comes from "somewhere" instead of the | |
112 | default stream page. | |
113 | """ | |
114 | self._fetch_token_from = location | |
115 | return self | |
116 | ||
117 | def post(self, string, data, headers={}, params={}, **kwargs): | |
118 | """This method posts data to session. | |
119 | Performs additional checks if needed. | |
120 | ||
121 | Example: | |
122 | To post to 'foo' one should call `post('foo', data={})`. | |
123 | ||
124 | :param string: URL to post without the pod's URL and slash eg. 'status_messages'. | |
125 | :type string: str | |
126 | :param data: Data to post. | |
127 | :param headers: Headers (optional). | |
128 | :type headers: dict | |
129 | :param params: Parameters (optional). | |
130 | :type params: dict | |
131 | """ | |
132 | string = '{0}/{1}'.format(self.pod, string) | |
133 | if 'X-CSRF-Token' not in headers: | |
134 | headers['X-CSRF-Token'] = self.get_token() | |
8bfe79c7 | 135 | if not kwargs: kwargs = self._requests_kwargs |
b9d087bb | 136 | request = self._session.post(string, data, headers=headers, params=params, **kwargs) |
d95ff94a C |
137 | return request |
138 | ||
139 | def put(self, string, data=None, headers={}, params={}, **kwargs): | |
140 | """This method PUTs to session. | |
141 | """ | |
142 | string = '{0}/{1}'.format(self.pod, string) | |
143 | if 'X-CSRF-Token' not in headers: | |
144 | headers['X-CSRF-Token'] = self.get_token() | |
8bfe79c7 | 145 | if not kwargs: kwargs = self._requests_kwargs |
d95ff94a | 146 | if data is not None: request = self._session.put(string, data, headers=headers, params=params, **kwargs) |
b9d087bb | 147 | else: request = self._session.put(string, headers=headers, params=params, **kwargs) |
d95ff94a C |
148 | return request |
149 | ||
150 | def delete(self, string, data = None, headers={}, **kwargs): | |
151 | """This method lets you send delete request to session. | |
152 | Performs additional checks if needed. | |
153 | ||
154 | :param string: URL to use. | |
155 | :type string: str | |
156 | :param data: Data to use. | |
157 | :param headers: Headers to use (optional). | |
158 | :type headers: dict | |
159 | """ | |
160 | string = '{0}/{1}'.format(self.pod, string) | |
161 | if 'X-CSRF-Token' not in headers: | |
162 | headers['X-CSRF-Token'] = self.get_token() | |
8bfe79c7 | 163 | if not kwargs: kwargs = self._requests_kwargs |
b9d087bb | 164 | request = self._session.delete(string, data=data, headers=headers, **kwargs) |
d95ff94a C |
165 | return request |
166 | ||
8bfe79c7 C |
167 | def _checkCamo(self): |
168 | response = self._session.head("{0}/camo/".format(self.pod), | |
169 | **self._requests_kwargs) | |
170 | if response.status_code == 200: self._camo_enabled = True | |
171 | else: self._camo_enabled = False | |
172 | ||
173 | def camo(self): return self._camo_enabled; | |
174 | ||
d95ff94a C |
175 | def _setlogin(self, username, password): |
176 | """This function is used to set data for login. | |
177 | ||
178 | .. note:: | |
179 | It should be called before _login() function. | |
180 | """ | |
181 | self._login_data = {'user[username]': username, | |
182 | 'user[password]': password, | |
183 | 'authenticity_token': self._fetchtoken()} | |
184 | ||
185 | def _login(self): | |
186 | """Handles actual login request. | |
187 | Raises LoginError if login failed. | |
188 | """ | |
189 | request = self.post('users/sign_in', | |
190 | data=self._login_data, | |
191 | allow_redirects=False) | |
192 | if request.status_code != 302: | |
193 | raise errors.LoginError('{0}: login failed'.format(request.status_code)) | |
8bfe79c7 | 194 | self._checkCamo() |
d95ff94a C |
195 | |
196 | def login(self, remember_me=1): | |
197 | """This function is used to log in to a pod. | |
198 | Will raise LoginError if password or username was not specified. | |
199 | """ | |
200 | if not self._login_data['user[username]'] or not self._login_data['user[password]']: | |
201 | raise errors.LoginError('username and/or password is not specified') | |
202 | self._login_data['user[remember_me]'] = remember_me | |
203 | status = self._login() | |
204 | self._login_data = {} | |
205 | return self | |
206 | ||
207 | def logout(self): | |
208 | """Logs out from a pod. | |
209 | When logged out you can't do anything. | |
210 | """ | |
211 | self.get('users/sign_out') | |
212 | self.token = '' | |
b9d087bb C |
213 | self._userdata = {} |
214 | self._diaspora_session = '' | |
215 | self._camo_enabled = False | |
d95ff94a | 216 | |
74dbe1b3 | 217 | def podswitch(self, pod, username, password, login=True): |
d95ff94a C |
218 | """Switches pod from current to another one. |
219 | """ | |
220 | self.pod = pod | |
221 | self._setlogin(username, password) | |
74dbe1b3 | 222 | if login: self._login() |
d95ff94a C |
223 | |
224 | def _fetchtoken(self): | |
225 | """This method tries to get token string needed for authentication on D*. | |
226 | ||
227 | :returns: token string | |
228 | """ | |
229 | request = self.get(self._fetch_token_from) | |
230 | token = self._token_regex.search(request.text) | |
231 | if token is None: token = self._token_regex_2.search(request.text) | |
232 | if token is not None: token = token.group(1) | |
233 | else: raise errors.TokenError('could not find valid CSRF token') | |
234 | self._token = token | |
235 | self._fetch_token_from = 'stream' | |
236 | return token | |
237 | ||
238 | def get_token(self, fetch=True): | |
239 | """This function returns a token needed for authentication in most cases. | |
240 | **Notice:** using repr() is recommended method for getting token. | |
241 | ||
242 | Each time it is run a _fetchtoken() is called and refreshed token is stored. | |
243 | ||
244 | It is more safe to use than _fetchtoken(). | |
245 | By setting new you can request new token or decide to get stored one. | |
246 | If no token is stored new one will be fetched anyway. | |
247 | ||
248 | :returns: string -- token used to authenticate | |
249 | """ | |
250 | try: | |
251 | if fetch or not self._token: self._fetchtoken() | |
252 | except requests.exceptions.ConnectionError as e: | |
253 | warnings.warn('{0} was cought: reusing old token'.format(e)) | |
254 | finally: | |
255 | if not self._token: raise errors.TokenError('cannot obtain token and no previous token found for reuse') | |
256 | return self._token | |
257 | ||
258 | def getSessionToken(self): | |
259 | """Returns session token string (_diaspora_session). | |
260 | """ | |
261 | return self._diaspora_session | |
262 | ||
8bfe79c7 | 263 | def userdata(self): return self._userdata |
d95ff94a C |
264 | |
265 | def getUserData(self): | |
266 | """Returns user data. | |
267 | """ | |
268 | request = self.get('bookmarklet') | |
269 | userdata = self._userinfo_regex.search(request.text) | |
270 | if userdata is None: userdata = self._userinfo_regex_2.search(request.text) | |
271 | if userdata is None: raise errors.DiaspyError('cannot find user data') | |
272 | userdata = userdata.group(1) | |
273 | self._userdata = json.loads(userdata) | |
274 | return self._userdata | |
275 | ||
276 | def set_verify_SSL(self, verify): | |
277 | """Sets whether there should be an error if a SSL-Certificate could not be verified. | |
278 | """ | |
279 | self._verify_SSL = verify | |
b9d087bb | 280 | self._requests_kwargs.update({'verify':verify}) |