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'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.preloads')
  28     _verify_SSL = True
  29
  30     def __init__(self, pod, username, password, schema='https'):
  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
  38         """
  39         self.pod = pod
  40         self._session = requests.Session()
  41         self._login_data = {'user[remember_me]': 1, 'utf8': '✓'}
  42         self._userdata = {}
  43         self._token = ''
  44         self._diaspora_session = ''
  45         self._cookies = self._fetchcookies()
  46         try:
  47             #self._setlogin(username, password)
  48             self._login_data = {'user[username]': username,
  49                                 'user[password]': password,
  50                                 'authenticity_token': self._fetchtoken()}
  51             success = True
  52         except requests.exceptions.MissingSchema:
  53             self.pod = '{0}://{1}'.format(schema, self.pod)
  54             warnings.warn('schema was missing')
  55             success = False
  56         finally:
  57             pass
  58         try:
  59             if not success:
  60                 self._login_data = {'user[username]': username,
  61                                     'user[password]': password,
  62                                     'authenticity_token': self._fetchtoken()}
  63         except Exception as e:
  64             raise errors.LoginError('cannot create login data (caused by: {0})'.format(e))
  65
  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
  78     def get(self, string, headers={}, params={}, direct=False, **kwargs):
  79         """This method gets data from session.
  80         Performs additional checks if needed.
  81
  82         Example:
  83             To obtain 'foo' from pod one should call `get('foo')`.
  84
  85         :param string: URL to get without the pod's URL and slash eg. 'stream'.
  86         :type string: str
  87         :param direct: if passed as True it will not be expanded
  88         :type direct: bool
  89         """
  90         if not direct: url = '{0}/{1}'.format(self.pod, string)
  91         else: url = string
  92         return self._session.get(url, params=params, headers=headers, verify=self._verify_SSL, **kwargs)
  93
  94     def post(self, string, data, headers={}, params={}, **kwargs):
  95         """This method posts data to session.
  96         Performs additional checks if needed.
  97
  98         Example:
  99             To post to 'foo' one should call `post('foo', data={})`.
 100
 101         :param string: URL to post without the pod's URL and slash eg. 'status_messages'.
 102         :type string: str
 103         :param data: Data to post.
 104         :param headers: Headers (optional).
 105         :type headers: dict
 106         :param params: Parameters (optional).
 107         :type params: dict
 108         """
 109         string = '{0}/{1}'.format(self.pod, string)
 110         request = self._session.post(string, data, headers=headers, params=params, verify=self._verify_SSL, **kwargs)
 111         return request
 112
 113     def put(self, string, data=None, headers={}, params={}, **kwargs):
 114         """This method PUTs to session.
 115         """
 116         string = '{0}/{1}'.format(self.pod, string)
 117         if data is not None: request = self._session.put(string, data, headers=headers, params=params, **kwargs)
 118         else: request = self._session.put(string, headers=headers, params=params, verify=self._verify_SSL, **kwargs)
 119         return request
 120
 121     def delete(self, string, data, headers={}, **kwargs):
 122         """This method lets you send delete request to session.
 123         Performs additional checks if needed.
 124
 125         :param string: URL to use.
 126         :type string: str
 127         :param data: Data to use.
 128         :param headers: Headers to use (optional).
 129         :type headers: dict
 130         """
 131         string = '{0}/{1}'.format(self.pod, string)
 132         request = self._session.delete(string, data=data, headers=headers, **kwargs)
 133         return request
 134
 135     def _setlogin(self, username, password):
 136         """This function is used to set data for login.
 137
 138         .. note::
 139             It should be called before _login() function.
 140         """
 141         self._login_data = {'user[username]': username,
 142                             'user[password]': password,
 143                             'authenticity_token': self._fetchtoken()}
 144
 145     def _login(self):
 146         """Handles actual login request.
 147         Raises LoginError if login failed.
 148         """
 149         request = self.post('users/sign_in',
 150                             data=self._login_data,
 151                             allow_redirects=False)
 152         if request.status_code != 302:
 153             raise errors.LoginError('{0}: login failed'.format(request.status_code))
 154
 155     def login(self, remember_me=1):
 156         """This function is used to log in to a pod.
 157         Will raise LoginError if password or username was not specified.
 158         """
 159         if not self._login_data['user[username]'] or not self._login_data['user[password]']:
 160             raise errors.LoginError('username and/or password is not specified')
 161         self._login_data['user[remember_me]'] = remember_me
 162         status = self._login()
 163         self._login_data = {}
 164         return self
 165
 166     def logout(self):
 167         """Logs out from a pod.
 168         When logged out you can't do anything.
 169         """
 170         self.get('users/sign_out')
 171         self.token = ''
 172
 173     def podswitch(self, pod, username, password):
 174         """Switches pod from current to another one.
 175         """
 176         self.pod = pod
 177         self._setlogin(username, password)
 178         self._login()
 179
 180     def _fetchtoken(self):
 181         """This method tries to get token string needed for authentication on D*.
 182
 183         :returns: token string
 184         """
 185         request = self.get('stream')
 186         token = self._token_regex.search(request.text)
 187         if token is None: self._token_regex_2.search(request.text)
 188         if token is not None: token = token.group(1)
 189         else: raise errors.TokenError('could not find valid CSRF token')
 190         self._token = token
 191         return token
 192
 193     def get_token(self, fetch=True):
 194         """This function returns a token needed for authentication in most cases.
 195         **Notice:** using repr() is recommended method for getting token.
 196
 197         Each time it is run a _fetchtoken() is called and refreshed token is stored.
 198
 199         It is more safe to use than _fetchtoken().
 200         By setting new you can request new token or decide to get stored one.
 201         If no token is stored new one will be fetched anyway.
 202
 203         :returns: string -- token used to authenticate
 204         """
 205         try:
 206             if fetch or not self._token: self._fetchtoken()
 207         except requests.exceptions.ConnectionError as e:
 208             warnings.warn('{0} was cought: reusing old token'.format(e))
 209         finally:
 210             if not self._token: raise errors.TokenError('cannot obtain token and no previous token found for reuse')
 211         return self._token
 212
 213     def getSessionToken(self):
 214         """Returns session token string (_diaspora_session).
 215         """
 216         return self._diaspora_session
 217
 218     def getUserData(self):
 219         """Returns user data.
 220         """
 221         request = self.get('bookmarklet')
 222         userdata = self._userinfo_regex.search(request.text)
 223         if userdata is None: userdata = self._userinfo_regex_2.search(request.text)
 224         if userdata is None: raise errors.DiaspyError('cannot find user data')
 225         userdata = userdata.group(1)
 226         return json.loads(userdata)
 227
 228     def set_verify_SSL(self, verify):
 229         """Sets whether there should be an error if a SSL-Certificate could not be verified.
 230         """
 231         self._verify_SSL = verify