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.')
  28         _verify_SSL = True
  29
  30         def __init__(self, pod, username, password, schema='https', **requestsKwargs):
  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                 :param requestsKwargs: default kwargs for requests (proxy, timeout, etc)
  39                 :type requestsKwargs: keyworded arguments
  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 = ''
  47                 self._fetch_token_from = 'stream'
  48                 self._requests_kwargs = {'verify':self._verify_SSL}
  49                 if requestsKwargs: self._requests_kwargs.update(requestsKwargs)
  50
  51                 self._camo_enabled = False
  52                 try: self._setlogin(username, password)
  53                 except requests.exceptions.MissingSchema:
  54                         self.pod = '{0}://{1}'.format(schema, self.pod)
  55                         warnings.warn('schema was missing')
  56
  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()
  61
  62         def __bool__(self):
  63                 if self._token: return True
  64                 return False
  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 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
  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
 102                 if not kwargs: kwargs = self._requests_kwargs
 103                 return self._session.get(url, params=params, headers=headers, **kwargs)
 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()
 135                 if not kwargs: kwargs = self._requests_kwargs
 136                 request = self._session.post(string, data, headers=headers, params=params, **kwargs)
 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()
 145                 if not kwargs: kwargs = self._requests_kwargs
 146                 if data is not None: request = self._session.put(string, data, headers=headers, params=params, **kwargs)
 147                 else: request = self._session.put(string, headers=headers, params=params, **kwargs)
 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()
 163                 if not kwargs: kwargs = self._requests_kwargs
 164                 request = self._session.delete(string, data=data, headers=headers, **kwargs)
 165                 return request
 166
 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
 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))
 194                 self._checkCamo()
 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 = ''
 213                 self._userdata = {}
 214                 self._diaspora_session = ''
 215                 self._camo_enabled = False
 216
 217         def podswitch(self, pod, username, password, login=True):
 218                 """Switches pod from current to another one.
 219                 """
 220                 self.pod = pod
 221                 self._setlogin(username, password)
 222                 if login: self._login()
 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
 263         def userdata(self): return self._userdata
 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
 280                 self._requests_kwargs.update({'verify':verify})