#!/usr/bin/env python import re import requests import json import warnings """This module abstracts connection to pod. """ class LoginError(Exception): pass class TokenError(Exception): pass class Connection(): """Object representing connection with the pod. """ _token_regex = re.compile(r'content="(.*?)"\s+name="csrf-token') _userinfo_regex = re.compile(r'window.current_user_attributes = ({.*})') def __init__(self, pod, username='', password='', schema='https'): """ :param pod: The complete url of the diaspora pod to use. :type pod: str :param username: The username used to log in. :type username: str :param password: The password used to log in. :type password: str """ self.pod = pod self.session = requests.Session() self.login_data = {} self.token = '' try: self._setlogin(username, password) except requests.exceptions.MissingSchema: self.pod = '{0}://{1}'.format(schema, self.pod) warnings.warn('schema was missing') finally: pass try: self._setlogin(username, password) except Exception as e: raise LoginError('cannot create login data (caused by: {0})'.format(e)) def __repr__(self): """Returns token string. It will be easier to change backend if programs will just use: repr(connection) instead of calling a specified method. """ return self.get_token() def get(self, string, headers={}, params={}): """This method gets data from session. Performs additional checks if needed. Example: To obtain 'foo' from pod one should call `get('foo')`. :param string: URL to get without the pod's URL and slash eg. 'stream'. :type string: str """ return self.session.get('{0}/{1}'.format(self.pod, string), params=params, headers=headers) def post(self, string, data, headers={}, params={}): """This method posts data to session. Performs additional checks if needed. Example: To post to 'foo' one should call `post('foo', data={})`. :param string: URL to post without the pod's URL and slash eg. 'status_messages'. :type string: str :param data: Data to post. :param headers: Headers (optional). :type headers: dict :param params: Parameters (optional). :type params: dict """ string = '{0}/{1}'.format(self.pod, string) request = self.session.post(string, data, headers=headers, params=params) return request def put(self, string, data=None, headers={}, params={}): """This method PUTs to session. """ string = '{0}/{1}'.format(self.pod, string) if data is not None: request = self.session.put(string, data, headers=headers, params=params) else: request = self.session.put(string, headers=headers, params=params) return request def delete(self, string, data, headers={}): """This method lets you send delete request to session. Performs additional checks if needed. :param string: URL to use. :type string: str :param data: Data to use. :param headers: Headers to use (optional). :type headers: dict """ string = '{0}/{1}'.format(self.pod, string) request = self.session.delete(string, data=data, headers=headers) return request def _setlogin(self, username, password): """This function is used to set data for login. .. note:: It should be called before _login() function. """ self.username, self.password = username, password self.login_data = {'user[username]': self.username, 'user[password]': self.password, 'authenticity_token': self._fetchtoken()} def _login(self): """Handles actual login request. Raises LoginError if login failed. """ request = self.post('users/sign_in', data=self.login_data, headers={'accept': 'application/json'}) if request.status_code != 201: raise LoginError('{0}: login failed'.format(request.status_code)) def login(self, username='', password=''): """This function is used to log in to a pod. Will raise LoginError if password or username was not specified. """ if username and password: self._setlogin(username, password) if not self.username or not self.password: raise LoginError('password or username not specified') self._login() def logout(self): """Logs out from a pod. When logged out you can't do anything. """ self.get('users/sign_out') self.username = '' self.token = '' self.password = '' def podswitch(self, pod): """Switches pod from current to another one. """ self.pod = pod self._login() def getUserInfo(self): """This function returns the current user's attributes. :returns: dict -- json formatted user info. """ request = self.get('bookmarklet') try: userdata = json.loads(self._userinfo_regex.search(request.text).group(1)) except AttributeError: raise errors.DiaspyError('cannot find user data') return userdata def _fetchtoken(self): """This method tries to get token string needed for authentication on D*. :returns: token string """ request = self.get('stream') token = self._token_regex.search(request.text).group(1) self.token = token return token def get_token(self, fetch=False): """This function returns a token needed for authentication in most cases. Each time it is run a _fetchtoken() is called and refreshed token is stored. It is more safe to use than _fetchtoken(). By setting new you can request new token or decide to get stored one. If no token is stored new one will be fatched anyway. :returns: string -- token used to authenticate """ try: if fetch: self._fetchtoken() if not self.token: self._fetchtoken() except requests.exceptions.ConnectionError as e: warnings.warn('{0} was cought: reusing old token'.format(e)) finally: if not self.token: raise TokenError('cannot obtain token and no previous token found for reuse') return self.token