diaspy/connection.py

   1 #!/usr/bin/env python
   2
   3
   4 """This module abstracts connection to pod.
   5 """
   6
   7
   8 import json
   9 import re
  10 import requests
  11 import warnings
  12
  13 from diaspy import errors
  14
  15
  16 DEBUG = True
  17
  18
  19 class Connection():
  20     """Object representing connection with the pod.
  21     """
  22     _token_regex = re.compile(r'content="(.*?)"\s+name="csrf-token')
  23     _userinfo_regex = re.compile(r'window.current_user_attributes = ({.*})')
  24     _userinfo_regex_2 = re.compile(r'gon.user=({.*});gon.preloads')
  25
  26     def __init__(self, pod, username, password, schema='https'):
  27         """
  28         :param pod: The complete url of the diaspora pod to use.
  29         :type pod: str
  30         :param username: The username used to log in.
  31         :type username: str
  32         :param password: The password used to log in.
  33         :type password: str
  34         """
  35         self.pod = pod
  36         self._session = requests.Session()
  37         self._login_data = {'user[remember_me]': 1, 'utf8': '✓'}
  38         self._userdata = {}
  39         self._token = ''
  40         self._diaspora_session = ''
  41         self._cookies = self._fetchcookies()
  42         try:
  43             #self._setlogin(username, password)
  44             self._login_data = {'user[username]': username,
  45                                 'user[password]': password,
  46                                 'authenticity_token': self._fetchtoken()}
  47             success = True
  48         except requests.exceptions.MissingSchema:
  49             self.pod = '{0}://{1}'.format(schema, self.pod)
  50             warnings.warn('schema was missing')
  51             success = False
  52         finally:
  53             pass
  54         try:
  55             if not success:
  56                 self._login_data = {'user[username]': username,
  57                                     'user[password]': password,
  58                                     'authenticity_token': self._fetchtoken()}
  59         except Exception as e:
  60             raise errors.LoginError('cannot create login data (caused by: {0})'.format(e))
  61
  62     def _fetchcookies(self):
  63         request = self.get('stream')
  64         return request.cookies
  65
  66     def __repr__(self):
  67         """Returns token string.
  68         It will be easier to change backend if programs will just use:
  69             repr(connection)
  70         instead of calling a specified method.
  71         """
  72         return self._fetchtoken()
  73
  74     def get(self, string, headers={}, params={}, direct=False, **kwargs):
  75         """This method gets data from session.
  76         Performs additional checks if needed.
  77
  78         Example:
  79             To obtain 'foo' from pod one should call `get('foo')`.
  80
  81         :param string: URL to get without the pod's URL and slash eg. 'stream'.
  82         :type string: str
  83         :param direct: if passed as True it will not be expanded
  84         :type direct: bool
  85         """
  86         if not direct: url = '{0}/{1}'.format(self.pod, string)
  87         else: url = string
  88         return self._session.get(url, params=params, headers=headers, **kwargs)
  89
  90     def post(self, string, data, headers={}, params={}, **kwargs):
  91         """This method posts data to session.
  92         Performs additional checks if needed.
  93
  94         Example:
  95             To post to 'foo' one should call `post('foo', data={})`.
  96
  97         :param string: URL to post without the pod's URL and slash eg. 'status_messages'.
  98         :type string: str
  99         :param data: Data to post.
 100         :param headers: Headers (optional).
 101         :type headers: dict
 102         :param params: Parameters (optional).
 103         :type params: dict
 104         """
 105         string = '{0}/{1}'.format(self.pod, string)
 106         request = self._session.post(string, data, headers=headers, params=params, **kwargs)
 107         return request
 108
 109     def put(self, string, data=None, headers={}, params={}, **kwargs):
 110         """This method PUTs to session.
 111         """
 112         string = '{0}/{1}'.format(self.pod, string)
 113         if data is not None: request = self._session.put(string, data, headers=headers, params=params, **kwargs)
 114         else: request = self._session.put(string, headers=headers, params=params, **kwargs)
 115         return request
 116
 117     def delete(self, string, data, headers={}, **kwargs):
 118         """This method lets you send delete request to session.
 119         Performs additional checks if needed.
 120
 121         :param string: URL to use.
 122         :type string: str
 123         :param data: Data to use.
 124         :param headers: Headers to use (optional).
 125         :type headers: dict
 126         """
 127         string = '{0}/{1}'.format(self.pod, string)
 128         request = self._session.delete(string, data=data, headers=headers, **kwargs)
 129         return request
 130
 131     def _setlogin(self, username, password):
 132         """This function is used to set data for login.
 133
 134         .. note::
 135             It should be called before _login() function.
 136         """
 137         self._login_data = {'user[username]': username,
 138                             'user[password]': password,
 139                             'authenticity_token': self._fetchtoken()}
 140
 141     def _login(self):
 142         """Handles actual login request.
 143         Raises LoginError if login failed.
 144         """
 145         request = self.post('users/sign_in',
 146                             data=self._login_data,
 147                             allow_redirects=False)
 148         if request.status_code != 302:
 149             raise errors.LoginError('{0}: login failed'.format(request.status_code))
 150         return request.status_code
 151
 152     def login(self, remember_me=1):
 153         """This function is used to log in to a pod.
 154         Will raise LoginError if password or username was not specified.
 155         """
 156         if not self._login_data['user[username]'] or not self._login_data['user[password]']:
 157             raise errors.LoginError('username and/or password is not specified')
 158         self._login_data['user[remember_me]'] = remember_me
 159         status = self._login()
 160         self._login_data = {}
 161         return status
 162
 163     def logout(self):
 164         """Logs out from a pod.
 165         When logged out you can't do anything.
 166         """
 167         self.get('users/sign_out')
 168         self.token = ''
 169
 170     def podswitch(self, pod, username, password):
 171         """Switches pod from current to another one.
 172         """
 173         self.pod = pod
 174         self._setlogin(username, password)
 175         self._login()
 176
 177     def getUserInfo(self, fetch=False):
 178         """This function returns the current user's attributes.
 179
 180         :returns: dict -- json formatted user info.
 181         """
 182         request = self.get('bookmarklet')
 183         userdata = self._userinfo_regex.search(request.text)
 184         if userdata is None: userdata = self._userinfo_regex_2.search(request.text)
 185         if userdata is None: raise errors.DiaspyError('cannot find user data')
 186         userdata = userdata.group(1)
 187         return json.loads(userdata)
 188
 189     def _fetchtoken(self):
 190         """This method tries to get token string needed for authentication on D*.
 191
 192         :returns: token string
 193         """
 194         request = self.get('stream')
 195         token = self._token_regex.search(request.text).group(1)
 196         self._token = token
 197         return token
 198
 199     def get_token(self, fetch=True):
 200         """This function returns a token needed for authentication in most cases.
 201         **Notice:** using repr() is recommended method for getting token.
 202
 203         Each time it is run a _fetchtoken() is called and refreshed token is stored.
 204
 205         It is more safe to use than _fetchtoken().
 206         By setting new you can request new token or decide to get stored one.
 207         If no token is stored new one will be fetched anyway.
 208
 209         :returns: string -- token used to authenticate
 210         """
 211         try:
 212             if fetch or not self._token: self._fetchtoken()
 213         except requests.exceptions.ConnectionError as e:
 214             warnings.warn('{0} was cought: reusing old token'.format(e))
 215         finally:
 216             if not self._token: raise errors.TokenError('cannot obtain token and no previous token found for reuse')
 217         return self._token
 218
 219     def getSessionToken(self):
 220         """Returns session token string (_diaspora_session).
 221         """
 222         return self._diaspora_session