diaspy/connection.py

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